mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-17 21:25:46 +02:00
293 lines
11 KiB
Python
293 lines
11 KiB
Python
"""Synchronous shared-stream fan-out controller for v3 thread streaming."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import logging
|
|
import random
|
|
import threading
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from queue import Queue as _Queue
|
|
from typing import Any
|
|
|
|
from langchain_protocol import Event, SubscribeParams
|
|
|
|
from langgraph_sdk.stream.subscription import compute_union_filter, filter_covers
|
|
from langgraph_sdk.stream.transport.sync_http import (
|
|
SyncEventStreamHandle,
|
|
SyncProtocolSseTransport,
|
|
)
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class _SyncSubscription:
|
|
id: int
|
|
params: SubscribeParams
|
|
queue: _Queue[Event | None] = field(default_factory=_Queue)
|
|
# Why: using `queue.Queue` in the annotation causes ty to resolve `queue`
|
|
# as the field being defined (name shadowing), not the stdlib module.
|
|
|
|
|
|
_DEFAULT_RUN_START_TIMEOUT: float = 30.0
|
|
|
|
|
|
class SyncStreamController:
|
|
"""Owns the sync shared SSE handle, subscription registry, and fan-out thread."""
|
|
|
|
def __init__(
|
|
self,
|
|
transport: SyncProtocolSseTransport,
|
|
*,
|
|
run_start_gate: threading.Event | None = None,
|
|
run_start_timeout: float = _DEFAULT_RUN_START_TIMEOUT,
|
|
max_reconnect_attempts: int = 5,
|
|
reconnect_backoff_base: float = 0.1,
|
|
reconnect_backoff_cap: float = 10.0,
|
|
) -> None:
|
|
self._transport = transport
|
|
self._next_subscription_id = 1
|
|
self._subscriptions: dict[int, _SyncSubscription] = {}
|
|
self._seen_event_ids: set[str] = set()
|
|
self._shared_stream: SyncEventStreamHandle | None = None
|
|
self._shared_stream_filter: dict[str, Any] | None = None
|
|
self._fanout_thread: threading.Thread | None = None
|
|
self._closed = False
|
|
self._lock = threading.RLock()
|
|
self._cursor: int | None = None
|
|
# When None, no gate is applied and reconcile_stream proceeds immediately.
|
|
# SyncThreadStream passes an un-set Event so subscriptions wait until
|
|
# run.start completes.
|
|
self._run_start_gate = run_start_gate
|
|
self._run_start_timeout = run_start_timeout
|
|
self._max_reconnect_attempts = max_reconnect_attempts
|
|
self._reconnect_backoff_base = reconnect_backoff_base
|
|
self._reconnect_backoff_cap = reconnect_backoff_cap
|
|
self._drain_threads: set[threading.Thread] = set()
|
|
|
|
def register_subscription(self, params: SubscribeParams) -> _SyncSubscription:
|
|
with self._lock:
|
|
sub = _SyncSubscription(id=self._next_subscription_id, params=params)
|
|
self._next_subscription_id += 1
|
|
self._subscriptions[sub.id] = sub
|
|
return sub
|
|
|
|
def unregister_subscription(self, subscription_id: int) -> None:
|
|
with self._lock:
|
|
self._subscriptions.pop(subscription_id, None)
|
|
|
|
def reconcile_stream(self, candidate_filter: SubscribeParams) -> None:
|
|
if self._run_start_gate is not None and not self._run_start_gate.wait(
|
|
timeout=self._run_start_timeout
|
|
):
|
|
raise TimeoutError("Sync run.start gate timeout.")
|
|
with self._lock:
|
|
if (
|
|
self._shared_stream is not None
|
|
and self._shared_stream_filter is not None
|
|
and filter_covers(self._shared_stream_filter, dict(candidate_filter))
|
|
):
|
|
return
|
|
new_filter = self._compute_current_union(extra=candidate_filter)
|
|
old_stream = self._shared_stream
|
|
self._shared_stream = self._transport.open_event_stream(
|
|
self._filter_with_since(new_filter)
|
|
)
|
|
self._shared_stream_filter = new_filter
|
|
if old_stream is not None:
|
|
drain_thread = threading.Thread(
|
|
target=self._drain_and_close,
|
|
args=(old_stream,),
|
|
daemon=True,
|
|
name="langgraph-sdk-sync-rotation-drain",
|
|
)
|
|
self._drain_threads.add(drain_thread)
|
|
drain_thread.start()
|
|
|
|
def ensure_fanout_running(self) -> None:
|
|
with self._lock:
|
|
if self._fanout_thread is not None and self._fanout_thread.is_alive():
|
|
return
|
|
self._fanout_thread = threading.Thread(
|
|
target=self._fanout,
|
|
name="langgraph-sdk-sync-stream-fanout",
|
|
daemon=True,
|
|
)
|
|
self._fanout_thread.start()
|
|
|
|
def _fanout(self) -> None:
|
|
from langgraph_sdk.stream.subscription import matches_subscription
|
|
|
|
while True:
|
|
with self._lock:
|
|
if self._closed:
|
|
return
|
|
shared = self._shared_stream
|
|
if shared is None:
|
|
return
|
|
try:
|
|
for event in self._dedup_iter(shared.events):
|
|
with self._lock:
|
|
if self._closed:
|
|
break
|
|
subscriptions = list(self._subscriptions.values())
|
|
for sub in subscriptions:
|
|
if matches_subscription(event, sub.params):
|
|
sub.queue.put(event)
|
|
except Exception:
|
|
pass # transport drop — attempt reconnect below
|
|
|
|
with self._lock:
|
|
if self._shared_stream is not shared:
|
|
continue # rotation happened; pick up new stream
|
|
# No rotation — check if this was a transport drop
|
|
err = shared.error()
|
|
if err is not None and not self._closed:
|
|
if self._reconnect_shared_stream():
|
|
continue
|
|
break
|
|
|
|
with self._lock:
|
|
for sub in self._subscriptions.values():
|
|
sub.queue.put(None)
|
|
|
|
def _reconnect_shared_stream(self) -> bool:
|
|
with self._lock:
|
|
base_filter = self._shared_stream_filter
|
|
if base_filter is None:
|
|
return False
|
|
for _ in range(self._max_reconnect_attempts):
|
|
with self._lock:
|
|
if self._closed:
|
|
return False
|
|
params = self._filter_with_since(base_filter)
|
|
try:
|
|
new_stream = self._transport.open_event_stream(params)
|
|
except Exception:
|
|
continue
|
|
with self._lock:
|
|
self._shared_stream = new_stream
|
|
return True
|
|
return False
|
|
|
|
def _compute_current_union(
|
|
self, extra: SubscribeParams | None = None
|
|
) -> dict[str, Any]:
|
|
filters = [dict(sub.params) for sub in self._subscriptions.values()]
|
|
if extra is not None:
|
|
filters.append(dict(extra))
|
|
return compute_union_filter(filters)
|
|
|
|
def observe_applied_through_seq(self, seq: Any) -> None:
|
|
"""Advance the reconnect cursor from a command response meta sequence."""
|
|
with self._lock:
|
|
self._observe_seq(seq)
|
|
|
|
def _observe_event(self, event: Event) -> None:
|
|
with self._lock:
|
|
self._observe_seq(event.get("seq"))
|
|
|
|
def _observe_seq(self, seq: Any) -> None:
|
|
if isinstance(seq, int) and (self._cursor is None or seq > self._cursor):
|
|
self._cursor = seq
|
|
|
|
def _filter_with_since(self, params: dict[str, Any]) -> dict[str, Any]:
|
|
out = dict(params)
|
|
if self._cursor is not None:
|
|
out["since"] = self._cursor
|
|
return out
|
|
|
|
def _dedup_iter(self, source: Any) -> Any:
|
|
for event in source:
|
|
event_id = event.get("event_id")
|
|
if event_id is not None:
|
|
if event_id in self._seen_event_ids:
|
|
continue
|
|
self._seen_event_ids.add(event_id)
|
|
self._observe_event(event)
|
|
yield event
|
|
|
|
def _drain_and_close(self, handle: SyncEventStreamHandle) -> None:
|
|
"""Drain remaining events from an old handle before closing it.
|
|
|
|
Runs in a background thread spawned by `reconcile_stream` on rotation
|
|
so buffered events are not lost when the shared stream is replaced.
|
|
Events are dispatched to subscribers regardless of `_closed` so that
|
|
already-buffered events reach consumers before the handle is closed.
|
|
"""
|
|
from langgraph_sdk.stream.subscription import matches_subscription
|
|
|
|
try:
|
|
for event in self._dedup_iter(handle.events):
|
|
with self._lock:
|
|
subscriptions = list(self._subscriptions.values())
|
|
for sub in subscriptions:
|
|
if matches_subscription(event, sub.params):
|
|
sub.queue.put(event)
|
|
except Exception as err:
|
|
_logger.debug("rotation drain exception: %r", err)
|
|
finally:
|
|
with contextlib.suppress(Exception):
|
|
handle.close()
|
|
with self._lock:
|
|
self._drain_threads.discard(threading.current_thread())
|
|
|
|
def _reconnect_sleep(self, attempt: int) -> None:
|
|
"""Sleep with exponential backoff + jitter before a reconnect attempt."""
|
|
base = self._reconnect_backoff_base
|
|
cap = self._reconnect_backoff_cap
|
|
delay = min(cap, base * (2**attempt))
|
|
jitter = random.uniform(0, delay * 0.25)
|
|
time.sleep(delay + jitter)
|
|
|
|
def _reconnect_shared_stream(self) -> bool:
|
|
"""Attempt to reopen the shared stream after a transport drop.
|
|
|
|
Returns True if a new stream was successfully opened, False if all
|
|
reconnect attempts were exhausted or the controller was closed.
|
|
"""
|
|
base_filter = self._shared_stream_filter
|
|
if base_filter is None:
|
|
return False
|
|
for attempt in range(self._max_reconnect_attempts):
|
|
if self._closed:
|
|
return False
|
|
if attempt > 0:
|
|
self._reconnect_sleep(attempt - 1)
|
|
try:
|
|
new_handle = self._transport.open_event_stream(
|
|
self._filter_with_since(base_filter)
|
|
)
|
|
old = self._shared_stream
|
|
self._shared_stream = new_handle
|
|
if old is not None:
|
|
with contextlib.suppress(Exception):
|
|
old.close()
|
|
return True
|
|
except Exception as err:
|
|
_logger.debug("sync reconnect attempt %d failed: %r", attempt, err)
|
|
return False
|
|
|
|
def close(self) -> None:
|
|
with self._lock:
|
|
if self._closed:
|
|
return
|
|
self._closed = True
|
|
shared = self._shared_stream
|
|
for sub in self._subscriptions.values():
|
|
sub.queue.put(None)
|
|
if shared is not None:
|
|
shared.close()
|
|
thread = self._fanout_thread
|
|
if thread is not None and thread.is_alive():
|
|
with contextlib.suppress(RuntimeError):
|
|
thread.join(timeout=1.0)
|
|
with self._lock:
|
|
drain_threads = set(self._drain_threads)
|
|
for drain in drain_threads:
|
|
if drain.is_alive():
|
|
with contextlib.suppress(RuntimeError):
|
|
drain.join(timeout=1.0)
|