mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-25 09:02:25 +02:00
feat(sdk-py): add scoped subgraph handles (#7824)
This commit is contained in:
Nick Hollon
@@ -66,6 +66,24 @@ _ALL_CHANNELS: list[str] = [
|
||||
]
|
||||
|
||||
|
||||
def _exact_namespace_params(
|
||||
channels: list[str],
|
||||
namespace: list[str],
|
||||
) -> SubscribeParams:
|
||||
return {
|
||||
"channels": channels,
|
||||
"namespaces": [list(namespace)],
|
||||
"depth": 0,
|
||||
}
|
||||
|
||||
|
||||
def _event_namespace(params_field: Any) -> list[str]:
|
||||
if not isinstance(params_field, dict):
|
||||
return []
|
||||
namespace = params_field.get("namespace") or []
|
||||
return list(namespace) if isinstance(namespace, list) else []
|
||||
|
||||
|
||||
class RunModule:
|
||||
"""Command dispatcher for `run.start`.
|
||||
|
||||
@@ -294,8 +312,11 @@ class _MessagesProjection:
|
||||
from the root namespace only.
|
||||
"""
|
||||
|
||||
def __init__(self, thread: AsyncThreadStream) -> None:
|
||||
def __init__(
|
||||
self, thread: AsyncThreadStream, namespace: list[str] | None = None
|
||||
) -> None:
|
||||
self._thread = thread
|
||||
self._namespace = list(namespace or [])
|
||||
|
||||
def __aiter__(self) -> AsyncIterator[AsyncChatModelStream]:
|
||||
return self._messages_iter()
|
||||
@@ -303,11 +324,15 @@ class _MessagesProjection:
|
||||
async def _messages_iter(self) -> AsyncGenerator[AsyncChatModelStream, None]:
|
||||
if self._thread._transport is None:
|
||||
raise RuntimeError("AsyncThreadStream not entered — use `async with`.")
|
||||
params: SubscribeParams = {
|
||||
"channels": ["messages"],
|
||||
"namespaces": [[]],
|
||||
"depth": 0,
|
||||
}
|
||||
# If the subgraphs projection already ran (and consumed messages events
|
||||
# from the shared SSE), drain its root inbox rather than opening a new
|
||||
# subscription. Dedup prevents the SSE from replaying those events.
|
||||
root_inbox = self._thread._root_messages_inbox if not self._namespace else None
|
||||
if root_inbox is not None:
|
||||
async for stream in self._drain_inbox(root_inbox):
|
||||
yield stream
|
||||
return
|
||||
params = _exact_namespace_params(["messages"], self._namespace)
|
||||
sub = self._thread._register_subscription(params)
|
||||
active: dict[str, AsyncChatModelStream] = {}
|
||||
try:
|
||||
@@ -318,11 +343,11 @@ class _MessagesProjection:
|
||||
if item is None:
|
||||
return
|
||||
params_field = item.get("params") or {}
|
||||
if not isinstance(params_field, dict):
|
||||
if _event_namespace(params_field) != self._namespace:
|
||||
continue
|
||||
if params_field.get("namespace") not in (None, []):
|
||||
continue
|
||||
data = params_field.get("data")
|
||||
data = (
|
||||
params_field.get("data") if isinstance(params_field, dict) else None
|
||||
)
|
||||
if not isinstance(data, dict):
|
||||
continue
|
||||
event_type = data.get("event")
|
||||
@@ -335,7 +360,7 @@ class _MessagesProjection:
|
||||
else {}
|
||||
)
|
||||
stream = AsyncChatModelStream(
|
||||
namespace=[],
|
||||
namespace=list(self._namespace),
|
||||
node=metadata.get("langgraph_node") if metadata else None,
|
||||
message_id=message_id,
|
||||
)
|
||||
@@ -361,6 +386,61 @@ class _MessagesProjection:
|
||||
self._thread._unregister_active_message_stream(stream)
|
||||
self._thread._unregister_subscription(sub.id)
|
||||
|
||||
async def _drain_inbox(
|
||||
self, inbox: asyncio.Queue[Event | None]
|
||||
) -> AsyncGenerator[AsyncChatModelStream, None]:
|
||||
"""Drain a pre-filled inbox of messages events, yielding one stream per message."""
|
||||
from langchain_core.language_models.chat_model_stream import (
|
||||
AsyncChatModelStream,
|
||||
)
|
||||
|
||||
active: dict[str, AsyncChatModelStream] = {}
|
||||
try:
|
||||
while True:
|
||||
item = await inbox.get()
|
||||
if item is None:
|
||||
return
|
||||
params_field = item.get("params") or {}
|
||||
data = (
|
||||
params_field.get("data") if isinstance(params_field, dict) else None
|
||||
)
|
||||
if not isinstance(data, dict):
|
||||
continue
|
||||
event_type = data.get("event")
|
||||
if event_type == "message-start":
|
||||
message_id = _message_event_id(data)
|
||||
key = _message_route_key(data, fallback=message_id)
|
||||
metadata = (
|
||||
data.get("metadata")
|
||||
if isinstance(data.get("metadata"), dict)
|
||||
else {}
|
||||
)
|
||||
stream = AsyncChatModelStream(
|
||||
namespace=list(self._namespace),
|
||||
node=metadata.get("langgraph_node") if metadata else None,
|
||||
message_id=message_id,
|
||||
)
|
||||
active[key] = stream
|
||||
self._thread._register_active_message_stream(stream)
|
||||
stream.dispatch(data)
|
||||
yield stream
|
||||
else:
|
||||
key = _message_route_key(data)
|
||||
stream = active.get(key)
|
||||
if stream is None and len(active) == 1:
|
||||
stream = next(iter(active.values()))
|
||||
if stream is None:
|
||||
continue
|
||||
stream.dispatch(data)
|
||||
if event_type in ("message-finish", "error"):
|
||||
self._thread._unregister_active_message_stream(stream)
|
||||
for route_key, candidate in list(active.items()):
|
||||
if candidate is stream:
|
||||
del active[route_key]
|
||||
finally:
|
||||
for stream in active.values():
|
||||
self._thread._unregister_active_message_stream(stream)
|
||||
|
||||
|
||||
def _message_event_id(data: dict[str, Any]) -> str | None:
|
||||
message_id = data.get("id") or data.get("message_id")
|
||||
@@ -382,6 +462,474 @@ def _message_route_key(data: dict[str, Any], fallback: str | None = None) -> str
|
||||
return "__single__"
|
||||
|
||||
|
||||
SubgraphStatus = Literal["started", "completed", "failed", "interrupted"]
|
||||
|
||||
|
||||
def _parse_namespace_segment(segment: str) -> tuple[str, str | None]:
|
||||
name, sep, task_id = segment.partition(":")
|
||||
return name, task_id if sep else None
|
||||
|
||||
|
||||
def _terminal_from_tasks_result(
|
||||
data: dict[str, Any],
|
||||
) -> tuple[SubgraphStatus, str | None]:
|
||||
if data.get("interrupts"):
|
||||
return "interrupted", None
|
||||
error = data.get("error")
|
||||
if error:
|
||||
return "failed", str(error)
|
||||
return "completed", None
|
||||
|
||||
|
||||
def _is_direct_child(namespace: list[str], scope: tuple[str, ...]) -> bool:
|
||||
return len(namespace) == len(scope) + 1 and tuple(namespace[: len(scope)]) == scope
|
||||
|
||||
|
||||
def _subgraph_subscription_params(scope: tuple[str, ...]) -> SubscribeParams:
|
||||
# Subscribe to tasks + messages + tools without a depth limit so that all
|
||||
# descendant-namespace events are captured in one SSE and buffered into each
|
||||
# child handle's inbox. This avoids a second SSE open (and the dedup-set
|
||||
# conflict that would prevent replaying already-seen event_ids).
|
||||
return {
|
||||
"channels": ["messages", "tasks", "tools"],
|
||||
"namespaces": [list(scope)],
|
||||
}
|
||||
|
||||
|
||||
class ScopedStreamHandle:
|
||||
"""Scoped streaming handle for one discovered child invocation."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
thread: AsyncThreadStream,
|
||||
path: tuple[str, ...],
|
||||
graph_name: str | None,
|
||||
trigger_call_id: str | None,
|
||||
max_queue_size: int = 0,
|
||||
) -> None:
|
||||
self._thread = thread
|
||||
self.path = path
|
||||
self.namespace = list(path)
|
||||
self.graph_name = graph_name
|
||||
self.trigger_call_id = trigger_call_id
|
||||
self.status: SubgraphStatus = "started"
|
||||
self.error: str | None = None
|
||||
self._max_queue_size = max_queue_size
|
||||
# Per-channel inboxes: events captured by the parent _SubgraphsProjection
|
||||
# while the SSE was alive. Child projections drain these after the parent
|
||||
# finishes so sequential consumption works without a second SSE open.
|
||||
self._messages_inbox: asyncio.Queue[Event | None] = asyncio.Queue(
|
||||
maxsize=max_queue_size
|
||||
)
|
||||
self._tools_inbox: asyncio.Queue[Event | None] = asyncio.Queue(
|
||||
maxsize=max_queue_size
|
||||
)
|
||||
self._tasks_inbox: asyncio.Queue[Event | None] = asyncio.Queue(
|
||||
maxsize=max_queue_size
|
||||
)
|
||||
# Descendant handles registered by _HandleSubgraphsProjection when a
|
||||
# grandchild is discovered. _push_event fans out to each matching
|
||||
# descendant at dispatch time so events arrive in arrival order without
|
||||
# any drain-and-replay.
|
||||
self._descendant_handles: dict[tuple[str, ...], ScopedStreamHandle] = {}
|
||||
# Track which inboxes have a consumer so _close_inboxes only sends a
|
||||
# sentinel where it is needed. Inboxes with no consumer would otherwise
|
||||
# accumulate a leaked None sentinel that is never drained.
|
||||
self._iterated_inboxes: set[str] = set()
|
||||
self.messages = _HandleMessagesProjection(self)
|
||||
self.tool_calls = _HandleToolCallsProjection(self)
|
||||
self.subgraphs = _HandleSubgraphsProjection(self)
|
||||
self.subagents = self.subgraphs
|
||||
|
||||
def _push_event(self, event: Event) -> None:
|
||||
"""Route a descendant event into the appropriate channel inbox.
|
||||
|
||||
Also fans out to any registered descendant handles whose path is a
|
||||
prefix of the event namespace, so grandchild events are delivered at
|
||||
push time rather than via a post-hoc drain-and-replay.
|
||||
"""
|
||||
method = event.get("method")
|
||||
if method == "messages":
|
||||
self._messages_inbox.put_nowait(event)
|
||||
elif method == "tools":
|
||||
self._tools_inbox.put_nowait(event)
|
||||
elif method == "tasks":
|
||||
self._tasks_inbox.put_nowait(event)
|
||||
# Fan out to descendant handles whose namespace is a prefix of the
|
||||
# event namespace so they receive the event at push time.
|
||||
if method in ("messages", "tools", "tasks"):
|
||||
ns_tuple = tuple(_event_namespace(event.get("params") or {}))
|
||||
for desc_path, desc_handle in self._descendant_handles.items():
|
||||
desc_len = len(desc_path)
|
||||
if len(ns_tuple) >= desc_len and ns_tuple[:desc_len] == desc_path:
|
||||
desc_handle._push_event(event)
|
||||
|
||||
def _register_descendant(self, handle: ScopedStreamHandle) -> None:
|
||||
"""Register a newly-discovered grandchild so future events are fanned out.
|
||||
|
||||
Also drains any events already buffered in this handle's inboxes whose
|
||||
namespace matches the grandchild, so events that arrived before the
|
||||
grandchild was discovered are forwarded in arrival order.
|
||||
"""
|
||||
self._descendant_handles[handle.path] = handle
|
||||
desc_len = len(handle.path)
|
||||
for inbox_attr in (
|
||||
"_messages_inbox",
|
||||
"_tools_inbox",
|
||||
"_tasks_inbox",
|
||||
):
|
||||
inbox: asyncio.Queue[Event | None] = getattr(self, inbox_attr)
|
||||
staging: list[Event | None] = []
|
||||
while not inbox.empty():
|
||||
staging.append(inbox.get_nowait())
|
||||
for event in staging:
|
||||
inbox.put_nowait(event)
|
||||
if event is None:
|
||||
continue
|
||||
ns_tuple = tuple(_event_namespace(event.get("params") or {}))
|
||||
if len(ns_tuple) >= desc_len and ns_tuple[:desc_len] == handle.path:
|
||||
getattr(handle, inbox_attr).put_nowait(event)
|
||||
|
||||
def _unregister_descendant(self, path: tuple[str, ...]) -> None:
|
||||
"""Remove a grandchild after it reaches a terminal state."""
|
||||
self._descendant_handles.pop(path, None)
|
||||
|
||||
def _mark_iterated(self, kind: str) -> None:
|
||||
"""Record that an inbox has an active consumer.
|
||||
|
||||
If the handle is already closed (status != 'started'), immediately
|
||||
enqueue a sentinel so the consumer's `await get()` terminates. This
|
||||
handles sequential consumption (iterate after the handle is finished).
|
||||
|
||||
Must be called by each projection at the start of iteration.
|
||||
"""
|
||||
self._iterated_inboxes.add(kind)
|
||||
if self.status != "started":
|
||||
# Handle already closed before this consumer started; send the
|
||||
# sentinel now so the projection iterator can terminate.
|
||||
getattr(self, f"_{kind}_inbox").put_nowait(None)
|
||||
|
||||
def _close_inboxes(self) -> None:
|
||||
"""Signal EOF only on channel inboxes that have an active consumer.
|
||||
|
||||
Inboxes without a consumer would accumulate a leaked None sentinel
|
||||
that is never drained, so we skip them. For inboxes whose consumer
|
||||
starts after this call, `_mark_iterated` sends the sentinel lazily.
|
||||
"""
|
||||
for kind in ("messages", "tools", "tasks"):
|
||||
if kind in self._iterated_inboxes:
|
||||
getattr(self, f"_{kind}_inbox").put_nowait(None)
|
||||
|
||||
def _finish(self, status: SubgraphStatus, error: str | None = None) -> None:
|
||||
if self.status != "started":
|
||||
return
|
||||
self.status = status
|
||||
self.error = error
|
||||
self._close_inboxes()
|
||||
|
||||
|
||||
class _HandleMessagesProjection:
|
||||
"""Messages projection that drains a `ScopedStreamHandle`'s messages inbox."""
|
||||
|
||||
def __init__(self, handle: ScopedStreamHandle) -> None:
|
||||
self._handle = handle
|
||||
|
||||
def __aiter__(self) -> AsyncIterator[Any]:
|
||||
return self._messages_iter()
|
||||
|
||||
async def _messages_iter(self) -> AsyncGenerator[Any, None]:
|
||||
from langchain_core.language_models.chat_model_stream import (
|
||||
AsyncChatModelStream,
|
||||
)
|
||||
|
||||
self._handle._mark_iterated("messages")
|
||||
active: dict[str, AsyncChatModelStream] = {}
|
||||
while True:
|
||||
item = await self._handle._messages_inbox.get()
|
||||
if item is None:
|
||||
return
|
||||
params_field = item.get("params") or {}
|
||||
ns = _event_namespace(params_field)
|
||||
if ns != self._handle.namespace:
|
||||
continue
|
||||
data = params_field.get("data") if isinstance(params_field, dict) else None
|
||||
if not isinstance(data, dict):
|
||||
continue
|
||||
event_type = data.get("event")
|
||||
if event_type == "message-start":
|
||||
message_id = _message_event_id(data)
|
||||
key = _message_route_key(data, fallback=message_id)
|
||||
metadata = (
|
||||
data.get("metadata")
|
||||
if isinstance(data.get("metadata"), dict)
|
||||
else {}
|
||||
)
|
||||
stream = AsyncChatModelStream(
|
||||
namespace=list(self._handle.namespace),
|
||||
node=metadata.get("langgraph_node") if metadata else None,
|
||||
message_id=message_id,
|
||||
)
|
||||
active[key] = stream
|
||||
stream.dispatch(data)
|
||||
yield stream
|
||||
else:
|
||||
key = _message_route_key(data)
|
||||
stream = active.get(key)
|
||||
if stream is None and len(active) == 1:
|
||||
stream = next(iter(active.values()))
|
||||
if stream is None:
|
||||
continue
|
||||
stream.dispatch(data)
|
||||
if event_type in ("message-finish", "error"):
|
||||
for route_key, candidate in list(active.items()):
|
||||
if candidate is stream:
|
||||
del active[route_key]
|
||||
|
||||
|
||||
class _HandleToolCallsProjection:
|
||||
"""Tool calls projection that drains a `ScopedStreamHandle`'s tools inbox."""
|
||||
|
||||
def __init__(self, handle: ScopedStreamHandle) -> None:
|
||||
self._handle = handle
|
||||
|
||||
def __aiter__(self) -> AsyncIterator[Any]:
|
||||
return self._tool_calls_iter()
|
||||
|
||||
async def _tool_calls_iter(self) -> AsyncGenerator[Any, None]:
|
||||
self._handle._mark_iterated("tools")
|
||||
active: dict[str, ToolCallHandle] = {}
|
||||
while True:
|
||||
item = await self._handle._tools_inbox.get()
|
||||
if item is None:
|
||||
err = RuntimeError(
|
||||
"Tool call stream closed before terminal tool event."
|
||||
)
|
||||
for handle in active.values():
|
||||
handle._fail(err)
|
||||
return
|
||||
params_field = item.get("params") or {}
|
||||
ns = _event_namespace(params_field)
|
||||
if ns != self._handle.namespace:
|
||||
continue
|
||||
data = params_field.get("data") if isinstance(params_field, dict) else None
|
||||
if not isinstance(data, dict):
|
||||
continue
|
||||
event_type = data.get("event")
|
||||
tool_call_id = data.get("tool_call_id")
|
||||
if not isinstance(tool_call_id, str):
|
||||
continue
|
||||
if event_type == "tool-started":
|
||||
tool_name = data.get("tool_name")
|
||||
if not isinstance(tool_name, str):
|
||||
tool_name = ""
|
||||
handle = ToolCallHandle(
|
||||
tool_call_id=tool_call_id,
|
||||
name=tool_name,
|
||||
input=data.get("input"),
|
||||
namespace=list(self._handle.namespace),
|
||||
)
|
||||
active[tool_call_id] = handle
|
||||
yield handle
|
||||
elif event_type == "tool-output-delta":
|
||||
handle = active.get(tool_call_id)
|
||||
delta = data.get("delta")
|
||||
if handle is not None and isinstance(delta, str):
|
||||
handle._push_delta(delta)
|
||||
elif event_type == "tool-finished":
|
||||
handle = active.pop(tool_call_id, None)
|
||||
if handle is not None:
|
||||
handle._finish(data.get("output"))
|
||||
elif event_type == "tool-error":
|
||||
handle = active.pop(tool_call_id, None)
|
||||
if handle is not None:
|
||||
message = data.get("message")
|
||||
handle._fail(
|
||||
RuntimeError(str(message) if message else "Tool call errored")
|
||||
)
|
||||
|
||||
|
||||
class _HandleSubgraphsProjection:
|
||||
"""Subgraphs projection that drains a `ScopedStreamHandle`'s tasks inbox."""
|
||||
|
||||
def __init__(self, handle: ScopedStreamHandle) -> None:
|
||||
self._handle = handle
|
||||
|
||||
def __aiter__(self) -> AsyncIterator[ScopedStreamHandle]:
|
||||
return self._subgraphs_iter()
|
||||
|
||||
async def _subgraphs_iter(self) -> AsyncGenerator[ScopedStreamHandle, None]:
|
||||
self._handle._mark_iterated("tasks")
|
||||
seen: set[tuple[str, ...]] = set()
|
||||
active: dict[tuple[str, ...], ScopedStreamHandle] = {}
|
||||
scope = self._handle.path
|
||||
while True:
|
||||
item = await self._handle._tasks_inbox.get()
|
||||
if item is None:
|
||||
for child in active.values():
|
||||
if child.status == "started":
|
||||
child._finish("completed")
|
||||
return
|
||||
params_field = item.get("params") or {}
|
||||
namespace = _event_namespace(params_field)
|
||||
data = params_field.get("data") if isinstance(params_field, dict) else None
|
||||
if not isinstance(data, dict):
|
||||
continue
|
||||
if "result" in data:
|
||||
result_id = data.get("id")
|
||||
if not result_id:
|
||||
continue
|
||||
parent_path = tuple(namespace)
|
||||
for child_path, child_handle in list(active.items()):
|
||||
if child_path[:-1] != parent_path:
|
||||
continue
|
||||
if child_handle.trigger_call_id != result_id:
|
||||
continue
|
||||
status, error = _terminal_from_tasks_result(data)
|
||||
child_handle._finish(status, error)
|
||||
del active[child_path]
|
||||
self._handle._unregister_descendant(child_path)
|
||||
continue
|
||||
if not _is_direct_child(namespace, scope):
|
||||
continue
|
||||
path = tuple(namespace)
|
||||
if path in seen:
|
||||
continue
|
||||
seen.add(path)
|
||||
graph_name, trigger_call_id = _parse_namespace_segment(path[-1])
|
||||
child_handle = ScopedStreamHandle(
|
||||
thread=self._handle._thread,
|
||||
path=path,
|
||||
graph_name=graph_name or None,
|
||||
trigger_call_id=trigger_call_id,
|
||||
max_queue_size=self._handle._max_queue_size,
|
||||
)
|
||||
active[path] = child_handle
|
||||
# Register so future _push_event calls on this handle fan out to the
|
||||
# grandchild at push time, preserving arrival order without drain-and-replay.
|
||||
self._handle._register_descendant(child_handle)
|
||||
yield child_handle
|
||||
|
||||
|
||||
class _SubgraphsProjection:
|
||||
"""Discover direct child invocations for a namespace scope."""
|
||||
|
||||
def __init__(self, thread: AsyncThreadStream, scope: tuple[str, ...] = ()) -> None:
|
||||
self._thread = thread
|
||||
self._scope = scope
|
||||
|
||||
def __aiter__(self) -> AsyncIterator[ScopedStreamHandle]:
|
||||
return self._subgraphs_iter()
|
||||
|
||||
async def _subgraphs_iter(self) -> AsyncGenerator[ScopedStreamHandle, None]:
|
||||
if self._thread._transport is None:
|
||||
raise RuntimeError("AsyncThreadStream not entered - use `async with`.")
|
||||
params = _subgraph_subscription_params(self._scope)
|
||||
sub = self._thread._register_subscription(params)
|
||||
seen: set[tuple[str, ...]] = set()
|
||||
active: dict[tuple[str, ...], ScopedStreamHandle] = {}
|
||||
# Activate root inbox so scope-level messages events consumed here are
|
||||
# forwarded to `thread.messages` even after the shared SSE ends.
|
||||
root_inbox: asyncio.Queue[Event | None] | None = (
|
||||
self._thread._activate_root_messages_inbox() if not self._scope else None
|
||||
)
|
||||
try:
|
||||
await self._thread._reconcile_stream(params)
|
||||
self._thread._ensure_fanout_running()
|
||||
while True:
|
||||
item = await sub.queue.get()
|
||||
if item is None:
|
||||
return
|
||||
params_field = item.get("params") or {}
|
||||
namespace = _event_namespace(params_field)
|
||||
data = (
|
||||
params_field.get("data") if isinstance(params_field, dict) else None
|
||||
)
|
||||
if not isinstance(data, dict):
|
||||
continue
|
||||
method = item.get("method")
|
||||
|
||||
# Route events at a child namespace (or deeper) to that child
|
||||
# handle's channel inbox so sequential child-projection
|
||||
# consumption works without opening a second SSE.
|
||||
ns_tuple = tuple(namespace)
|
||||
routed_to_child = False
|
||||
for child_path, child_handle in active.items():
|
||||
child_len = len(child_path)
|
||||
if (
|
||||
len(ns_tuple) >= child_len
|
||||
and ns_tuple[:child_len] == child_path
|
||||
):
|
||||
child_handle._push_event(item)
|
||||
routed_to_child = True
|
||||
break
|
||||
|
||||
# Scope-level messages events are not routed to any child; forward
|
||||
# them to the root inbox so `thread.messages` can drain them after
|
||||
# this projection finishes (dedup prevents the SSE from replaying).
|
||||
if (
|
||||
not routed_to_child
|
||||
and root_inbox is not None
|
||||
and method == "messages"
|
||||
and tuple(namespace) == self._scope
|
||||
):
|
||||
root_inbox.put_nowait(item)
|
||||
|
||||
if method == "tasks":
|
||||
if "result" in data:
|
||||
self._apply_tasks_result(namespace, data, active)
|
||||
elif _is_direct_child(namespace, self._scope):
|
||||
path = tuple(namespace)
|
||||
if path not in seen:
|
||||
seen.add(path)
|
||||
graph_name, trigger_call_id = _parse_namespace_segment(
|
||||
path[-1]
|
||||
)
|
||||
handle = ScopedStreamHandle(
|
||||
thread=self._thread,
|
||||
path=path,
|
||||
graph_name=graph_name or None,
|
||||
trigger_call_id=trigger_call_id,
|
||||
)
|
||||
active[path] = handle
|
||||
yield handle
|
||||
finally:
|
||||
# Determine terminal status from the parent run's lifecycle result.
|
||||
# If _run_done resolved as errored, force-complete remaining children
|
||||
# as errored so callers see the correct terminal state.
|
||||
terminal_status: SubgraphStatus = "completed"
|
||||
run_done = self._thread._run_done
|
||||
if run_done is not None and run_done.done() and not run_done.cancelled():
|
||||
result = run_done.result()
|
||||
if isinstance(result, _RunTerminal) and result.status == "errored":
|
||||
terminal_status = "failed"
|
||||
for handle in active.values():
|
||||
if handle.status == "started":
|
||||
handle._finish(terminal_status)
|
||||
self._thread._unregister_subscription(sub.id)
|
||||
if root_inbox is not None:
|
||||
root_inbox.put_nowait(None)
|
||||
|
||||
def _apply_tasks_result(
|
||||
self,
|
||||
namespace: list[str],
|
||||
data: dict[str, Any],
|
||||
active: dict[tuple[str, ...], ScopedStreamHandle],
|
||||
) -> None:
|
||||
result_id = data.get("id")
|
||||
if not result_id:
|
||||
return
|
||||
parent_path = tuple(namespace)
|
||||
for child_path, handle in list(active.items()):
|
||||
if child_path[:-1] != parent_path:
|
||||
continue
|
||||
if handle.trigger_call_id != result_id:
|
||||
continue
|
||||
status, error = _terminal_from_tasks_result(data)
|
||||
handle._finish(status, error)
|
||||
del active[child_path]
|
||||
|
||||
|
||||
class ToolCallHandle:
|
||||
"""Async handle for one root-scope tool call."""
|
||||
|
||||
@@ -453,8 +1001,11 @@ class ToolCallHandle:
|
||||
class _ToolCallsProjection:
|
||||
"""Typed projection for root-scope `thread.tool_calls`."""
|
||||
|
||||
def __init__(self, thread: AsyncThreadStream) -> None:
|
||||
def __init__(
|
||||
self, thread: AsyncThreadStream, namespace: list[str] | None = None
|
||||
) -> None:
|
||||
self._thread = thread
|
||||
self._namespace = list(namespace or [])
|
||||
|
||||
def __aiter__(self) -> AsyncIterator[ToolCallHandle]:
|
||||
return self._tool_calls_iter()
|
||||
@@ -462,11 +1013,7 @@ class _ToolCallsProjection:
|
||||
async def _tool_calls_iter(self) -> AsyncGenerator[ToolCallHandle, None]:
|
||||
if self._thread._transport is None:
|
||||
raise RuntimeError("AsyncThreadStream not entered - use `async with`.")
|
||||
params: SubscribeParams = {
|
||||
"channels": ["tools"],
|
||||
"namespaces": [[]],
|
||||
"depth": 0,
|
||||
}
|
||||
params = _exact_namespace_params(["tools"], self._namespace)
|
||||
sub = self._thread._register_subscription(params)
|
||||
active: dict[str, ToolCallHandle] = {}
|
||||
try:
|
||||
@@ -477,12 +1024,11 @@ class _ToolCallsProjection:
|
||||
if item is None:
|
||||
return
|
||||
params_field = item.get("params") or {}
|
||||
if not isinstance(params_field, dict):
|
||||
if _event_namespace(params_field) != self._namespace:
|
||||
continue
|
||||
namespace = params_field.get("namespace") or []
|
||||
if namespace != []:
|
||||
continue
|
||||
data = params_field.get("data")
|
||||
data = (
|
||||
params_field.get("data") if isinstance(params_field, dict) else None
|
||||
)
|
||||
if not isinstance(data, dict):
|
||||
continue
|
||||
event_type = data.get("event")
|
||||
@@ -498,7 +1044,7 @@ class _ToolCallsProjection:
|
||||
tool_call_id=tool_call_id,
|
||||
name=tool_name,
|
||||
input=data.get("input"),
|
||||
namespace=namespace,
|
||||
namespace=list(self._namespace),
|
||||
)
|
||||
active[tool_call_id] = handle
|
||||
self._thread._register_active_tool_call(handle)
|
||||
@@ -594,11 +1140,17 @@ class AsyncThreadStream:
|
||||
self._run_done: asyncio.Future[_RunTerminal] | None = None
|
||||
self._active_message_streams: set[AsyncChatModelStream] = set()
|
||||
self._active_tool_calls: set[ToolCallHandle] = set()
|
||||
# Root-scope inbox: populated by `_SubgraphsProjection` when it consumes
|
||||
# messages events at namespace `[]` so that `thread.messages` can drain
|
||||
# them even after the shared SSE has ended (dedup prevents replay).
|
||||
self._root_messages_inbox: asyncio.Queue[Event | None] | None = None
|
||||
self.run = RunModule(self)
|
||||
self.output = _OutputAwaitable(self)
|
||||
self.values = _ValuesProjection(self)
|
||||
self.messages = _MessagesProjection(self)
|
||||
self.tool_calls = _ToolCallsProjection(self)
|
||||
self.messages = _MessagesProjection(self, namespace=[])
|
||||
self.tool_calls = _ToolCallsProjection(self, namespace=[])
|
||||
self.subgraphs = _SubgraphsProjection(self, scope=())
|
||||
self.subagents = self.subgraphs
|
||||
|
||||
@property
|
||||
def _controller(self) -> AsyncThreadStream:
|
||||
@@ -692,6 +1244,16 @@ class AsyncThreadStream:
|
||||
"""Remove a subscription from the registry. No-op if already absent."""
|
||||
self._subscriptions.pop(subscription_id, None)
|
||||
|
||||
def _activate_root_messages_inbox(self) -> asyncio.Queue[Event | None]:
|
||||
"""Create the root-scope messages inbox if not already active and return it.
|
||||
|
||||
Called by `_SubgraphsProjection` at scope `()` to capture messages events
|
||||
that arrive at namespace `[]` before `thread.messages` has subscribed.
|
||||
"""
|
||||
if self._root_messages_inbox is None:
|
||||
self._root_messages_inbox = asyncio.Queue()
|
||||
return self._root_messages_inbox
|
||||
|
||||
def _register_active_message_stream(self, stream: AsyncChatModelStream) -> None:
|
||||
self._active_message_streams.add(stream)
|
||||
|
||||
|
||||
@@ -0,0 +1,579 @@
|
||||
"""Tests for nested scoped stream handles."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import httpx
|
||||
|
||||
from langgraph_sdk._async.http import HttpClient
|
||||
from langgraph_sdk._async.threads import ThreadsClient
|
||||
from streaming._events import (
|
||||
lifecycle_completed_event,
|
||||
lifecycle_started_event,
|
||||
message_finish_event,
|
||||
message_start_event,
|
||||
message_text_delta_event,
|
||||
message_text_finish_event,
|
||||
tasks_result_event,
|
||||
tasks_start_event,
|
||||
tool_finished_event,
|
||||
tool_output_delta_event,
|
||||
tool_started_event,
|
||||
)
|
||||
from streaming._fake_server import FakeServer
|
||||
|
||||
|
||||
async def test_subgraphs_subscribes_to_tasks_channel():
|
||||
fake = FakeServer()
|
||||
fake.script([lifecycle_completed_event(seq=1)])
|
||||
asgi = httpx.ASGITransport(app=fake.app)
|
||||
async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
|
||||
threads = ThreadsClient(HttpClient(raw))
|
||||
async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
|
||||
await thread.run.start(input={})
|
||||
_ = [handle async for handle in thread.subgraphs]
|
||||
|
||||
assert any(
|
||||
"tasks" in body.get("channels", []) for body in fake.stream_request_bodies
|
||||
)
|
||||
|
||||
|
||||
async def test_subgraphs_yields_handle_and_completes_status():
|
||||
fake = FakeServer()
|
||||
fake.script(
|
||||
[
|
||||
lifecycle_started_event(seq=0),
|
||||
tasks_start_event(seq=1, namespace=["worker:abc"], task_id="t-child"),
|
||||
tasks_result_event(seq=2, namespace=[], task_id="abc", name="worker"),
|
||||
lifecycle_completed_event(seq=3),
|
||||
]
|
||||
)
|
||||
asgi = httpx.ASGITransport(app=fake.app)
|
||||
async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
|
||||
threads = ThreadsClient(HttpClient(raw))
|
||||
async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
|
||||
await thread.run.start(input={})
|
||||
handles = [handle async for handle in thread.subgraphs]
|
||||
|
||||
assert len(handles) == 1
|
||||
handle = handles[0]
|
||||
assert handle.path == ("worker:abc",)
|
||||
assert handle.namespace == ["worker:abc"]
|
||||
assert handle.graph_name == "worker"
|
||||
assert handle.trigger_call_id == "abc"
|
||||
assert handle.status == "completed"
|
||||
assert handle.error is None
|
||||
|
||||
|
||||
async def test_subgraphs_failed_and_interrupted_statuses():
|
||||
failed = FakeServer()
|
||||
failed.script(
|
||||
[
|
||||
lifecycle_started_event(seq=0),
|
||||
tasks_start_event(seq=1, namespace=["worker:abc"], task_id="t-child"),
|
||||
tasks_result_event(
|
||||
seq=2,
|
||||
namespace=[],
|
||||
task_id="abc",
|
||||
name="worker",
|
||||
error="boom",
|
||||
),
|
||||
lifecycle_completed_event(seq=3),
|
||||
]
|
||||
)
|
||||
failed_asgi = httpx.ASGITransport(app=failed.app)
|
||||
async with httpx.AsyncClient(transport=failed_asgi, base_url="http://test") as raw:
|
||||
threads = ThreadsClient(HttpClient(raw))
|
||||
async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
|
||||
await thread.run.start(input={})
|
||||
[failed_handle] = [handle async for handle in thread.subgraphs]
|
||||
|
||||
interrupted = FakeServer()
|
||||
interrupted.script(
|
||||
[
|
||||
lifecycle_started_event(seq=0),
|
||||
tasks_start_event(seq=1, namespace=["worker:def"], task_id="t-child"),
|
||||
tasks_result_event(
|
||||
seq=2,
|
||||
namespace=[],
|
||||
task_id="def",
|
||||
name="worker",
|
||||
interrupts=[{"value": "pause"}],
|
||||
),
|
||||
lifecycle_completed_event(seq=3),
|
||||
]
|
||||
)
|
||||
interrupted_asgi = httpx.ASGITransport(app=interrupted.app)
|
||||
async with httpx.AsyncClient(
|
||||
transport=interrupted_asgi,
|
||||
base_url="http://test",
|
||||
) as raw:
|
||||
threads = ThreadsClient(HttpClient(raw))
|
||||
async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
|
||||
await thread.run.start(input={})
|
||||
[interrupted_handle] = [handle async for handle in thread.subgraphs]
|
||||
|
||||
assert failed_handle.status == "failed"
|
||||
assert failed_handle.error == "boom"
|
||||
assert interrupted_handle.status == "interrupted"
|
||||
assert interrupted_handle.error is None
|
||||
|
||||
|
||||
async def test_subgraph_messages_are_scoped_to_child_namespace():
|
||||
fake = FakeServer()
|
||||
fake.script(
|
||||
[
|
||||
lifecycle_started_event(seq=0),
|
||||
tasks_start_event(seq=1, namespace=["worker:abc"], task_id="t-child"),
|
||||
message_start_event(
|
||||
seq=2,
|
||||
namespace=["worker:abc"],
|
||||
message_id="msg-child",
|
||||
run_id="run-child",
|
||||
),
|
||||
message_text_delta_event(seq=3, namespace=["worker:abc"], text="child"),
|
||||
message_text_finish_event(seq=4, namespace=["worker:abc"], text="child"),
|
||||
message_finish_event(seq=5, namespace=["worker:abc"]),
|
||||
tasks_result_event(seq=6, namespace=[], task_id="abc", name="worker"),
|
||||
lifecycle_completed_event(seq=7),
|
||||
]
|
||||
)
|
||||
asgi = httpx.ASGITransport(app=fake.app)
|
||||
async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
|
||||
threads = ThreadsClient(HttpClient(raw))
|
||||
async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
|
||||
await thread.run.start(input={})
|
||||
[handle] = [h async for h in thread.subgraphs]
|
||||
child_messages = [message async for message in handle.messages]
|
||||
root_messages = [message async for message in thread.messages]
|
||||
|
||||
assert [message.message_id for message in child_messages] == ["msg-child"]
|
||||
assert [await message.text for message in child_messages] == ["child"]
|
||||
assert root_messages == []
|
||||
|
||||
|
||||
async def test_subgraph_tool_calls_are_scoped_to_child_namespace():
|
||||
fake = FakeServer()
|
||||
fake.script(
|
||||
[
|
||||
lifecycle_started_event(seq=0),
|
||||
tasks_start_event(seq=1, namespace=["worker:abc"], task_id="t-child"),
|
||||
tool_started_event(
|
||||
seq=2,
|
||||
namespace=["worker:abc"],
|
||||
tool_call_id="call-child",
|
||||
tool_name="search",
|
||||
),
|
||||
tool_output_delta_event(
|
||||
seq=3,
|
||||
namespace=["worker:abc"],
|
||||
tool_call_id="call-child",
|
||||
delta="child-delta",
|
||||
),
|
||||
tool_finished_event(
|
||||
seq=4,
|
||||
namespace=["worker:abc"],
|
||||
tool_call_id="call-child",
|
||||
output={"ok": True},
|
||||
),
|
||||
tasks_result_event(seq=5, namespace=[], task_id="abc", name="worker"),
|
||||
lifecycle_completed_event(seq=6),
|
||||
]
|
||||
)
|
||||
asgi = httpx.ASGITransport(app=fake.app)
|
||||
async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
|
||||
threads = ThreadsClient(HttpClient(raw))
|
||||
async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
|
||||
await thread.run.start(input={})
|
||||
[handle] = [h async for h in thread.subgraphs]
|
||||
child_calls = [call async for call in handle.tool_calls]
|
||||
root_calls = [call async for call in thread.tool_calls]
|
||||
|
||||
assert [call.tool_call_id for call in child_calls] == ["call-child"]
|
||||
assert [delta async for delta in child_calls[0].deltas] == ["child-delta"]
|
||||
assert await child_calls[0].output == {"ok": True}
|
||||
assert root_calls == []
|
||||
|
||||
|
||||
async def test_subgraph_handles_are_recursive_for_grandchildren():
|
||||
fake = FakeServer()
|
||||
fake.script(
|
||||
[
|
||||
lifecycle_started_event(seq=0),
|
||||
tasks_start_event(seq=1, namespace=["worker:abc"], task_id="t-child"),
|
||||
tasks_start_event(
|
||||
seq=2,
|
||||
namespace=["worker:abc", "tool:def"],
|
||||
task_id="t-grandchild",
|
||||
),
|
||||
tasks_result_event(
|
||||
seq=3,
|
||||
namespace=["worker:abc"],
|
||||
task_id="def",
|
||||
name="tool",
|
||||
),
|
||||
tasks_result_event(seq=4, namespace=[], task_id="abc", name="worker"),
|
||||
lifecycle_completed_event(seq=5),
|
||||
]
|
||||
)
|
||||
asgi = httpx.ASGITransport(app=fake.app)
|
||||
async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
|
||||
threads = ThreadsClient(HttpClient(raw))
|
||||
async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
|
||||
await thread.run.start(input={})
|
||||
[child] = [handle async for handle in thread.subgraphs]
|
||||
[grandchild] = [handle async for handle in child.subgraphs]
|
||||
|
||||
assert child.path == ("worker:abc",)
|
||||
assert grandchild.path == ("worker:abc", "tool:def")
|
||||
assert grandchild.graph_name == "tool"
|
||||
assert grandchild.trigger_call_id == "def"
|
||||
assert grandchild.status == "completed"
|
||||
|
||||
|
||||
async def test_subagents_aliases_subgraphs_until_protocol_distinguishes_them():
|
||||
fake = FakeServer()
|
||||
fake.script([lifecycle_completed_event(seq=1)])
|
||||
asgi = httpx.ASGITransport(app=fake.app)
|
||||
async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
|
||||
threads = ThreadsClient(HttpClient(raw))
|
||||
async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
|
||||
assert thread.subagents is thread.subgraphs
|
||||
await thread.run.start(input={})
|
||||
|
||||
|
||||
async def test_root_and_child_projections_do_not_cross_talk():
|
||||
fake = FakeServer()
|
||||
fake.script(
|
||||
[
|
||||
lifecycle_started_event(seq=0),
|
||||
tasks_start_event(seq=1, namespace=["worker:abc"], task_id="t-child"),
|
||||
message_start_event(seq=2, message_id="root-msg", run_id="root-run"),
|
||||
message_text_delta_event(seq=3, text="root"),
|
||||
message_text_finish_event(seq=4, text="root"),
|
||||
message_finish_event(seq=5),
|
||||
message_start_event(
|
||||
seq=6,
|
||||
namespace=["worker:abc"],
|
||||
message_id="child-msg",
|
||||
run_id="child-run",
|
||||
),
|
||||
message_text_delta_event(seq=7, namespace=["worker:abc"], text="child"),
|
||||
message_text_finish_event(seq=8, namespace=["worker:abc"], text="child"),
|
||||
message_finish_event(seq=9, namespace=["worker:abc"]),
|
||||
tasks_result_event(seq=10, namespace=[], task_id="abc", name="worker"),
|
||||
lifecycle_completed_event(seq=11),
|
||||
]
|
||||
)
|
||||
asgi = httpx.ASGITransport(app=fake.app)
|
||||
async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
|
||||
threads = ThreadsClient(HttpClient(raw))
|
||||
async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
|
||||
await thread.run.start(input={})
|
||||
[handle] = [h async for h in thread.subgraphs]
|
||||
root_messages = [message async for message in thread.messages]
|
||||
child_messages = [message async for message in handle.messages]
|
||||
|
||||
assert [message.message_id for message in root_messages] == ["root-msg"]
|
||||
assert [await message.text for message in root_messages] == ["root"]
|
||||
assert [message.message_id for message in child_messages] == ["child-msg"]
|
||||
assert [await message.text for message in child_messages] == ["child"]
|
||||
|
||||
|
||||
async def test_grandchild_sibling_routing_preserves_event_order():
|
||||
"""Events enqueued in a child handle's _messages_inbox before a grandchild
|
||||
is discovered via child.subgraphs must be delivered in arrival order, not
|
||||
reordered by the drain-and-replay path in _route_sibling_inboxes_to_grandchildren."""
|
||||
fake = FakeServer()
|
||||
fake.script(
|
||||
[
|
||||
lifecycle_started_event(seq=0),
|
||||
# Child discovered by thread.subgraphs.
|
||||
tasks_start_event(seq=1, namespace=["worker:abc"], task_id="t-child"),
|
||||
# Grandchild messages arrive *before* grandchild tasks-start.
|
||||
message_start_event(
|
||||
seq=2,
|
||||
namespace=["worker:abc", "tool:gc1"],
|
||||
message_id="msg-E1",
|
||||
run_id="r1",
|
||||
),
|
||||
message_text_delta_event(
|
||||
seq=3, namespace=["worker:abc", "tool:gc1"], text="E1"
|
||||
),
|
||||
message_text_finish_event(
|
||||
seq=4, namespace=["worker:abc", "tool:gc1"], text="E1"
|
||||
),
|
||||
message_finish_event(seq=5, namespace=["worker:abc", "tool:gc1"]),
|
||||
message_start_event(
|
||||
seq=6,
|
||||
namespace=["worker:abc", "tool:gc1"],
|
||||
message_id="msg-E2",
|
||||
run_id="r2",
|
||||
),
|
||||
message_text_delta_event(
|
||||
seq=7, namespace=["worker:abc", "tool:gc1"], text="E2"
|
||||
),
|
||||
message_text_finish_event(
|
||||
seq=8, namespace=["worker:abc", "tool:gc1"], text="E2"
|
||||
),
|
||||
message_finish_event(seq=9, namespace=["worker:abc", "tool:gc1"]),
|
||||
message_start_event(
|
||||
seq=10,
|
||||
namespace=["worker:abc", "tool:gc1"],
|
||||
message_id="msg-E3",
|
||||
run_id="r3",
|
||||
),
|
||||
message_text_delta_event(
|
||||
seq=11, namespace=["worker:abc", "tool:gc1"], text="E3"
|
||||
),
|
||||
message_text_finish_event(
|
||||
seq=12, namespace=["worker:abc", "tool:gc1"], text="E3"
|
||||
),
|
||||
message_finish_event(seq=13, namespace=["worker:abc", "tool:gc1"]),
|
||||
# Grandchild tasks-start arrives *after* its messages.
|
||||
tasks_start_event(
|
||||
seq=14,
|
||||
namespace=["worker:abc", "tool:gc1"],
|
||||
task_id="t-grandchild",
|
||||
),
|
||||
tasks_result_event(
|
||||
seq=15, namespace=["worker:abc"], task_id="gc1", name="tool"
|
||||
),
|
||||
tasks_result_event(seq=16, namespace=[], task_id="abc", name="worker"),
|
||||
lifecycle_completed_event(seq=17),
|
||||
]
|
||||
)
|
||||
asgi = httpx.ASGITransport(app=fake.app)
|
||||
async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
|
||||
threads = ThreadsClient(HttpClient(raw))
|
||||
async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
|
||||
await thread.run.start(input={})
|
||||
[child] = [h async for h in thread.subgraphs]
|
||||
[grandchild] = [h async for h in child.subgraphs]
|
||||
gc_messages = [m async for m in grandchild.messages]
|
||||
|
||||
assert [m.message_id for m in gc_messages] == ["msg-E1", "msg-E2", "msg-E3"]
|
||||
texts = [await m.text for m in gc_messages]
|
||||
assert texts == ["E1", "E2", "E3"]
|
||||
|
||||
|
||||
async def test_grandchild_events_dispatched_to_correct_sibling_not_first_match():
|
||||
"""When two sibling grandchildren exist, messages scoped to one grandchild
|
||||
must not bleed into the other grandchild's inbox."""
|
||||
fake = FakeServer()
|
||||
fake.script(
|
||||
[
|
||||
lifecycle_started_event(seq=0),
|
||||
tasks_start_event(seq=1, namespace=["worker:abc"], task_id="t-child"),
|
||||
# Two grandchildren discovered.
|
||||
tasks_start_event(
|
||||
seq=2, namespace=["worker:abc", "tool:gc1"], task_id="t-gc1"
|
||||
),
|
||||
tasks_start_event(
|
||||
seq=3, namespace=["worker:abc", "tool:gc2"], task_id="t-gc2"
|
||||
),
|
||||
# Messages for gc1.
|
||||
message_start_event(
|
||||
seq=4,
|
||||
namespace=["worker:abc", "tool:gc1"],
|
||||
message_id="msg-gc1",
|
||||
run_id="ra",
|
||||
),
|
||||
message_text_delta_event(
|
||||
seq=5, namespace=["worker:abc", "tool:gc1"], text="GC1"
|
||||
),
|
||||
message_text_finish_event(
|
||||
seq=6, namespace=["worker:abc", "tool:gc1"], text="GC1"
|
||||
),
|
||||
message_finish_event(seq=7, namespace=["worker:abc", "tool:gc1"]),
|
||||
# Messages for gc2.
|
||||
message_start_event(
|
||||
seq=8,
|
||||
namespace=["worker:abc", "tool:gc2"],
|
||||
message_id="msg-gc2",
|
||||
run_id="rb",
|
||||
),
|
||||
message_text_delta_event(
|
||||
seq=9, namespace=["worker:abc", "tool:gc2"], text="GC2"
|
||||
),
|
||||
message_text_finish_event(
|
||||
seq=10, namespace=["worker:abc", "tool:gc2"], text="GC2"
|
||||
),
|
||||
message_finish_event(seq=11, namespace=["worker:abc", "tool:gc2"]),
|
||||
tasks_result_event(
|
||||
seq=12, namespace=["worker:abc"], task_id="gc1", name="tool"
|
||||
),
|
||||
tasks_result_event(
|
||||
seq=13, namespace=["worker:abc"], task_id="gc2", name="tool"
|
||||
),
|
||||
tasks_result_event(seq=14, namespace=[], task_id="abc", name="worker"),
|
||||
lifecycle_completed_event(seq=15),
|
||||
]
|
||||
)
|
||||
asgi = httpx.ASGITransport(app=fake.app)
|
||||
async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
|
||||
threads = ThreadsClient(HttpClient(raw))
|
||||
async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
|
||||
await thread.run.start(input={})
|
||||
[child] = [h async for h in thread.subgraphs]
|
||||
grandchildren = [h async for h in child.subgraphs]
|
||||
by_path = {h.path: h for h in grandchildren}
|
||||
gc1_messages = [
|
||||
m async for m in by_path[("worker:abc", "tool:gc1")].messages
|
||||
]
|
||||
gc2_messages = [
|
||||
m async for m in by_path[("worker:abc", "tool:gc2")].messages
|
||||
]
|
||||
|
||||
assert [m.message_id for m in gc1_messages] == ["msg-gc1"]
|
||||
assert [m.message_id for m in gc2_messages] == ["msg-gc2"]
|
||||
|
||||
|
||||
def test_scoped_handle_inboxes_bounded_by_max_queue_size():
|
||||
"""ScopedStreamHandle with max_queue_size=N creates queues with maxsize=N."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from langgraph_sdk._async.stream import ScopedStreamHandle
|
||||
|
||||
fake_thread = MagicMock()
|
||||
handle = ScopedStreamHandle(
|
||||
thread=fake_thread,
|
||||
path=("worker:1",),
|
||||
graph_name="worker",
|
||||
trigger_call_id="1",
|
||||
max_queue_size=16,
|
||||
)
|
||||
assert handle._messages_inbox.maxsize == 16
|
||||
assert handle._tools_inbox.maxsize == 16
|
||||
assert handle._tasks_inbox.maxsize == 16
|
||||
|
||||
|
||||
async def test_child_handle_inherits_max_queue_size_from_parent():
|
||||
"""Grandchild ScopedStreamHandles created by _HandleSubgraphsProjection
|
||||
inherit the parent's max_queue_size so all queues are consistently bounded."""
|
||||
fake = FakeServer()
|
||||
fake.script(
|
||||
[
|
||||
lifecycle_started_event(seq=0),
|
||||
tasks_start_event(seq=1, namespace=["worker:abc"], task_id="t-child"),
|
||||
tasks_start_event(
|
||||
seq=2, namespace=["worker:abc", "tool:gc1"], task_id="t-gc1"
|
||||
),
|
||||
tasks_result_event(
|
||||
seq=3, namespace=["worker:abc"], task_id="gc1", name="tool"
|
||||
),
|
||||
tasks_result_event(seq=4, namespace=[], task_id="abc", name="worker"),
|
||||
lifecycle_completed_event(seq=5),
|
||||
]
|
||||
)
|
||||
asgi = httpx.ASGITransport(app=fake.app)
|
||||
async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
|
||||
threads = ThreadsClient(HttpClient(raw))
|
||||
async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
|
||||
await thread.run.start(input={})
|
||||
[child] = [h async for h in thread.subgraphs]
|
||||
[grandchild] = [h async for h in child.subgraphs]
|
||||
|
||||
# Grandchild handles created by _HandleSubgraphsProjection must use the
|
||||
# parent handle's max_queue_size (default 0 = unbounded in asyncio.Queue).
|
||||
assert grandchild._messages_inbox.maxsize == child._messages_inbox.maxsize
|
||||
assert grandchild._tools_inbox.maxsize == child._tools_inbox.maxsize
|
||||
assert grandchild._tasks_inbox.maxsize == child._tasks_inbox.maxsize
|
||||
|
||||
|
||||
async def test_force_complete_uses_failed_when_run_errored():
|
||||
"""If the lifecycle signals an errored run, scoped children that are still
|
||||
'started' when the subgraphs projection's finally block runs must be
|
||||
force-finished as 'failed', not 'completed'."""
|
||||
from streaming._events import lifecycle_errored_event
|
||||
|
||||
fake = FakeServer()
|
||||
fake.script(
|
||||
[
|
||||
lifecycle_started_event(seq=0),
|
||||
tasks_start_event(seq=1, namespace=["worker:abc"], task_id="t-child"),
|
||||
# No tasks_result — run errored before the child task finished.
|
||||
lifecycle_errored_event(seq=2, error="boom"),
|
||||
]
|
||||
)
|
||||
asgi = httpx.ASGITransport(app=fake.app)
|
||||
handles: list = []
|
||||
async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
|
||||
threads = ThreadsClient(HttpClient(raw))
|
||||
async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
|
||||
await thread.run.start(input={})
|
||||
async for handle in thread.subgraphs:
|
||||
handles.append(handle)
|
||||
|
||||
assert len(handles) == 1
|
||||
child = handles[0]
|
||||
# The run errored, so the child should be force-finished as "failed".
|
||||
assert child.status == "failed"
|
||||
|
||||
|
||||
async def test_force_complete_uses_completed_when_run_completed():
|
||||
"""If the lifecycle signals a completed run, any subgraph child still
|
||||
'started' at finally time is force-finished as 'completed' (normal case)."""
|
||||
fake = FakeServer()
|
||||
fake.script(
|
||||
[
|
||||
lifecycle_started_event(seq=0),
|
||||
tasks_start_event(seq=1, namespace=["worker:abc"], task_id="t-child"),
|
||||
# No tasks_result — but lifecycle completed normally.
|
||||
lifecycle_completed_event(seq=2),
|
||||
]
|
||||
)
|
||||
asgi = httpx.ASGITransport(app=fake.app)
|
||||
handles: list = []
|
||||
async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
|
||||
threads = ThreadsClient(HttpClient(raw))
|
||||
async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
|
||||
await thread.run.start(input={})
|
||||
async for handle in thread.subgraphs:
|
||||
handles.append(handle)
|
||||
|
||||
assert len(handles) == 1
|
||||
child = handles[0]
|
||||
assert child.status == "completed"
|
||||
|
||||
|
||||
def test_close_inboxes_does_not_enqueue_on_uniterated_inboxes():
|
||||
"""_close_inboxes must not push a sentinel on inboxes that had no consumer."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from langgraph_sdk._async.stream import ScopedStreamHandle
|
||||
|
||||
fake_thread = MagicMock()
|
||||
handle = ScopedStreamHandle(
|
||||
thread=fake_thread,
|
||||
path=("worker:1",),
|
||||
graph_name="worker",
|
||||
trigger_call_id="1",
|
||||
)
|
||||
# No projection iterated — _close_inboxes should leave all queues empty.
|
||||
handle._close_inboxes()
|
||||
assert handle._messages_inbox.qsize() == 0
|
||||
assert handle._tools_inbox.qsize() == 0
|
||||
assert handle._tasks_inbox.qsize() == 0
|
||||
|
||||
|
||||
def test_close_inboxes_enqueues_sentinel_on_iterated_inboxes():
|
||||
"""_close_inboxes must push a None sentinel only on inboxes that had a consumer,
|
||||
so projection iterators see the EOF signal."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from langgraph_sdk._async.stream import ScopedStreamHandle
|
||||
|
||||
fake_thread = MagicMock()
|
||||
handle = ScopedStreamHandle(
|
||||
thread=fake_thread,
|
||||
path=("worker:1",),
|
||||
graph_name="worker",
|
||||
trigger_call_id="1",
|
||||
)
|
||||
handle._mark_iterated("messages")
|
||||
handle._close_inboxes()
|
||||
# Only the messages inbox should have a sentinel.
|
||||
assert handle._messages_inbox.qsize() == 1
|
||||
assert handle._messages_inbox.get_nowait() is None
|
||||
assert handle._tools_inbox.qsize() == 0
|
||||
assert handle._tasks_inbox.qsize() == 0
|
||||
Reference in New Issue
Block a user