feat(sdk-py): extract stream decoders and add interleave_projections (#7935)

## 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)
This commit is contained in:
Nick Hollon
2026-05-29 17:14:14 -04:00
committed by GitHub
parent 68fa011fc9
commit 1fcb768182
8 changed files with 1951 additions and 466 deletions
+230 -207
View File
@@ -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:<id>``. 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]:
+246 -258
View File
@@ -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:
@@ -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:<run_id>` slot.
"""
message_id = _message_event_id(data)
if message_id is not None:
return f"message:{message_id}"
if fallback is not None:
return f"message:{fallback}"
return "__single__"
SubgraphStatus = Literal["started", "completed", "failed", "interrupted"]
def _parse_namespace_segment(segment: str) -> tuple[str, str | None]:
name, sep, task_id = segment.partition(":")
return name, task_id if sep else None
def _terminal_from_tasks_result(
data: dict[str, Any],
) -> tuple[SubgraphStatus, str | None]:
if data.get("interrupts"):
return "interrupted", None
error = data.get("error")
if error:
return "failed", str(error)
return "completed", None
def _is_direct_child(namespace: list[str], scope: tuple[str, ...]) -> bool:
return len(namespace) == len(scope) + 1 and tuple(namespace[: len(scope)]) == scope
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
+12
View File
@@ -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]:
@@ -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="")
@@ -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"]
@@ -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
@@ -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