From 30fea6468733ec89fa6631a30688a6d5646de4e9 Mon Sep 17 00:00:00 2001 From: Nick Hollon Date: Wed, 27 May 2026 12:11:06 -0400 Subject: [PATCH] feat(sdk-py): add messages and tool call projections (#7823) --- libs/langgraph/uv.lock | 2 + libs/prebuilt/uv.lock | 2 + libs/sdk-py/langgraph_sdk/_async/stream.py | 302 +++++++++++++++++- libs/sdk-py/langgraph_sdk/stream/__init__.py | 15 + libs/sdk-py/pyproject.toml | 7 +- libs/sdk-py/tests/streaming/_events.py | 215 +++++++++++++ .../streaming/test_messages_projection.py | 209 ++++++++++++ .../tests/streaming/test_subscription.py | 26 ++ .../streaming/test_tool_calls_projection.py | 272 ++++++++++++++++ libs/sdk-py/uv.lock | 2 + 10 files changed, 1042 insertions(+), 10 deletions(-) create mode 100644 libs/sdk-py/tests/streaming/test_messages_projection.py create mode 100644 libs/sdk-py/tests/streaming/test_tool_calls_projection.py diff --git a/libs/langgraph/uv.lock b/libs/langgraph/uv.lock index d3ead6815..bd9769ca8 100644 --- a/libs/langgraph/uv.lock +++ b/libs/langgraph/uv.lock @@ -1828,6 +1828,7 @@ name = "langgraph-sdk" source = { editable = "../sdk-py" } dependencies = [ { name = "httpx" }, + { name = "langchain-core" }, { name = "langchain-protocol" }, { name = "orjson" }, ] @@ -1835,6 +1836,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "httpx", specifier = ">=0.25.2" }, + { name = "langchain-core", specifier = ">=1.4.0,<2" }, { name = "langchain-protocol", specifier = ">=0.0.15" }, { name = "orjson", specifier = ">=3.11.5" }, ] diff --git a/libs/prebuilt/uv.lock b/libs/prebuilt/uv.lock index 2b76c0d2b..8ba5930b3 100644 --- a/libs/prebuilt/uv.lock +++ b/libs/prebuilt/uv.lock @@ -597,6 +597,7 @@ name = "langgraph-sdk" source = { editable = "../sdk-py" } dependencies = [ { name = "httpx" }, + { name = "langchain-core" }, { name = "langchain-protocol" }, { name = "orjson" }, ] @@ -604,6 +605,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "httpx", specifier = ">=0.25.2" }, + { name = "langchain-core", specifier = ">=1.4.0,<2" }, { name = "langchain-protocol", specifier = ">=0.0.15" }, { name = "orjson", specifier = ">=3.11.5" }, ] diff --git a/libs/sdk-py/langgraph_sdk/_async/stream.py b/libs/sdk-py/langgraph_sdk/_async/stream.py index 543659ceb..d19983157 100644 --- a/libs/sdk-py/langgraph_sdk/_async/stream.py +++ b/libs/sdk-py/langgraph_sdk/_async/stream.py @@ -18,6 +18,7 @@ from collections.abc import AsyncGenerator, AsyncIterator, Generator, Mapping from dataclasses import dataclass, field from typing import Any, Literal, TypedDict +from langchain_core.language_models.chat_model_stream import AsyncChatModelStream from langchain_protocol import Event, SubscribeParams from langgraph_sdk._async.http import HttpClient @@ -285,6 +286,265 @@ class _ValuesProjection: self._thread._unregister_subscription(sub.id) +class _MessagesProjection: + """Typed projection for root-scope `thread.messages`. + + Iterating yields one `AsyncChatModelStream` per message-start event. + Each iterator owns its own `messages` subscription and routes events + from the root namespace only. + """ + + def __init__(self, thread: AsyncThreadStream) -> None: + self._thread = thread + + def __aiter__(self) -> AsyncIterator[AsyncChatModelStream]: + return self._messages_iter() + + 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, + } + sub = self._thread._register_subscription(params) + active: dict[str, AsyncChatModelStream] = {} + 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 {} + if not isinstance(params_field, dict): + continue + if params_field.get("namespace") not in (None, []): + continue + data = params_field.get("data") + 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=[], + 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: + # 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"): + 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) + self._thread._unregister_subscription(sub.id) + + +def _message_event_id(data: dict[str, Any]) -> str | None: + message_id = data.get("id") or data.get("message_id") + return str(message_id) if message_id is not None else None + + +def _message_route_key(data: dict[str, Any], fallback: str | None = None) -> str: + """Return the routing key for a message-channel event in `active`. + + Keys on `message_id` when available so concurrent messages that share the + same `run_id` (two AI turns in one agent step) route to independent streams + rather than colliding on a shared `run:` slot. + """ + message_id = _message_event_id(data) + if message_id is not None: + return f"message:{message_id}" + if fallback is not None: + return f"message:{fallback}" + return "__single__" + + +class ToolCallHandle: + """Async handle for one root-scope tool call.""" + + def __init__( + self, + *, + tool_call_id: str, + name: str, + input: Any = None, + namespace: list[str] | None = None, + max_queue_size: int = 1024, + ) -> None: + self.tool_call_id = tool_call_id + self.name = name + self.input = input + self.namespace = list(namespace or []) + self.done = False + self.error: BaseException | None = None + loop = asyncio.get_running_loop() + self.output: asyncio.Future[Any] = loop.create_future() + self._deltas: asyncio.Queue[str | None] = asyncio.Queue(maxsize=max_queue_size) + self._deltas_consumed = False + + @property + def deltas(self) -> AsyncIterator[str]: + """Stream tool output deltas emitted before the terminal event. + + Raises: + RuntimeError: if called more than once — the underlying queue is + single-consumer and cannot be fanned out safely. + """ + if self._deltas_consumed: + raise RuntimeError( + "ToolCallHandle.deltas can only be iterated by a single consumer." + ) + self._deltas_consumed = True + return self._delta_iter() + + async def _delta_iter(self) -> AsyncGenerator[str, None]: + while True: + item = await self._deltas.get() + if item is None: + return # errors surface via output, not deltas + yield item + + def _push_delta(self, delta: str) -> None: + if self.done: + return + self._deltas.put_nowait(delta) + + def _finish(self, output: Any) -> None: + if self.done: + return + self.done = True + if not self.output.done(): + self.output.set_result(output) + self._deltas.put_nowait(None) + + def _fail(self, err: BaseException) -> None: + if self.done: + return + self.done = True + self.error = err + if not self.output.done(): + self.output.set_exception(err) + self._deltas.put_nowait(None) + + +class _ToolCallsProjection: + """Typed projection for root-scope `thread.tool_calls`.""" + + def __init__(self, thread: AsyncThreadStream) -> None: + self._thread = thread + + def __aiter__(self) -> AsyncIterator[ToolCallHandle]: + return self._tool_calls_iter() + + 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, + } + sub = self._thread._register_subscription(params) + active: dict[str, ToolCallHandle] = {} + 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 {} + if not isinstance(params_field, dict): + continue + namespace = params_field.get("namespace") or [] + if namespace != []: + continue + data = params_field.get("data") + 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=namespace, + ) + active[tool_call_id] = handle + self._thread._register_active_tool_call(handle) + yield handle + elif event_type == "tool-output-delta": + handle = active.get(tool_call_id) + delta = data.get("delta") + if handle is not None and isinstance(delta, str): + handle._push_delta(delta) + elif event_type == "tool-finished": + handle = active.pop(tool_call_id, None) + if handle is not None: + self._thread._unregister_active_tool_call(handle) + handle._finish(data.get("output")) + elif event_type == "tool-error": + handle = active.pop(tool_call_id, None) + if handle is not None: + self._thread._unregister_active_tool_call(handle) + message = data.get("message") + handle._fail( + RuntimeError( + str(message) if message else "Tool call errored" + ) + ) + finally: + # Read terminal error from _run_done if it is already resolved. + # We do NOT block here: callers who need a terminal observation + # should await `thread.output` directly. Blocking in iterator + # teardown would stall every early break or exception exit for + # up to the full shield-wait timeout (previously 1 s). + run_done = self._thread._run_done + terminal_err: BaseException | None = None + if run_done is not None and run_done.done() and not run_done.cancelled(): + terminal = run_done.result() + terminal_err = terminal.error + err = ( + terminal_err + if terminal_err is not None + else RuntimeError("Tool call stream closed before terminal tool event.") + ) + for handle in active.values(): + self._thread._unregister_active_tool_call(handle) + handle._fail(err) + self._thread._unregister_subscription(sub.id) + + class AsyncThreadStream: """Async context manager for one thread's v3 streaming session. @@ -332,9 +592,13 @@ class AsyncThreadStream: self._run_start_ready: asyncio.Future[None] | None = None self._run_seen: bool = False self._run_done: asyncio.Future[_RunTerminal] | None = None + self._active_message_streams: set[AsyncChatModelStream] = set() + self._active_tool_calls: set[ToolCallHandle] = set() self.run = RunModule(self) self.output = _OutputAwaitable(self) self.values = _ValuesProjection(self) + self.messages = _MessagesProjection(self) + self.tool_calls = _ToolCallsProjection(self) @property def _controller(self) -> AsyncThreadStream: @@ -402,6 +666,8 @@ class AsyncThreadStream: await self._lifecycle_watcher_task if self._lifecycle_watcher_handle is not None: await self._lifecycle_watcher_handle.close() + self._fail_active_message_streams(asyncio.CancelledError()) + self._fail_active_tool_calls(asyncio.CancelledError()) if self._fanout_task is not None: self._fanout_task.cancel() with contextlib.suppress(Exception, asyncio.CancelledError): @@ -426,6 +692,28 @@ class AsyncThreadStream: """Remove a subscription from the registry. No-op if already absent.""" self._subscriptions.pop(subscription_id, None) + def _register_active_message_stream(self, stream: AsyncChatModelStream) -> None: + self._active_message_streams.add(stream) + + def _unregister_active_message_stream(self, stream: AsyncChatModelStream) -> None: + self._active_message_streams.discard(stream) + + def _fail_active_message_streams(self, err: BaseException) -> None: + for stream in list(self._active_message_streams): + stream.fail(err) + self._active_message_streams.clear() + + def _register_active_tool_call(self, handle: ToolCallHandle) -> None: + self._active_tool_calls.add(handle) + + def _unregister_active_tool_call(self, handle: ToolCallHandle) -> None: + self._active_tool_calls.discard(handle) + + def _fail_active_tool_calls(self, err: BaseException) -> None: + for handle in list(self._active_tool_calls): + handle._fail(err) + self._active_tool_calls.clear() + def subscribe( self, channels: list[str], @@ -750,15 +1038,11 @@ class AsyncThreadStream: error_msg = ( data.get("error") if isinstance(data, dict) else None ) - run_done.set_result( - _RunTerminal( - status="errored", - error=RuntimeError( - f"Run errored: {error_msg}" - if error_msg - else "Run errored" - ), - ) + error = RuntimeError( + f"Run errored: {error_msg}" if error_msg else "Run errored" ) + self._fail_active_message_streams(error) + self._fail_active_tool_calls(error) + run_done.set_result(_RunTerminal(status="errored", error=error)) else: run_done.set_result(_RunTerminal(status="completed")) diff --git a/libs/sdk-py/langgraph_sdk/stream/__init__.py b/libs/sdk-py/langgraph_sdk/stream/__init__.py index e69de29bb..a1fe5ed76 100644 --- a/libs/sdk-py/langgraph_sdk/stream/__init__.py +++ b/libs/sdk-py/langgraph_sdk/stream/__init__.py @@ -0,0 +1,15 @@ +"""Stream module for LangGraph SDK v3.""" + +from langchain_protocol import ( + Channel, + Event, + Namespace, + SubscribeParams, +) + +__all__ = [ + "Channel", + "Event", + "Namespace", + "SubscribeParams", +] diff --git a/libs/sdk-py/pyproject.toml b/libs/sdk-py/pyproject.toml index b47f0618a..c0698ac03 100644 --- a/libs/sdk-py/pyproject.toml +++ b/libs/sdk-py/pyproject.toml @@ -11,7 +11,12 @@ requires-python = ">=3.10" readme = "README.md" license = "MIT" license-files = ['LICENSE'] -dependencies = ["httpx>=0.25.2", "orjson>=3.11.5", "langchain-protocol>=0.0.15"] +dependencies = [ + "httpx>=0.25.2", + "orjson>=3.11.5", + "langchain-protocol>=0.0.15", + "langchain-core>=1.4.0,<2", +] [tool.hatch.version] path = "langgraph_sdk/__init__.py" diff --git a/libs/sdk-py/tests/streaming/_events.py b/libs/sdk-py/tests/streaming/_events.py index 86eb9fc82..1ef35f7c4 100644 --- a/libs/sdk-py/tests/streaming/_events.py +++ b/libs/sdk-py/tests/streaming/_events.py @@ -69,3 +69,218 @@ def input_requested_event( seq: int = 0, namespace: list[str] | None = None ) -> dict[str, Any]: return _base(seq, "input.requested", namespace or [], {"interrupt_id": "i-1"}) + + +def message_start_event( + seq: int = 0, + namespace: list[str] | None = None, + *, + message_id: str = "msg-1", + role: str = "ai", + run_id: str = "run-1", + node: str = "agent", +) -> dict[str, Any]: + return _base( + seq, + "messages", + namespace or [], + { + "event": "message-start", + "id": message_id, + "role": role, + "metadata": {"run_id": run_id, "langgraph_node": node}, + }, + ) + + +def message_text_delta_event( + seq: int = 0, + namespace: list[str] | None = None, + *, + text: str, + index: int = 0, + message_id: str | None = None, +) -> dict[str, Any]: + data: dict[str, Any] = { + "event": "content-block-delta", + "index": index, + "delta": {"type": "text-delta", "text": text}, + } + if message_id is not None: + data["id"] = message_id + return _base(seq, "messages", namespace or [], data) + + +def message_text_finish_event( + seq: int = 0, + namespace: list[str] | None = None, + *, + text: str, + index: int = 0, + message_id: str | None = None, +) -> dict[str, Any]: + data: dict[str, Any] = { + "event": "content-block-finish", + "index": index, + "content": {"type": "text", "text": text}, + } + if message_id is not None: + data["id"] = message_id + return _base(seq, "messages", namespace or [], data) + + +def message_finish_event( + seq: int = 0, + namespace: list[str] | None = None, + *, + input_tokens: int = 1, + output_tokens: int = 1, + message_id: str | None = None, +) -> dict[str, Any]: + data: dict[str, Any] = { + "event": "message-finish", + "usage": { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens, + }, + } + if message_id is not None: + data["id"] = message_id + return _base(seq, "messages", namespace or [], data) + + +def message_error_event( + seq: int = 0, + namespace: list[str] | None = None, + *, + message: str = "model failed", + code: str = "provider_error", + message_id: str | None = None, +) -> dict[str, Any]: + data: dict[str, Any] = {"event": "error", "message": message, "code": code} + if message_id is not None: + data["id"] = message_id + return _base(seq, "messages", namespace or [], data) + + +def tool_started_event( + seq: int = 0, + namespace: list[str] | None = None, + *, + tool_call_id: str = "call-1", + tool_name: str = "search", + input: Any = None, +) -> dict[str, Any]: + payload: dict[str, Any] = { + "event": "tool-started", + "tool_call_id": tool_call_id, + "tool_name": tool_name, + } + if input is not None: + payload["input"] = input + return _base(seq, "tools", namespace or [], payload) + + +def tool_output_delta_event( + seq: int = 0, + namespace: list[str] | None = None, + *, + tool_call_id: str = "call-1", + delta: str = "", +) -> dict[str, Any]: + return _base( + seq, + "tools", + namespace or [], + { + "event": "tool-output-delta", + "tool_call_id": tool_call_id, + "delta": delta, + }, + ) + + +def tool_finished_event( + seq: int = 0, + namespace: list[str] | None = None, + *, + tool_call_id: str = "call-1", + output: Any = None, +) -> dict[str, Any]: + return _base( + seq, + "tools", + namespace or [], + { + "event": "tool-finished", + "tool_call_id": tool_call_id, + "output": output, + }, + ) + + +def tool_error_event( + seq: int = 0, + namespace: list[str] | None = None, + *, + tool_call_id: str = "call-1", + message: str = "tool failed", + code: str = "tool_error", +) -> dict[str, Any]: + return _base( + seq, + "tools", + namespace or [], + { + "event": "tool-error", + "tool_call_id": tool_call_id, + "message": message, + "code": code, + }, + ) + + +def tasks_start_event( + seq: int = 0, + namespace: list[str] | None = None, + *, + task_id: str = "task-1", + name: str = "node", + input: Any = None, +) -> dict[str, Any]: + return _base( + seq, + "tasks", + namespace or [], + { + "id": task_id, + "name": name, + "input": input, + "triggers": [], + }, + ) + + +def tasks_result_event( + seq: int = 0, + namespace: list[str] | None = None, + *, + task_id: str = "task-1", + name: str = "node", + result: Any = None, + error: str | None = None, + interrupts: list[dict[str, Any]] | None = None, +) -> dict[str, Any]: + return _base( + seq, + "tasks", + namespace or [], + { + "id": task_id, + "name": name, + "result": result if result is not None else {}, + "error": error, + "interrupts": interrupts or [], + }, + ) diff --git a/libs/sdk-py/tests/streaming/test_messages_projection.py b/libs/sdk-py/tests/streaming/test_messages_projection.py new file mode 100644 index 000000000..35b359094 --- /dev/null +++ b/libs/sdk-py/tests/streaming/test_messages_projection.py @@ -0,0 +1,209 @@ +"""Tests for `thread.messages` - typed async message projection.""" + +from __future__ import annotations + +import httpx +import pytest +from langchain_core.language_models.chat_model_stream import AsyncChatModelStream +from langchain_core.messages import AIMessage + +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_error_event, + message_finish_event, + message_start_event, + message_text_delta_event, + message_text_finish_event, +) +from streaming._fake_server import FakeServer + + +async def test_messages_subscribes_to_messages_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={}) + _ = [message async for message in thread.messages] + + assert any( + "messages" in body.get("channels", []) for body in fake.stream_request_bodies + ) + + +async def test_messages_yields_async_chat_model_stream_and_text_deltas(): + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + message_start_event(seq=1, message_id="msg-1", run_id="run-1"), + message_text_delta_event(seq=2, text="hel", message_id="msg-1"), + message_text_delta_event(seq=3, text="lo", message_id="msg-1"), + message_text_finish_event(seq=4, text="hello", message_id="msg-1"), + message_finish_event( + seq=5, input_tokens=2, output_tokens=3, message_id="msg-1" + ), + 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={}) + streams = [message async for message in thread.messages] + + assert len(streams) == 1 + message = streams[0] + assert isinstance(message, AsyncChatModelStream) + assert message.message_id == "msg-1" + assert [delta async for delta in message.text] == ["hel", "lo"] + assert await message.text == "hello" + output = await message.output + assert isinstance(output, AIMessage) + assert output.id == "msg-1" + assert output.content == [{"type": "text", "text": "hello", "index": 0}] + assert output.usage_metadata == { + "input_tokens": 2, + "output_tokens": 3, + "total_tokens": 5, + } + + +async def test_messages_multiple_messages_are_distinct_streams(): + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + message_start_event(seq=1, message_id="msg-1", run_id="run-1"), + message_text_delta_event(seq=2, text="one", message_id="msg-1"), + message_text_finish_event(seq=3, text="one", message_id="msg-1"), + message_finish_event(seq=4, message_id="msg-1"), + message_start_event(seq=5, message_id="msg-2", run_id="run-2"), + message_text_delta_event(seq=6, text="two", message_id="msg-2"), + message_text_finish_event(seq=7, text="two", message_id="msg-2"), + message_finish_event(seq=8, message_id="msg-2"), + lifecycle_completed_event(seq=9), + ] + ) + 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={}) + streams = [message async for message in thread.messages] + + assert [stream.message_id for stream in streams] == ["msg-1", "msg-2"] + assert [await stream.text for stream in streams] == ["one", "two"] + + +async def test_messages_ignores_nested_namespace_for_root_projection(): + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + message_start_event(seq=1, namespace=["child:1"], message_id="nested"), + message_text_delta_event(seq=2, namespace=["child:1"], text="nested"), + message_finish_event(seq=3, namespace=["child:1"]), + lifecycle_completed_event(seq=4), + ] + ) + 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={}) + streams = [message async for message in thread.messages] + + assert streams == [] + + +async def test_messages_error_event_fails_active_stream(): + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + message_start_event(seq=1, message_id="msg-1"), + message_error_event(seq=2, message="model failed", message_id="msg-1"), + 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={}) + streams = [message async for message in thread.messages] + + assert len(streams) == 1 + with pytest.raises(RuntimeError, match="model failed"): + await streams[0].output + + +async def test_messages_concurrent_same_run_id_route_independently(): + """Two messages sharing a run_id must route to independent streams. + + The old `_message_route_key` keyed on `run_id` when present, so both + message-start events mapped to the same `active` slot and the second + overwrote the first. Subsequent deltas and finish events all routed to + the wrong (or missing) stream. + """ + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + # Both messages share run_id="run-1" (same agent turn) + message_start_event(seq=1, message_id="msg-A", run_id="run-1"), + message_start_event(seq=2, message_id="msg-B", run_id="run-1"), + message_text_delta_event(seq=3, text="alpha", message_id="msg-A"), + message_text_delta_event(seq=4, text="beta", message_id="msg-B"), + message_finish_event(seq=5, message_id="msg-A"), + message_finish_event(seq=6, message_id="msg-B"), + 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={}) + streams = [msg async for msg in thread.messages] + + assert [s.message_id for s in streams] == ["msg-A", "msg-B"] + assert [await s.text for s in streams] == ["alpha", "beta"] + + +async def test_messages_orphan_delta_without_matching_key_is_dropped(): + """A delta whose message_id doesn't match any active stream must be dropped. + + The old code fell back to routing the event to the only active stream when + `len(active) == 1`, causing orphan/mismatched deltas to silently corrupt + an unrelated stream's content. + """ + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + message_start_event(seq=1, message_id="msg-A"), + # Delta with a mismatched message_id — must be dropped, not routed to msg-A. + message_text_delta_event(seq=2, text="orphan", message_id="msg-UNKNOWN"), + message_text_delta_event(seq=3, text="real", message_id="msg-A"), + message_finish_event(seq=4, message_id="msg-A"), + 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={}) + streams = [msg async for msg in thread.messages] + + assert len(streams) == 1 + # Only the correctly-keyed delta "real" must appear; "orphan" must be dropped. + assert await streams[0].text == "real" diff --git a/libs/sdk-py/tests/streaming/test_subscription.py b/libs/sdk-py/tests/streaming/test_subscription.py index 9e8bc92ae..3d3878d9f 100644 --- a/libs/sdk-py/tests/streaming/test_subscription.py +++ b/libs/sdk-py/tests/streaming/test_subscription.py @@ -1,7 +1,25 @@ from __future__ import annotations import pytest +from langchain_protocol import ( + Channel as ProtocolChannel, +) +from langchain_protocol import ( + Event as ProtocolEvent, +) +from langchain_protocol import ( + Namespace as ProtocolNamespace, +) +from langchain_protocol import ( + SubscribeParams as ProtocolSubscribeParams, +) +from langgraph_sdk.stream import ( + Channel, + Event, + Namespace, + SubscribeParams, +) from langgraph_sdk.stream.subscription import ( compute_union_filter, filter_covers, @@ -236,3 +254,11 @@ def test_filter_covers_bounded_coverer_does_not_cover_unbounded_target(): coverer = {"channels": ["values"], "depth": 2} target = {"channels": ["values"]} # unbounded assert filter_covers(coverer, target) is False + + +def test_protocol_types_are_importable_from_stream_module(): + """Test that v3 protocol types are re-exported from langgraph_sdk.stream.""" + assert Channel is ProtocolChannel + assert Event is ProtocolEvent + assert Namespace is ProtocolNamespace + assert SubscribeParams is ProtocolSubscribeParams diff --git a/libs/sdk-py/tests/streaming/test_tool_calls_projection.py b/libs/sdk-py/tests/streaming/test_tool_calls_projection.py new file mode 100644 index 000000000..f01e03d81 --- /dev/null +++ b/libs/sdk-py/tests/streaming/test_tool_calls_projection.py @@ -0,0 +1,272 @@ +"""Tests for `thread.tool_calls` - typed async tool-call projection.""" + +from __future__ import annotations + +import time + +import httpx +import pytest + +from langgraph_sdk._async.http import HttpClient +from langgraph_sdk._async.threads import ThreadsClient +from streaming._events import ( + lifecycle_completed_event, + lifecycle_errored_event, + lifecycle_started_event, + tool_error_event, + tool_finished_event, + tool_output_delta_event, + tool_started_event, +) +from streaming._fake_server import FakeServer + + +async def test_tool_calls_subscribes_to_tools_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={}) + _ = [call async for call in thread.tool_calls] + + assert any( + "tools" in body.get("channels", []) for body in fake.stream_request_bodies + ) + + +async def test_tool_calls_yields_handle_deltas_and_output(): + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + tool_started_event( + seq=1, + tool_call_id="call-1", + tool_name="search", + input={"query": "sf weather"}, + ), + tool_output_delta_event(seq=2, tool_call_id="call-1", delta="part "), + tool_output_delta_event(seq=3, tool_call_id="call-1", delta="two"), + tool_finished_event( + seq=4, + tool_call_id="call-1", + output={"temperature": 68}, + ), + 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={}) + calls = [call async for call in thread.tool_calls] + + assert len(calls) == 1 + call = calls[0] + assert call.tool_call_id == "call-1" + assert call.name == "search" + assert call.input == {"query": "sf weather"} + assert call.namespace == [] + assert call.done is True + assert [delta async for delta in call.deltas] == ["part ", "two"] + assert await call.output == {"temperature": 68} + + +async def test_tool_calls_multiple_concurrent_calls_route_by_id(): + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + tool_started_event(seq=1, tool_call_id="call-a", tool_name="alpha"), + tool_started_event(seq=2, tool_call_id="call-b", tool_name="beta"), + tool_output_delta_event(seq=3, tool_call_id="call-b", delta="b1"), + tool_output_delta_event(seq=4, tool_call_id="call-a", delta="a1"), + tool_finished_event(seq=5, tool_call_id="call-a", output="A"), + tool_finished_event(seq=6, tool_call_id="call-b", output="B"), + 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={}) + calls = [call async for call in thread.tool_calls] + + by_id = {call.tool_call_id: call for call in calls} + assert set(by_id) == {"call-a", "call-b"} + assert [delta async for delta in by_id["call-a"].deltas] == ["a1"] + assert [delta async for delta in by_id["call-b"].deltas] == ["b1"] + assert await by_id["call-a"].output == "A" + assert await by_id["call-b"].output == "B" + + +async def test_tool_calls_ignores_nested_namespace_for_root_projection(): + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + tool_started_event(seq=1, namespace=["child:1"], tool_call_id="nested"), + tool_finished_event(seq=2, namespace=["child:1"], tool_call_id="nested"), + 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={}) + calls = [call async for call in thread.tool_calls] + + assert calls == [] + + +async def test_tool_calls_error_event_fails_output_and_deltas(): + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + tool_started_event(seq=1, tool_call_id="call-1"), + tool_output_delta_event(seq=2, tool_call_id="call-1", delta="before"), + tool_error_event(seq=3, tool_call_id="call-1", message="boom"), + lifecycle_completed_event(seq=4), + ] + ) + 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={}) + calls = [call async for call in thread.tool_calls] + + assert len(calls) == 1 + assert [delta async for delta in calls[0].deltas] == ["before"] + with pytest.raises(RuntimeError, match="boom"): + await calls[0].output + + +async def test_tool_calls_run_error_fails_active_handle(): + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + tool_started_event(seq=1, tool_call_id="call-1"), + lifecycle_errored_event(seq=2, error="run failed"), + ] + ) + 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={}) + calls = [call async for call in thread.tool_calls] + + assert len(calls) == 1 + with pytest.raises(RuntimeError, match="Run errored: run failed"): + await calls[0].output + + +async def test_tool_calls_stream_end_fails_active_handle(): + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + tool_started_event(seq=1, tool_call_id="call-1"), + lifecycle_completed_event(seq=2), + ] + ) + 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={}) + calls = [call async for call in thread.tool_calls] + + assert len(calls) == 1 + with pytest.raises(RuntimeError, match="closed before terminal tool event"): + await calls[0].output + + +async def test_tool_calls_explicit_aclose_does_not_block_1s(): + """Explicitly closing the tool_calls iterator must return in <500ms. + + The old finally block did `await asyncio.wait_for(asyncio.shield(run_done), + timeout=1.0)` unconditionally. When the caller explicitly calls aclose() on + the generator before any lifecycle terminal event arrives, this caused a + mandatory 1-second stall per iterator close. + """ + fake = FakeServer() + # Script has a started lifecycle and one tool, but NO terminal lifecycle. + # If the shield-wait is present, aclose() will block for 1s. + fake.script( + [ + lifecycle_started_event(seq=0), + tool_started_event(seq=1, tool_call_id="call-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={}) + # _tool_calls_iter() is an AsyncGenerator; cast so the type checker + # knows aclose() is available without a bare AsyncIterator protocol. + from collections.abc import AsyncGenerator + + gen: AsyncGenerator = thread.tool_calls._tool_calls_iter() + _call = await gen.__anext__() # receive the one tool-started handle + start = time.monotonic() + await gen.aclose() # explicitly close — must not stall 1s + elapsed = time.monotonic() - start + assert elapsed < 0.5, f"tool_calls aclose() took {elapsed:.3f}s (expected <0.5s)" + + +def test_tool_call_handle_deltas_queue_is_bounded(): + """ToolCallHandle._deltas must be constructed with a bounded asyncio.Queue. + + Unbounded queues allow producers to enqueue indefinitely, causing memory + growth when consumers are slow. + """ + import asyncio + + # We need a running loop to create the Future inside ToolCallHandle.__init__. + async def _make() -> None: + from langgraph_sdk._async.stream import ToolCallHandle + + handle_default = ToolCallHandle(tool_call_id="tc1", name="foo") + assert handle_default._deltas.maxsize > 0, ( + "default maxsize must be positive (bounded)" + ) + + handle_custom = ToolCallHandle(tool_call_id="tc2", name="bar", max_queue_size=8) + assert handle_custom._deltas.maxsize == 8 + + asyncio.run(_make()) + + +def test_tool_call_handle_deltas_single_consumer_guard(): + """Accessing `handle.deltas` a second time must raise immediately. + + `_deltas` is a single-consumer queue; fanning out to multiple consumers + would cause each consumer to miss events already consumed by the other. + The property must raise before returning the iterator so the caller + sees the error even without iterating. + """ + import asyncio + + async def _run() -> None: + from langgraph_sdk._async.stream import ToolCallHandle + + handle = ToolCallHandle(tool_call_id="tc1", name="foo") + + # First access: fine — returns the iterator. + _iter_1 = handle.deltas + + # Second access: must raise immediately (before any iteration). + with pytest.raises(RuntimeError, match="single consumer"): + _ = handle.deltas + + asyncio.run(_run()) diff --git a/libs/sdk-py/uv.lock b/libs/sdk-py/uv.lock index dfb4020e6..7dfbfb629 100644 --- a/libs/sdk-py/uv.lock +++ b/libs/sdk-py/uv.lock @@ -484,6 +484,7 @@ name = "langgraph-sdk" source = { editable = "." } dependencies = [ { name = "httpx" }, + { name = "langchain-core" }, { name = "langchain-protocol" }, { name = "orjson" }, ] @@ -519,6 +520,7 @@ test = [ [package.metadata] requires-dist = [ { name = "httpx", specifier = ">=0.25.2" }, + { name = "langchain-core", specifier = ">=1.4.0,<2" }, { name = "langchain-protocol", specifier = ">=0.0.15" }, { name = "orjson", specifier = ">=3.11.5" }, ]