mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-17 21:25:46 +02:00
feat(sdk-py): add sync scoped subgraphs (#7828)
This commit is contained in:
Nick Hollon
@@ -76,17 +76,49 @@ def _event_namespace(params_field: Any) -> list[str]:
|
||||
return list(namespace) if isinstance(namespace, list) else []
|
||||
|
||||
|
||||
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:
|
||||
return {
|
||||
"channels": ["messages", "tasks", "tools"],
|
||||
"namespaces": [list(scope)],
|
||||
}
|
||||
|
||||
|
||||
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:
|
||||
metadata = data.get("metadata") if isinstance(data.get("metadata"), dict) else {}
|
||||
run_id = metadata.get("run_id") if metadata else None
|
||||
"""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 run_id is not None:
|
||||
return f"run:{run_id}"
|
||||
if message_id is not None:
|
||||
return f"message:{message_id}"
|
||||
if fallback is not None:
|
||||
@@ -324,6 +356,12 @@ class _SyncMessagesProjection:
|
||||
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"):
|
||||
@@ -335,9 +373,9 @@ class _SyncMessagesProjection:
|
||||
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:
|
||||
# No active stream matches this event's key. Drop rather
|
||||
# than silently misroute to the only remaining stream.
|
||||
continue
|
||||
stream.dispatch(data)
|
||||
if event_type in ("message-finish", "error"):
|
||||
@@ -407,6 +445,8 @@ def _drain_messages_inbox(
|
||||
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"):
|
||||
@@ -418,9 +458,9 @@ def _drain_messages_inbox(
|
||||
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:
|
||||
# No active stream matches this event's key. Drop rather
|
||||
# than silently misroute to the only remaining stream.
|
||||
continue
|
||||
stream.dispatch(data)
|
||||
if event_type in ("message-finish", "error"):
|
||||
@@ -639,6 +679,457 @@ class _SyncToolCallsProjection:
|
||||
self._thread._unregister_subscription(sub.id)
|
||||
|
||||
|
||||
class SyncScopedStreamHandle:
|
||||
"""Scoped streaming handle for one discovered child invocation."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
thread: SyncThreadStream,
|
||||
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
|
||||
self._finish_lock = threading.Lock()
|
||||
self._messages_inbox: queue.Queue[Event | None] = queue.Queue(
|
||||
maxsize=max_queue_size
|
||||
)
|
||||
self._tools_inbox: queue.Queue[Event | None] = queue.Queue(
|
||||
maxsize=max_queue_size
|
||||
)
|
||||
self._tasks_inbox: queue.Queue[Event | None] = queue.Queue(
|
||||
maxsize=max_queue_size
|
||||
)
|
||||
# Descendant handles registered by _SyncHandleSubgraphsProjection 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, ...], SyncScopedStreamHandle] = {}
|
||||
# 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 = _SyncHandleMessagesProjection(self)
|
||||
self.tool_calls = _SyncHandleToolCallsProjection(self)
|
||||
self.subgraphs = _SyncHandleSubgraphsProjection(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 list(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: SyncScopedStreamHandle) -> 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: queue.Queue[Event | None] = getattr(self, inbox_attr)
|
||||
staging: list[Event | None] = []
|
||||
while True:
|
||||
try:
|
||||
staging.append(inbox.get_nowait())
|
||||
except queue.Empty:
|
||||
break
|
||||
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 `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:
|
||||
with self._finish_lock:
|
||||
if self.status != "started":
|
||||
return
|
||||
self.status = status
|
||||
self.error = error
|
||||
self._close_inboxes()
|
||||
|
||||
|
||||
class _SyncHandleMessagesProjection:
|
||||
"""Messages projection that drains a `SyncScopedStreamHandle`'s messages inbox."""
|
||||
|
||||
def __init__(self, handle: SyncScopedStreamHandle) -> None:
|
||||
self._handle = handle
|
||||
|
||||
def __iter__(self) -> Iterator[ChatModelStream]:
|
||||
return self._messages_iter()
|
||||
|
||||
def _messages_iter(self) -> Iterator[ChatModelStream]:
|
||||
self._handle._mark_iterated("messages")
|
||||
active: dict[str, ChatModelStream] = {}
|
||||
namespace = self._handle.namespace
|
||||
inbox = self._handle._messages_inbox
|
||||
try:
|
||||
while True:
|
||||
item = inbox.get()
|
||||
if item is None:
|
||||
return
|
||||
params_field = item.get("params") or {}
|
||||
ns = _event_namespace(params_field)
|
||||
if ns != 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(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]
|
||||
finally:
|
||||
pass
|
||||
|
||||
|
||||
class _SyncHandleToolCallsProjection:
|
||||
"""Tool calls projection that drains a `SyncScopedStreamHandle`'s tools inbox."""
|
||||
|
||||
def __init__(self, handle: SyncScopedStreamHandle) -> None:
|
||||
self._handle = handle
|
||||
|
||||
def __iter__(self) -> Iterator[SyncToolCallHandle]:
|
||||
return self._tool_calls_iter()
|
||||
|
||||
def _tool_calls_iter(self) -> Iterator[SyncToolCallHandle]:
|
||||
self._handle._mark_iterated("tools")
|
||||
active: dict[str, SyncToolCallHandle] = {}
|
||||
namespace = self._handle.namespace
|
||||
inbox = self._handle._tools_inbox
|
||||
while True:
|
||||
item = inbox.get()
|
||||
if item is None:
|
||||
err = RuntimeError(
|
||||
"Tool call stream closed before terminal tool event."
|
||||
)
|
||||
for h in active.values():
|
||||
h._fail(err)
|
||||
return
|
||||
params_field = item.get("params") or {}
|
||||
ns = _event_namespace(params_field)
|
||||
if ns != 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(namespace),
|
||||
)
|
||||
active[tool_call_id] = handle
|
||||
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:
|
||||
h._finish(data.get("output"))
|
||||
elif event_type == "tool-error":
|
||||
h = active.pop(tool_call_id, None)
|
||||
if h is not None:
|
||||
message = data.get("message")
|
||||
h._fail(
|
||||
RuntimeError(str(message) if message else "Tool call errored")
|
||||
)
|
||||
|
||||
|
||||
class _SyncHandleSubgraphsProjection:
|
||||
"""Subgraphs projection that drains a `SyncScopedStreamHandle`'s tasks inbox."""
|
||||
|
||||
def __init__(self, handle: SyncScopedStreamHandle) -> None:
|
||||
self._handle = handle
|
||||
|
||||
def __iter__(self) -> Iterator[SyncScopedStreamHandle]:
|
||||
return self._subgraphs_iter()
|
||||
|
||||
def _subgraphs_iter(self) -> Iterator[SyncScopedStreamHandle]:
|
||||
self._handle._mark_iterated("tasks")
|
||||
seen: set[tuple[str, ...]] = set()
|
||||
active: dict[tuple[str, ...], SyncScopedStreamHandle] = {}
|
||||
scope = self._handle.path
|
||||
while True:
|
||||
item = self._handle._tasks_inbox.get()
|
||||
if item is None:
|
||||
# Determine terminal status from the parent run's lifecycle result.
|
||||
# If _run_done resolved as errored, force-complete remaining children
|
||||
# as failed so callers see the correct terminal state.
|
||||
terminal_status: SubgraphStatus = "completed"
|
||||
run_done = self._handle._thread._run_done
|
||||
if run_done is not None and run_done.done():
|
||||
try:
|
||||
result = run_done.result(timeout=0)
|
||||
if (
|
||||
isinstance(result, _RunTerminal)
|
||||
and result.status == "errored"
|
||||
):
|
||||
terminal_status = "failed"
|
||||
except Exception:
|
||||
pass
|
||||
for child in active.values():
|
||||
if child.status == "started":
|
||||
child._finish(terminal_status)
|
||||
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 = SyncScopedStreamHandle(
|
||||
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 _SyncSubgraphsProjection:
|
||||
"""Discover direct child invocations for a namespace scope."""
|
||||
|
||||
def __init__(self, thread: SyncThreadStream, scope: tuple[str, ...] = ()) -> None:
|
||||
self._thread = thread
|
||||
self._scope = scope
|
||||
|
||||
def __iter__(self) -> Iterator[SyncScopedStreamHandle]:
|
||||
return self._subgraphs_iter()
|
||||
|
||||
def _subgraphs_iter(self) -> Iterator[SyncScopedStreamHandle]:
|
||||
if self._thread._transport is None:
|
||||
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] = {}
|
||||
root_inbox: queue.Queue[Event | None] | None = (
|
||||
self._thread._activate_root_messages_inbox() if not self._scope else None
|
||||
)
|
||||
try:
|
||||
self._thread._reconcile_stream(params)
|
||||
self._thread._ensure_fanout_running()
|
||||
while True:
|
||||
item = 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")
|
||||
|
||||
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.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
|
||||
finally:
|
||||
# Determine terminal status from the run's lifecycle result.
|
||||
# If _run_done resolved as errored, force-complete remaining children
|
||||
# as failed 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():
|
||||
try:
|
||||
result = run_done.result(timeout=0)
|
||||
if isinstance(result, _RunTerminal) and result.status == "errored":
|
||||
terminal_status = "failed"
|
||||
except Exception:
|
||||
pass
|
||||
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, ...], 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 SyncThreadStream:
|
||||
"""Synchronous context manager for one thread's v3 streaming session.
|
||||
|
||||
@@ -680,6 +1171,8 @@ class SyncThreadStream:
|
||||
self.values = _SyncValuesProjection(self)
|
||||
self.messages = _SyncMessagesProjection(self, namespace=[])
|
||||
self.tool_calls = _SyncToolCallsProjection(self, namespace=[])
|
||||
self.subgraphs = _SyncSubgraphsProjection(self, scope=())
|
||||
self.subagents = self.subgraphs
|
||||
|
||||
def __enter__(self) -> SyncThreadStream:
|
||||
if self._closed:
|
||||
|
||||
@@ -20,6 +20,8 @@ from streaming._events import (
|
||||
message_start_event,
|
||||
message_text_delta_event,
|
||||
message_text_finish_event,
|
||||
tasks_result_event,
|
||||
tasks_start_event,
|
||||
tool_error_event,
|
||||
tool_finished_event,
|
||||
tool_output_delta_event,
|
||||
@@ -171,6 +173,60 @@ def test_sync_messages_error_event_fails_active_stream():
|
||||
_ = streams[0].output
|
||||
|
||||
|
||||
def test_sync_subgraph_scoped_messages_and_tool_calls():
|
||||
fake = SyncFakeServer()
|
||||
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"]),
|
||||
tool_started_event(
|
||||
seq=6,
|
||||
namespace=["worker:abc"],
|
||||
tool_call_id="call-child",
|
||||
tool_name="search",
|
||||
),
|
||||
tool_output_delta_event(
|
||||
seq=7,
|
||||
namespace=["worker:abc"],
|
||||
tool_call_id="call-child",
|
||||
delta="delta",
|
||||
),
|
||||
tool_finished_event(
|
||||
seq=8,
|
||||
namespace=["worker:abc"],
|
||||
tool_call_id="call-child",
|
||||
output={"child": True},
|
||||
),
|
||||
tasks_result_event(seq=9, namespace=[], task_id="abc", name="worker"),
|
||||
lifecycle_completed_event(seq=10),
|
||||
]
|
||||
)
|
||||
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] = list(thread.subgraphs)
|
||||
child_messages = list(handle.messages)
|
||||
child_calls = list(handle.tool_calls)
|
||||
|
||||
assert handle.status == "completed"
|
||||
assert handle.path == ("worker:abc",)
|
||||
assert [message.message_id for message in child_messages] == ["msg-child"]
|
||||
assert [str(message.text) for message in child_messages] == ["child"]
|
||||
assert [call.tool_call_id for call in child_calls] == ["call-child"]
|
||||
assert list(child_calls[0].deltas) == ["delta"]
|
||||
assert child_calls[0].output == {"child": True}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Task 10.7 — comprehensive sync tool_calls projection tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -0,0 +1,760 @@
|
||||
"""Sync mirror of test_scoped_handles.py for SyncScopedStreamHandle."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
from concurrent.futures import ThreadPoolExecutor, wait
|
||||
|
||||
import httpx
|
||||
|
||||
from langgraph_sdk._sync.http import SyncHttpClient
|
||||
from langgraph_sdk._sync.stream import SyncScopedStreamHandle
|
||||
from langgraph_sdk._sync.threads import SyncThreadsClient
|
||||
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._sync_fake_server import SyncFakeServer
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Task 11.1: _finish idempotency — double-finish must not double-close inboxes
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_sync_scoped_handle_finish_is_idempotent():
|
||||
"""Calling _finish twice must not enqueue a second sentinel on each inbox."""
|
||||
handle = SyncScopedStreamHandle(
|
||||
thread=None, # ty: ignore[invalid-argument-type]
|
||||
path=("worker:abc",),
|
||||
graph_name="worker",
|
||||
trigger_call_id="abc",
|
||||
)
|
||||
# Mark all three inboxes so _close_inboxes sends sentinels to each.
|
||||
handle._mark_iterated("messages")
|
||||
handle._mark_iterated("tools")
|
||||
handle._mark_iterated("tasks")
|
||||
handle._finish("completed")
|
||||
handle._finish("completed") # second call must be a no-op
|
||||
|
||||
# Each inbox should have exactly one None sentinel from the first _finish.
|
||||
assert handle._messages_inbox.qsize() == 1
|
||||
assert handle._tools_inbox.qsize() == 1
|
||||
assert handle._tasks_inbox.qsize() == 1
|
||||
assert handle._messages_inbox.get_nowait() is None
|
||||
assert handle._tools_inbox.get_nowait() is None
|
||||
assert handle._tasks_inbox.get_nowait() is None
|
||||
|
||||
|
||||
def test_sync_scoped_handle_finish_concurrent_only_one_wins():
|
||||
"""Concurrent _finish calls from two threads: only the first must close inboxes."""
|
||||
handle = SyncScopedStreamHandle(
|
||||
thread=None, # ty: ignore[invalid-argument-type]
|
||||
path=("worker:abc",),
|
||||
graph_name="worker",
|
||||
trigger_call_id="abc",
|
||||
)
|
||||
# Mark all three inboxes so _close_inboxes sends sentinels to each.
|
||||
handle._mark_iterated("messages")
|
||||
handle._mark_iterated("tools")
|
||||
handle._mark_iterated("tasks")
|
||||
barrier = threading.Barrier(2)
|
||||
errors: list[Exception] = []
|
||||
|
||||
def _call_finish(status: str) -> None:
|
||||
try:
|
||||
barrier.wait()
|
||||
handle._finish(status) # ty: ignore[invalid-argument-type]
|
||||
except Exception as exc:
|
||||
errors.append(exc)
|
||||
|
||||
t1 = threading.Thread(target=_call_finish, args=("completed",))
|
||||
t2 = threading.Thread(target=_call_finish, args=("failed",))
|
||||
t1.start()
|
||||
t2.start()
|
||||
t1.join()
|
||||
t2.join()
|
||||
|
||||
assert not errors
|
||||
# Each inbox must have exactly one None sentinel regardless of which thread won.
|
||||
assert handle._messages_inbox.qsize() == 1
|
||||
assert handle._tools_inbox.qsize() == 1
|
||||
assert handle._tasks_inbox.qsize() == 1
|
||||
assert handle.status in ("completed", "failed")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Task 11.2: subgraphs yields handle with correct metadata and completion status
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_sync_subgraphs_yields_handle_and_completes_status():
|
||||
fake = SyncFakeServer()
|
||||
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),
|
||||
]
|
||||
)
|
||||
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={})
|
||||
handles = list(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
|
||||
|
||||
|
||||
def test_sync_subgraphs_failed_and_interrupted_statuses():
|
||||
failed_fake = SyncFakeServer()
|
||||
failed_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", error="boom"
|
||||
),
|
||||
lifecycle_completed_event(seq=3),
|
||||
]
|
||||
)
|
||||
with httpx.Client(transport=failed_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={})
|
||||
[failed_handle] = list(thread.subgraphs)
|
||||
|
||||
interrupted_fake = SyncFakeServer()
|
||||
interrupted_fake.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),
|
||||
]
|
||||
)
|
||||
with httpx.Client(
|
||||
transport=interrupted_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={})
|
||||
[interrupted_handle] = list(thread.subgraphs)
|
||||
|
||||
assert failed_handle.status == "failed"
|
||||
assert failed_handle.error == "boom"
|
||||
assert interrupted_handle.status == "interrupted"
|
||||
assert interrupted_handle.error is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Task 11.3: grandchild routing via _route_sibling_inboxes_to_grandchildren
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_sync_subgraph_handles_are_recursive_for_grandchildren():
|
||||
fake = SyncFakeServer()
|
||||
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),
|
||||
]
|
||||
)
|
||||
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] = list(thread.subgraphs)
|
||||
[grandchild] = list(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"
|
||||
|
||||
|
||||
def test_sync_subgraph_messages_are_scoped_to_child_namespace():
|
||||
fake = SyncFakeServer()
|
||||
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),
|
||||
]
|
||||
)
|
||||
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] = list(thread.subgraphs)
|
||||
child_messages = list(handle.messages)
|
||||
root_messages = list(thread.messages)
|
||||
|
||||
assert [m.message_id for m in child_messages] == ["msg-child"]
|
||||
assert [str(m.text) for m in child_messages] == ["child"]
|
||||
assert root_messages == []
|
||||
|
||||
|
||||
def test_sync_subgraph_tool_calls_are_scoped_to_child_namespace():
|
||||
fake = SyncFakeServer()
|
||||
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),
|
||||
]
|
||||
)
|
||||
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] = list(thread.subgraphs)
|
||||
child_calls = list(handle.tool_calls)
|
||||
root_calls = list(thread.tool_calls)
|
||||
|
||||
assert [call.tool_call_id for call in child_calls] == ["call-child"]
|
||||
assert list(child_calls[0].deltas) == ["child-delta"]
|
||||
assert child_calls[0].output == {"ok": True}
|
||||
assert root_calls == []
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Task 11.4: root/child cross-talk regression
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_sync_root_and_child_projections_do_not_cross_talk():
|
||||
fake = SyncFakeServer()
|
||||
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),
|
||||
]
|
||||
)
|
||||
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] = list(thread.subgraphs)
|
||||
root_messages = list(thread.messages)
|
||||
child_messages = list(handle.messages)
|
||||
|
||||
assert [m.message_id for m in root_messages] == ["root-msg"]
|
||||
assert [str(m.text) for m in root_messages] == ["root"]
|
||||
assert [m.message_id for m in child_messages] == ["child-msg"]
|
||||
assert [str(m.text) for m in child_messages] == ["child"]
|
||||
|
||||
|
||||
def test_sync_subagents_aliases_subgraphs():
|
||||
fake = SyncFakeServer()
|
||||
fake.script([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:
|
||||
assert thread.subagents is thread.subgraphs
|
||||
thread.run.start(input={})
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fix A: sibling routing dispatches at push time, preserves order, fans out
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_sync_scoped_handle_has_descendant_handles_dict():
|
||||
"""SyncScopedStreamHandle must expose _descendant_handles so push-time
|
||||
fan-out can be registered without drain-and-replay."""
|
||||
handle = SyncScopedStreamHandle(
|
||||
thread=None, # ty: ignore[invalid-argument-type]
|
||||
path=("worker:abc",),
|
||||
graph_name="worker",
|
||||
trigger_call_id="abc",
|
||||
)
|
||||
# Must exist and be empty at construction.
|
||||
assert hasattr(handle, "_descendant_handles")
|
||||
assert isinstance(handle._descendant_handles, dict)
|
||||
assert len(handle._descendant_handles) == 0
|
||||
|
||||
|
||||
def test_sync_register_descendant_forwards_buffered_events_in_order():
|
||||
"""_register_descendant must drain already-buffered events whose namespace
|
||||
matches the new grandchild, push them into the grandchild, and preserve
|
||||
the original arrival order in the parent inbox."""
|
||||
from langchain_protocol import Event
|
||||
|
||||
parent = SyncScopedStreamHandle(
|
||||
thread=None, # ty: ignore[invalid-argument-type]
|
||||
path=("worker:abc",),
|
||||
graph_name="worker",
|
||||
trigger_call_id="abc",
|
||||
)
|
||||
child_path = ("worker:abc", "tool:gc1")
|
||||
grandchild = SyncScopedStreamHandle(
|
||||
thread=None, # ty: ignore[invalid-argument-type]
|
||||
path=child_path,
|
||||
graph_name="tool",
|
||||
trigger_call_id="gc1",
|
||||
)
|
||||
# Put two grandchild-scoped events and one parent-scoped event into parent inbox.
|
||||
evt_gc1: Event = message_start_event(
|
||||
seq=1, namespace=list(child_path), message_id="m1", run_id="r1"
|
||||
) # ty: ignore[invalid-assignment]
|
||||
evt_gc2: Event = message_start_event(
|
||||
seq=2, namespace=list(child_path), message_id="m2", run_id="r2"
|
||||
) # ty: ignore[invalid-assignment]
|
||||
evt_parent: Event = message_start_event(
|
||||
seq=3, namespace=list(parent.path), message_id="m3", run_id="r3"
|
||||
) # ty: ignore[invalid-assignment]
|
||||
parent._messages_inbox.put_nowait(evt_gc1)
|
||||
parent._messages_inbox.put_nowait(evt_parent)
|
||||
parent._messages_inbox.put_nowait(evt_gc2)
|
||||
|
||||
parent._register_descendant(grandchild)
|
||||
|
||||
# Grandchild inbox must have exactly the two grandchild events.
|
||||
assert grandchild._messages_inbox.qsize() == 2
|
||||
first = grandchild._messages_inbox.get_nowait()
|
||||
second = grandchild._messages_inbox.get_nowait()
|
||||
assert isinstance(first, dict) and isinstance(second, dict)
|
||||
assert (first.get("params") or {}).get("data", {}).get("id") == "m1"
|
||||
assert (second.get("params") or {}).get("data", {}).get("id") == "m2"
|
||||
|
||||
# Parent inbox must still contain all 3 events in original order.
|
||||
assert parent._messages_inbox.qsize() == 3
|
||||
items = [parent._messages_inbox.get_nowait() for _ in range(3)]
|
||||
ids = [
|
||||
(i.get("params") or {}).get("data", {}).get("id")
|
||||
if isinstance(i, dict)
|
||||
else None
|
||||
for i in items
|
||||
]
|
||||
assert ids == ["m1", "m3", "m2"]
|
||||
|
||||
|
||||
def test_sync_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."""
|
||||
fake = SyncFakeServer()
|
||||
fake.script(
|
||||
[
|
||||
lifecycle_started_event(seq=0),
|
||||
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),
|
||||
]
|
||||
)
|
||||
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] = list(thread.subgraphs)
|
||||
[grandchild] = list(child.subgraphs)
|
||||
gc_messages = list(grandchild.messages)
|
||||
|
||||
assert [m.message_id for m in gc_messages] == ["msg-E1", "msg-E2", "msg-E3"]
|
||||
texts = [str(m.text) for m in gc_messages]
|
||||
assert texts == ["E1", "E2", "E3"]
|
||||
|
||||
|
||||
def test_sync_grandchild_events_dispatched_to_correct_sibling():
|
||||
"""When two sibling grandchildren exist, messages scoped to one grandchild
|
||||
must not bleed into the other grandchild's inbox."""
|
||||
fake = SyncFakeServer()
|
||||
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_start_event(
|
||||
seq=3, namespace=["worker:abc", "tool:gc2"], task_id="t-gc2"
|
||||
),
|
||||
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"]),
|
||||
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),
|
||||
]
|
||||
)
|
||||
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] = list(thread.subgraphs)
|
||||
grandchildren = list(child.subgraphs)
|
||||
by_path = {h.path: h for h in grandchildren}
|
||||
gc1_messages = list(by_path[("worker:abc", "tool:gc1")].messages)
|
||||
gc2_messages = list(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"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Task 11.5: _finish thread-safe status transition via threading.Lock
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _finish_one(
|
||||
idx: int,
|
||||
handle: SyncScopedStreamHandle,
|
||||
barrier: threading.Barrier,
|
||||
errors: list[Exception],
|
||||
statuses: list[str],
|
||||
) -> None:
|
||||
try:
|
||||
barrier.wait()
|
||||
status = statuses[idx % len(statuses)]
|
||||
handle._finish(status) # ty: ignore[invalid-argument-type]
|
||||
except Exception as exc:
|
||||
errors.append(exc)
|
||||
|
||||
|
||||
def test_sync_scoped_handle_finish_thread_safe_with_20_concurrent_calls():
|
||||
"""20 concurrent _finish calls must yield exactly one terminal status and one
|
||||
sentinel per inbox (deterministic even under high contention).
|
||||
"""
|
||||
n_workers = 20
|
||||
statuses = ["completed", "failed", "interrupted"]
|
||||
|
||||
for _ in range(50): # repeat to increase chance of catching races
|
||||
handle = SyncScopedStreamHandle(
|
||||
thread=None, # ty: ignore[invalid-argument-type]
|
||||
path=("worker:abc",),
|
||||
graph_name="worker",
|
||||
trigger_call_id="abc",
|
||||
)
|
||||
# Mark all inboxes so _close_inboxes sends sentinels to each.
|
||||
handle._mark_iterated("messages")
|
||||
handle._mark_iterated("tools")
|
||||
handle._mark_iterated("tasks")
|
||||
barrier = threading.Barrier(n_workers)
|
||||
errors: list[Exception] = []
|
||||
|
||||
with ThreadPoolExecutor(max_workers=n_workers) as pool:
|
||||
futures = [
|
||||
pool.submit(_finish_one, i, handle, barrier, errors, statuses)
|
||||
for i in range(n_workers)
|
||||
]
|
||||
wait(futures)
|
||||
|
||||
assert not errors, f"Unexpected exception(s): {errors}"
|
||||
# Status must be one of the valid terminal values (not "started").
|
||||
assert handle.status in ("completed", "failed", "interrupted"), (
|
||||
f"Unexpected status: {handle.status!r}"
|
||||
)
|
||||
# Each inbox must have exactly one None sentinel — the lock must ensure
|
||||
# that _close_inboxes() is called exactly once.
|
||||
assert handle._messages_inbox.qsize() == 1, (
|
||||
f"messages_inbox has {handle._messages_inbox.qsize()} items (expected 1)"
|
||||
)
|
||||
assert handle._tools_inbox.qsize() == 1, (
|
||||
f"tools_inbox has {handle._tools_inbox.qsize()} items (expected 1)"
|
||||
)
|
||||
assert handle._tasks_inbox.qsize() == 1, (
|
||||
f"tasks_inbox has {handle._tasks_inbox.qsize()} items (expected 1)"
|
||||
)
|
||||
assert handle._messages_inbox.get_nowait() is None
|
||||
assert handle._tools_inbox.get_nowait() is None
|
||||
assert handle._tasks_inbox.get_nowait() is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fix B: bound SyncScopedStreamHandle inboxes via max_queue_size
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_sync_scoped_handle_inboxes_bounded_by_max_queue_size():
|
||||
"""SyncScopedStreamHandle with max_queue_size=N creates queues with maxsize=N."""
|
||||
handle = SyncScopedStreamHandle(
|
||||
thread=None, # ty: ignore[invalid-argument-type]
|
||||
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
|
||||
|
||||
|
||||
def test_sync_child_handle_inherits_max_queue_size_from_parent():
|
||||
"""Grandchild SyncScopedStreamHandles created by _SyncHandleSubgraphsProjection
|
||||
inherit the parent's max_queue_size so all queues are consistently bounded."""
|
||||
fake = SyncFakeServer()
|
||||
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),
|
||||
]
|
||||
)
|
||||
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] = list(thread.subgraphs)
|
||||
[grandchild] = list(child.subgraphs)
|
||||
|
||||
# Grandchild queues must have the same maxsize as the parent queues.
|
||||
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
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fix C: force-complete uses parent terminal status
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_sync_force_complete_uses_failed_when_run_errored():
|
||||
"""If the lifecycle signals an errored run, scoped children that are still
|
||||
'started' when the subgraphs iterator's finally block runs must be
|
||||
force-finished as 'failed', not 'completed'."""
|
||||
from streaming._events import lifecycle_errored_event
|
||||
|
||||
fake = SyncFakeServer()
|
||||
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"),
|
||||
]
|
||||
)
|
||||
handles: list[SyncScopedStreamHandle] = []
|
||||
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 handle in thread.subgraphs:
|
||||
handles.append(handle)
|
||||
|
||||
assert len(handles) == 1
|
||||
child = handles[0]
|
||||
assert child.status == "failed"
|
||||
|
||||
|
||||
def test_sync_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 = SyncFakeServer()
|
||||
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),
|
||||
]
|
||||
)
|
||||
handles: list[SyncScopedStreamHandle] = []
|
||||
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 handle in thread.subgraphs:
|
||||
handles.append(handle)
|
||||
|
||||
assert len(handles) == 1
|
||||
child = handles[0]
|
||||
assert child.status == "completed"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fix D: only enqueue close sentinel on inboxes that had a consumer
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_sync_close_inboxes_does_not_enqueue_on_uniterated_inboxes():
|
||||
"""_close_inboxes must not push a sentinel on inboxes that had no consumer."""
|
||||
handle = SyncScopedStreamHandle(
|
||||
thread=None, # ty: ignore[invalid-argument-type]
|
||||
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_sync_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."""
|
||||
handle = SyncScopedStreamHandle(
|
||||
thread=None, # ty: ignore[invalid-argument-type]
|
||||
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