mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-17 21:25:46 +02:00
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:
@@ -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]:
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user