"""Async thread-centric streaming surface for the v3 protocol. `AsyncThreadStream` is an async context manager that owns a `ProtocolSseTransport` for one thread, dispatches commands (`run.start`, `run.respond`), exposes typed subscriptions over a single shared SSE (`subscribe`, `events`), surfaces lifecycle state (`interrupted`, `interrupts`) via an always-on lifecycle watcher SSE, and provides typed projections (`thread.values`, `thread.messages`, `thread.tool_calls`, `thread.extensions`). Direct port of `libs/sdk/src/client/stream/index.ts`. """ from __future__ import annotations import asyncio import contextlib import random from collections.abc import AsyncGenerator, AsyncIterator, Generator, Mapping from dataclasses import dataclass, field from typing import Any, Literal, TypedDict, cast from langchain_core.language_models.chat_model_stream import AsyncChatModelStream from langchain_protocol import Event, SubscribeParams from langgraph_sdk._async.http import HttpClient from langgraph_sdk.schema import QueryParamTypes from langgraph_sdk.stream.controller import _SeenEventIds from langgraph_sdk.stream.decoders import ( DataDecoder, Decoder, ExtensionsDecoder, MessagesDecoder, SubgraphsDecoder, ToolCallsDecoder, validate_interleave_channels, ) from langgraph_sdk.stream.subscription import compute_union_filter, infer_channel from langgraph_sdk.stream.transport import ( AsyncProtocolTransport, EventStreamHandle, ProtocolSseTransport, ProtocolWebSocketTransport, ) class InterruptPayload(TypedDict): """Payload surfaced when the server requests human input for a thread.""" interrupt_id: str value: Any namespace: list[str] @dataclass class _RunTerminal: """Terminal state record resolved into `_run_done` on lifecycle completion.""" status: Literal["completed", "errored"] error: BaseException | None = None @dataclass class _Subscription: """Internal record for one active subscription on an `AsyncThreadStream`.""" id: int params: SubscribeParams queue: asyncio.Queue = field(default_factory=asyncio.Queue) # Why: asyncio.Queue[Event | None] as a subscript in the field annotation # causes a type error with ty; bare asyncio.Queue is accepted. # All public protocol channels used by the raw `events`/`subscribe` surface. # Typed projections open narrower channel filters on the shared SSE. _ALL_CHANNELS: list[str] = [ "values", "updates", "messages", "tools", "lifecycle", "input", "checkpoints", "tasks", "custom", ] def _exact_namespace_params( channels: list[str], namespace: list[str], ) -> SubscribeParams: return { "channels": channels, "namespaces": [list(namespace)], "depth": 0, } def _event_namespace(params_field: Any) -> list[str]: if not isinstance(params_field, dict): return [] namespace = params_field.get("namespace") or [] return list(namespace) if isinstance(namespace, list) else [] _ROOT_TERMINAL_LIFECYCLE_EVENTS = frozenset({"completed", "failed"}) def _is_root_terminal_lifecycle(event: Any) -> bool: """Return True for a root-namespace lifecycle event marking run end. Matches the wire shape ``{method: "lifecycle", params: {namespace: [], data: {event: "completed" | "failed"}}}``. Subgraph lifecycle events (non-empty namespace) do not terminate the parent run. """ if not isinstance(event, dict): return False if event.get("method") != "lifecycle": return False params = event.get("params") or {} if not isinstance(params, dict): return False if params.get("namespace") or []: return False data = params.get("data") or {} if not isinstance(data, dict): return False return data.get("event") in _ROOT_TERMINAL_LIFECYCLE_EVENTS class _AgentModule: """Assistant graph helpers scoped to one thread stream.""" def __init__(self, owner: AsyncThreadStream) -> None: self._owner = owner async def get_tree( self, *, xray: int | bool = False, headers: Mapping[str, str] | None = None, params: QueryParamTypes | None = None, ) -> dict[str, list[dict[str, Any]]]: if self._owner._closed: raise RuntimeError("AsyncThreadStream is closed.") query_params: dict[str, Any] = {} if xray: query_params["xray"] = xray if params: query_params.update(dict(params)) request_headers = {**self._owner._headers, **dict(headers or {})} return await self._owner._http.get( f"/assistants/{self._owner.assistant_id}/graph", params=query_params, headers=request_headers or None, ) class RunModule: """Command dispatcher for `run.start`. Bound to one `AsyncThreadStream`; accesses its transport and id allocator. """ def __init__(self, owner: AsyncThreadStream) -> None: self._owner = owner async def start( self, *, input: Any = None, config: dict[str, Any] | None = None, metadata: dict[str, Any] | None = None, ) -> dict[str, Any]: """Send `run.start` to the server. Returns the result (`{"run_id": ...}`).""" params: dict[str, Any] = {"assistant_id": self._owner.assistant_id} if input is not None: params["input"] = input if config is not None: params["config"] = config if metadata is not None: params["metadata"] = metadata loop = asyncio.get_running_loop() gate: asyncio.Future[None] = loop.create_future() self._owner._run_start_ready = gate try: result = await self._owner._send_command("run.start", params) if not gate.done(): gate.set_result(None) self._owner._run_seen = True return result except BaseException as err: # Why: gate MUST reject on any exit type, including CancelledError, # so awaiters see the failure rather than hanging indefinitely. if not gate.done(): gate.set_exception(err) raise finally: # Why: concurrent run.start calls (multitask_strategy="enqueue") # can replace _run_start_ready before our finally fires. # Identity-check before clearing so the later call's gate isn't stomped. if self._owner._run_start_ready is gate: self._owner._run_start_ready = None # Why: if the gate stored an exception that no awaiter consumed, # retrieve it here to suppress asyncio's GC warning. The exception # is already propagated to our caller via the `raise` above. if gate.done() and not gate.cancelled(): gate.exception() async def respond( self, response: Any, *, interrupt_id: str | None = None, ) -> dict[str, Any]: """Reply to a server-side interrupt and resume the run. Args: response: the response value forwarded as `params.response` on the wire (protocol field name). interrupt_id: optional explicit id. When omitted, requires exactly one outstanding interrupt and uses its id. Raises: RuntimeError: no outstanding interrupts; `interrupt_id` is None but multiple interrupts are outstanding; or the explicit `interrupt_id` doesn't match any outstanding interrupt. """ # Why: take the `interrupts` snapshot AND dispatch the command under # `_interrupts_lock`, so the lifecycle watcher's terminal-clear path # cannot wipe `interrupts` between the snapshot and the wire send. async with self._owner._interrupts_lock: outstanding = list(self._owner.interrupts) if interrupt_id is None: if len(outstanding) == 0: raise RuntimeError( "thread.run.respond: no outstanding interrupt. Provide " "an explicit `interrupt_id` or wait for " "`thread.interrupted`." ) if len(outstanding) > 1: ids = [p["interrupt_id"] for p in outstanding] raise RuntimeError( f"thread.run.respond: ambiguous — {len(outstanding)} " f"outstanding interrupts ({ids!r}). Provide an explicit " "`interrupt_id`." ) match = outstanding[0] else: match = next( (p for p in outstanding if p["interrupt_id"] == interrupt_id), None, ) if match is None: raise RuntimeError( f"thread.run.respond: interrupt_id {interrupt_id!r} does " "not match any outstanding interrupt in " "`thread.interrupts`." ) params = { "interrupt_id": match["interrupt_id"], "namespace": match["namespace"], "response": response, } return await self._owner._send_command("input.respond", params) async def _close_after(handle: EventStreamHandle, *, delay: float = 0.0) -> None: """Close a handle, optionally after a brief delay. Used to detach closing the old stream from the synchronous rotation step so the new stream can absorb server-side replayed events first. """ if delay: await asyncio.sleep(delay) await handle.close() class _OutputAwaitable: """Awaitable that waits for lifecycle completion then fetches durable thread state. Multiple awaiters share one underlying task (idempotent task caching). Call `with_timeout(seconds)` to bound the wait on the lifecycle terminal. """ def __init__(self, thread: AsyncThreadStream) -> None: self._thread = thread self._task: asyncio.Task[Any] | None = None self._timeout: float | None = None def __await__(self): # type: ignore[override] return self._get_task().__await__() def with_timeout(self, timeout: float) -> _OutputAwaitable: """Return a new awaitable that raises `asyncio.TimeoutError` after `timeout` seconds. Bounds the wait for the lifecycle terminal (and only that wait); the subsequent REST GET for terminal state is not bounded. Returns a fresh `_OutputAwaitable` so the original `thread.output` is unaffected. """ bounded = _OutputAwaitable(self._thread) bounded._timeout = timeout return bounded def _get_task(self) -> asyncio.Task[Any]: """Return the shared fetch task, creating it on first call. A cancelled task is intentionally NOT respawned: subsequent awaiters receive `asyncio.CancelledError` from the shared task instead of triggering a fresh REST GET. This preserves "one fetch per `thread.output`" semantics even when callers wrap awaits with `asyncio.wait_for` (which cancels the underlying task on timeout). """ if self._task is None: self._task = asyncio.create_task(self._fetch()) return self._task async def _fetch(self) -> Any: """Fetch terminal thread state, waiting for the lifecycle if needed.""" # Fast path: explicit thread_id with no run in flight — state may already # be terminal so we can skip the lifecycle wait entirely. if self._thread._can_return_existing_state_immediately(): state = await self._thread._fetch_state() if self._thread._state_is_terminal(state): return state["values"] # Normal path: wait for the lifecycle terminal signal. if self._timeout is not None: terminal = await asyncio.wait_for( self._thread._wait_for_run_done(), timeout=self._timeout ) else: terminal = await self._thread._wait_for_run_done() if terminal.error is not None: raise terminal.error state = await self._thread._fetch_state() return state["values"] class _ValuesProjection: """Typed projection for `thread.values` — yields state snapshots as they arrive. Supports both `async for` (live stream of state snapshots) and `await` (delegates to `thread.output` for the terminal state value). """ def __init__(self, thread: AsyncThreadStream) -> None: self._thread = thread def __await__(self) -> Generator[Any, None, Any]: return self._thread.output.__await__() def __aiter__(self) -> AsyncIterator[Any]: return self._values_iter() async def _values_iter(self) -> AsyncGenerator[Any, None]: if self._thread._transport is None: raise RuntimeError("AsyncThreadStream not entered — use `async with`.") params: SubscribeParams = {"channels": ["values"]} sub = self._thread._register_subscription(params) decoder = DataDecoder("values") try: await self._thread._reconcile_stream(params) self._thread._ensure_fanout_running() state = await self._thread._fetch_state() yield state["values"] while True: item = await sub.queue.get() if item is None: return for out in decoder.feed(item): yield out finally: self._thread._unregister_subscription(sub.id) class _MessagesProjection: """Typed projection for root-scope `thread.messages`. Iterating yields one `AsyncChatModelStream` per message-start event. Each iterator owns its own `messages` subscription and routes events from the root namespace only. """ def __init__( self, thread: AsyncThreadStream, namespace: list[str] | None = None ) -> None: self._thread = thread self._namespace = list(namespace or []) def __aiter__(self) -> AsyncIterator[AsyncChatModelStream]: return self._messages_iter() async def _messages_iter(self) -> AsyncGenerator[AsyncChatModelStream, None]: if self._thread._transport is None: raise RuntimeError("AsyncThreadStream not entered — use `async with`.") # If the subgraphs projection already ran (and consumed messages events # from the shared SSE), drain its root inbox rather than opening a new # subscription. Dedup prevents the SSE from replaying those events. root_inbox = self._thread._root_messages_inbox if not self._namespace else None if root_inbox is not None: async for stream in self._drain_inbox(root_inbox): yield stream return params = _exact_namespace_params(["messages"], self._namespace) sub = self._thread._register_subscription(params) decoder = MessagesDecoder( namespace=self._namespace, stream_factory=lambda *, namespace, node, message_id: AsyncChatModelStream( namespace=namespace, node=node, message_id=message_id ), ) registered: list[AsyncChatModelStream] = [] try: await self._thread._reconcile_stream(params) self._thread._ensure_fanout_running() while True: item = await sub.queue.get() if item is None: return for stream in decoder.feed(item): self._thread._register_active_message_stream(stream) registered.append(stream) yield stream finally: for stream in registered: self._thread._unregister_active_message_stream(stream) self._thread._unregister_subscription(sub.id) async def _drain_inbox( self, inbox: asyncio.Queue[Event | None] ) -> AsyncGenerator[AsyncChatModelStream, None]: """Drain a pre-filled inbox of messages events, yielding one stream per message.""" from langchain_core.language_models.chat_model_stream import ( AsyncChatModelStream, ) active: dict[str, AsyncChatModelStream] = {} try: while True: item = await inbox.get() if item is None: return params_field = item.get("params") or {} data = ( params_field.get("data") if isinstance(params_field, dict) else None ) if not isinstance(data, dict): continue event_type = data.get("event") if event_type == "message-start": message_id = _message_event_id(data) key = _message_route_key(data, fallback=message_id) metadata = ( data.get("metadata") if isinstance(data.get("metadata"), dict) else {} ) stream = AsyncChatModelStream( namespace=list(self._namespace), node=metadata.get("langgraph_node") if metadata else None, message_id=message_id, ) active[key] = stream self._thread._register_active_message_stream(stream) stream.dispatch(data) yield stream else: key = _message_route_key(data) stream = active.get(key) if stream is None and len(active) == 1: stream = next(iter(active.values())) if stream is None: continue stream.dispatch(data) if event_type in ("message-finish", "error"): self._thread._unregister_active_message_stream(stream) for route_key, candidate in list(active.items()): if candidate is stream: del active[route_key] finally: for stream in active.values(): self._thread._unregister_active_message_stream(stream) def _message_event_id(data: dict[str, Any]) -> str | None: message_id = data.get("id") or data.get("message_id") return str(message_id) if message_id is not None else None def _message_route_key(data: dict[str, Any], fallback: str | None = None) -> str: """Return the routing key for a message-channel event in `active`. Keys on `message_id` when available so concurrent messages that share the same `run_id` (two AI turns in one agent step) route to independent streams rather than colliding on a shared `run:` slot. """ message_id = _message_event_id(data) if message_id is not None: return f"message:{message_id}" if fallback is not None: return f"message:{fallback}" return "__single__" SubgraphStatus = Literal["started", "completed", "failed", "interrupted"] def _parse_namespace_segment(segment: str) -> tuple[str, str | None]: name, sep, task_id = segment.partition(":") return name, task_id if sep else None def _terminal_from_tasks_result( data: dict[str, Any], ) -> tuple[SubgraphStatus, str | None]: if data.get("interrupts"): return "interrupted", None error = data.get("error") if error: return "failed", str(error) return "completed", None def _is_direct_child(namespace: list[str], scope: tuple[str, ...]) -> bool: return len(namespace) == len(scope) + 1 and tuple(namespace[: len(scope)]) == scope def _subgraph_subscription_params(scope: tuple[str, ...]) -> SubscribeParams: # Subscribe to tasks + messages + tools + lifecycle without a depth limit # so all descendant-namespace events are captured in one SSE and buffered # into each child handle's inbox. ``lifecycle`` is included so child- # namespace ``started`` events (the canonical signal for # ``create_deep_agent``-style subagent discovery, matching JS behavior) # reach ``_subgraphs_iter``; servers that surface child invocations via # ``tasks`` events instead are also handled via the existing ``method == # "tasks"`` branch. return { "channels": ["messages", "tasks", "tools", "lifecycle"], "namespaces": [list(scope)], } class ScopedStreamHandle: """Scoped streaming handle for one discovered child invocation.""" def __init__( self, *, thread: AsyncThreadStream, path: tuple[str, ...], graph_name: str | None, trigger_call_id: str | None, max_queue_size: int = 0, ) -> None: self._thread = thread self.path = path self.namespace = list(path) self.graph_name = graph_name self.trigger_call_id = trigger_call_id self.status: SubgraphStatus = "started" self.error: str | None = None self._max_queue_size = max_queue_size # Per-channel inboxes: events captured by the parent _SubgraphsProjection # while the SSE was alive. Child projections drain these after the parent # finishes so sequential consumption works without a second SSE open. self._messages_inbox: asyncio.Queue[Event | None] = asyncio.Queue( maxsize=max_queue_size ) self._tools_inbox: asyncio.Queue[Event | None] = asyncio.Queue( maxsize=max_queue_size ) self._tasks_inbox: asyncio.Queue[Event | None] = asyncio.Queue( maxsize=max_queue_size ) # Descendant handles registered by _HandleSubgraphsProjection when a # grandchild is discovered. _push_event fans out to each matching # descendant at dispatch time so events arrive in arrival order without # any drain-and-replay. self._descendant_handles: dict[tuple[str, ...], ScopedStreamHandle] = {} # Track which inboxes have a consumer so _close_inboxes only sends a # sentinel where it is needed. Inboxes with no consumer would otherwise # accumulate a leaked None sentinel that is never drained. self._iterated_inboxes: set[str] = set() self.messages = _HandleMessagesProjection(self) self.tool_calls = _HandleToolCallsProjection(self) self.subgraphs = _HandleSubgraphsProjection(self) self.subagents = self.subgraphs self.extensions = _ExtensionsProjection(thread, namespace=list(path)) def _push_event(self, event: Event) -> None: """Route a descendant event into the appropriate channel inbox. Also fans out to any registered descendant handles whose path is a prefix of the event namespace, so grandchild events are delivered at push time rather than via a post-hoc drain-and-replay. """ method = event.get("method") if method == "messages": self._messages_inbox.put_nowait(event) elif method == "tools": self._tools_inbox.put_nowait(event) elif method == "tasks": self._tasks_inbox.put_nowait(event) # Fan out to descendant handles whose namespace is a prefix of the # event namespace so they receive the event at push time. if method in ("messages", "tools", "tasks"): ns_tuple = tuple(_event_namespace(event.get("params") or {})) for desc_path, desc_handle in self._descendant_handles.items(): desc_len = len(desc_path) if len(ns_tuple) >= desc_len and ns_tuple[:desc_len] == desc_path: desc_handle._push_event(event) def _register_descendant(self, handle: ScopedStreamHandle) -> None: """Register a newly-discovered grandchild so future events are fanned out. Also drains any events already buffered in this handle's inboxes whose namespace matches the grandchild, so events that arrived before the grandchild was discovered are forwarded in arrival order. """ self._descendant_handles[handle.path] = handle desc_len = len(handle.path) for inbox_attr in ( "_messages_inbox", "_tools_inbox", "_tasks_inbox", ): inbox: asyncio.Queue[Event | None] = getattr(self, inbox_attr) staging: list[Event | None] = [] while not inbox.empty(): staging.append(inbox.get_nowait()) for event in staging: inbox.put_nowait(event) if event is None: continue ns_tuple = tuple(_event_namespace(event.get("params") or {})) if len(ns_tuple) >= desc_len and ns_tuple[:desc_len] == handle.path: getattr(handle, inbox_attr).put_nowait(event) def _unregister_descendant(self, path: tuple[str, ...]) -> None: """Remove a grandchild after it reaches a terminal state.""" self._descendant_handles.pop(path, None) def _mark_iterated(self, kind: str) -> None: """Record that an inbox has an active consumer. If the handle is already closed (status != 'started'), immediately enqueue a sentinel so the consumer's `await get()` terminates. This handles sequential consumption (iterate after the handle is finished). Must be called by each projection at the start of iteration. """ self._iterated_inboxes.add(kind) if self.status != "started": # Handle already closed before this consumer started; send the # sentinel now so the projection iterator can terminate. getattr(self, f"_{kind}_inbox").put_nowait(None) def _close_inboxes(self) -> None: """Signal EOF only on channel inboxes that have an active consumer. Inboxes without a consumer would accumulate a leaked None sentinel that is never drained, so we skip them. For inboxes whose consumer starts after this call, `_mark_iterated` sends the sentinel lazily. """ for kind in ("messages", "tools", "tasks"): if kind in self._iterated_inboxes: getattr(self, f"_{kind}_inbox").put_nowait(None) def _finish(self, status: SubgraphStatus, error: str | None = None) -> None: if self.status != "started": return self.status = status self.error = error self._close_inboxes() class _HandleMessagesProjection: """Messages projection that drains a `ScopedStreamHandle`'s messages inbox.""" def __init__(self, handle: ScopedStreamHandle) -> None: self._handle = handle def __aiter__(self) -> AsyncIterator[Any]: return self._messages_iter() async def _messages_iter(self) -> AsyncGenerator[Any, None]: from langchain_core.language_models.chat_model_stream import ( AsyncChatModelStream, ) self._handle._mark_iterated("messages") active: dict[str, AsyncChatModelStream] = {} while True: item = await self._handle._messages_inbox.get() if item is None: return params_field = item.get("params") or {} ns = _event_namespace(params_field) if ns != self._handle.namespace: continue data = params_field.get("data") if isinstance(params_field, dict) else None if not isinstance(data, dict): continue event_type = data.get("event") if event_type == "message-start": message_id = _message_event_id(data) key = _message_route_key(data, fallback=message_id) metadata = ( data.get("metadata") if isinstance(data.get("metadata"), dict) else {} ) stream = AsyncChatModelStream( namespace=list(self._handle.namespace), node=metadata.get("langgraph_node") if metadata else None, message_id=message_id, ) active[key] = stream stream.dispatch(data) yield stream else: key = _message_route_key(data) stream = active.get(key) if stream is None and len(active) == 1: stream = next(iter(active.values())) if stream is None: continue stream.dispatch(data) if event_type in ("message-finish", "error"): for route_key, candidate in list(active.items()): if candidate is stream: del active[route_key] class _HandleToolCallsProjection: """Tool calls projection that drains a `ScopedStreamHandle`'s tools inbox.""" def __init__(self, handle: ScopedStreamHandle) -> None: self._handle = handle def __aiter__(self) -> AsyncIterator[Any]: return self._tool_calls_iter() async def _tool_calls_iter(self) -> AsyncGenerator[Any, None]: self._handle._mark_iterated("tools") active: dict[str, ToolCallHandle] = {} while True: item = await self._handle._tools_inbox.get() if item is None: err = RuntimeError( "Tool call stream closed before terminal tool event." ) for handle in active.values(): handle._fail(err) return params_field = item.get("params") or {} ns = _event_namespace(params_field) if ns != self._handle.namespace: continue data = params_field.get("data") if isinstance(params_field, dict) else None if not isinstance(data, dict): continue event_type = data.get("event") tool_call_id = data.get("tool_call_id") if not isinstance(tool_call_id, str): continue if event_type == "tool-started": tool_name = data.get("tool_name") if not isinstance(tool_name, str): tool_name = "" handle = ToolCallHandle( tool_call_id=tool_call_id, name=tool_name, input=data.get("input"), namespace=list(self._handle.namespace), ) active[tool_call_id] = handle yield handle elif event_type == "tool-output-delta": handle = active.get(tool_call_id) delta = data.get("delta") if handle is not None and isinstance(delta, str): handle._push_delta(delta) elif event_type == "tool-finished": handle = active.pop(tool_call_id, None) if handle is not None: handle._finish(data.get("output")) elif event_type == "tool-error": handle = active.pop(tool_call_id, None) if handle is not None: message = data.get("message") handle._fail( RuntimeError(str(message) if message else "Tool call errored") ) class _HandleSubgraphsProjection: """Subgraphs projection that drains a `ScopedStreamHandle`'s tasks inbox.""" def __init__(self, handle: ScopedStreamHandle) -> None: self._handle = handle def __aiter__(self) -> AsyncIterator[ScopedStreamHandle]: return self._subgraphs_iter() def _route_sibling_inboxes_to_grandchildren( self, active: dict[tuple[str, ...], ScopedStreamHandle], ) -> None: """Drain non-blocking events from parent's messages/tools inboxes to grandchildren. Called after each tasks event so grandchild handles receive events that were enqueued in the parent handle's inboxes before (or just after) the grandchild was discovered. """ parent = self._handle for inbox_attr, grandchild_attr in ( ("_messages_inbox", "_messages_inbox"), ("_tools_inbox", "_tools_inbox"), ): inbox: asyncio.Queue[Event | None] = getattr(parent, inbox_attr) staging: list[Event | None] = [] # Drain without blocking. while not inbox.empty(): staging.append(inbox.get_nowait()) for event in staging: if event is None: # Re-queue the EOF sentinel — it belongs to the parent inbox consumer. inbox.put_nowait(None) continue event_params = event.get("params") or {} ns_tuple = tuple(_event_namespace(event_params)) routed = False for _child_path, grandchild in active.items(): grandchild_len = len(grandchild.path) if ( len(ns_tuple) >= grandchild_len and ns_tuple[:grandchild_len] == grandchild.path ): gc_inbox: asyncio.Queue[Event | None] = getattr( grandchild, grandchild_attr ) gc_inbox.put_nowait(event) routed = True break if not routed: # Not a grandchild event — put it back for the handle projection. inbox.put_nowait(event) async def _subgraphs_iter(self) -> AsyncGenerator[ScopedStreamHandle, None]: self._handle._mark_iterated("tasks") seen: set[tuple[str, ...]] = set() active: dict[tuple[str, ...], ScopedStreamHandle] = {} scope = self._handle.path while True: item = await self._handle._tasks_inbox.get() if item is None: for child in active.values(): if child.status == "started": child._finish("completed") return params_field = item.get("params") or {} namespace = _event_namespace(params_field) data = params_field.get("data") if isinstance(params_field, dict) else None if not isinstance(data, dict): continue if "result" in data: result_id = data.get("id") if not result_id: continue parent_path = tuple(namespace) for child_path, child_handle in list(active.items()): if child_path[:-1] != parent_path: continue if child_handle.trigger_call_id != result_id: continue status, error = _terminal_from_tasks_result(data) child_handle._finish(status, error) del active[child_path] self._handle._unregister_descendant(child_path) continue if not _is_direct_child(namespace, scope): continue path = tuple(namespace) if path in seen: continue seen.add(path) graph_name, trigger_call_id = _parse_namespace_segment(path[-1]) child_handle = ScopedStreamHandle( thread=self._handle._thread, path=path, graph_name=graph_name or None, trigger_call_id=trigger_call_id, max_queue_size=self._handle._max_queue_size, ) active[path] = child_handle # Register so future _push_event calls on this handle fan out to the # grandchild at push time, preserving arrival order without drain-and-replay. self._handle._register_descendant(child_handle) yield child_handle class _SubgraphsProjection: """Discover direct child invocations for a namespace scope.""" def __init__(self, thread: AsyncThreadStream, scope: tuple[str, ...] = ()) -> None: self._thread = thread self._scope = scope def __aiter__(self) -> AsyncIterator[ScopedStreamHandle]: return self._subgraphs_iter() async def _subgraphs_iter(self) -> AsyncGenerator[ScopedStreamHandle, None]: if self._thread._transport is None: raise RuntimeError("AsyncThreadStream not entered - use `async with`.") params = _subgraph_subscription_params(self._scope) sub = self._thread._register_subscription(params) decoder = SubgraphsDecoder( scope=self._scope, handle_factory=lambda *, path, graph_name, trigger_call_id: ( ScopedStreamHandle( thread=self._thread, path=path, graph_name=graph_name, trigger_call_id=trigger_call_id, ) ), ) root_inbox: asyncio.Queue[Event | None] | None = ( self._thread._activate_root_messages_inbox() if not self._scope else None ) try: await self._thread._reconcile_stream(params) self._thread._ensure_fanout_running() while True: item = await sub.queue.get() if item is None: return params_field = item.get("params") or {} if ( root_inbox is not None and item.get("method") == "messages" and tuple(_event_namespace(params_field)) == self._scope ): root_inbox.put_nowait(item) for handle in decoder.feed(item): yield handle finally: # Determine terminal status from the parent run's lifecycle result. # If _run_done resolved as errored, force-complete remaining children # as errored so callers see the correct terminal state. terminal_status: SubgraphStatus = "completed" run_done = self._thread._run_done if run_done is not None and run_done.done() and not run_done.cancelled(): result = run_done.result() if isinstance(result, _RunTerminal) and result.status == "errored": terminal_status = "failed" for handle in decoder._active.values(): if handle.status == "started": handle._finish(terminal_status) self._thread._unregister_subscription(sub.id) if root_inbox is not None: root_inbox.put_nowait(None) class ToolCallHandle: """Async handle for one root-scope tool call.""" def __init__( self, *, tool_call_id: str, name: str, input: Any = None, namespace: list[str] | None = None, max_queue_size: int = 1024, ) -> None: self.tool_call_id = tool_call_id self.name = name self.input = input self.namespace = list(namespace or []) self.done = False self.error: BaseException | None = None loop = asyncio.get_running_loop() self.output: asyncio.Future[Any] = loop.create_future() self._deltas: asyncio.Queue[str | None] = asyncio.Queue(maxsize=max_queue_size) self._deltas_consumed = False @property def deltas(self) -> AsyncIterator[str]: """Stream tool output deltas emitted before the terminal event. Raises: RuntimeError: if called more than once — the underlying queue is single-consumer and cannot be fanned out safely. """ if self._deltas_consumed: raise RuntimeError( "ToolCallHandle.deltas can only be iterated by a single consumer." ) self._deltas_consumed = True return self._delta_iter() async def _delta_iter(self) -> AsyncGenerator[str, None]: while True: item = await self._deltas.get() if item is None: return # errors surface via output, not deltas yield item def _push_delta(self, delta: str) -> None: if self.done: return self._deltas.put_nowait(delta) def _finish(self, output: Any) -> None: if self.done: return self.done = True if not self.output.done(): self.output.set_result(output) self._deltas.put_nowait(None) def _fail(self, err: BaseException) -> None: if self.done: return self.done = True self.error = err if not self.output.done(): self.output.set_exception(err) self._deltas.put_nowait(None) class _ToolCallsProjection: """Typed projection for root-scope `thread.tool_calls`.""" def __init__( self, thread: AsyncThreadStream, namespace: list[str] | None = None ) -> None: self._thread = thread self._namespace = list(namespace or []) def __aiter__(self) -> AsyncIterator[ToolCallHandle]: return self._tool_calls_iter() async def _tool_calls_iter(self) -> AsyncGenerator[ToolCallHandle, None]: if self._thread._transport is None: raise RuntimeError("AsyncThreadStream not entered - use `async with`.") params = _exact_namespace_params(["tools"], self._namespace) sub = self._thread._register_subscription(params) decoder = ToolCallsDecoder( namespace=self._namespace, handle_factory=lambda *, tool_call_id, name, input, namespace: ( ToolCallHandle( tool_call_id=tool_call_id, name=name, input=input, namespace=namespace, ) ), ) registered: list[ToolCallHandle] = [] try: await self._thread._reconcile_stream(params) self._thread._ensure_fanout_running() while True: item = await sub.queue.get() if item is None: return for handle in decoder.feed(item): self._thread._register_active_tool_call(handle) registered.append(handle) yield handle finally: # Read terminal error from _run_done if it is already resolved. # We do NOT block here: callers who need a terminal observation # should await `thread.output` directly. Blocking in iterator # teardown would stall every early break or exception exit for # up to the full shield-wait timeout (previously 1 s). run_done = self._thread._run_done terminal_err: BaseException | None = None if run_done is not None and run_done.done() and not run_done.cancelled(): terminal = run_done.result() terminal_err = terminal.error err: BaseException = ( terminal_err if terminal_err is not None else RuntimeError("Tool call stream closed before terminal tool event.") ) for handle in list(decoder._active.values()): handle._fail(err) for handle in registered: self._thread._unregister_active_tool_call(handle) self._thread._unregister_subscription(sub.id) class _ExtensionsProjection: """Mapping from extension name to custom event payload stream. Repeated access for the same `name` returns the cached projection so that callers receive the same subscription handle across multiple references to `thread.extensions["foo"]` within one session. """ def __init__(self, thread: AsyncThreadStream, namespace: list[str]) -> None: self._thread = thread self._namespace = namespace self._cache: dict[str, _ExtensionProjection] = {} def __getitem__(self, name: str) -> _ExtensionProjection: if not name: raise ValueError("extension name must be non-empty.") if name not in self._cache: self._cache[name] = _ExtensionProjection( self._thread, name=name, namespace=self._namespace ) return self._cache[name] class _ExtensionProjection: def __init__( self, thread: AsyncThreadStream, *, name: str, namespace: list[str], ) -> None: self._thread = thread self._name = name self._namespace = namespace def __aiter__(self) -> AsyncIterator[dict[str, Any]]: return self._iter() async def _iter(self) -> AsyncGenerator[dict[str, Any], None]: params: SubscribeParams = {"channels": [f"custom:{self._name}"]} if self._namespace: params["namespaces"] = [self._namespace] sub = self._thread._register_subscription(params) decoder = ExtensionsDecoder(name=self._name) try: if self._thread._closed: return await self._thread._reconcile_stream(params) self._thread._ensure_fanout_running() while True: item = await sub.queue.get() if item is None: return for out in decoder.feed(item): yield out finally: self._thread._unregister_subscription(sub.id) class AsyncThreadStream: """Async context manager for one thread's v3 streaming session. Construct via `client.threads.stream(thread_id=None, *, assistant_id, ...)` rather than instantiating directly. """ def __init__( self, *, http: HttpClient, thread_id: str, assistant_id: str, headers: Mapping[str, str] | None = None, max_queue_size: int = 1024, run_start_timeout: float | None = None, explicit_thread_id: bool = False, transport_kind: Literal["sse", "websocket"] = "sse", ) -> None: self._http = http self._headers = dict(headers or {}) self.thread_id = thread_id self.assistant_id = assistant_id self._max_queue_size = max_queue_size self._run_start_timeout = run_start_timeout self._explicit_thread_id = explicit_thread_id self._transport_kind = transport_kind self._closed = False self._transport: AsyncProtocolTransport | None = None self._open_handles: list[EventStreamHandle] = [] self._next_command_id = 1 self._next_subscription_id = 1 self._subscriptions: dict[int, _Subscription] = {} self._seen_event_ids = _SeenEventIds() self._shared_stream: EventStreamHandle | None = None self._shared_stream_filter: dict[str, Any] | None = None self._fanout_task: asyncio.Task[None] | None = None self.interrupted: bool = False self.interrupts: list[InterruptPayload] = [] # Why: serialize the `interrupts` snapshot in `run.respond` with the # terminal-clear path in `_apply_lifecycle_event`, so a respond() in # flight cannot send a stale `interrupt_id` after the lifecycle watcher # observes a `completed`/`errored` event. self._interrupts_lock = asyncio.Lock() self._lifecycle_watcher_task: asyncio.Task[None] | None = None self._lifecycle_watcher_handle: EventStreamHandle | None = None self._lifecycle_cursor: int | None = None self._lifecycle_max_reconnect_attempts = 5 # Shared-stream reconnect knobs: applied by `_fanout` after a post-ready # transport drop so subscribers (messages/tools/tasks/values projections, # subgraph child handles) survive a brief SSE disconnect without losing # buffered events. Cursor (`_cursor`) is replayed as `since` so the # server resumes from where the prior stream left off; per-event # `event_id` dedup in `_dedup_iter` drops any overlap on the new stream. self._shared_max_reconnect_attempts = 5 self._shared_reconnect_backoff_base = 0.1 self._shared_reconnect_backoff_cap = 2.0 self._run_start_ready: asyncio.Future[None] | None = None self._run_seen: bool = False self._run_done: asyncio.Future[_RunTerminal] | None = None self._cursor: int | None = None self._active_message_streams: set[AsyncChatModelStream] = set() self._active_tool_calls: set[ToolCallHandle] = set() # Root-scope inbox: populated by `_SubgraphsProjection` when it consumes # messages events at namespace `[]` so that `thread.messages` can drain # them even after the shared SSE has ended (dedup prevents replay). self._root_messages_inbox: asyncio.Queue[Event | None] | None = None self.run = RunModule(self) self.agent = _AgentModule(self) self.output = _OutputAwaitable(self) self.values = _ValuesProjection(self) self.messages = _MessagesProjection(self, namespace=[]) self.tool_calls = _ToolCallsProjection(self, namespace=[]) self.subgraphs = _SubgraphsProjection(self, scope=()) self.subagents = self.subgraphs self.extensions = _ExtensionsProjection(self, namespace=[]) @property def _controller(self) -> AsyncThreadStream: """Return self as the subscription controller (duck-type compatible with StreamController). Exposes `_subscriptions` so tests can verify subscription counts via `thread._controller._subscriptions` without requiring a separate controller object. """ return self async def __aenter__(self) -> AsyncThreadStream: if self._closed: raise RuntimeError("AsyncThreadStream is closed and cannot be re-entered.") transport_cls = ( ProtocolWebSocketTransport if self._transport_kind == "websocket" else ProtocolSseTransport ) self._transport = transport_cls( client=self._http.client, thread_id=self.thread_id, headers=self._headers, max_queue_size=self._max_queue_size, ) # Create the run-done future here (async context guarantees a running loop). self._run_done = asyncio.get_running_loop().create_future() # Start the lifecycle watcher immediately so reattach and thread.output # work without a preceding run.start call. self._ensure_lifecycle_watcher_running() return self async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None: try: await self.close() except BaseException as close_err: if exc is None: raise # Original exception takes precedence; chain close error as context. close_err.__context__ = exc @property def events(self) -> AsyncIterator[Event]: """Return a fresh subscription to ALL channels. Each property access opens a new subscription; callers iterating twice will see two independent streams (both filtered by the same channel union). Terminates when the stream closes (server hangup, `__aexit__`, or transport-level close). """ if self._transport is None: raise RuntimeError("AsyncThreadStream not entered — use `async with`.") handle = self._transport.open_event_stream({"channels": _ALL_CHANNELS}) self._open_handles.append(handle) return handle.events async def close(self) -> None: """Tear down the thread stream. Idempotent.""" if self._closed: return self._closed = True for handle in self._open_handles: await handle.close() # Cancel _run_done so thread.output doesn't wait forever on close. run_done = self._run_done if run_done is not None and not run_done.done(): run_done.cancel() if self._lifecycle_watcher_task is not None: self._lifecycle_watcher_task.cancel() with contextlib.suppress(Exception, asyncio.CancelledError): await self._lifecycle_watcher_task if self._lifecycle_watcher_handle is not None: await self._lifecycle_watcher_handle.close() self._fail_active_message_streams(asyncio.CancelledError()) self._fail_active_tool_calls(asyncio.CancelledError()) if self._fanout_task is not None: self._fanout_task.cancel() with contextlib.suppress(Exception, asyncio.CancelledError): await self._fanout_task if self._shared_stream is not None: await self._shared_stream.close() if self._transport is not None: await self._transport.close() def _register_subscription(self, params: SubscribeParams) -> _Subscription: """Allocate a subscription id and add it to the registry.""" sub = _Subscription( id=self._next_subscription_id, params=params, queue=asyncio.Queue(maxsize=self._max_queue_size), ) self._next_subscription_id += 1 self._subscriptions[sub.id] = sub return sub def _unregister_subscription(self, subscription_id: int) -> None: """Remove a subscription from the registry. No-op if already absent.""" self._subscriptions.pop(subscription_id, None) def _activate_root_messages_inbox(self) -> asyncio.Queue[Event | None]: """Create the root-scope messages inbox if not already active and return it. Called by `_SubgraphsProjection` at scope `()` to capture messages events that arrive at namespace `[]` before `thread.messages` has subscribed. """ if self._root_messages_inbox is None: self._root_messages_inbox = asyncio.Queue() return self._root_messages_inbox def _register_active_message_stream(self, stream: AsyncChatModelStream) -> None: self._active_message_streams.add(stream) def _unregister_active_message_stream(self, stream: AsyncChatModelStream) -> None: self._active_message_streams.discard(stream) def _fail_active_message_streams(self, err: BaseException) -> None: for stream in list(self._active_message_streams): stream.fail(err) self._active_message_streams.clear() def _register_active_tool_call(self, handle: ToolCallHandle) -> None: self._active_tool_calls.add(handle) def _unregister_active_tool_call(self, handle: ToolCallHandle) -> None: self._active_tool_calls.discard(handle) def _fail_active_tool_calls(self, err: BaseException) -> None: for handle in list(self._active_tool_calls): handle._fail(err) self._active_tool_calls.clear() def _signal_paused(self) -> None: """Wake every active projection iterator on interrupt / run end. Pushes the terminal sentinel (`None`) into every subscription queue. Iterators see `None` and return; the shared SSE keeps running so re-iteration after `run.respond(...)` registers a fresh subscription and resumes. `root_messages_inbox` is intentionally NOT signaled here: the subgraphs projection that populates it is responsible for pushing the terminal `None` in its own `finally` block, so any message events it redirected to the inbox land before the sentinel. Signaling root_inbox here would race the redirection and could drop messages. """ # On a saturated queue the consumer is already behind; the iterator # will still terminate when it drains to this point. for sub in list(self._subscriptions.values()): with contextlib.suppress(asyncio.QueueFull): sub.queue.put_nowait(None) def observe_applied_through_seq(self, seq: Any) -> None: """Advance the reconnect cursor from a command response meta sequence.""" if isinstance(seq, int) and (self._cursor is None or seq > self._cursor): self._cursor = seq def _observe_event(self, event: Event) -> None: seq = event.get("seq") if isinstance(seq, int) and (self._cursor is None or seq > self._cursor): self._cursor = seq def subscribe( self, channels: list[str], *, namespaces: list[list[str]] | None = None, depth: int | None = None, ) -> AsyncIterator[Event]: """Open a typed subscription against the shared SSE. Returns an async iterator that yields raw `Event` dicts matching the given filter. Multiple concurrent subscribes share one HTTP connection whose union expands or rotates as subscriptions come and go. """ if self._transport is None: raise RuntimeError("AsyncThreadStream not entered — use `async with`.") params: SubscribeParams = {"channels": list(channels)} if namespaces is not None: params["namespaces"] = namespaces if depth is not None: params["depth"] = depth return self._subscription_iter(params) async def interleave_projections( self, channels: list[str] ) -> AsyncIterator[tuple[str, Any]]: """Yield `(channel_name, item)` tuples across multiple projections. One shared subscription drives all per-channel decoders; items arrive in server-emit order (the SDK analog of `GraphRunStream.interleave`). Args: channels: Flat list of `"values"`, `"messages"`, `"tool_calls"`, `"subgraphs"`, and/or extension names. Built-ins yield their typed item (snapshot dict / `AsyncChatModelStream` / `ToolCallHandle` / `ScopedStreamHandle`); an extension yields its payload dict, keyed by the bare extension name. Note: Handles and streams are yielded eagerly (before their sub-stream completes), so items arrive interleaved in real time. To receive a fully-resolved handle (output already populated), use the dedicated `thread.tool_calls` / `thread.messages` projections instead. """ validate_interleave_channels(channels) if self._transport is None: raise RuntimeError("AsyncThreadStream not entered — use `async with`.") decoders: dict[str, Decoder] = {} sub_params: list[SubscribeParams] = [] for ch in channels: if ch == "values": decoders[ch] = DataDecoder("values") sub_params.append({"channels": ["values"]}) elif ch in ("updates", "checkpoints", "tasks"): # Plain payload channels (local Updates/Checkpoints/Tasks # analog). Root-scope filter is load-bearing: a co-requested # unscoped `values` widens the merged subscription to all # namespaces, so the decoder itself keeps subgraph payloads out. decoders[ch] = DataDecoder(ch, namespace=[]) sub_params.append(_exact_namespace_params([ch], [])) elif ch == "messages": decoders[ch] = MessagesDecoder( namespace=[], stream_factory=lambda *, namespace, node, message_id: ( AsyncChatModelStream( namespace=namespace, node=node, message_id=message_id ) ), ) sub_params.append(_exact_namespace_params(["messages"], [])) elif ch == "tool_calls": decoders[ch] = ToolCallsDecoder( namespace=[], handle_factory=lambda *, tool_call_id, name, input, namespace: ( ToolCallHandle( tool_call_id=tool_call_id, name=name, input=input, namespace=namespace, ) ), ) sub_params.append(_exact_namespace_params(["tools"], [])) elif ch == "subgraphs": decoders[ch] = SubgraphsDecoder( scope=(), handle_factory=lambda *, path, graph_name, trigger_call_id: ( ScopedStreamHandle( thread=self, path=path, graph_name=graph_name, trigger_call_id=trigger_call_id, ) ), ) sub_params.append(_subgraph_subscription_params(())) else: decoders[ch] = ExtensionsDecoder(name=ch) sub_params.append({"channels": [f"custom:{ch}"]}) if not sub_params: return merged = cast( SubscribeParams, compute_union_filter(cast(list[dict[str, Any]], sub_params)), ) subgraphs = decoders.get("subgraphs") # Track decoder-created handles so teardown can finalize anything still # in flight; otherwise an awaiting `handle.output` / `handle.messages` # would hang after an early break or run termination. registered_tool_calls: list[ToolCallHandle] = [] registered_message_streams: list[AsyncChatModelStream] = [] try: async for event in self._subscription_iter(merged): if subgraphs is not None: for item in subgraphs.feed(event): yield ("subgraphs", item) wire = infer_channel(event) public = self._interleave_public_name(wire) # subgraphs is driven separately above (it consumes all events); never dispatch it here. if public is not None and public != "subgraphs": decoder = decoders.get(public) if decoder is not None: for item in decoder.feed(event): if public == "tool_calls": self._register_active_tool_call(item) registered_tool_calls.append(item) elif public == "messages": self._register_active_message_stream(item) registered_message_streams.append(item) yield (public, item) finally: self._finalize_interleave_decoders( decoders.get("tool_calls"), subgraphs, registered_tool_calls, registered_message_streams, ) @staticmethod def _interleave_public_name(wire: str | None) -> str | None: """Map a wire channel name to the public channel name used in interleave tuples.""" if wire is None: return None if wire == "tools": return "tool_calls" if wire.startswith("custom:"): return wire[len("custom:") :] return wire # values, messages (tasks/lifecycle pass through with no decoder match) def _finalize_interleave_decoders( self, tool_calls: Decoder | None, subgraphs: Decoder | None, registered_tool_calls: list[ToolCallHandle], registered_message_streams: list[AsyncChatModelStream], ) -> None: """Finalize in-flight handles when `interleave_projections` tears down. Mirrors the terminal handling of the dedicated `_ToolCallsProjection` / `_SubgraphsProjection`: in-flight tool calls are failed (so awaiting `handle.output` can't hang) and discovered subgraph children are force-completed with the run's terminal status. """ run_done = self._run_done resolved = ( run_done.result() if run_done is not None and run_done.done() and not run_done.cancelled() else None ) if isinstance(tool_calls, ToolCallsDecoder): err: BaseException = ( resolved.error if resolved is not None and resolved.error is not None else RuntimeError("Tool call stream closed before terminal tool event.") ) for handle in list(tool_calls._active.values()): handle._fail(err) for handle in registered_tool_calls: self._unregister_active_tool_call(handle) for stream in registered_message_streams: self._unregister_active_message_stream(stream) if isinstance(subgraphs, SubgraphsDecoder): terminal_status: SubgraphStatus = ( "failed" if isinstance(resolved, _RunTerminal) and resolved.status == "errored" else "completed" ) for child in subgraphs._active.values(): if child.status == "started": child._finish(terminal_status) async def _subscription_iter( self, params: SubscribeParams ) -> AsyncGenerator[Event, None]: sub = self._register_subscription(params) try: if self._closed: return await self._reconcile_stream(params) self._ensure_fanout_running() while True: item = await sub.queue.get() if item is None: return yield item finally: self._unregister_subscription(sub.id) def _ensure_fanout_running(self) -> None: if self._fanout_task is None or self._fanout_task.done(): self._fanout_task = asyncio.create_task(self._fanout()) async def _fanout(self) -> None: """Single consumer of the shared SSE; routes events to subscriptions. Why: rotation in `_reconcile_stream` replaces `_shared_stream` mid-loop. Re-read `self._shared_stream` on each outer iteration so we always consume from the current handle. The old handle's iterator exhausts naturally after `_close_after` closes it. On a post-ready transport drop (non-cancelled error in `shared.done`), attempts to reconnect up to `_shared_max_reconnect_attempts` times so scoped projections (subgraph child handles, message streams) survive without losing buffered events. The reconnect replays `since=` and `_dedup_iter` drops any overlap. """ from langgraph_sdk.stream.subscription import matches_subscription while not self._closed: shared = self._shared_stream if shared is None: return try: async for event in self._dedup_iter(shared.events): if self._closed: break self._observe_event(event) for sub in list(self._subscriptions.values()): if matches_subscription(event, sub.params): sub.queue.put_nowait(event) # On root-terminal lifecycle, push the `None` sentinel # into all subscription queues so projection iterators # exit when the run ends naturally. Runs on the shared # SSE so the terminal is processed in seq order with # the projection events -- any in-flight values / # tools / messages events for this run are already # queued before None. if _is_root_terminal_lifecycle(event): self._signal_paused() except Exception: # Pump errored — fall through to error-handling/reconnect. pass if self._shared_stream is shared: # No rotation happened; the stream genuinely ended. Check # `shared.done` for a post-ready drop and, if so, attempt to # reconnect with `since=` so subscribers don't lose # buffered events on a transient transport failure. err = await shared.done if ( err is not None and not isinstance(err, asyncio.CancelledError) and not self._closed ): with contextlib.suppress(Exception): await shared.close() if await self._reconnect_shared_stream(): continue break # Rotation: loop again to pick up the new _shared_stream. # Terminate consumers cleanly on shutdown / stream-end. for sub in self._subscriptions.values(): sub.queue.put_nowait(None) async def _reconnect_sleep(self, attempt: int) -> None: """Sleep with exponential backoff and jitter for reconnect attempt `attempt`.""" base = self._shared_reconnect_backoff_base cap = self._shared_reconnect_backoff_cap delay = min(cap, base * (2**attempt)) jitter = random.uniform(0, delay * 0.25) await asyncio.sleep(delay + jitter) async def _reconnect_shared_stream(self) -> bool: """Attempt to reopen the shared stream after a post-ready transport drop. Returns: `True` if a new stream was opened (caller should resume fanout), `False` if all reconnect attempts were exhausted or the controller was closed in the meantime. """ if self._transport is None: return False # Use the current shared-stream filter (latest computed union); if # subscriptions changed during the drop, this picks up the new shape. base_filter = self._shared_stream_filter if base_filter is None: return False for attempt in range(self._shared_max_reconnect_attempts): if self._closed: return False stream_params: dict[str, Any] = dict(base_filter) if self._cursor is not None: stream_params["since"] = self._cursor try: new_stream = self._transport.open_event_stream(stream_params) await new_stream.ready except asyncio.CancelledError: raise except Exception: await self._reconnect_sleep(attempt) continue self._shared_stream = new_stream return True return False async def _reconcile_stream(self, candidate_filter: SubscribeParams) -> None: """Ensure the shared SSE covers `candidate_filter`. Rotate if not. Open-new-before-close-old: any events buffered server-side between the two opens are replayed on the new SSE, and the per-thread `_seen_event_ids` set dedupes the overlap. Awaits `new_stream.ready` so the HTTP connection is established before returning, guaranteeing that both old and new streams are simultaneously connected during rotation (enabling correct peak-count tracking and dedup correctness). """ await self._await_run_start_gate(timeout=self._run_start_timeout) from langgraph_sdk.stream.subscription import filter_covers if self._transport is None: raise RuntimeError("AsyncThreadStream not entered — use `async with`.") 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 # Existing stream is sufficient. new_filter = self._compute_current_union(extra=candidate_filter) stream_params: dict[str, Any] = dict(new_filter) if self._cursor is not None: stream_params["since"] = self._cursor new_stream = self._transport.open_event_stream(stream_params) old_stream = self._shared_stream self._shared_stream = new_stream self._shared_stream_filter = new_filter # Await the new stream's ready future so the HTTP connection is # established before we schedule the old stream's close. This ensures # old and new are simultaneously open during the rotation window. await new_stream.ready if old_stream is not None: # Schedule the old stream's close as a separate task so the # caller doesn't pay close() latency in the rotation hot path. asyncio.create_task(_close_after(old_stream)) # noqa: RUF006 def _compute_current_union( self, extra: SubscribeParams | None = None ) -> dict[str, Any]: from langgraph_sdk.stream.subscription import compute_union_filter filters: list[dict[str, Any]] = [ dict(sub.params) for sub in self._subscriptions.values() ] if extra is not None: filters.append(dict(extra)) # Always include lifecycle in the shared SSE so the fanout consumer # sees root-terminal events in seq order with the projection events. # See `_is_root_terminal_lifecycle` -- the fanout uses it to push # the `None` sentinel into sub queues when the run ends naturally, # which is what makes projection iterators exit on a long-lived # SSE that doesn't EOF after the run. Per-subscription filtering # (`matches_subscription`) drops lifecycle events for any # subscription that didn't ask for them, so user-visible queues # don't see leaked events. filters.append({"channels": ["lifecycle"]}) return compute_union_filter(filters) async def _dedup_iter(self, source: AsyncIterator[Event]) -> AsyncIterator[Event]: async 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) yield event async def _send_command( self, method: str, params: dict[str, Any] ) -> dict[str, Any]: """Send a protocol command and return the `result` payload. Returns `{}` for 202/204 responses (no body). Raises `RuntimeError` with the protocol code/message when the server returns an error envelope (`{"type": "error", ...}`). """ if self._transport is None: raise RuntimeError("AsyncThreadStream not entered — use `async with`.") command_id = self._next_command_id self._next_command_id += 1 response = await self._transport.send_command( {"id": command_id, "method": method, "params": params} ) if response is None: # 202/204 — no body. Caller gets an empty result. return {} if response.get("type") == "error": code = response.get("error", "unknown") message = response.get("message", "") raise RuntimeError(f"Protocol error [{code}]: {message}") meta = response.get("meta") if isinstance(meta, dict): applied_through_seq = meta.get("applied_through_seq") if self._controller is not None: self._controller.observe_applied_through_seq(applied_through_seq) return response.get("result", {}) async def _await_run_start_gate(self, *, timeout: float | None = None) -> None: """Wait for the current run.start to commit the thread server-side. No-op when no run.start is in flight. Re-raises if run.start failed. Raises `asyncio.TimeoutError` if `timeout` is set and the gate does not resolve in time; the gate itself is left intact for later callers. """ gate = self._run_start_ready if gate is None or gate.done(): return if timeout is None: await gate else: await asyncio.wait_for(asyncio.shield(gate), timeout=timeout) def _ensure_lifecycle_watcher_running(self) -> None: if self._lifecycle_watcher_task is not None: return self._lifecycle_watcher_task = asyncio.create_task( self._run_lifecycle_watcher() ) def _observe_lifecycle_event(self, event: Event) -> None: seq = event.get("seq") if isinstance(seq, int) and ( self._lifecycle_cursor is None or seq > self._lifecycle_cursor ): self._lifecycle_cursor = seq def _lifecycle_stream_params(self) -> dict[str, Any]: params: dict[str, Any] = {"channels": ["lifecycle", "input"]} if self._lifecycle_cursor is not None: params["since"] = self._lifecycle_cursor return params async def _run_lifecycle_watcher(self) -> None: """Always-on SSE consuming lifecycle + input channels.""" if self._transport is None: return reconnect_attempts = 0 while not self._closed: try: handle = self._transport.open_event_stream( self._lifecycle_stream_params() ) self._lifecycle_watcher_handle = handle await asyncio.wait_for(handle.ready, timeout=5.0) async for event in handle.events: if self._closed: return self._observe_lifecycle_event(event) await self._apply_lifecycle_event(event) err = await handle.done if err is None or isinstance(err, asyncio.CancelledError): # Clean EOF: stream ended without a terminal lifecycle # event. Resolve `_run_done` as errored so awaiters of # `thread.output` don't hang. if err is None: run_done = self._run_done if run_done is not None and not run_done.done(): run_done.set_result( _RunTerminal( status="errored", error=RuntimeError( "lifecycle stream ended before terminal event" ), ) ) return reconnect_attempts += 1 if reconnect_attempts > self._lifecycle_max_reconnect_attempts: raise err await asyncio.sleep(0.05) except asyncio.CancelledError: raise except Exception as exc: reconnect_attempts += 1 if reconnect_attempts <= self._lifecycle_max_reconnect_attempts: await asyncio.sleep(0.05) continue run_done = self._run_done if run_done is not None and not run_done.done(): run_done.set_result( _RunTerminal( status="errored", error=RuntimeError(f"Lifecycle transport failed: {exc}"), ) ) return async def _fetch_state(self) -> dict[str, Any]: """Fetch the current thread state from the REST endpoint.""" return await self._http.get( f"/threads/{self.thread_id}/state", headers=self._headers or None, ) def _state_is_terminal(self, state: dict[str, Any]) -> bool: """Return `True` if the thread state has no pending tasks or next nodes.""" return not state.get("next") and not state.get("tasks") def _can_return_existing_state_immediately(self) -> bool: """Return `True` if we can try the REST state before waiting on the lifecycle. True only when the caller passed an explicit `thread_id` (not a minted UUID) and no run has been seen yet, indicating a potential reattach to an already-terminal thread. """ return self._explicit_thread_id and not self._run_seen async def _wait_for_run_done(self) -> _RunTerminal: """Await `_run_done`, raising if the stream was never entered or no run exists. Raises: RuntimeError: stream not entered, or no run started and no explicit thread_id was provided. """ if self._run_done is None: raise RuntimeError("AsyncThreadStream not entered — use async with") if not self._run_seen and not self._explicit_thread_id: raise RuntimeError( "thread.output: no run has been started and no explicit thread_id " "was provided. Call thread.run.start() first." ) return await self._run_done async def _apply_lifecycle_event(self, event: Event) -> None: """Update `interrupted` / `interrupts` / `_run_done` from a lifecycle or input event.""" method = event.get("method") if method == "input.requested": params = event.get("params") or {} data = params.get("data") if isinstance(params, dict) else None interrupt_id = data.get("interrupt_id") if isinstance(data, dict) else None if isinstance(interrupt_id, str): payload: InterruptPayload = { "interrupt_id": interrupt_id, "value": data.get("value") if isinstance(data, dict) else None, "namespace": params.get("namespace") or [] if isinstance(params, dict) else [], } async with self._interrupts_lock: was_interrupted = self.interrupted self.interrupts.append(payload) self.interrupted = True # On the rising edge of `interrupted`, push the terminal # sentinel into every active projection subscription so their # iterators exit cleanly. The run is paused — not done — so # the shared SSE and fanout keep running; a subsequent # `async for snap in thread.values:` (or any other # projection) registers a fresh subscription and resumes # iteration once the consumer calls `run.respond(...)`. if not was_interrupted: self._signal_paused() elif method == "lifecycle": params = event.get("params") or {} data = params.get("data") if isinstance(params, dict) else None phase = data.get("event") if isinstance(data, dict) else None if phase in ("started", "running"): # Mark that we have observed an active run so thread.output # knows a run exists (handles reattach without run.start). self._run_seen = True elif phase in ("completed", "failed"): # Why: interrupts describe current-run state. Clear on terminal # lifecycle so a subsequent run.respond() can't fire against a # stale prior-run interrupt_id. Acquire `_interrupts_lock` so # any in-flight `run.respond` either completes against the # pre-clear snapshot or sees the cleared state — never both. async with self._interrupts_lock: self.interrupted = False self.interrupts = [] run_done = self._run_done if run_done is not None and not run_done.done(): if phase == "failed": error_msg = ( data.get("error") if isinstance(data, dict) else None ) error = RuntimeError( f"Run errored: {error_msg}" if error_msg else "Run errored" ) self._fail_active_message_streams(error) self._fail_active_tool_calls(error) run_done.set_result(_RunTerminal(status="errored", error=error)) else: run_done.set_result(_RunTerminal(status="completed"))