diff --git a/libs/sdk-py/langgraph_sdk/_async/stream.py b/libs/sdk-py/langgraph_sdk/_async/stream.py index d19983157..e653fea06 100644 --- a/libs/sdk-py/langgraph_sdk/_async/stream.py +++ b/libs/sdk-py/langgraph_sdk/_async/stream.py @@ -66,6 +66,24 @@ _ALL_CHANNELS: list[str] = [ ] +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`. @@ -294,8 +312,11 @@ class _MessagesProjection: from the root namespace only. """ - def __init__(self, thread: AsyncThreadStream) -> None: + 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() @@ -303,11 +324,15 @@ class _MessagesProjection: async def _messages_iter(self) -> AsyncGenerator[AsyncChatModelStream, None]: if self._thread._transport is None: raise RuntimeError("AsyncThreadStream not entered — use `async with`.") - params: SubscribeParams = { - "channels": ["messages"], - "namespaces": [[]], - "depth": 0, - } + # 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: @@ -318,11 +343,11 @@ class _MessagesProjection: if item is None: return params_field = item.get("params") or {} - if not isinstance(params_field, dict): + if _event_namespace(params_field) != self._namespace: continue - if params_field.get("namespace") not in (None, []): - continue - data = params_field.get("data") + data = ( + params_field.get("data") if isinstance(params_field, dict) else None + ) if not isinstance(data, dict): continue event_type = data.get("event") @@ -335,7 +360,7 @@ class _MessagesProjection: else {} ) stream = AsyncChatModelStream( - namespace=[], + namespace=list(self._namespace), node=metadata.get("langgraph_node") if metadata else None, message_id=message_id, ) @@ -361,6 +386,61 @@ class _MessagesProjection: 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") @@ -382,6 +462,474 @@ def _message_route_key(data: dict[str, Any], fallback: str | None = None) -> str 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() + + 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.""" @@ -453,8 +1001,11 @@ class ToolCallHandle: class _ToolCallsProjection: """Typed projection for root-scope `thread.tool_calls`.""" - def __init__(self, thread: AsyncThreadStream) -> None: + 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() @@ -462,11 +1013,7 @@ class _ToolCallsProjection: async def _tool_calls_iter(self) -> AsyncGenerator[ToolCallHandle, None]: if self._thread._transport is None: raise RuntimeError("AsyncThreadStream not entered - use `async with`.") - params: SubscribeParams = { - "channels": ["tools"], - "namespaces": [[]], - "depth": 0, - } + params = _exact_namespace_params(["tools"], self._namespace) sub = self._thread._register_subscription(params) active: dict[str, ToolCallHandle] = {} try: @@ -477,12 +1024,11 @@ class _ToolCallsProjection: if item is None: return params_field = item.get("params") or {} - if not isinstance(params_field, dict): + if _event_namespace(params_field) != self._namespace: continue - namespace = params_field.get("namespace") or [] - if namespace != []: - continue - data = params_field.get("data") + data = ( + params_field.get("data") if isinstance(params_field, dict) else None + ) if not isinstance(data, dict): continue event_type = data.get("event") @@ -498,7 +1044,7 @@ class _ToolCallsProjection: tool_call_id=tool_call_id, name=tool_name, input=data.get("input"), - namespace=namespace, + namespace=list(self._namespace), ) active[tool_call_id] = handle self._thread._register_active_tool_call(handle) @@ -594,11 +1140,17 @@ class AsyncThreadStream: self._run_done: asyncio.Future[_RunTerminal] | 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) - self.tool_calls = _ToolCallsProjection(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: @@ -692,6 +1244,16 @@ class AsyncThreadStream: """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) diff --git a/libs/sdk-py/tests/streaming/test_scoped_handles.py b/libs/sdk-py/tests/streaming/test_scoped_handles.py new file mode 100644 index 000000000..c5614e0f4 --- /dev/null +++ b/libs/sdk-py/tests/streaming/test_scoped_handles.py @@ -0,0 +1,579 @@ +"""Tests for nested scoped stream handles.""" + +from __future__ import annotations + +import httpx + +from langgraph_sdk._async.http import HttpClient +from langgraph_sdk._async.threads import ThreadsClient +from streaming._events import ( + lifecycle_completed_event, + lifecycle_started_event, + message_finish_event, + message_start_event, + message_text_delta_event, + message_text_finish_event, + tasks_result_event, + tasks_start_event, + tool_finished_event, + tool_output_delta_event, + tool_started_event, +) +from streaming._fake_server import FakeServer + + +async def test_subgraphs_subscribes_to_tasks_channel(): + fake = FakeServer() + fake.script([lifecycle_completed_event(seq=1)]) + asgi = httpx.ASGITransport(app=fake.app) + async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw: + threads = ThreadsClient(HttpClient(raw)) + async with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + await thread.run.start(input={}) + _ = [handle async for handle in thread.subgraphs] + + assert any( + "tasks" in body.get("channels", []) for body in fake.stream_request_bodies + ) + + +async def test_subgraphs_yields_handle_and_completes_status(): + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + tasks_start_event(seq=1, namespace=["worker:abc"], task_id="t-child"), + tasks_result_event(seq=2, namespace=[], task_id="abc", name="worker"), + lifecycle_completed_event(seq=3), + ] + ) + asgi = httpx.ASGITransport(app=fake.app) + async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw: + threads = ThreadsClient(HttpClient(raw)) + async with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + await thread.run.start(input={}) + handles = [handle async for handle in thread.subgraphs] + + assert len(handles) == 1 + handle = handles[0] + assert handle.path == ("worker:abc",) + assert handle.namespace == ["worker:abc"] + assert handle.graph_name == "worker" + assert handle.trigger_call_id == "abc" + assert handle.status == "completed" + assert handle.error is None + + +async def test_subgraphs_failed_and_interrupted_statuses(): + failed = FakeServer() + failed.script( + [ + lifecycle_started_event(seq=0), + tasks_start_event(seq=1, namespace=["worker:abc"], task_id="t-child"), + tasks_result_event( + seq=2, + namespace=[], + task_id="abc", + name="worker", + error="boom", + ), + lifecycle_completed_event(seq=3), + ] + ) + failed_asgi = httpx.ASGITransport(app=failed.app) + async with httpx.AsyncClient(transport=failed_asgi, base_url="http://test") as raw: + threads = ThreadsClient(HttpClient(raw)) + async with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + await thread.run.start(input={}) + [failed_handle] = [handle async for handle in thread.subgraphs] + + interrupted = FakeServer() + interrupted.script( + [ + lifecycle_started_event(seq=0), + tasks_start_event(seq=1, namespace=["worker:def"], task_id="t-child"), + tasks_result_event( + seq=2, + namespace=[], + task_id="def", + name="worker", + interrupts=[{"value": "pause"}], + ), + lifecycle_completed_event(seq=3), + ] + ) + interrupted_asgi = httpx.ASGITransport(app=interrupted.app) + async with httpx.AsyncClient( + transport=interrupted_asgi, + base_url="http://test", + ) as raw: + threads = ThreadsClient(HttpClient(raw)) + async with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + await thread.run.start(input={}) + [interrupted_handle] = [handle async for handle in thread.subgraphs] + + assert failed_handle.status == "failed" + assert failed_handle.error == "boom" + assert interrupted_handle.status == "interrupted" + assert interrupted_handle.error is None + + +async def test_subgraph_messages_are_scoped_to_child_namespace(): + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + tasks_start_event(seq=1, namespace=["worker:abc"], task_id="t-child"), + message_start_event( + seq=2, + namespace=["worker:abc"], + message_id="msg-child", + run_id="run-child", + ), + message_text_delta_event(seq=3, namespace=["worker:abc"], text="child"), + message_text_finish_event(seq=4, namespace=["worker:abc"], text="child"), + message_finish_event(seq=5, namespace=["worker:abc"]), + tasks_result_event(seq=6, namespace=[], task_id="abc", name="worker"), + lifecycle_completed_event(seq=7), + ] + ) + asgi = httpx.ASGITransport(app=fake.app) + async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw: + threads = ThreadsClient(HttpClient(raw)) + async with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + await thread.run.start(input={}) + [handle] = [h async for h in thread.subgraphs] + child_messages = [message async for message in handle.messages] + root_messages = [message async for message in thread.messages] + + assert [message.message_id for message in child_messages] == ["msg-child"] + assert [await message.text for message in child_messages] == ["child"] + assert root_messages == [] + + +async def test_subgraph_tool_calls_are_scoped_to_child_namespace(): + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + tasks_start_event(seq=1, namespace=["worker:abc"], task_id="t-child"), + tool_started_event( + seq=2, + namespace=["worker:abc"], + tool_call_id="call-child", + tool_name="search", + ), + tool_output_delta_event( + seq=3, + namespace=["worker:abc"], + tool_call_id="call-child", + delta="child-delta", + ), + tool_finished_event( + seq=4, + namespace=["worker:abc"], + tool_call_id="call-child", + output={"ok": True}, + ), + tasks_result_event(seq=5, namespace=[], task_id="abc", name="worker"), + lifecycle_completed_event(seq=6), + ] + ) + asgi = httpx.ASGITransport(app=fake.app) + async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw: + threads = ThreadsClient(HttpClient(raw)) + async with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + await thread.run.start(input={}) + [handle] = [h async for h in thread.subgraphs] + child_calls = [call async for call in handle.tool_calls] + root_calls = [call async for call in thread.tool_calls] + + assert [call.tool_call_id for call in child_calls] == ["call-child"] + assert [delta async for delta in child_calls[0].deltas] == ["child-delta"] + assert await child_calls[0].output == {"ok": True} + assert root_calls == [] + + +async def test_subgraph_handles_are_recursive_for_grandchildren(): + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + tasks_start_event(seq=1, namespace=["worker:abc"], task_id="t-child"), + tasks_start_event( + seq=2, + namespace=["worker:abc", "tool:def"], + task_id="t-grandchild", + ), + tasks_result_event( + seq=3, + namespace=["worker:abc"], + task_id="def", + name="tool", + ), + tasks_result_event(seq=4, namespace=[], task_id="abc", name="worker"), + lifecycle_completed_event(seq=5), + ] + ) + asgi = httpx.ASGITransport(app=fake.app) + async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw: + threads = ThreadsClient(HttpClient(raw)) + async with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + await thread.run.start(input={}) + [child] = [handle async for handle in thread.subgraphs] + [grandchild] = [handle async for handle in child.subgraphs] + + assert child.path == ("worker:abc",) + assert grandchild.path == ("worker:abc", "tool:def") + assert grandchild.graph_name == "tool" + assert grandchild.trigger_call_id == "def" + assert grandchild.status == "completed" + + +async def test_subagents_aliases_subgraphs_until_protocol_distinguishes_them(): + fake = FakeServer() + fake.script([lifecycle_completed_event(seq=1)]) + asgi = httpx.ASGITransport(app=fake.app) + async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw: + threads = ThreadsClient(HttpClient(raw)) + async with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + assert thread.subagents is thread.subgraphs + await thread.run.start(input={}) + + +async def test_root_and_child_projections_do_not_cross_talk(): + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + tasks_start_event(seq=1, namespace=["worker:abc"], task_id="t-child"), + message_start_event(seq=2, message_id="root-msg", run_id="root-run"), + message_text_delta_event(seq=3, text="root"), + message_text_finish_event(seq=4, text="root"), + message_finish_event(seq=5), + message_start_event( + seq=6, + namespace=["worker:abc"], + message_id="child-msg", + run_id="child-run", + ), + message_text_delta_event(seq=7, namespace=["worker:abc"], text="child"), + message_text_finish_event(seq=8, namespace=["worker:abc"], text="child"), + message_finish_event(seq=9, namespace=["worker:abc"]), + tasks_result_event(seq=10, namespace=[], task_id="abc", name="worker"), + lifecycle_completed_event(seq=11), + ] + ) + asgi = httpx.ASGITransport(app=fake.app) + async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw: + threads = ThreadsClient(HttpClient(raw)) + async with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + await thread.run.start(input={}) + [handle] = [h async for h in thread.subgraphs] + root_messages = [message async for message in thread.messages] + child_messages = [message async for message in handle.messages] + + assert [message.message_id for message in root_messages] == ["root-msg"] + assert [await message.text for message in root_messages] == ["root"] + assert [message.message_id for message in child_messages] == ["child-msg"] + assert [await message.text for message in child_messages] == ["child"] + + +async def test_grandchild_sibling_routing_preserves_event_order(): + """Events enqueued in a child handle's _messages_inbox before a grandchild + is discovered via child.subgraphs must be delivered in arrival order, not + reordered by the drain-and-replay path in _route_sibling_inboxes_to_grandchildren.""" + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + # Child discovered by thread.subgraphs. + tasks_start_event(seq=1, namespace=["worker:abc"], task_id="t-child"), + # Grandchild messages arrive *before* grandchild tasks-start. + message_start_event( + seq=2, + namespace=["worker:abc", "tool:gc1"], + message_id="msg-E1", + run_id="r1", + ), + message_text_delta_event( + seq=3, namespace=["worker:abc", "tool:gc1"], text="E1" + ), + message_text_finish_event( + seq=4, namespace=["worker:abc", "tool:gc1"], text="E1" + ), + message_finish_event(seq=5, namespace=["worker:abc", "tool:gc1"]), + message_start_event( + seq=6, + namespace=["worker:abc", "tool:gc1"], + message_id="msg-E2", + run_id="r2", + ), + message_text_delta_event( + seq=7, namespace=["worker:abc", "tool:gc1"], text="E2" + ), + message_text_finish_event( + seq=8, namespace=["worker:abc", "tool:gc1"], text="E2" + ), + message_finish_event(seq=9, namespace=["worker:abc", "tool:gc1"]), + message_start_event( + seq=10, + namespace=["worker:abc", "tool:gc1"], + message_id="msg-E3", + run_id="r3", + ), + message_text_delta_event( + seq=11, namespace=["worker:abc", "tool:gc1"], text="E3" + ), + message_text_finish_event( + seq=12, namespace=["worker:abc", "tool:gc1"], text="E3" + ), + message_finish_event(seq=13, namespace=["worker:abc", "tool:gc1"]), + # Grandchild tasks-start arrives *after* its messages. + tasks_start_event( + seq=14, + namespace=["worker:abc", "tool:gc1"], + task_id="t-grandchild", + ), + tasks_result_event( + seq=15, namespace=["worker:abc"], task_id="gc1", name="tool" + ), + tasks_result_event(seq=16, namespace=[], task_id="abc", name="worker"), + lifecycle_completed_event(seq=17), + ] + ) + asgi = httpx.ASGITransport(app=fake.app) + async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw: + threads = ThreadsClient(HttpClient(raw)) + async with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + await thread.run.start(input={}) + [child] = [h async for h in thread.subgraphs] + [grandchild] = [h async for h in child.subgraphs] + gc_messages = [m async for m in grandchild.messages] + + assert [m.message_id for m in gc_messages] == ["msg-E1", "msg-E2", "msg-E3"] + texts = [await m.text for m in gc_messages] + assert texts == ["E1", "E2", "E3"] + + +async def test_grandchild_events_dispatched_to_correct_sibling_not_first_match(): + """When two sibling grandchildren exist, messages scoped to one grandchild + must not bleed into the other grandchild's inbox.""" + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + tasks_start_event(seq=1, namespace=["worker:abc"], task_id="t-child"), + # Two grandchildren discovered. + tasks_start_event( + seq=2, namespace=["worker:abc", "tool:gc1"], task_id="t-gc1" + ), + tasks_start_event( + seq=3, namespace=["worker:abc", "tool:gc2"], task_id="t-gc2" + ), + # Messages for gc1. + message_start_event( + seq=4, + namespace=["worker:abc", "tool:gc1"], + message_id="msg-gc1", + run_id="ra", + ), + message_text_delta_event( + seq=5, namespace=["worker:abc", "tool:gc1"], text="GC1" + ), + message_text_finish_event( + seq=6, namespace=["worker:abc", "tool:gc1"], text="GC1" + ), + message_finish_event(seq=7, namespace=["worker:abc", "tool:gc1"]), + # Messages for gc2. + message_start_event( + seq=8, + namespace=["worker:abc", "tool:gc2"], + message_id="msg-gc2", + run_id="rb", + ), + message_text_delta_event( + seq=9, namespace=["worker:abc", "tool:gc2"], text="GC2" + ), + message_text_finish_event( + seq=10, namespace=["worker:abc", "tool:gc2"], text="GC2" + ), + message_finish_event(seq=11, namespace=["worker:abc", "tool:gc2"]), + tasks_result_event( + seq=12, namespace=["worker:abc"], task_id="gc1", name="tool" + ), + tasks_result_event( + seq=13, namespace=["worker:abc"], task_id="gc2", name="tool" + ), + tasks_result_event(seq=14, namespace=[], task_id="abc", name="worker"), + lifecycle_completed_event(seq=15), + ] + ) + asgi = httpx.ASGITransport(app=fake.app) + async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw: + threads = ThreadsClient(HttpClient(raw)) + async with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + await thread.run.start(input={}) + [child] = [h async for h in thread.subgraphs] + grandchildren = [h async for h in child.subgraphs] + by_path = {h.path: h for h in grandchildren} + gc1_messages = [ + m async for m in by_path[("worker:abc", "tool:gc1")].messages + ] + gc2_messages = [ + m async for m in by_path[("worker:abc", "tool:gc2")].messages + ] + + assert [m.message_id for m in gc1_messages] == ["msg-gc1"] + assert [m.message_id for m in gc2_messages] == ["msg-gc2"] + + +def test_scoped_handle_inboxes_bounded_by_max_queue_size(): + """ScopedStreamHandle with max_queue_size=N creates queues with maxsize=N.""" + from unittest.mock import MagicMock + + from langgraph_sdk._async.stream import ScopedStreamHandle + + fake_thread = MagicMock() + handle = ScopedStreamHandle( + thread=fake_thread, + path=("worker:1",), + graph_name="worker", + trigger_call_id="1", + max_queue_size=16, + ) + assert handle._messages_inbox.maxsize == 16 + assert handle._tools_inbox.maxsize == 16 + assert handle._tasks_inbox.maxsize == 16 + + +async def test_child_handle_inherits_max_queue_size_from_parent(): + """Grandchild ScopedStreamHandles created by _HandleSubgraphsProjection + inherit the parent's max_queue_size so all queues are consistently bounded.""" + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + tasks_start_event(seq=1, namespace=["worker:abc"], task_id="t-child"), + tasks_start_event( + seq=2, namespace=["worker:abc", "tool:gc1"], task_id="t-gc1" + ), + tasks_result_event( + seq=3, namespace=["worker:abc"], task_id="gc1", name="tool" + ), + tasks_result_event(seq=4, namespace=[], task_id="abc", name="worker"), + lifecycle_completed_event(seq=5), + ] + ) + asgi = httpx.ASGITransport(app=fake.app) + async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw: + threads = ThreadsClient(HttpClient(raw)) + async with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + await thread.run.start(input={}) + [child] = [h async for h in thread.subgraphs] + [grandchild] = [h async for h in child.subgraphs] + + # Grandchild handles created by _HandleSubgraphsProjection must use the + # parent handle's max_queue_size (default 0 = unbounded in asyncio.Queue). + assert grandchild._messages_inbox.maxsize == child._messages_inbox.maxsize + assert grandchild._tools_inbox.maxsize == child._tools_inbox.maxsize + assert grandchild._tasks_inbox.maxsize == child._tasks_inbox.maxsize + + +async def test_force_complete_uses_failed_when_run_errored(): + """If the lifecycle signals an errored run, scoped children that are still + 'started' when the subgraphs projection's finally block runs must be + force-finished as 'failed', not 'completed'.""" + from streaming._events import lifecycle_errored_event + + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + tasks_start_event(seq=1, namespace=["worker:abc"], task_id="t-child"), + # No tasks_result — run errored before the child task finished. + lifecycle_errored_event(seq=2, error="boom"), + ] + ) + asgi = httpx.ASGITransport(app=fake.app) + handles: list = [] + async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw: + threads = ThreadsClient(HttpClient(raw)) + async with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + await thread.run.start(input={}) + async for handle in thread.subgraphs: + handles.append(handle) + + assert len(handles) == 1 + child = handles[0] + # The run errored, so the child should be force-finished as "failed". + assert child.status == "failed" + + +async def test_force_complete_uses_completed_when_run_completed(): + """If the lifecycle signals a completed run, any subgraph child still + 'started' at finally time is force-finished as 'completed' (normal case).""" + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + tasks_start_event(seq=1, namespace=["worker:abc"], task_id="t-child"), + # No tasks_result — but lifecycle completed normally. + lifecycle_completed_event(seq=2), + ] + ) + asgi = httpx.ASGITransport(app=fake.app) + handles: list = [] + async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw: + threads = ThreadsClient(HttpClient(raw)) + async with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + await thread.run.start(input={}) + async for handle in thread.subgraphs: + handles.append(handle) + + assert len(handles) == 1 + child = handles[0] + assert child.status == "completed" + + +def test_close_inboxes_does_not_enqueue_on_uniterated_inboxes(): + """_close_inboxes must not push a sentinel on inboxes that had no consumer.""" + from unittest.mock import MagicMock + + from langgraph_sdk._async.stream import ScopedStreamHandle + + fake_thread = MagicMock() + handle = ScopedStreamHandle( + thread=fake_thread, + path=("worker:1",), + graph_name="worker", + trigger_call_id="1", + ) + # No projection iterated — _close_inboxes should leave all queues empty. + handle._close_inboxes() + assert handle._messages_inbox.qsize() == 0 + assert handle._tools_inbox.qsize() == 0 + assert handle._tasks_inbox.qsize() == 0 + + +def test_close_inboxes_enqueues_sentinel_on_iterated_inboxes(): + """_close_inboxes must push a None sentinel only on inboxes that had a consumer, + so projection iterators see the EOF signal.""" + from unittest.mock import MagicMock + + from langgraph_sdk._async.stream import ScopedStreamHandle + + fake_thread = MagicMock() + handle = ScopedStreamHandle( + thread=fake_thread, + path=("worker:1",), + graph_name="worker", + trigger_call_id="1", + ) + handle._mark_iterated("messages") + handle._close_inboxes() + # Only the messages inbox should have a sentinel. + assert handle._messages_inbox.qsize() == 1 + assert handle._messages_inbox.get_nowait() is None + assert handle._tools_inbox.qsize() == 0 + assert handle._tasks_inbox.qsize() == 0