From 1fcb7681825faf0372fc1556910b0337ea10d761 Mon Sep 17 00:00:00 2001 From: Nick Hollon Date: Fri, 29 May 2026 17:14:14 -0400 Subject: [PATCH] feat(sdk-py): extract stream decoders and add interleave_projections (#7935) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Refactors the four (now five) sdk-py streaming projections into reusable, transport-agnostic `Decoder` classes and adds a new `interleave_projections(channels)` method to `AsyncThreadStream` and `SyncThreadStream` that drives multiple decoders from one shared subscription, yielding `(channel_name, item)` tuples in arrival order (the SDK analog of local `GraphRunStream.interleave`). - **New `langgraph_sdk/stream/decoders.py`**: pure `feed(event) -> Iterable[item]` state machines — `ValuesDecoder`, `MessagesDecoder`, `ToolCallsDecoder`, `SubgraphsDecoder`, `ExtensionsDecoder` — behind a `Decoder` Protocol. No subscription/queue/thread access. - **Projection migration (async + sync)**: the five existing projections now delegate their per-event logic to the decoders. Behavior-preserving — the existing test suite is the regression net. Thread-coupled side effects (active-stream registration, terminal-error-on-close, root-inbox forwarding) stay in the projection wrappers; sync messages/tool_calls keep their pre-dispatch contract (handle/stream resolved on yield) via a FIFO-head buffer. - **`interleave_projections`**: flat-namespace channel list (built-ins + extension names), `tool_calls`↔`tools` wire mapping, subgraphs fed every event, extensions keyed by bare name. ### Notable - Fixes a latent sync bug surfaced by the refactor: two tool calls / messages whose events interleave previously dropped the second; both now surface. Locked in with a regression test. - `Decoder.feed` takes `Mapping[str, Any]` (read-only), so the Protocol is load-bearing in both stream files. ### Deferred (not in this PR) - Wiring `RemoteGraph._RemoteGraphRunStream.interleave` to `interleave_projections` (gated on #7927). - Migrating the handle-scoped projections (`_Handle*Projection`) to the decoders — hence the small, verified-identical helper duplication between `decoders.py` and `_async/stream.py`. - `interleave_projections` handles aren't registered for thread-close cleanup (additive follow-up). ## Test Plan - [x] `make test` in `libs/sdk-py` — 464 passed, 0 failures - [x] `make format` / `make lint` (ruff + ty) clean - [x] Per-decoder unit tests in `tests/streaming/test_decoders.py` - [x] `interleave_projections` tests (single-channel, multi-channel arrival order, builtin+extension mix, tool_calls public-name, subgraphs discovery) async + sync - [x] Existing projection suites pass unchanged (regression net for the migration) --- libs/sdk-py/langgraph_sdk/_async/stream.py | 437 ++++++++------- libs/sdk-py/langgraph_sdk/_sync/stream.py | 504 +++++++++--------- libs/sdk-py/langgraph_sdk/stream/decoders.py | 359 +++++++++++++ libs/sdk-py/tests/streaming/_events.py | 12 + libs/sdk-py/tests/streaming/test_decoders.py | 460 ++++++++++++++++ .../tests/streaming/test_sync_projections.py | 78 ++- .../streaming/test_sync_thread_stream.py | 307 +++++++++++ .../tests/streaming/test_thread_stream.py | 260 +++++++++ 8 files changed, 1951 insertions(+), 466 deletions(-) create mode 100644 libs/sdk-py/langgraph_sdk/stream/decoders.py create mode 100644 libs/sdk-py/tests/streaming/test_decoders.py diff --git a/libs/sdk-py/langgraph_sdk/_async/stream.py b/libs/sdk-py/langgraph_sdk/_async/stream.py index d574f5794..870ce2556 100644 --- a/libs/sdk-py/langgraph_sdk/_async/stream.py +++ b/libs/sdk-py/langgraph_sdk/_async/stream.py @@ -18,7 +18,7 @@ import contextlib import random from collections.abc import AsyncGenerator, AsyncIterator, Generator, Mapping from dataclasses import dataclass, field -from typing import Any, Literal, TypedDict +from typing import Any, Literal, TypedDict, cast from langchain_core.language_models.chat_model_stream import AsyncChatModelStream from langchain_protocol import Event, SubscribeParams @@ -26,6 +26,16 @@ from langchain_protocol import Event, SubscribeParams from langgraph_sdk._async.http import HttpClient from langgraph_sdk.schema import QueryParamTypes from langgraph_sdk.stream.controller import _SeenEventIds +from langgraph_sdk.stream.decoders import ( + DataDecoder, + Decoder, + ExtensionsDecoder, + MessagesDecoder, + SubgraphsDecoder, + ToolCallsDecoder, + validate_interleave_channels, +) +from langgraph_sdk.stream.subscription import compute_union_filter, infer_channel from langgraph_sdk.stream.transport import ( AsyncProtocolTransport, EventStreamHandle, @@ -348,6 +358,7 @@ class _ValuesProjection: raise RuntimeError("AsyncThreadStream not entered — use `async with`.") params: SubscribeParams = {"channels": ["values"]} sub = self._thread._register_subscription(params) + decoder = DataDecoder("values") try: await self._thread._reconcile_stream(params) self._thread._ensure_fanout_running() @@ -357,12 +368,8 @@ class _ValuesProjection: 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 + for out in decoder.feed(item): + yield out finally: self._thread._unregister_subscription(sub.id) @@ -397,7 +404,13 @@ class _MessagesProjection: return params = _exact_namespace_params(["messages"], self._namespace) sub = self._thread._register_subscription(params) - active: dict[str, AsyncChatModelStream] = {} + decoder = MessagesDecoder( + namespace=self._namespace, + stream_factory=lambda *, namespace, node, message_id: AsyncChatModelStream( + namespace=namespace, node=node, message_id=message_id + ), + ) + registered: list[AsyncChatModelStream] = [] try: await self._thread._reconcile_stream(params) self._thread._ensure_fanout_running() @@ -405,59 +418,12 @@ class _MessagesProjection: 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 + for stream in decoder.feed(item): self._thread._register_active_message_stream(stream) - stream.dispatch(data) + registered.append(stream) yield stream - else: - key = _message_route_key(data) - stream = active.get(key) - if stream is None and key == "__single__" and len(active) == 1: - # Content-block events (content-block-start / - # content-block-delta / content-block-finish / - # message-finish) don't carry the message ``id`` - # on the wire, so ``_message_route_key`` returns - # ``__single__`` while the active stream was - # registered under ``message:``. When exactly - # one stream is active, that mismatch is - # unambiguous -- the events belong to it. - # Events that DO carry an explicit id which - # doesn't match any active stream are still - # dropped (orphan-delta safety, see - # ``test_messages_orphan_delta_without_matching_key_is_dropped``). - 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(): + for stream in registered: self._thread._unregister_active_message_stream(stream) self._thread._unregister_subscription(sub.id) @@ -950,10 +916,17 @@ class _SubgraphsProjection: 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. + decoder = SubgraphsDecoder( + scope=self._scope, + handle_factory=lambda *, path, graph_name, trigger_call_id: ( + ScopedStreamHandle( + thread=self._thread, + path=path, + graph_name=graph_name, + trigger_call_id=trigger_call_id, + ) + ), + ) root_inbox: asyncio.Queue[Event | None] | None = ( self._thread._activate_root_messages_inbox() if not self._scope else None ) @@ -965,80 +938,14 @@ class _SubgraphsProjection: 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 is not None + and item.get("method") == "messages" + and tuple(_event_namespace(params_field)) == 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 - elif ( - method == "lifecycle" - and data.get("event") == "started" - and _is_direct_child(namespace, self._scope) - ): - # ``create_deep_agent`` and similar surfaces signal - # subagent invocation via a child-namespace - # ``lifecycle: started`` event rather than a ``tasks`` - # event. JS does the same (see ``langgraphjs`` - # ``stream/handles/subgraphs.ts``). - 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 + for handle in decoder.feed(item): + yield handle finally: # Determine terminal status from the parent run's lifecycle result. # If _run_done resolved as errored, force-complete remaining children @@ -1049,32 +956,13 @@ class _SubgraphsProjection: result = run_done.result() if isinstance(result, _RunTerminal) and result.status == "errored": terminal_status = "failed" - for handle in active.values(): + for handle in decoder._active.values(): if handle.status == "started": handle._finish(terminal_status) self._thread._unregister_subscription(sub.id) if root_inbox is not None: root_inbox.put_nowait(None) - 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.""" @@ -1161,7 +1049,18 @@ class _ToolCallsProjection: 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] = {} + decoder = ToolCallsDecoder( + namespace=self._namespace, + handle_factory=lambda *, tool_call_id, name, input, namespace: ( + ToolCallHandle( + tool_call_id=tool_call_id, + name=name, + input=input, + namespace=namespace, + ) + ), + ) + registered: list[ToolCallHandle] = [] try: await self._thread._reconcile_stream(params) self._thread._ensure_fanout_running() @@ -1169,52 +1068,10 @@ class _ToolCallsProjection: 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 + for handle in decoder.feed(item): self._thread._register_active_tool_call(handle) + registered.append(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 @@ -1226,14 +1083,15 @@ class _ToolCallsProjection: if run_done is not None and run_done.done() and not run_done.cancelled(): terminal = run_done.result() terminal_err = terminal.error - err = ( + err: BaseException = ( 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) + for handle in list(decoder._active.values()): handle._fail(err) + for handle in registered: + self._thread._unregister_active_tool_call(handle) self._thread._unregister_subscription(sub.id) @@ -1280,6 +1138,7 @@ class _ExtensionProjection: if self._namespace: params["namespaces"] = [self._namespace] sub = self._thread._register_subscription(params) + decoder = ExtensionsDecoder(name=self._name) try: if self._thread._closed: return @@ -1289,12 +1148,8 @@ class _ExtensionProjection: item = await sub.queue.get() if item is None: return - event_params = item.get("params") or {} - data = ( - event_params.get("data") if isinstance(event_params, dict) else None - ) - if isinstance(data, dict): - yield data + for out in decoder.feed(item): + yield out finally: self._thread._unregister_subscription(sub.id) @@ -1558,6 +1413,174 @@ class AsyncThreadStream: params["depth"] = depth return self._subscription_iter(params) + async def interleave_projections( + self, channels: list[str] + ) -> AsyncIterator[tuple[str, Any]]: + """Yield `(channel_name, item)` tuples across multiple projections. + + One shared subscription drives all per-channel decoders; items arrive + in server-emit order (the SDK analog of `GraphRunStream.interleave`). + + Args: + channels: Flat list of `"values"`, `"messages"`, `"tool_calls"`, + `"subgraphs"`, and/or extension names. Built-ins yield their + typed item (snapshot dict / `AsyncChatModelStream` / + `ToolCallHandle` / `ScopedStreamHandle`); an extension yields + its payload dict, keyed by the bare extension name. + + Note: + Handles and streams are yielded eagerly (before their sub-stream + completes), so items arrive interleaved in real time. To receive a + fully-resolved handle (output already populated), use the dedicated + `thread.tool_calls` / `thread.messages` projections instead. + """ + validate_interleave_channels(channels) + if self._transport is None: + raise RuntimeError("AsyncThreadStream not entered — use `async with`.") + decoders: dict[str, Decoder] = {} + sub_params: list[SubscribeParams] = [] + for ch in channels: + if ch == "values": + decoders[ch] = DataDecoder("values") + sub_params.append({"channels": ["values"]}) + elif ch in ("updates", "checkpoints", "tasks"): + # Plain payload channels (local Updates/Checkpoints/Tasks + # analog). Root-scope filter is load-bearing: a co-requested + # unscoped `values` widens the merged subscription to all + # namespaces, so the decoder itself keeps subgraph payloads out. + decoders[ch] = DataDecoder(ch, namespace=[]) + sub_params.append(_exact_namespace_params([ch], [])) + elif ch == "messages": + decoders[ch] = MessagesDecoder( + namespace=[], + stream_factory=lambda *, namespace, node, message_id: ( + AsyncChatModelStream( + namespace=namespace, node=node, message_id=message_id + ) + ), + ) + sub_params.append(_exact_namespace_params(["messages"], [])) + elif ch == "tool_calls": + decoders[ch] = ToolCallsDecoder( + namespace=[], + handle_factory=lambda *, tool_call_id, name, input, namespace: ( + ToolCallHandle( + tool_call_id=tool_call_id, + name=name, + input=input, + namespace=namespace, + ) + ), + ) + sub_params.append(_exact_namespace_params(["tools"], [])) + elif ch == "subgraphs": + decoders[ch] = SubgraphsDecoder( + scope=(), + handle_factory=lambda *, path, graph_name, trigger_call_id: ( + ScopedStreamHandle( + thread=self, + path=path, + graph_name=graph_name, + trigger_call_id=trigger_call_id, + ) + ), + ) + sub_params.append(_subgraph_subscription_params(())) + else: + decoders[ch] = ExtensionsDecoder(name=ch) + sub_params.append({"channels": [f"custom:{ch}"]}) + if not sub_params: + return + merged = cast( + SubscribeParams, + compute_union_filter(cast(list[dict[str, Any]], sub_params)), + ) + subgraphs = decoders.get("subgraphs") + # Track decoder-created handles so teardown can finalize anything still + # in flight; otherwise an awaiting `handle.output` / `handle.messages` + # would hang after an early break or run termination. + registered_tool_calls: list[ToolCallHandle] = [] + registered_message_streams: list[AsyncChatModelStream] = [] + try: + async for event in self._subscription_iter(merged): + if subgraphs is not None: + for item in subgraphs.feed(event): + yield ("subgraphs", item) + wire = infer_channel(event) + public = self._interleave_public_name(wire) + # subgraphs is driven separately above (it consumes all events); never dispatch it here. + if public is not None and public != "subgraphs": + decoder = decoders.get(public) + if decoder is not None: + for item in decoder.feed(event): + if public == "tool_calls": + self._register_active_tool_call(item) + registered_tool_calls.append(item) + elif public == "messages": + self._register_active_message_stream(item) + registered_message_streams.append(item) + yield (public, item) + finally: + self._finalize_interleave_decoders( + decoders.get("tool_calls"), + subgraphs, + registered_tool_calls, + registered_message_streams, + ) + + @staticmethod + def _interleave_public_name(wire: str | None) -> str | None: + """Map a wire channel name to the public channel name used in interleave tuples.""" + if wire is None: + return None + if wire == "tools": + return "tool_calls" + if wire.startswith("custom:"): + return wire[len("custom:") :] + return wire # values, messages (tasks/lifecycle pass through with no decoder match) + + def _finalize_interleave_decoders( + self, + tool_calls: Decoder | None, + subgraphs: Decoder | None, + registered_tool_calls: list[ToolCallHandle], + registered_message_streams: list[AsyncChatModelStream], + ) -> None: + """Finalize in-flight handles when `interleave_projections` tears down. + + Mirrors the terminal handling of the dedicated `_ToolCallsProjection` / + `_SubgraphsProjection`: in-flight tool calls are failed (so awaiting + `handle.output` can't hang) and discovered subgraph children are + force-completed with the run's terminal status. + """ + run_done = self._run_done + resolved = ( + run_done.result() + if run_done is not None and run_done.done() and not run_done.cancelled() + else None + ) + if isinstance(tool_calls, ToolCallsDecoder): + err: BaseException = ( + resolved.error + if resolved is not None and resolved.error is not None + else RuntimeError("Tool call stream closed before terminal tool event.") + ) + for handle in list(tool_calls._active.values()): + handle._fail(err) + for handle in registered_tool_calls: + self._unregister_active_tool_call(handle) + for stream in registered_message_streams: + self._unregister_active_message_stream(stream) + if isinstance(subgraphs, SubgraphsDecoder): + terminal_status: SubgraphStatus = ( + "failed" + if isinstance(resolved, _RunTerminal) and resolved.status == "errored" + else "completed" + ) + for child in subgraphs._active.values(): + if child.status == "started": + child._finish(terminal_status) + async def _subscription_iter( self, params: SubscribeParams ) -> AsyncGenerator[Event, None]: diff --git a/libs/sdk-py/langgraph_sdk/_sync/stream.py b/libs/sdk-py/langgraph_sdk/_sync/stream.py index f987e794b..9405e7f83 100644 --- a/libs/sdk-py/langgraph_sdk/_sync/stream.py +++ b/libs/sdk-py/langgraph_sdk/_sync/stream.py @@ -17,13 +17,23 @@ import queue import threading from collections.abc import Iterator, Mapping from dataclasses import dataclass -from typing import Any, Literal, TypedDict +from typing import Any, Literal, TypedDict, cast from langchain_core.language_models.chat_model_stream import ChatModelStream from langchain_protocol import Event, SubscribeParams from langgraph_sdk._sync.http import SyncHttpClient from langgraph_sdk.schema import QueryParamTypes +from langgraph_sdk.stream.decoders import ( + DataDecoder, + Decoder, + ExtensionsDecoder, + MessagesDecoder, + SubgraphsDecoder, + ToolCallsDecoder, + validate_interleave_channels, +) +from langgraph_sdk.stream.subscription import compute_union_filter, infer_channel from langgraph_sdk.stream.sync_controller import SyncStreamController, _SyncSubscription from langgraph_sdk.stream.transport import ( SyncEventStreamHandle, @@ -295,6 +305,7 @@ class _SyncValuesProjection: raise RuntimeError("SyncThreadStream not entered — use `with`.") params: SubscribeParams = {"channels": ["values"]} sub = self._thread._register_subscription(params) + decoder = DataDecoder("values") try: self._thread._reconcile_stream(params) self._thread._ensure_fanout_running() @@ -304,12 +315,7 @@ class _SyncValuesProjection: item = 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 + yield from decoder.feed(cast(dict[str, Any], item)) finally: self._thread._unregister_subscription(sub.id) @@ -340,87 +346,35 @@ class _SyncMessagesProjection: return params = _exact_namespace_params(["messages"], self._namespace) sub = self._thread._register_subscription(params) - active: dict[str, ChatModelStream] = {} + decoder = MessagesDecoder( + namespace=self._namespace, + stream_factory=lambda *, namespace, node, message_id: ChatModelStream( + namespace=namespace, node=node, message_id=message_id + ), + ) + registered: list[ChatModelStream] = [] + pending: list[ChatModelStream] = [] try: self._thread._reconcile_stream(params) self._thread._ensure_fanout_running() while True: item = sub.queue.get() if item is None: + # EOF: surface remaining streams (possibly incomplete) in start order. + while pending: + yield pending.pop(0) 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 = ChatModelStream( - namespace=list(self._namespace), - node=metadata.get("langgraph_node") if metadata else None, - message_id=message_id, - ) - active[key] = stream + for stream in decoder.feed(cast(dict[str, Any], item)): self._thread._register_active_message_stream(stream) - stream.dispatch(data) - # Pre-dispatch all remaining events for this message so the - # caller can access str(message.text) inside a for loop. - while not stream._done: - next_item = sub.queue.get() - if next_item is None: - sub.queue.put(None) - break - next_params = next_item.get("params") or {} - next_data = ( - next_params.get("data") - if isinstance(next_params, dict) - else None - ) - if not isinstance(next_data, dict): - continue - next_event_type = next_data.get("event") - next_key = _message_route_key(next_data) - target = active.get(next_key) - if ( - target is None - and next_key == "__single__" - and len(active) == 1 - ): - target = next(iter(active.values())) - if target is not None: - target.dispatch(next_data) - if next_event_type in ("message-finish", "error"): - self._thread._unregister_active_message_stream(target) - for rk, cand in list(active.items()): - if cand is target: - del active[rk] - yield stream - else: - key = _message_route_key(data) - stream = active.get(key) - if stream is None and key == "__single__" 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] + registered.append(stream) + pending.append(stream) + # Surface in start order once each (and all earlier) streams are done, + # so `str(stream.text)` is ready on yield (sync pre-dispatch contract). + while pending and pending[0]._done: + yield pending.pop(0) finally: - for s in active.values(): - self._thread._unregister_active_message_stream(s) + for stream in registered: + self._thread._unregister_active_message_stream(stream) self._thread._unregister_subscription(sub.id) @@ -595,100 +549,37 @@ class _SyncToolCallsProjection: raise RuntimeError("SyncThreadStream not entered — use `with`.") params = _exact_namespace_params(["tools"], self._namespace) sub = self._thread._register_subscription(params) - active: dict[str, SyncToolCallHandle] = {} + decoder = ToolCallsDecoder( + namespace=self._namespace, + handle_factory=lambda *, tool_call_id, name, input, namespace: ( + SyncToolCallHandle( + tool_call_id=tool_call_id, + name=name, + input=input, + namespace=namespace, + ) + ), + ) + registered: list[SyncToolCallHandle] = [] + pending: list[SyncToolCallHandle] = [] try: self._thread._reconcile_stream(params) self._thread._ensure_fanout_running() while True: item = sub.queue.get() if item is None: + # EOF: surface remaining handles (possibly incomplete) in start order. + while pending: + yield pending.pop(0) 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 = SyncToolCallHandle( - tool_call_id=tool_call_id, - name=tool_name, - input=data.get("input"), - namespace=list(self._namespace), - ) - active[tool_call_id] = handle + for handle in decoder.feed(cast(dict[str, Any], item)): self._thread._register_active_tool_call(handle) - # Pre-dispatch events until this tool call completes so that - # `call.output` is resolved when the caller receives the handle. - while not handle.done: - next_item = sub.queue.get() - if next_item is None: - sub.queue.put(None) - break - next_params = next_item.get("params") or {} - if _event_namespace(next_params) != self._namespace: - continue - next_data = ( - next_params.get("data") - if isinstance(next_params, dict) - else None - ) - if not isinstance(next_data, dict): - continue - next_event_type = next_data.get("event") - next_tcid = next_data.get("tool_call_id") - if not isinstance(next_tcid, str): - continue - if next_event_type == "tool-output-delta": - h = active.get(next_tcid) - delta = next_data.get("delta") - if h is not None and isinstance(delta, str): - h._push_delta(delta) - elif next_event_type == "tool-finished": - h = active.pop(next_tcid, None) - if h is not None: - self._thread._unregister_active_tool_call(h) - h._finish(next_data.get("output")) - elif next_event_type == "tool-error": - h = active.pop(next_tcid, None) - if h is not None: - self._thread._unregister_active_tool_call(h) - message = next_data.get("message") - h._fail( - RuntimeError( - str(message) if message else "Tool call errored" - ) - ) - yield handle - elif event_type == "tool-output-delta": - h = active.get(tool_call_id) - delta = data.get("delta") - if h is not None and isinstance(delta, str): - h._push_delta(delta) - elif event_type == "tool-finished": - h = active.pop(tool_call_id, None) - if h is not None: - self._thread._unregister_active_tool_call(h) - h._finish(data.get("output")) - elif event_type == "tool-error": - h = active.pop(tool_call_id, None) - if h is not None: - self._thread._unregister_active_tool_call(h) - message = data.get("message") - h._fail( - RuntimeError( - str(message) if message else "Tool call errored" - ) - ) + registered.append(handle) + pending.append(handle) + # Surface in start order once each (and all earlier) handles are done, + # so `call.output` is resolved when the caller receives the handle. + while pending and pending[0].done: + yield pending.pop(0) finally: # Read terminal error from _run_done if it is already resolved. # We do NOT block here: callers who need a terminal observation @@ -708,9 +599,10 @@ class _SyncToolCallsProjection: if terminal_err is not None else RuntimeError("Tool call stream closed before terminal tool event.") ) - for h in active.values(): - self._thread._unregister_active_tool_call(h) - h._fail(err) + for handle in list(decoder._active.values()): + handle._fail(err) + for handle in registered: + self._thread._unregister_active_tool_call(handle) self._thread._unregister_subscription(sub.id) @@ -1067,8 +959,17 @@ class _SyncSubgraphsProjection: raise RuntimeError("SyncThreadStream not entered — use `with`.") params = _subgraph_subscription_params(self._scope) sub = self._thread._register_subscription(params) - seen: set[tuple[str, ...]] = set() - active: dict[tuple[str, ...], SyncScopedStreamHandle] = {} + decoder = SubgraphsDecoder( + scope=self._scope, + handle_factory=lambda *, path, graph_name, trigger_call_id: ( + SyncScopedStreamHandle( + thread=self._thread, + path=path, + graph_name=graph_name, + trigger_call_id=trigger_call_id, + ) + ), + ) root_inbox: queue.Queue[Event | None] | None = ( self._thread._activate_root_messages_inbox() if not self._scope else None ) @@ -1080,71 +981,14 @@ class _SyncSubgraphsProjection: 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") - - 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 - if ( - not routed_to_child - and root_inbox is not None - and method == "messages" - and tuple(namespace) == self._scope + root_inbox is not None + and item.get("method") == "messages" + and tuple(_event_namespace(params_field)) == 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 = SyncScopedStreamHandle( - thread=self._thread, - path=path, - graph_name=graph_name or None, - trigger_call_id=trigger_call_id, - ) - active[path] = handle - yield handle - elif ( - method == "lifecycle" - and data.get("event") == "started" - and _is_direct_child(namespace, self._scope) - ): - # ``create_deep_agent`` subagent discovery: child- - # namespace ``lifecycle: started`` rather than ``tasks``. - path = tuple(namespace) - if path not in seen: - seen.add(path) - graph_name, trigger_call_id = _parse_namespace_segment(path[-1]) - handle = SyncScopedStreamHandle( - thread=self._thread, - path=path, - graph_name=graph_name or None, - trigger_call_id=trigger_call_id, - ) - active[path] = handle - yield handle + for handle in decoder.feed(cast(dict[str, Any], item)): + yield handle finally: # Determine terminal status from the run's lifecycle result. # If _run_done resolved as errored, force-complete remaining children @@ -1158,32 +1002,13 @@ class _SyncSubgraphsProjection: terminal_status = "failed" except Exception: pass - for handle in active.values(): + for handle in decoder._active.values(): if handle.status == "started": handle._finish(terminal_status) self._thread._unregister_subscription(sub.id) if root_inbox is not None: root_inbox.put_nowait(None) - def _apply_tasks_result( - self, - namespace: list[str], - data: dict[str, Any], - active: dict[tuple[str, ...], SyncScopedStreamHandle], - ) -> 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 _SyncExtensionsProjection: """Mapping from extension name to custom event payload stream. @@ -1230,6 +1055,7 @@ class _SyncExtensionProjection: if self._namespace: params["namespaces"] = [self._namespace] sub = self._thread._register_subscription(params) + decoder = ExtensionsDecoder(name=self._name) try: if self._thread._closed: return @@ -1239,12 +1065,7 @@ class _SyncExtensionProjection: item = sub.queue.get() if item is None: return - event_params = item.get("params") or {} - data = ( - event_params.get("data") if isinstance(event_params, dict) else None - ) - if isinstance(data, dict): - yield data + yield from decoder.feed(cast(dict[str, Any], item)) finally: self._thread._unregister_subscription(sub.id) @@ -1461,6 +1282,173 @@ class SyncThreadStream: params["depth"] = depth return self._subscription_iter(params) + def interleave_projections(self, channels: list[str]) -> Iterator[tuple[str, Any]]: + """Yield `(channel_name, item)` tuples across multiple projections. + + One shared subscription drives all per-channel decoders; items arrive + in server-emit order (the SDK analog of `GraphRunStream.interleave`). + + Args: + channels: Flat list of `"values"`, `"messages"`, `"tool_calls"`, + `"subgraphs"`, and/or extension names. Built-ins yield their + typed item (snapshot dict / `ChatModelStream` / + `SyncToolCallHandle` / `SyncScopedStreamHandle`); an extension + yields its payload dict, keyed by the bare extension name. + + Note: + Handles and streams are yielded eagerly (before their sub-stream + completes), so items arrive interleaved in real time. To receive a + fully-resolved handle (output already populated), use the dedicated + `thread.tool_calls` / `thread.messages` projections instead. + """ + validate_interleave_channels(channels) + if self._transport is None: + raise RuntimeError("SyncThreadStream not entered — use `with`.") + decoders: dict[str, Decoder] = {} + sub_params: list[dict[str, Any]] = [] + for ch in channels: + if ch == "values": + decoders[ch] = DataDecoder("values") + sub_params.append({"channels": ["values"]}) + elif ch in ("updates", "checkpoints", "tasks"): + # Plain payload channels (local Updates/Checkpoints/Tasks + # analog). Root-scope filter is load-bearing: a co-requested + # unscoped `values` widens the merged subscription to all + # namespaces, so the decoder itself keeps subgraph payloads out. + decoders[ch] = DataDecoder(ch, namespace=[]) + sub_params.append(dict(_exact_namespace_params([ch], []))) + elif ch == "messages": + decoders[ch] = MessagesDecoder( + namespace=[], + stream_factory=lambda *, namespace, node, message_id: ( + ChatModelStream( + namespace=namespace, node=node, message_id=message_id + ) + ), + ) + sub_params.append(dict(_exact_namespace_params(["messages"], []))) + elif ch == "tool_calls": + decoders[ch] = ToolCallsDecoder( + namespace=[], + handle_factory=lambda *, tool_call_id, name, input, namespace: ( + SyncToolCallHandle( + tool_call_id=tool_call_id, + name=name, + input=input, + namespace=namespace, + ) + ), + ) + sub_params.append(dict(_exact_namespace_params(["tools"], []))) + elif ch == "subgraphs": + decoders[ch] = SubgraphsDecoder( + scope=(), + handle_factory=lambda *, path, graph_name, trigger_call_id: ( + SyncScopedStreamHandle( + thread=self, + path=path, + graph_name=graph_name, + trigger_call_id=trigger_call_id, + ) + ), + ) + sub_params.append(dict(_subgraph_subscription_params(()))) + else: + decoders[ch] = ExtensionsDecoder(name=ch) + sub_params.append({"channels": [f"custom:{ch}"]}) + if not sub_params: + return + merged = cast( + SubscribeParams, + compute_union_filter(cast("list[dict[str, Any]]", sub_params)), + ) + subgraphs = decoders.get("subgraphs") + # Track decoder-created handles so teardown can finalize anything still + # in flight; otherwise an awaiting `handle.output` / `handle.messages` + # would block after an early break or run termination. + registered_tool_calls: list[SyncToolCallHandle] = [] + registered_message_streams: list[ChatModelStream] = [] + try: + for event in self._subscription_iter(merged): + if subgraphs is not None: + for item in subgraphs.feed(event): + yield ("subgraphs", item) + wire = infer_channel(event) + public = self._interleave_public_name(wire) + # subgraphs is driven separately above (it consumes all events); never dispatch it here. + if public is not None and public != "subgraphs": + decoder = decoders.get(public) + if decoder is not None: + for item in decoder.feed(event): + if public == "tool_calls": + self._register_active_tool_call(item) + registered_tool_calls.append(item) + elif public == "messages": + self._register_active_message_stream(item) + registered_message_streams.append(item) + yield (public, item) + finally: + self._finalize_interleave_decoders( + decoders.get("tool_calls"), + subgraphs, + registered_tool_calls, + registered_message_streams, + ) + + @staticmethod + def _interleave_public_name(wire: str | None) -> str | None: + """Map a wire channel name to the public channel name used in interleave tuples.""" + if wire is None: + return None + if wire == "tools": + return "tool_calls" + if wire.startswith("custom:"): + return wire[len("custom:") :] + return wire # values, messages (tasks/lifecycle pass through with no decoder match) + + def _finalize_interleave_decoders( + self, + tool_calls: Decoder | None, + subgraphs: Decoder | None, + registered_tool_calls: list[SyncToolCallHandle], + registered_message_streams: list[ChatModelStream], + ) -> None: + """Finalize in-flight handles when `interleave_projections` tears down. + + Mirrors the terminal handling of the dedicated `_SyncToolCallsProjection` + / `_SyncSubgraphsProjection`: in-flight tool calls are failed (so a + blocking `handle.output` can't hang) and discovered subgraph children + are force-completed with the run's terminal status. + """ + run_done = self._run_done + resolved: _RunTerminal | None = None + if run_done is not None and run_done.done(): + try: + resolved = run_done.result(timeout=0) + except Exception: + resolved = None + if isinstance(tool_calls, ToolCallsDecoder): + err: BaseException = ( + resolved.error + if resolved is not None and resolved.error is not None + else RuntimeError("Tool call stream closed before terminal tool event.") + ) + for handle in list(tool_calls._active.values()): + handle._fail(err) + for handle in registered_tool_calls: + self._unregister_active_tool_call(handle) + for stream in registered_message_streams: + self._unregister_active_message_stream(stream) + if isinstance(subgraphs, SubgraphsDecoder): + terminal_status: SubgraphStatus = ( + "failed" + if resolved is not None and resolved.status == "errored" + else "completed" + ) + for child in subgraphs._active.values(): + if child.status == "started": + child._finish(terminal_status) + def _subscription_iter(self, params: SubscribeParams) -> Iterator[Event]: sub = self._register_subscription(params) try: diff --git a/libs/sdk-py/langgraph_sdk/stream/decoders.py b/libs/sdk-py/langgraph_sdk/stream/decoders.py new file mode 100644 index 000000000..45959bf46 --- /dev/null +++ b/libs/sdk-py/langgraph_sdk/stream/decoders.py @@ -0,0 +1,359 @@ +"""Per-channel event → items state machines. + +Used both by the projection iterators (`_ValuesProjection`, +`_MessagesProjection`, `_ToolCallsProjection`, `_SubgraphsProjection`) on +`AsyncThreadStream` / `SyncThreadStream`, and by `interleave_projections`, +which drives multiple decoders from one shared subscription. +""" + +from __future__ import annotations + +from collections.abc import Callable, Iterable, Mapping +from typing import Any, Literal, Protocol + +#: Channel names the public ``interleave_projections`` API accepts as built-ins. +SUPPORTED_INTERLEAVE_CHANNELS = ( + "values", + "messages", + "tool_calls", + "subgraphs", + "updates", + "checkpoints", + "tasks", +) + +#: Channel names that ``infer_channel`` recognizes as first-class protocol +#: methods but that ``interleave_projections`` has no decoder for. Routing them +#: to the extension/``custom:`` fallback would subscribe to a channel that never +#: matches and silently yield nothing, so they are rejected up front (fail +#: closed). ``lifecycle`` is control-plane (drives run output/interrupt); ``tools`` +#: is the wire alias for the public ``tool_calls`` channel. +RESERVED_INTERLEAVE_CHANNELS = frozenset({"lifecycle", "tools", "input"}) + + +def validate_interleave_channels(channels: list[str]) -> None: + """Reject reserved protocol channel names before they hit the fallback. + + Genuine extension names pass through untouched; only names that + ``infer_channel`` treats as built-in methods without an interleave decoder + are rejected, so a typo'd or unsupported protocol channel surfaces an error + instead of an empty stream. + """ + for ch in channels: + if ch in RESERVED_INTERLEAVE_CHANNELS: + hint = ' (use "tool_calls")' if ch == "tools" else "" + raise ValueError( + f"{ch!r} is not a valid interleave_projections channel{hint}. " + f"Supported channels: {', '.join(SUPPORTED_INTERLEAVE_CHANNELS)}, " + "or an extension name." + ) + + +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 [] + + +def _message_event_id(data: dict[str, Any]) -> str | None: + message_id = data.get("id") or data.get("message_id") + return str(message_id) if message_id is not None else None + + +def _message_route_key(data: dict[str, Any], fallback: str | None = None) -> str: + """Return the routing key for a message-channel event in `active`. + + Keys on `message_id` when available so concurrent messages that share the + same `run_id` (two AI turns in one agent step) route to independent streams + rather than colliding on a shared `run:` slot. + """ + message_id = _message_event_id(data) + if message_id is not None: + return f"message:{message_id}" + if fallback is not None: + return f"message:{fallback}" + return "__single__" + + +SubgraphStatus = Literal["started", "completed", "failed", "interrupted"] + + +def _parse_namespace_segment(segment: str) -> tuple[str, str | None]: + name, sep, task_id = segment.partition(":") + return name, task_id if sep else None + + +def _terminal_from_tasks_result( + data: dict[str, Any], +) -> tuple[SubgraphStatus, str | None]: + if data.get("interrupts"): + return "interrupted", None + error = data.get("error") + if error: + return "failed", str(error) + return "completed", None + + +def _is_direct_child(namespace: list[str], scope: tuple[str, ...]) -> bool: + return len(namespace) == len(scope) + 1 and tuple(namespace[: len(scope)]) == scope + + +class Decoder(Protocol): + def feed(self, event: Mapping[str, Any]) -> Iterable[Any]: ... + + +class DataDecoder: + """Yields `params.data` from events of a single `method`. + + Covers the channels whose projection is just "emit the payload": `values`, + `updates`, `checkpoints`, `tasks` — the SDK analog of local's + `Values`/`Updates`/`Checkpoints`/`TasksTransformer`, all of which push + `params["data"]` unchanged. The REST-state seeding for `values` stays at + the projection layer; it is a one-shot pre-stream fetch, not part of the + event state machine. + + Args: + method: The protocol `method` this decoder consumes. + namespace: When not `None`, events whose namespace differs are ignored + (scope filter, mirroring the local transformers' `namespace != scope` + check). `None` consumes every namespace — the historical `values` + projection behavior, where subscription scoping is handled upstream. + """ + + def __init__(self, method: str, namespace: list[str] | None = None): + self._method = method + self._namespace = list(namespace) if namespace is not None else None + + def feed(self, event: Mapping[str, Any]) -> Iterable[Any]: + if event.get("method") != self._method: + return + params = event.get("params") or {} + if self._namespace is not None and _event_namespace(params) != self._namespace: + return + data = params.get("data") + if data is not None: + yield data + + +class MessagesDecoder: + """Yields one chat-model stream per `message-start` event. + + Subsequent events route to the matching stream via `stream.dispatch(data)`. + Mirrors the per-event body of `_MessagesProjection._messages_iter` + (`_async/stream.py:404-458`). The subscription open/close and the + `_root_messages_inbox` drain branch stay at the projection layer. + + Args: + namespace: Events whose namespace differs are ignored (scope filter). + stream_factory: Keyword-only `(namespace, node, message_id) -> stream`. + Sync binds `ChatModelStream`; async binds `AsyncChatModelStream`. + """ + + def __init__( + self, + namespace: list[str], + stream_factory: Callable[..., Any], + ): + self._namespace = list(namespace) + self._stream_factory = stream_factory + self._active: dict[str, Any] = {} # route_key -> stream + + def feed(self, event: Mapping[str, Any]) -> Iterable[Any]: + if event.get("method") != "messages": + return + params = event.get("params") or {} + if _event_namespace(params) != self._namespace: + return + data = params.get("data") + if not isinstance(data, dict): + return + if data.get("event") == "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 = self._stream_factory( + namespace=list(self._namespace), + node=metadata.get("langgraph_node") if metadata else None, + message_id=message_id, + ) + self._active[key] = stream + stream.dispatch(data) + yield stream + else: + key = _message_route_key(data) + stream = self._active.get(key) + if stream is None and key == "__single__" and len(self._active) == 1: + stream = next(iter(self._active.values())) + if stream is None: + return + stream.dispatch(data) + if data.get("event") in ("message-finish", "error"): + for route_key, candidate in list(self._active.items()): + if candidate is stream: + del self._active[route_key] + + +class ToolCallsDecoder: + """Yields one tool-call handle per `tool-started` event. + + Mirrors the per-event body of `_ToolCallsProjection._tool_calls_iter` + (`_async/stream.py:1168-1217`). The thread register/unregister and the + terminal-error-on-close finally stay at the projection / wrapper layer. + + Args: + namespace: Events whose namespace differs are ignored. + handle_factory: Keyword-only `(tool_call_id, name, input, namespace) -> handle`. + """ + + def __init__(self, namespace: list[str], handle_factory: Callable[..., Any]): + self._namespace = list(namespace) + self._handle_factory = handle_factory + self._active: dict[str, Any] = {} + + def feed(self, event: Mapping[str, Any]) -> Iterable[Any]: + if event.get("method") != "tools": + return + params = event.get("params") or {} + if _event_namespace(params) != self._namespace: + return + data = params.get("data") + if not isinstance(data, dict): + return + tool_call_id = data.get("tool_call_id") + if not isinstance(tool_call_id, str): + return + event_type = data.get("event") + if event_type == "tool-started": + name = data.get("tool_name") + handle = self._handle_factory( + tool_call_id=tool_call_id, + name=name if isinstance(name, str) else "", + input=data.get("input"), + namespace=list(self._namespace), + ) + self._active[tool_call_id] = handle + yield handle + elif event_type == "tool-output-delta": + handle = self._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 = self._active.pop(tool_call_id, None) + if handle is not None: + handle._finish(data.get("output")) + elif event_type == "tool-error": + handle = self._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 SubgraphsDecoder: + """Discovers child subgraph handles and fans out events to active ones. + + Mirrors the per-event body of `_SubgraphsProjection._subgraphs_iter` + (`_async/stream.py:963-1041`) plus `_apply_tasks_result`. Root-inbox + forwarding and terminal-status-on-close stay at the projection / wrapper + layer. + + Args: + scope: Tuple-form namespace of this decoder's parent. `()` for root. + handle_factory: Keyword-only `(path, graph_name, trigger_call_id) -> handle`. + """ + + def __init__(self, scope: tuple[str, ...], handle_factory: Callable[..., Any]): + self._scope = scope + self._handle_factory = handle_factory + self._active: dict[tuple[str, ...], Any] = {} + self._seen: set[tuple[str, ...]] = set() + + def feed(self, event: Mapping[str, Any]) -> Iterable[Any]: + params = event.get("params") or {} + namespace = _event_namespace(params) + data = params.get("data") + if not isinstance(data, dict): + return + method = event.get("method") + + # 1. Fanout: first active child whose path prefixes this namespace. + ns_tuple = tuple(namespace) + for child_path, child_handle in self._active.items(): + child_len = len(child_path) + if len(ns_tuple) >= child_len and ns_tuple[:child_len] == child_path: + child_handle._push_event(event) + break + + # 2 + 3. Discovery / status from tasks; discovery from lifecycle. + if method == "tasks": + if "result" in data: + self._apply_tasks_result(namespace, data) + elif _is_direct_child(namespace, self._scope): + yield from self._discover(namespace) + elif ( + method == "lifecycle" + and data.get("event") == "started" + and _is_direct_child(namespace, self._scope) + ): + yield from self._discover(namespace) + + def _discover(self, namespace: list[str]) -> Iterable[Any]: + path = tuple(namespace) + if path in self._seen: + return + self._seen.add(path) + graph_name, trigger_call_id = _parse_namespace_segment(path[-1]) + handle = self._handle_factory( + path=path, + graph_name=graph_name or None, + trigger_call_id=trigger_call_id, + ) + self._active[path] = handle + yield handle + + def _apply_tasks_result(self, namespace: list[str], data: dict[str, Any]) -> None: + result_id = data.get("id") + if not result_id: + return + parent_path = tuple(namespace) + for child_path, handle in list(self._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 self._active[child_path] + + +class ExtensionsDecoder: + """Yields `params.data` from one named custom channel. + + Mirrors `_ExtensionProjection._iter` (`_async/stream.py:1278-1299`), with + an added name filter so it can share one subscription in interleave. + + Args: + name: The extension name. Only `custom` events whose `data["name"]` + matches are consumed. + """ + + def __init__(self, name: str): + if not name: + raise ValueError("extension name must be non-empty.") + self._name = name + + def feed(self, event: Mapping[str, Any]) -> Iterable[Any]: + if event.get("method") != "custom": + return + params = event.get("params") or {} + data = params.get("data") + if not isinstance(data, dict): + return + if data.get("name") != self._name: + return + yield data diff --git a/libs/sdk-py/tests/streaming/_events.py b/libs/sdk-py/tests/streaming/_events.py index ba0d752e2..5db352df6 100644 --- a/libs/sdk-py/tests/streaming/_events.py +++ b/libs/sdk-py/tests/streaming/_events.py @@ -75,6 +75,18 @@ def values_event( return _base(seq, "values", namespace or [], data or {"values": {}}) +def updates_event( + seq: int = 0, namespace: list[str] | None = None, **data: Any +) -> dict[str, Any]: + return _base(seq, "updates", namespace or [], data or {}) + + +def checkpoints_event( + seq: int = 0, namespace: list[str] | None = None, **data: Any +) -> dict[str, Any]: + return _base(seq, "checkpoints", namespace or [], data or {}) + + def custom_event( seq: int = 0, name: str = "ext", namespace: list[str] | None = None, **data: Any ) -> dict[str, Any]: diff --git a/libs/sdk-py/tests/streaming/test_decoders.py b/libs/sdk-py/tests/streaming/test_decoders.py new file mode 100644 index 000000000..c1ebcd98d --- /dev/null +++ b/libs/sdk-py/tests/streaming/test_decoders.py @@ -0,0 +1,460 @@ +"""Unit tests for the per-channel Decoders. + +Each test drives a single decoder with synthetic events from `_events` and +asserts the items the decoder yields. +""" + +from __future__ import annotations + +from typing import Any + +from langgraph_sdk.stream.decoders import ( + DataDecoder, + ExtensionsDecoder, + MessagesDecoder, + SubgraphsDecoder, + ToolCallsDecoder, +) +from streaming._events import ( + checkpoints_event, + custom_event, + lifecycle_completed_event, + lifecycle_started_event, + message_error_event, + message_finish_event, + message_start_event, + message_text_delta_event, + tasks_result_event, + tasks_start_event, + tool_error_event, + tool_finished_event, + tool_output_delta_event, + tool_started_event, + updates_event, + values_event, +) + + +def test_data_decoder_yields_params_data(): + decoder = DataDecoder("values") + assert list(decoder.feed(values_event(seq=1, x=1))) == [{"x": 1}] + assert list(decoder.feed(values_event(seq=2, x=2, y=3))) == [{"x": 2, "y": 3}] + + +def test_data_decoder_ignores_other_methods(): + decoder = DataDecoder("values") + assert list(decoder.feed(lifecycle_completed_event(seq=1))) == [] + assert list(decoder.feed(updates_event(seq=2, foo=1))) == [] + + +def test_data_decoder_handles_updates_checkpoints_tasks_methods(): + assert list(DataDecoder("updates").feed(updates_event(seq=1, node={"x": 1}))) == [ + {"node": {"x": 1}} + ] + assert list( + DataDecoder("checkpoints").feed(checkpoints_event(seq=2, ts="t", v=4)) + ) == [{"ts": "t", "v": 4}] + # tasks payloads pass through verbatim as data dicts + [item] = list(DataDecoder("tasks").feed(tasks_start_event(seq=3, task_id="a"))) + assert item["id"] == "a" + + +def test_data_decoder_namespace_none_yields_regardless_of_namespace(): + decoder = DataDecoder("checkpoints", namespace=None) + assert list(decoder.feed(checkpoints_event(seq=1, namespace=["child"], v=1))) == [ + {"v": 1} + ] + + +def test_data_decoder_namespace_filter_drops_non_matching_namespace(): + decoder = DataDecoder("checkpoints", namespace=[]) + # root-namespace event is yielded; child-namespace event is filtered out + assert list(decoder.feed(checkpoints_event(seq=1, v=1))) == [{"v": 1}] + assert list(decoder.feed(checkpoints_event(seq=2, namespace=["child"], v=2))) == [] + + +class _FakeStream: + """Stand-in for AsyncChatModelStream/ChatModelStream in decoder tests.""" + + def __init__(self, *, namespace, node, message_id): + self.namespace = namespace + self.node = node + self.message_id = message_id + self.dispatched: list[dict] = [] + + def dispatch(self, data): + self.dispatched.append(data) + + +def _factory(*, namespace, node, message_id): + return _FakeStream(namespace=namespace, node=node, message_id=message_id) + + +def test_messages_decoder_yields_stream_on_message_start(): + decoder = MessagesDecoder(namespace=[], stream_factory=_factory) + streams = list( + decoder.feed(message_start_event(seq=1, message_id="m-1", node="agent")) + ) + assert len(streams) == 1 + assert streams[0].message_id == "m-1" + assert streams[0].node == "agent" + # The start event is dispatched into the stream too (matches stream.py:432). + assert ( + streams[0].dispatched and streams[0].dispatched[0]["event"] == "message-start" + ) + + +def test_messages_decoder_dispatches_delta_to_active_stream(): + decoder = MessagesDecoder(namespace=[], stream_factory=_factory) + [stream] = list(decoder.feed(message_start_event(seq=1, message_id="m-1"))) + delta = message_text_delta_event(seq=2, message_id="m-1", text="hi") + assert list(decoder.feed(delta)) == [] + assert stream.dispatched[-1]["event"] == "content-block-delta" + + +def test_messages_decoder_finish_retires_stream(): + decoder = MessagesDecoder(namespace=[], stream_factory=_factory) + [stream] = list(decoder.feed(message_start_event(seq=1, message_id="m-1"))) + list(decoder.feed(message_finish_event(seq=2, message_id="m-1"))) + assert stream.dispatched[-1]["event"] == "message-finish" + assert all(s is not stream for s in decoder._active.values()) + [again] = list(decoder.feed(message_start_event(seq=3, message_id="m-1"))) + assert again is not stream + + +def test_messages_decoder_error_retires_stream(): + decoder = MessagesDecoder(namespace=[], stream_factory=_factory) + [stream] = list(decoder.feed(message_start_event(seq=1, message_id="m-1"))) + list(decoder.feed(message_error_event(seq=2, message_id="m-1", message="boom"))) + assert stream.dispatched[-1]["event"] == "error" + assert all(s is not stream for s in decoder._active.values()) + + +def test_messages_decoder_single_fallback_routes_idless_events(): + decoder = MessagesDecoder(namespace=[], stream_factory=_factory) + [stream] = list(decoder.feed(message_start_event(seq=1, message_id="m-1"))) + list(decoder.feed(message_text_delta_event(seq=2, text="x"))) # no message_id + assert stream.dispatched[-1]["event"] == "content-block-delta" + + +def test_messages_decoder_drops_idful_events_for_unknown_stream(): + decoder = MessagesDecoder(namespace=[], stream_factory=_factory) + [stream] = list(decoder.feed(message_start_event(seq=1, message_id="m-1"))) + list(decoder.feed(message_text_delta_event(seq=2, message_id="ghost", text="x"))) + assert all(d["event"] != "content-block-delta" for d in stream.dispatched) + + +def test_messages_decoder_ignores_other_namespaces(): + decoder = MessagesDecoder(namespace=[], stream_factory=_factory) + assert ( + list( + decoder.feed( + message_start_event(seq=1, namespace=["child"], message_id="m-1") + ) + ) + == [] + ) + + +def test_messages_decoder_drops_idless_delta_when_multiple_active(): + decoder = MessagesDecoder(namespace=[], stream_factory=_factory) + [a] = list(decoder.feed(message_start_event(seq=1, message_id="m-1"))) + [b] = list(decoder.feed(message_start_event(seq=2, message_id="m-2"))) + # id-less delta is ambiguous with two active streams -> dropped, routed to neither + list(decoder.feed(message_text_delta_event(seq=3, text="x"))) + assert all(d["event"] != "content-block-delta" for d in a.dispatched) + assert all(d["event"] != "content-block-delta" for d in b.dispatched) + + +class _FakeToolHandle: + def __init__(self, *, tool_call_id, name, input, namespace): + self.tool_call_id = tool_call_id + self.name = name + self.input = input + self.namespace = namespace + self.deltas: list[str] = [] + self.finished_output: Any = None + self.finished = False + self.error: BaseException | None = None + + def _push_delta(self, delta): + self.deltas.append(delta) + + def _finish(self, output): + self.finished = True + self.finished_output = output + + def _fail(self, exc): + self.error = exc + + +def _tool_factory(*, tool_call_id, name, input, namespace): + return _FakeToolHandle( + tool_call_id=tool_call_id, name=name, input=input, namespace=namespace + ) + + +def test_tool_calls_decoder_yields_handle_on_start(): + decoder = ToolCallsDecoder(namespace=[], handle_factory=_tool_factory) + [handle] = list( + decoder.feed(tool_started_event(seq=1, tool_call_id="tc-1", tool_name="search")) + ) + assert handle.tool_call_id == "tc-1" + assert handle.name == "search" + + +def test_tool_calls_decoder_routes_delta_finish_and_error(): + decoder = ToolCallsDecoder(namespace=[], handle_factory=_tool_factory) + [h] = list( + decoder.feed(tool_started_event(seq=1, tool_call_id="tc-1", tool_name="search")) + ) + list(decoder.feed(tool_output_delta_event(seq=2, tool_call_id="tc-1", delta="x"))) + assert h.deltas == ["x"] + list( + decoder.feed( + tool_finished_event(seq=3, tool_call_id="tc-1", output={"ok": True}) + ) + ) + assert h.finished and h.finished_output == {"ok": True} + + decoder2 = ToolCallsDecoder(namespace=[], handle_factory=_tool_factory) + [h2] = list( + decoder2.feed( + tool_started_event(seq=1, tool_call_id="tc-2", tool_name="search") + ) + ) + list(decoder2.feed(tool_error_event(seq=2, tool_call_id="tc-2", message="boom"))) + assert isinstance(h2.error, RuntimeError) and "boom" in str(h2.error) + + +def test_tool_calls_decoder_drops_events_for_unknown_id(): + decoder = ToolCallsDecoder(namespace=[], handle_factory=_tool_factory) + assert ( + list( + decoder.feed( + tool_output_delta_event(seq=1, tool_call_id="ghost", delta="x") + ) + ) + == [] + ) + + +def test_tool_calls_decoder_finish_and_error_retire_handle(): + decoder = ToolCallsDecoder(namespace=[], handle_factory=_tool_factory) + [h] = list( + decoder.feed(tool_started_event(seq=1, tool_call_id="tc-1", tool_name="search")) + ) + list(decoder.feed(tool_finished_event(seq=2, tool_call_id="tc-1"))) + # retired: a late delta for the same id is now dropped + list( + decoder.feed(tool_output_delta_event(seq=3, tool_call_id="tc-1", delta="late")) + ) + assert h.deltas == [] + + +def test_tool_calls_decoder_ignores_other_namespaces(): + decoder = ToolCallsDecoder(namespace=[], handle_factory=_tool_factory) + assert ( + list( + decoder.feed( + tool_started_event(seq=1, namespace=["child"], tool_call_id="tc-1") + ) + ) + == [] + ) + + +def test_tool_calls_decoder_error_retires_handle(): + decoder = ToolCallsDecoder(namespace=[], handle_factory=_tool_factory) + [h] = list( + decoder.feed(tool_started_event(seq=1, tool_call_id="tc-1", tool_name="search")) + ) + list(decoder.feed(tool_error_event(seq=2, tool_call_id="tc-1", message="boom"))) + # retired: a late delta for the same id is now dropped + list( + decoder.feed(tool_output_delta_event(seq=3, tool_call_id="tc-1", delta="late")) + ) + assert h.deltas == [] + + +def test_tool_calls_decoder_skips_non_str_tool_call_id(): + decoder = ToolCallsDecoder(namespace=[], handle_factory=_tool_factory) + # Build a tools event whose data.tool_call_id is not a string. + bad = tool_started_event(seq=1, tool_call_id="tc-1", tool_name="search") + bad["params"]["data"]["tool_call_id"] = 123 + assert list(decoder.feed(bad)) == [] + + +def test_tool_calls_decoder_skips_non_str_delta(): + decoder = ToolCallsDecoder(namespace=[], handle_factory=_tool_factory) + [h] = list( + decoder.feed(tool_started_event(seq=1, tool_call_id="tc-1", tool_name="search")) + ) + evt = tool_output_delta_event(seq=2, tool_call_id="tc-1", delta="x") + evt["params"]["data"]["delta"] = 123 # non-str delta must be ignored + list(decoder.feed(evt)) + assert h.deltas == [] + + +def test_tool_calls_decoder_defaults_missing_tool_name_to_empty(): + decoder = ToolCallsDecoder(namespace=[], handle_factory=_tool_factory) + evt = tool_started_event(seq=1, tool_call_id="tc-1", tool_name="search") + del evt["params"]["data"]["tool_name"] # absent tool_name -> handle.name == "" + [h] = list(decoder.feed(evt)) + assert h.name == "" + + +def test_tool_calls_decoder_error_message_defaults_when_blank(): + decoder = ToolCallsDecoder(namespace=[], handle_factory=_tool_factory) + [h] = list( + decoder.feed(tool_started_event(seq=1, tool_call_id="tc-1", tool_name="search")) + ) + evt = tool_error_event(seq=2, tool_call_id="tc-1") + evt["params"]["data"]["message"] = "" # blank -> default message + list(decoder.feed(evt)) + assert str(h.error) == "Tool call errored" + + +class _FakeScopedHandle: + def __init__(self, *, path, graph_name, trigger_call_id): + self.path = path + self.graph_name = graph_name + self.trigger_call_id = trigger_call_id + self.status = "started" + self.error = None + self.events: list[dict] = [] + + def _push_event(self, event): + self.events.append(event) + + def _finish(self, status, error=None): + if self.status != "started": + return + self.status = status + self.error = error + + +def _scoped_factory(*, path, graph_name, trigger_call_id): + return _FakeScopedHandle( + path=path, graph_name=graph_name, trigger_call_id=trigger_call_id + ) + + +def test_subgraphs_decoder_discovers_on_lifecycle_started_once(): + decoder = SubgraphsDecoder(scope=(), handle_factory=_scoped_factory) + [h] = list(decoder.feed(lifecycle_started_event(seq=1, namespace=["child"]))) + assert h.path == ("child",) + assert list(decoder.feed(lifecycle_started_event(seq=2, namespace=["child"]))) == [] + + +def test_subgraphs_decoder_discovers_on_tasks_start_without_result(): + decoder = SubgraphsDecoder(scope=(), handle_factory=_scoped_factory) + [h] = list(decoder.feed(tasks_start_event(seq=1, namespace=["child"]))) + assert h.path == ("child",) + + +def test_subgraphs_decoder_parses_graph_name_and_trigger_from_segment(): + decoder = SubgraphsDecoder(scope=(), handle_factory=_scoped_factory) + [h] = list(decoder.feed(tasks_start_event(seq=1, namespace=["agent:call-1"]))) + assert h.graph_name == "agent" + assert h.trigger_call_id == "call-1" + + +def test_subgraphs_decoder_fans_out_events_to_active_child(): + decoder = SubgraphsDecoder(scope=(), handle_factory=_scoped_factory) + [h] = list(decoder.feed(lifecycle_started_event(seq=1, namespace=["child"]))) + inner = message_start_event(seq=2, namespace=["child"], message_id="m") + assert list(decoder.feed(inner)) == [] + assert inner in h.events # whole event pushed, not just data + + +def test_subgraphs_decoder_fans_out_grandchild_to_direct_child(): + decoder = SubgraphsDecoder(scope=(), handle_factory=_scoped_factory) + [h] = list(decoder.feed(lifecycle_started_event(seq=1, namespace=["child"]))) + grand = message_start_event(seq=2, namespace=["child", "grand"], message_id="m") + list(decoder.feed(grand)) + assert grand in h.events + + +def test_subgraphs_decoder_tasks_result_at_parent_finalizes_child(): + decoder = SubgraphsDecoder(scope=(), handle_factory=_scoped_factory) + # child discovered with a colon segment -> trigger_call_id == "call-1" + [h] = list(decoder.feed(tasks_start_event(seq=1, namespace=["agent:call-1"]))) + assert h.trigger_call_id == "call-1" + # the finalizing tasks-result is emitted at the PARENT (root) namespace, + # with id (task_id) matching the child's trigger_call_id + list( + decoder.feed( + tasks_result_event(seq=2, namespace=[], task_id="call-1", result={"ok": 1}) + ) + ) + assert h.status == "completed" + # finalized + removed from active: later child-namespace events no longer fan out + later = message_start_event(seq=3, namespace=["agent:call-1"], message_id="m") + list(decoder.feed(later)) + assert later not in h.events + + +def test_subgraphs_decoder_tasks_result_failed_status(): + decoder = SubgraphsDecoder(scope=(), handle_factory=_scoped_factory) + [h] = list(decoder.feed(tasks_start_event(seq=1, namespace=["agent:call-1"]))) + list( + decoder.feed( + tasks_result_event(seq=2, namespace=[], task_id="call-1", error="boom") + ) + ) + assert h.status == "failed" + assert h.error == "boom" + + +def test_subgraphs_decoder_tasks_result_interrupted_status(): + decoder = SubgraphsDecoder(scope=(), handle_factory=_scoped_factory) + [h] = list(decoder.feed(tasks_start_event(seq=1, namespace=["agent:call-1"]))) + list( + decoder.feed( + tasks_result_event( + seq=2, namespace=[], task_id="call-1", interrupts=[{"id": "i-1"}] + ) + ) + ) + assert h.status == "interrupted" + + +def test_subgraphs_decoder_ignores_unrelated_and_scope_itself(): + decoder = SubgraphsDecoder(scope=("root",), handle_factory=_scoped_factory) + # not a direct child of ("root",): wrong depth / wrong prefix + assert list(decoder.feed(lifecycle_started_event(seq=1, namespace=["other"]))) == [] + # the scope's own namespace is not a discovery + assert list(decoder.feed(lifecycle_started_event(seq=2, namespace=["root"]))) == [] + + +def test_extensions_decoder_yields_full_data_for_matching_name(): + decoder = ExtensionsDecoder(name="foo") + # custom_event(name="foo", x=1) -> params.data = {"name": "foo", "x": 1} + assert list(decoder.feed(custom_event(seq=1, name="foo", x=1))) == [ + {"name": "foo", "x": 1} + ] + + +def test_extensions_decoder_ignores_other_extension_names(): + decoder = ExtensionsDecoder(name="foo") + assert list(decoder.feed(custom_event(seq=1, name="bar", x=1))) == [] + + +def test_extensions_decoder_ignores_non_custom_methods(): + decoder = ExtensionsDecoder(name="foo") + assert list(decoder.feed(lifecycle_completed_event(seq=1))) == [] + + +def test_extensions_decoder_ignores_non_dict_data(): + decoder = ExtensionsDecoder(name="foo") + evt = custom_event(seq=1, name="foo", x=1) + evt["params"]["data"] = "not-a-dict" + assert list(decoder.feed(evt)) == [] + + +def test_extensions_decoder_rejects_empty_name(): + import pytest + + with pytest.raises(ValueError): + ExtensionsDecoder(name="") diff --git a/libs/sdk-py/tests/streaming/test_sync_projections.py b/libs/sdk-py/tests/streaming/test_sync_projections.py index c7ee2199b..07708793e 100644 --- a/libs/sdk-py/tests/streaming/test_sync_projections.py +++ b/libs/sdk-py/tests/streaming/test_sync_projections.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import cast +from typing import Any, cast import httpx import pytest @@ -562,3 +562,79 @@ def test_sync_tool_call_handle_deltas_single_consumer_guard(): # Second access: must raise immediately (before any iteration). with pytest.raises(RuntimeError, match="single consumer"): _ = handle.deltas + + +def test_sync_messages_subscription_pre_dispatches_before_yield(): + """Over a live subscription, str(stream.text) must work immediately on yield.""" + fake = SyncFakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + message_start_event(seq=1, message_id="msg-1"), + message_text_delta_event(seq=2, text="hello", message_id="msg-1"), + message_text_finish_event(seq=3, text="hello", message_id="msg-1"), + message_finish_event(seq=4, message_id="msg-1"), + lifecycle_completed_event(seq=5), + ] + ) + fake.set_state({}) + collected: list[str] = [] + with httpx.Client(transport=fake.transport, base_url="http://test") as raw: + threads = SyncThreadsClient(SyncHttpClient(raw)) + with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + thread.run.start(input={}) + for stream in thread.messages: + collected.append(str(stream.text)) + assert collected == ["hello"] + + +def test_sync_tool_calls_subscription_resolves_output_before_yield(): + """Over a live subscription, call.output is resolved when the handle is yielded.""" + fake = SyncFakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + tool_started_event(seq=1, tool_call_id="call-1", tool_name="search"), + tool_finished_event(seq=2, tool_call_id="call-1", output={"ok": True}), + lifecycle_completed_event(seq=3), + ] + ) + fake.set_state({}) + outputs: list[Any] = [] + with httpx.Client(transport=fake.transport, base_url="http://test") as raw: + threads = SyncThreadsClient(SyncHttpClient(raw)) + with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + thread.run.start(input={}) + for call in thread.tool_calls: + outputs.append(call.output) # resolved (blocking) on yield + assert outputs == [{"ok": True}] + + +# --------------------------------------------------------------------------- +# Regression: interleaved-concurrent tool calls must BOTH surface +# --------------------------------------------------------------------------- + + +def test_sync_tool_calls_interleaved_concurrent_calls_both_surface(): + """Two tool calls whose events interleave must BOTH be yielded (regression: + the pre-decoder read-ahead silently dropped the second concurrent call).""" + fake = SyncFakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + tool_started_event(seq=1, tool_call_id="call-a", tool_name="search"), + tool_started_event(seq=2, tool_call_id="call-b", tool_name="lookup"), + tool_finished_event(seq=3, tool_call_id="call-a", output={"a": 1}), + tool_finished_event(seq=4, tool_call_id="call-b", output={"b": 2}), + lifecycle_completed_event(seq=5), + ] + ) + fake.set_state({}) + seen = [] + with httpx.Client(transport=fake.transport, base_url="http://test") as raw: + threads = SyncThreadsClient(SyncHttpClient(raw)) + with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + thread.run.start(input={}) + for call in thread.tool_calls: + seen.append(call.tool_call_id) + assert sorted(seen) == ["call-a", "call-b"] diff --git a/libs/sdk-py/tests/streaming/test_sync_thread_stream.py b/libs/sdk-py/tests/streaming/test_sync_thread_stream.py index c594eddeb..72221d8d9 100644 --- a/libs/sdk-py/tests/streaming/test_sync_thread_stream.py +++ b/libs/sdk-py/tests/streaming/test_sync_thread_stream.py @@ -602,3 +602,310 @@ def test_v3_streaming_sync_surface_smoke(): assert tools_result[0].name == "search" # ty: ignore[unresolved-attribute] assert results["progress"] == [{"name": "progress", "step": 1}] assert final == {"final": True} + + +# --------------------------------------------------------------------------- +# interleave_projections tests +# --------------------------------------------------------------------------- + + +def test_interleave_projections_single_channel_values(): + from streaming._events import ( + lifecycle_completed_event, + lifecycle_started_event, + values_event, + ) + + fake = SyncFakeServer() + fake.set_state({"counter": 0}) + fake.script( + [ + lifecycle_started_event(seq=0), + values_event(seq=1, counter=1), + values_event(seq=2, counter=2), + lifecycle_completed_event(seq=3), + ] + ) + with httpx.Client(transport=fake.transport, base_url="http://test") as raw: + threads = SyncThreadsClient(SyncHttpClient(raw)) + with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + thread.run.start(input={}) + items = [] + for ch, item in thread.interleave_projections(["values"]): + items.append((ch, item)) + assert ("values", {"counter": 1}) in items + assert ("values", {"counter": 2}) in items + assert all(ch == "values" for ch, _ in items) + + +def test_interleave_projections_values_and_messages_arrival_order(): + from streaming._events import ( + lifecycle_completed_event, + lifecycle_started_event, + message_finish_event, + message_start_event, + values_event, + ) + + fake = SyncFakeServer() + fake.set_state({"counter": 0}) + fake.script( + [ + lifecycle_started_event(seq=0), + values_event(seq=1, counter=1), + message_start_event(seq=2, message_id="m-1"), + values_event(seq=3, counter=2), + message_finish_event(seq=4, message_id="m-1"), + lifecycle_completed_event(seq=5), + ] + ) + with httpx.Client(transport=fake.transport, base_url="http://test") as raw: + threads = SyncThreadsClient(SyncHttpClient(raw)) + with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + thread.run.start(input={}) + order = [] + for ch, _ in thread.interleave_projections(["values", "messages"]): + order.append(ch) + if len(order) >= 3: + break + assert order[:3] == ["values", "messages", "values"] + + +def test_interleave_projections_mixes_builtin_and_extension(): + from streaming._events import ( + custom_event, + lifecycle_completed_event, + lifecycle_started_event, + values_event, + ) + + fake = SyncFakeServer() + fake.set_state({"counter": 0}) + fake.script( + [ + lifecycle_started_event(seq=0), + values_event(seq=1, counter=1), + custom_event(seq=2, name="foo", hello="world"), + lifecycle_completed_event(seq=3), + ] + ) + with httpx.Client(transport=fake.transport, base_url="http://test") as raw: + threads = SyncThreadsClient(SyncHttpClient(raw)) + with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + thread.run.start(input={}) + items = [] + for ch, item in thread.interleave_projections(["values", "foo"]): + items.append((ch, item)) + assert ("values", {"counter": 1}) in items + assert ("foo", {"name": "foo", "hello": "world"}) in items + + +def test_interleave_projections_tool_calls_uses_public_name(): + from streaming._events import ( + lifecycle_completed_event, + lifecycle_started_event, + tool_finished_event, + tool_started_event, + ) + + fake = SyncFakeServer() + fake.set_state({}) + fake.script( + [ + lifecycle_started_event(seq=0), + tool_started_event(seq=1, tool_call_id="call-1", tool_name="search"), + tool_finished_event(seq=2, tool_call_id="call-1", output={"ok": True}), + lifecycle_completed_event(seq=3), + ] + ) + with httpx.Client(transport=fake.transport, base_url="http://test") as raw: + threads = SyncThreadsClient(SyncHttpClient(raw)) + with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + thread.run.start(input={}) + names = [] + handle = None + for ch, item in thread.interleave_projections(["tool_calls"]): + names.append(ch) + if handle is None: + handle = item + break + assert names == ["tool_calls"] + assert handle is not None + assert handle.tool_call_id == "call-1" + + +def test_interleave_projections_subgraphs_discovers_child(): + from streaming._events import ( + lifecycle_completed_event, + lifecycle_started_event, + ) + + fake = SyncFakeServer() + fake.set_state({}) + fake.script( + [ + lifecycle_started_event(seq=0), + lifecycle_started_event(seq=1, namespace=["child"]), + lifecycle_completed_event(seq=2), + ] + ) + with httpx.Client(transport=fake.transport, base_url="http://test") as raw: + threads = SyncThreadsClient(SyncHttpClient(raw)) + with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + thread.run.start(input={}) + discovered = [] + for ch, handle in thread.interleave_projections(["subgraphs"]): + discovered.append((ch, handle.path)) + assert ("subgraphs", ("child",)) in discovered + + +def test_interleave_projections_inflight_tool_call_failed_on_break(): + """A tool handle held past an early break is failed in teardown, never left hanging.""" + from streaming._events import ( + lifecycle_completed_event, + lifecycle_started_event, + tool_started_event, + ) + + fake = SyncFakeServer() + fake.set_state({}) + fake.script( + [ + lifecycle_started_event(seq=0), + tool_started_event(seq=1, tool_call_id="call-1", tool_name="search"), + # no tool-finished: the call is still in flight when the consumer breaks + lifecycle_completed_event(seq=2), + ] + ) + with httpx.Client(transport=fake.transport, base_url="http://test") as raw: + threads = SyncThreadsClient(SyncHttpClient(raw)) + with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + thread.run.start(input={}) + handle = None + for _, item in thread.interleave_projections(["tool_calls"]): + handle = item + break + assert handle is not None + # Without teardown finalization this blocks forever; the bounded + # timeout turns a regression into a TimeoutError, not a RuntimeError. + with pytest.raises(RuntimeError): + handle._result.result(timeout=2) + + +def test_interleave_projections_inflight_subgraph_finished_on_terminal(): + """A discovered subgraph child with no terminal tasks-result is force-completed.""" + from streaming._events import ( + lifecycle_completed_event, + lifecycle_started_event, + ) + + fake = SyncFakeServer() + fake.set_state({}) + fake.script( + [ + lifecycle_started_event(seq=0), + lifecycle_started_event(seq=1, namespace=["child"]), + # no tasks-result for the child: it is still "started" at run end + lifecycle_completed_event(seq=2), + ] + ) + with httpx.Client(transport=fake.transport, base_url="http://test") as raw: + threads = SyncThreadsClient(SyncHttpClient(raw)) + with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + thread.run.start(input={}) + child = None + for _, handle in thread.interleave_projections(["subgraphs"]): + child = handle + assert child is not None + assert child.status == "completed" + + +@pytest.mark.parametrize("channel", ["lifecycle", "tools", "input"]) +def test_interleave_projections_rejects_reserved_channel(channel): + """Reserved protocol channel names raise instead of silently no-op'ing. + + `infer_channel` treats these as first-class methods, but they have no + interleave decoder, so routing them to the extension/`custom:` fallback + would subscribe to a channel that never matches and yield nothing. Fail + closed. (`updates`/`checkpoints`/`tasks` are supported and tested below.) + """ + from streaming._events import ( + lifecycle_completed_event, + lifecycle_started_event, + ) + + fake = SyncFakeServer() + fake.set_state({}) + fake.script([lifecycle_started_event(seq=0), lifecycle_completed_event(seq=1)]) + with httpx.Client(transport=fake.transport, base_url="http://test") as raw: + threads = SyncThreadsClient(SyncHttpClient(raw)) + with ( + threads.stream(thread_id="t-1", assistant_id="agent") as thread, + pytest.raises(ValueError, match=channel), + ): + for _ in thread.interleave_projections([channel]): + pass + + +def test_interleave_projections_data_channels_yield_payloads(): + """`updates`/`checkpoints`/`tasks` yield their raw `params.data` payloads.""" + from streaming._events import ( + checkpoints_event, + lifecycle_completed_event, + lifecycle_started_event, + tasks_start_event, + updates_event, + ) + + fake = SyncFakeServer() + fake.set_state({}) + fake.script( + [ + lifecycle_started_event(seq=0), + updates_event(seq=1, node={"v": 1}), + checkpoints_event(seq=2, ts="t-0", v=4), + tasks_start_event(seq=3, task_id="task-9"), + lifecycle_completed_event(seq=4), + ] + ) + with httpx.Client(transport=fake.transport, base_url="http://test") as raw: + threads = SyncThreadsClient(SyncHttpClient(raw)) + with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + thread.run.start(input={}) + items = list( + thread.interleave_projections(["updates", "checkpoints", "tasks"]) + ) + assert ("updates", {"node": {"v": 1}}) in items + assert ("checkpoints", {"ts": "t-0", "v": 4}) in items + assert any(ch == "tasks" and item.get("id") == "task-9" for ch, item in items) + + +def test_interleave_projections_data_channel_scoped_to_root_namespace(): + """A child-namespace checkpoint must not leak into a root interleave.""" + from streaming._events import ( + checkpoints_event, + lifecycle_completed_event, + lifecycle_started_event, + ) + + fake = SyncFakeServer() + fake.set_state({"counter": 0}) + fake.script( + [ + lifecycle_started_event(seq=0), + checkpoints_event(seq=1, namespace=["child"], scope="child"), + checkpoints_event(seq=2, scope="root"), + lifecycle_completed_event(seq=3), + ] + ) + with httpx.Client(transport=fake.transport, base_url="http://test") as raw: + threads = SyncThreadsClient(SyncHttpClient(raw)) + with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + thread.run.start(input={}) + checkpoints = [ + item + for ch, item in thread.interleave_projections(["values", "checkpoints"]) + if ch == "checkpoints" + ] + assert {"scope": "root"} in checkpoints + assert {"scope": "child"} not in checkpoints diff --git a/libs/sdk-py/tests/streaming/test_thread_stream.py b/libs/sdk-py/tests/streaming/test_thread_stream.py index 959915b16..eb71b3dbc 100644 --- a/libs/sdk-py/tests/streaming/test_thread_stream.py +++ b/libs/sdk-py/tests/streaming/test_thread_stream.py @@ -17,9 +17,17 @@ from langgraph_sdk.stream.transport import ( ProtocolWebSocketTransport, ) from streaming._events import ( + checkpoints_event, + custom_event, lifecycle_completed_event, lifecycle_event, lifecycle_started_event, + message_finish_event, + message_start_event, + tasks_start_event, + tool_finished_event, + tool_started_event, + updates_event, values_event, ) from streaming._fake_server import FakeServer @@ -971,3 +979,255 @@ async def test_v3_streaming_async_surface_smoke(): assert tool_calls[0].name == "search" assert progress == [{"name": "progress", "step": 1}] assert final == {"final": True} + + +async def test_interleave_projections_single_channel_values(): + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + values_event(seq=1, counter=1), + values_event(seq=2, counter=2), + lifecycle_completed_event(seq=3), + ] + ) + fake.set_state({"counter": 0}) + 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={}) + items = [] + async for ch, item in thread.interleave_projections(["values"]): + items.append((ch, item)) + assert ("values", {"counter": 1}) in items + assert ("values", {"counter": 2}) in items + assert all(ch == "values" for ch, _ in items) + + +async def test_interleave_projections_values_and_messages_arrival_order(): + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + values_event(seq=1, counter=1), + message_start_event(seq=2, message_id="m-1"), + values_event(seq=3, counter=2), + message_finish_event(seq=4, message_id="m-1"), + lifecycle_completed_event(seq=5), + ] + ) + fake.set_state({"counter": 0}) + 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={}) + order = [] + async for ch, _ in thread.interleave_projections(["values", "messages"]): + order.append(ch) + if len(order) >= 3: + break + assert order[:3] == ["values", "messages", "values"] + + +async def test_interleave_projections_mixes_builtin_and_extension(): + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + values_event(seq=1, counter=1), + custom_event(seq=2, name="foo", hello="world"), + lifecycle_completed_event(seq=3), + ] + ) + fake.set_state({"counter": 0}) + 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={}) + items = [] + async for ch, item in thread.interleave_projections(["values", "foo"]): + items.append((ch, item)) + assert ("values", {"counter": 1}) in items + # extension payload is the whole params.data including "name"; tuple uses bare name "foo" + assert ("foo", {"name": "foo", "hello": "world"}) in items + + +async def test_interleave_projections_tool_calls_uses_public_name(): + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + tool_started_event(seq=1, tool_call_id="call-1", tool_name="search"), + tool_finished_event(seq=2, tool_call_id="call-1", output={"ok": True}), + lifecycle_completed_event(seq=3), + ] + ) + fake.set_state({}) + 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={}) + names = [] + async for ch, item in thread.interleave_projections(["tool_calls"]): + names.append(ch) + assert item.tool_call_id == "call-1" # real ToolCallHandle + # tuple uses the PUBLIC name "tool_calls", never the wire name "tools" + assert names == ["tool_calls"] + + +async def test_interleave_projections_subgraphs_discovers_child(): + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + lifecycle_started_event(seq=1, namespace=["child"]), + lifecycle_completed_event(seq=2), + ] + ) + fake.set_state({}) + 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={}) + discovered = [] + async for ch, handle in thread.interleave_projections(["subgraphs"]): + discovered.append((ch, handle.path)) + assert ("subgraphs", ("child",)) in discovered + + +async def test_interleave_projections_inflight_tool_call_failed_on_break(): + """A tool handle held past an early break is failed in teardown, never left hanging.""" + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + tool_started_event(seq=1, tool_call_id="call-1", tool_name="search"), + # no tool-finished: the call is still in flight when the consumer breaks + lifecycle_completed_event(seq=2), + ] + ) + fake.set_state({}) + 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 = None + async for _, item in thread.interleave_projections(["tool_calls"]): + handle = item + break + assert handle is not None + # Without teardown finalization this would hang forever; wait_for + # turns a regression into a TimeoutError rather than a RuntimeError. + with pytest.raises(RuntimeError): + await asyncio.wait_for(handle.output, timeout=2) + + +async def test_interleave_projections_inflight_subgraph_finished_on_terminal(): + """A discovered subgraph child with no terminal tasks-result is force-completed.""" + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + lifecycle_started_event(seq=1, namespace=["child"]), + # no tasks-result for the child: it is still "started" at run end + lifecycle_completed_event(seq=2), + ] + ) + fake.set_state({}) + 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 = None + async for _, handle in thread.interleave_projections(["subgraphs"]): + child = handle + assert child is not None + assert child.status == "completed" + + +@pytest.mark.parametrize("channel", ["lifecycle", "tools", "input"]) +async def test_interleave_projections_rejects_reserved_channel(channel): + """Reserved protocol channel names raise instead of silently no-op'ing. + + `infer_channel` treats these as first-class methods, but they have no + interleave decoder, so routing them to the extension/`custom:` fallback + would subscribe to a channel that never matches and yield nothing. Fail + closed. (`updates`/`checkpoints`/`tasks` are supported and tested below.) + """ + fake = FakeServer() + fake.script([lifecycle_started_event(seq=0), lifecycle_completed_event(seq=1)]) + fake.set_state({}) + 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: + with pytest.raises(ValueError, match=channel): + async for _ in thread.interleave_projections([channel]): + pass + + +async def test_interleave_projections_data_channels_yield_payloads(): + """`updates`/`checkpoints`/`tasks` yield their raw `params.data` payloads.""" + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + updates_event(seq=1, node={"v": 1}), + checkpoints_event(seq=2, ts="t-0", v=4), + tasks_start_event(seq=3, task_id="task-9"), + lifecycle_completed_event(seq=4), + ] + ) + fake.set_state({}) + 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={}) + items = [] + async for ch, item in thread.interleave_projections( + ["updates", "checkpoints", "tasks"] + ): + items.append((ch, item)) + assert ("updates", {"node": {"v": 1}}) in items + assert ("checkpoints", {"ts": "t-0", "v": 4}) in items + assert any(ch == "tasks" and item.get("id") == "task-9" for ch, item in items) + + +async def test_interleave_projections_data_channel_scoped_to_root_namespace(): + """A child-namespace checkpoint must not leak into a root interleave. + + `values` subscribes unscoped, so `compute_union_filter` widens the merged + subscription to all namespaces; the `DataDecoder` root filter is what keeps + a subgraph checkpoint out of the root projection (mirrors local scope). + """ + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + checkpoints_event(seq=1, namespace=["child"], scope="child"), + checkpoints_event(seq=2, scope="root"), + lifecycle_completed_event(seq=3), + ] + ) + fake.set_state({"counter": 0}) + 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={}) + checkpoints = [] + async for ch, item in thread.interleave_projections( + ["values", "checkpoints"] + ): + if ch == "checkpoints": + checkpoints.append(item) + assert {"scope": "root"} in checkpoints + assert {"scope": "child"} not in checkpoints