mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-26 19:45:00 +02:00
1675 lines
70 KiB
Python
1675 lines
70 KiB
Python
"""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`), and surfaces lifecycle state (`interrupted`,
|
|
`interrupts`) via an always-on lifecycle watcher SSE. Typed projections
|
|
(`thread.values`, `thread.messages`, etc.) mirror the v3 protocol surface.
|
|
|
|
Direct port of `libs/sdk/src/client/stream/index.ts`.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
from collections.abc import AsyncGenerator, AsyncIterator, Generator, Mapping
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Literal, TypedDict
|
|
|
|
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.stream.transport import EventStreamHandle, ProtocolSseTransport
|
|
|
|
|
|
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` surface.
|
|
_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 []
|
|
|
|
|
|
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)
|
|
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
|
|
params_field = item.get("params") or {}
|
|
data = (
|
|
params_field.get("data") if isinstance(params_field, dict) else None
|
|
)
|
|
if data is not None:
|
|
yield data
|
|
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)
|
|
active: dict[str, 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
|
|
params_field = item.get("params") or {}
|
|
if _event_namespace(params_field) != self._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._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:
|
|
# No active stream matches this event's key. Drop rather
|
|
# than silently misroute to the only remaining stream.
|
|
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)
|
|
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:<run_id>` 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 without a depth limit so that all
|
|
# descendant-namespace events are captured in one SSE and buffered into each
|
|
# child handle's inbox. This avoids a second SSE open (and the dedup-set
|
|
# conflict that would prevent replaying already-seen event_ids).
|
|
return {
|
|
"channels": ["messages", "tasks", "tools"],
|
|
"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
|
|
|
|
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)
|
|
seen: set[tuple[str, ...]] = set()
|
|
active: dict[tuple[str, ...], ScopedStreamHandle] = {}
|
|
# Activate root inbox so scope-level messages events consumed here are
|
|
# forwarded to `thread.messages` even after the shared SSE ends.
|
|
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 {}
|
|
namespace = _event_namespace(params_field)
|
|
data = (
|
|
params_field.get("data") if isinstance(params_field, dict) else None
|
|
)
|
|
if not isinstance(data, dict):
|
|
continue
|
|
method = item.get("method")
|
|
|
|
# Route events at a child namespace (or deeper) to that child
|
|
# handle's channel inbox so sequential child-projection
|
|
# consumption works without opening a second SSE.
|
|
ns_tuple = tuple(namespace)
|
|
routed_to_child = False
|
|
for child_path, child_handle in active.items():
|
|
child_len = len(child_path)
|
|
if (
|
|
len(ns_tuple) >= child_len
|
|
and ns_tuple[:child_len] == child_path
|
|
):
|
|
child_handle._push_event(item)
|
|
routed_to_child = True
|
|
break
|
|
|
|
# Scope-level messages events are not routed to any child; forward
|
|
# them to the root inbox so `thread.messages` can drain them after
|
|
# this projection finishes (dedup prevents the SSE from replaying).
|
|
if (
|
|
not routed_to_child
|
|
and root_inbox is not None
|
|
and method == "messages"
|
|
and tuple(namespace) == self._scope
|
|
):
|
|
root_inbox.put_nowait(item)
|
|
|
|
if method == "tasks":
|
|
if "result" in data:
|
|
self._apply_tasks_result(namespace, data, active)
|
|
elif _is_direct_child(namespace, self._scope):
|
|
path = tuple(namespace)
|
|
if path not in seen:
|
|
seen.add(path)
|
|
graph_name, trigger_call_id = _parse_namespace_segment(
|
|
path[-1]
|
|
)
|
|
handle = ScopedStreamHandle(
|
|
thread=self._thread,
|
|
path=path,
|
|
graph_name=graph_name or None,
|
|
trigger_call_id=trigger_call_id,
|
|
)
|
|
active[path] = handle
|
|
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 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)
|
|
|
|
def _apply_tasks_result(
|
|
self,
|
|
namespace: list[str],
|
|
data: dict[str, Any],
|
|
active: dict[tuple[str, ...], ScopedStreamHandle],
|
|
) -> None:
|
|
result_id = data.get("id")
|
|
if not result_id:
|
|
return
|
|
parent_path = tuple(namespace)
|
|
for child_path, handle in list(active.items()):
|
|
if child_path[:-1] != parent_path:
|
|
continue
|
|
if handle.trigger_call_id != result_id:
|
|
continue
|
|
status, error = _terminal_from_tasks_result(data)
|
|
handle._finish(status, error)
|
|
del active[child_path]
|
|
|
|
|
|
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)
|
|
active: dict[str, 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
|
|
params_field = item.get("params") or {}
|
|
if _event_namespace(params_field) != self._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._namespace),
|
|
)
|
|
active[tool_call_id] = handle
|
|
self._thread._register_active_tool_call(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:
|
|
self._thread._unregister_active_tool_call(handle)
|
|
handle._finish(data.get("output"))
|
|
elif event_type == "tool-error":
|
|
handle = active.pop(tool_call_id, None)
|
|
if handle is not None:
|
|
self._thread._unregister_active_tool_call(handle)
|
|
message = data.get("message")
|
|
handle._fail(
|
|
RuntimeError(
|
|
str(message) if message else "Tool call errored"
|
|
)
|
|
)
|
|
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 = (
|
|
terminal_err
|
|
if terminal_err is not None
|
|
else RuntimeError("Tool call stream closed before terminal tool event.")
|
|
)
|
|
for handle in active.values():
|
|
self._thread._unregister_active_tool_call(handle)
|
|
handle._fail(err)
|
|
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,
|
|
) -> 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._closed = False
|
|
self._transport: ProtocolSseTransport | 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: set[str] = set()
|
|
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._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.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
|
|
|
|
@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.")
|
|
self._transport = ProtocolSseTransport(
|
|
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 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 _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.
|
|
"""
|
|
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)
|
|
except Exception:
|
|
# Pump errored — close all subscription queues so consumers
|
|
# don't hang.
|
|
for sub in self._subscriptions.values():
|
|
sub.queue.put_nowait(None)
|
|
raise
|
|
if self._shared_stream is shared:
|
|
# No rotation happened; stream genuinely ended.
|
|
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 _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))
|
|
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:
|
|
# Why: this watcher is intentionally one-shot. If it crashes, it stays
|
|
# dead until the AsyncThreadStream is closed.
|
|
if self._lifecycle_watcher_task is not None:
|
|
return
|
|
self._lifecycle_watcher_task = asyncio.create_task(
|
|
self._run_lifecycle_watcher()
|
|
)
|
|
|
|
async def _run_lifecycle_watcher(self) -> None:
|
|
"""Always-on SSE consuming lifecycle + input channels.
|
|
|
|
Independent of the union-filter shared stream so that interrupts
|
|
surface even when no other subscription is active. Starts immediately
|
|
on session entry (before any run.start) so reattach and thread.output
|
|
work for existing runs.
|
|
"""
|
|
if self._transport is None:
|
|
return
|
|
try:
|
|
handle = self._transport.open_event_stream(
|
|
{"channels": ["lifecycle", "input"]}
|
|
)
|
|
self._lifecycle_watcher_handle = handle
|
|
await asyncio.wait_for(handle.ready, timeout=5.0)
|
|
async for event in handle.events:
|
|
if self._closed:
|
|
return
|
|
await self._apply_lifecycle_event(event)
|
|
# Why: iterator exhausted without `_run_done` being resolved by a
|
|
# terminal lifecycle event. Surface any transport error captured
|
|
# on `handle.done`, otherwise treat the clean EOF as errored so
|
|
# awaiters of `_run_done` (e.g. `thread.output`) don't hang.
|
|
err = await handle.done
|
|
run_done = self._run_done
|
|
if run_done is not None and not run_done.done():
|
|
if err is not None:
|
|
run_done.set_result(
|
|
_RunTerminal(
|
|
status="errored",
|
|
error=RuntimeError(f"Lifecycle transport failed: {err}"),
|
|
)
|
|
)
|
|
else:
|
|
run_done.set_result(
|
|
_RunTerminal(
|
|
status="errored",
|
|
error=RuntimeError(
|
|
"lifecycle stream ended before terminal event"
|
|
),
|
|
)
|
|
)
|
|
return
|
|
except (Exception, asyncio.CancelledError) as exc:
|
|
# Why: advisory-only watcher. Any error (HTTP failure, malformed
|
|
# event in `_apply_lifecycle_event`, cancellation on close) must
|
|
# not crash the caller; the watcher is one-shot best-effort.
|
|
# Resolve _run_done with an error so thread.output doesn't wait
|
|
# forever when the lifecycle transport fails.
|
|
run_done = self._run_done
|
|
if run_done is not None and not run_done.done():
|
|
if not isinstance(exc, asyncio.CancelledError):
|
|
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:
|
|
self.interrupts.append(payload)
|
|
self.interrupted = True
|
|
elif method == "lifecycle":
|
|
params = event.get("params") or {}
|
|
data = params.get("data") if isinstance(params, dict) else None
|
|
phase = data.get("phase") 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", "errored"):
|
|
# 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 == "errored":
|
|
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"))
|