From 22259558cc613d8bfced2679b9185f552afea25d Mon Sep 17 00:00:00 2001 From: Nick Hollon Date: Wed, 27 May 2026 10:41:50 -0400 Subject: [PATCH] feat(sdk-py): add async thread stream skeleton (#7819) --- libs/sdk-py/langgraph_sdk/_async/stream.py | 148 ++++++++ libs/sdk-py/langgraph_sdk/_async/threads.py | 33 ++ .../langgraph_sdk/stream/transport/http.py | 25 +- libs/sdk-py/tests/streaming/_fake_server.py | 8 +- .../streaming/assert_transport_replays.py | 73 ++++ .../streaming/test_replay_conformance.py | 24 ++ .../tests/streaming/test_subscription.py | 6 +- .../tests/streaming/test_thread_stream.py | 356 ++++++++++++++++++ .../tests/streaming/test_transport_http.py | 111 +++++- libs/sdk-py/tests/test_api_parity.py | 18 +- 10 files changed, 780 insertions(+), 22 deletions(-) create mode 100644 libs/sdk-py/langgraph_sdk/_async/stream.py create mode 100644 libs/sdk-py/tests/streaming/assert_transport_replays.py create mode 100644 libs/sdk-py/tests/streaming/test_replay_conformance.py create mode 100644 libs/sdk-py/tests/streaming/test_thread_stream.py diff --git a/libs/sdk-py/langgraph_sdk/_async/stream.py b/libs/sdk-py/langgraph_sdk/_async/stream.py new file mode 100644 index 000000000..7cb6b3a6d --- /dev/null +++ b/libs/sdk-py/langgraph_sdk/_async/stream.py @@ -0,0 +1,148 @@ +"""Async thread-centric streaming surface for the v3 protocol. + +`AsyncThreadStream` is an async context manager that owns a +`ProtocolSseTransport` for one thread, dispatches `run.start` commands, +and exposes a raw `events` async iterable. + +Direct port of `libs/sdk/src/client/stream/index.ts`. +""" + +from __future__ import annotations + +from collections.abc import AsyncIterator +from typing import Any + +import httpx +from langchain_protocol import Event + +from langgraph_sdk.stream.transport import EventStreamHandle, ProtocolSseTransport + +# All public protocol channels used by the raw `events` surface. +_ALL_CHANNELS: list[str] = [ + "values", + "updates", + "messages", + "tools", + "lifecycle", + "input", + "checkpoints", + "tasks", + "custom", +] + + +class RunModule: + """Command dispatcher for `run.start`. + + Bound to one `AsyncThreadStream`; accesses its transport and id allocator. + """ + + def __init__(self, owner: AsyncThreadStream) -> None: + self._owner = owner + + async def start( + self, + *, + input: Any = None, + config: dict[str, Any] | None = None, + metadata: dict[str, Any] | None = None, + ) -> dict[str, Any]: + """Send `run.start` to the server. Returns the result (`{"run_id": ...}`).""" + params: dict[str, Any] = {"assistant_id": self._owner.assistant_id} + if input is not None: + params["input"] = input + if config is not None: + params["config"] = config + if metadata is not None: + params["metadata"] = metadata + return await self._owner._send_command("run.start", params) + + +class AsyncThreadStream: + """Async context manager for one thread's v3 streaming session. + + Construct via `client.threads.stream(thread_id=None, *, assistant_id, ...)` + rather than instantiating directly. + """ + + def __init__( + self, + *, + client: httpx.AsyncClient, + thread_id: str, + assistant_id: str, + ) -> None: + self._http_client = client + self.thread_id = thread_id + self.assistant_id = assistant_id + self._closed = False + self._transport: ProtocolSseTransport | None = None + self._open_handles: list[EventStreamHandle] = [] + self._next_command_id = 1 + self.run = RunModule(self) + + async def __aenter__(self) -> AsyncThreadStream: + self._transport = ProtocolSseTransport( + client=self._http_client, + thread_id=self.thread_id, + ) + return self + + async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None: + try: + await self.close() + except BaseException: + if exc is None: + raise + # If we got here, a body exception is already in flight; swallow the + # close error so the body exception propagates. + + @property + def events(self) -> AsyncIterator[Event]: + """Return a fresh subscription to ALL channels. + + Each property access opens a new subscription; callers iterating twice + will see two independent streams (both filtered by the same channel union). + Terminates when the stream closes (server hangup, `__aexit__`, or + transport-level close). + """ + if self._transport is None: + raise RuntimeError("AsyncThreadStream not entered — use `async with`.") + handle = self._transport.open_event_stream({"channels": _ALL_CHANNELS}) + self._open_handles.append(handle) + return handle.events + + async def close(self) -> None: + """Tear down the thread stream. Idempotent.""" + if self._closed: + return + self._closed = True + for handle in self._open_handles: + await handle.close() + if self._transport is not None: + await self._transport.close() + + async def _send_command( + self, method: str, params: dict[str, Any] + ) -> dict[str, Any]: + """Send a protocol command and return the `result` payload. + + Returns `{}` for 202/204 responses (no body). Raises `RuntimeError` + with the protocol code/message when the server returns an error + envelope (`{"type": "error", ...}`). + """ + if self._transport is None: + raise RuntimeError("AsyncThreadStream not entered — use `async with`.") + command_id = self._next_command_id + self._next_command_id += 1 + response = await self._transport.send_command( + {"id": command_id, "method": method, "params": params} + ) + if response is None: + # 202/204 — no body. Caller gets an empty result. + return {} + if response.get("type") == "error": + code = response.get("error", "unknown") + message = response.get("message", "") + raise RuntimeError(f"Protocol error [{code}]: {message}") + return response.get("result", {}) diff --git a/libs/sdk-py/langgraph_sdk/_async/threads.py b/libs/sdk-py/langgraph_sdk/_async/threads.py index 7dbe003c9..96efb4bdc 100644 --- a/libs/sdk-py/langgraph_sdk/_async/threads.py +++ b/libs/sdk-py/langgraph_sdk/_async/threads.py @@ -2,10 +2,12 @@ from __future__ import annotations +import uuid from collections.abc import AsyncIterator, Mapping, Sequence from typing import Any, Literal, overload from langgraph_sdk._async.http import HttpClient +from langgraph_sdk._async.stream import AsyncThreadStream from langgraph_sdk._shared.utilities import _quote_path_param from langgraph_sdk.schema import ( Checkpoint, @@ -734,6 +736,37 @@ class ThreadsClient: params=params, ) + def stream( + self, + thread_id: str | None = None, + *, + assistant_id: str, + headers: Mapping[str, str] | None = None, # noqa: ARG002 + ) -> AsyncThreadStream: + """Open a v3 thread-centric streaming session. + + When `thread_id` is None, a fresh UUIDv4 is minted client-side and + included in the URL of subsequent `POST /threads/{thread_id}/...` + calls. The server creates the thread row lazily on the first + `run.start` via the run payload's `if_not_exists: "create"`. The + v3 protocol response carries only `run_id`, never `thread_id`. + + Args: + thread_id: optional explicit thread identifier. Defaults to a + fresh UUIDv4. + assistant_id: assistant the run will use. Required. + headers: optional per-request headers. Reserved; not currently + forwarded. + + Returns: + An `AsyncThreadStream` to use as an async context manager. + """ + return AsyncThreadStream( + client=self.http.client, + thread_id=thread_id if thread_id is not None else str(uuid.uuid4()), + assistant_id=assistant_id, + ) + async def join_stream( self, thread_id: str, diff --git a/libs/sdk-py/langgraph_sdk/stream/transport/http.py b/libs/sdk-py/langgraph_sdk/stream/transport/http.py index b7cc93854..705814c57 100644 --- a/libs/sdk-py/langgraph_sdk/stream/transport/http.py +++ b/libs/sdk-py/langgraph_sdk/stream/transport/http.py @@ -67,11 +67,11 @@ class ProtocolSseTransport: def __init__( self, + *, client: httpx.AsyncClient, thread_id: str, commands_path: str | None = None, stream_path: str | None = None, - *, max_queue_size: int = 1024, ) -> None: self._client = client @@ -80,6 +80,7 @@ class ProtocolSseTransport: self._stream_url = stream_path or f"/threads/{thread_id}/stream/events" self._max_queue_size = max_queue_size self._closed = False + self._event_streams: set[asyncio.Task[None]] = set() async def send_command(self, command: dict[str, Any]) -> dict[str, Any] | None: """POST a command. Returns the response JSON, or `None` for 202/204. @@ -107,7 +108,7 @@ class ProtocolSseTransport: raise RuntimeError( "Protocol command did not return a valid response." ) from err - if not isinstance(payload, dict) or "command_id" not in payload: + if not isinstance(payload, dict) or "id" not in payload: raise RuntimeError("Protocol command did not return a valid response.") return payload @@ -116,8 +117,8 @@ class ProtocolSseTransport: Posts `params` as a SubscribeParams body to `/threads/{thread_id}/stream/events`. Returns an `EventStreamHandle` whose `events` async iterator yields typed - `Event` dicts as the server emits them. `handle.ready` resolves when - response headers arrive (or rejects on early failure). + `Event` dicts as the server emits them. `handle.ready` resolves on a 2xx + response (rejects on HTTP error or transport failure before headers). Reconnect: pass `params["since"]` to filter outbound seqs server-side. The cursor goes in the request body, not as a `Last-Event-ID` header. @@ -182,6 +183,8 @@ class ProtocolSseTransport: await queue.put(None) # sentinel: end of stream task = asyncio.create_task(pump()) + self._event_streams.add(task) + task.add_done_callback(self._event_streams.discard) async def aiter() -> AsyncIterator[Event]: while True: @@ -192,12 +195,22 @@ class ProtocolSseTransport: async def close() -> None: cancel_event.set() + # Why: pump may be mid-`finally`; ensure consumer unblocks. + queue.put_nowait(None) task.cancel() - with contextlib.suppress(asyncio.CancelledError, BaseException): + with contextlib.suppress(asyncio.CancelledError, Exception): await task return EventStreamHandle(events=aiter(), ready=ready, done=done, close=close) async def close(self) -> None: - """Mark the transport closed. Idempotent.""" + """Cancel any open event streams and mark the transport closed. Idempotent.""" + if self._closed: + return self._closed = True + tasks = list(self._event_streams) + for task in tasks: + task.cancel() + if tasks: + with contextlib.suppress(Exception, asyncio.CancelledError): + await asyncio.gather(*tasks, return_exceptions=True) diff --git a/libs/sdk-py/tests/streaming/_fake_server.py b/libs/sdk-py/tests/streaming/_fake_server.py index c8307db69..7c47604bc 100644 --- a/libs/sdk-py/tests/streaming/_fake_server.py +++ b/libs/sdk-py/tests/streaming/_fake_server.py @@ -51,9 +51,13 @@ class FakeServer: async def commands(request: Request) -> Response: body = orjson.loads(await request.body()) self.received_commands.append(body) - command_id = body.get("command_id") + command_id = body.get("id") return JSONResponse( - {"command_id": command_id, "result": {"run_id": "run-1"}} + { + "type": "success", + "id": command_id, + "result": {"run_id": "run-1"}, + } ) async def stream_events(request: Request) -> Response: diff --git a/libs/sdk-py/tests/streaming/assert_transport_replays.py b/libs/sdk-py/tests/streaming/assert_transport_replays.py new file mode 100644 index 000000000..88e5429ad --- /dev/null +++ b/libs/sdk-py/tests/streaming/assert_transport_replays.py @@ -0,0 +1,73 @@ +"""Public conformance helper for the transport replay contract. + +Usage: + + from tests.streaming.assert_transport_replays import assert_transport_replays + + async def test_my_transport(): + async with my_transport_factory() as harness: + await assert_transport_replays(harness) + +The helper publishes a few events into a transport's underlying buffer +(via whatever side-channel the implementation exposes — typically by +scripting the fake server) and verifies that a fresh `open_event_stream` +yields them all before any new live events. +""" + +from __future__ import annotations + +import asyncio +from typing import Protocol + +from langchain_protocol import Event + +from langgraph_sdk.stream.transport.http import ProtocolSseTransport +from streaming._events import lifecycle_event + + +class _ReplayableHarness(Protocol): + transport: ProtocolSseTransport + + def script_buffered(self, events: list[dict]) -> None: ... + + +async def assert_transport_replays( + harness: _ReplayableHarness, + *, + buffered_count: int = 3, + timeout: float = 1.0, +) -> None: + """Assert that `harness.transport` replays buffered events on subscribe. + + Args: + harness: object exposing an open `ProtocolSseTransport` plus a + `script_buffered(events)` method that queues events as if they + were buffered server-side before the subscription opens. + buffered_count: how many synthetic events to script. + timeout: per-step await timeout in seconds. + + Raises: + AssertionError: when fewer than `buffered_count` events arrive (or + arrive out of order) on the fresh stream before it closes. + """ + events = [lifecycle_event(seq=i) for i in range(buffered_count)] + harness.script_buffered(events) + handle = harness.transport.open_event_stream({"channels": ["lifecycle"]}) + await asyncio.wait_for(handle.ready, timeout=timeout) + + received: list[Event] = [] + + async def drain() -> None: + async for event in handle.events: + received.append(event) + + try: + await asyncio.wait_for(drain(), timeout=timeout) + finally: + await handle.close() + + seqs = [e["seq"] for e in received] + assert seqs == list(range(buffered_count)), ( + f"transport did not replay buffered events: expected " + f"{list(range(buffered_count))}, got {seqs}" + ) diff --git a/libs/sdk-py/tests/streaming/test_replay_conformance.py b/libs/sdk-py/tests/streaming/test_replay_conformance.py new file mode 100644 index 000000000..167692fa3 --- /dev/null +++ b/libs/sdk-py/tests/streaming/test_replay_conformance.py @@ -0,0 +1,24 @@ +from __future__ import annotations + +import httpx + +from langgraph_sdk.stream.transport.http import ProtocolSseTransport +from streaming._fake_server import FakeServer +from streaming.assert_transport_replays import assert_transport_replays + + +class _Harness: + def __init__(self, fake: FakeServer, transport: ProtocolSseTransport) -> None: + self._fake = fake + self.transport = transport + + def script_buffered(self, events): + self._fake.script(events) + + +async def test_fake_server_replays_buffered_events(): + fake = FakeServer() + asgi = httpx.ASGITransport(app=fake.app) + async with httpx.AsyncClient(transport=asgi, base_url="http://test") as client: + transport = ProtocolSseTransport(client=client, thread_id="t-1") + await assert_transport_replays(_Harness(fake, transport)) diff --git a/libs/sdk-py/tests/streaming/test_subscription.py b/libs/sdk-py/tests/streaming/test_subscription.py index f8712b908..066348fea 100644 --- a/libs/sdk-py/tests/streaming/test_subscription.py +++ b/libs/sdk-py/tests/streaming/test_subscription.py @@ -71,10 +71,11 @@ def test_namespace_matches_with_depth_limit(): ) def test_infer_channel_for_each_method(method, expected_channel): event = { + "type": "event", "method": method, "params": {"namespace": [], "data": {}}, "seq": 0, - "id": "e", + "event_id": "e", } assert infer_channel(event) == expected_channel # ty: ignore[invalid-argument-type] @@ -91,10 +92,11 @@ def test_infer_channel_unknown_method_returns_none(): assert ( infer_channel( { + "type": "event", "method": "unknown", "params": {"namespace": [], "data": {}}, "seq": 0, - "id": "e", + "event_id": "e", } # ty:ignore[invalid-argument-type] ) is None diff --git a/libs/sdk-py/tests/streaming/test_thread_stream.py b/libs/sdk-py/tests/streaming/test_thread_stream.py new file mode 100644 index 000000000..16eca70e2 --- /dev/null +++ b/libs/sdk-py/tests/streaming/test_thread_stream.py @@ -0,0 +1,356 @@ +from __future__ import annotations + +import re +import uuid + +import httpx +import pytest + +from langgraph_sdk._async.http import HttpClient +from langgraph_sdk._async.stream import AsyncThreadStream +from langgraph_sdk._async.threads import ThreadsClient +from streaming._events import lifecycle_event, values_event +from streaming._fake_server import FakeServer + + +async def test_thread_stream_stores_thread_id_and_assistant_id(): + async with httpx.AsyncClient(base_url="http://test") as client: + stream = AsyncThreadStream( + client=client, + thread_id="t-1", + assistant_id="agent", + ) + assert stream.thread_id == "t-1" + assert stream.assistant_id == "agent" + + +async def test_aenter_returns_self(): + async with httpx.AsyncClient(base_url="http://test") as client: + stream = AsyncThreadStream(client=client, thread_id="t-1", assistant_id="agent") + async with stream as entered: + assert entered is stream + + +async def test_aexit_marks_closed(): + async with httpx.AsyncClient(base_url="http://test") as client: + stream = AsyncThreadStream(client=client, thread_id="t-1", assistant_id="agent") + async with stream: + assert stream._closed is False + assert stream._closed is True + + +async def test_close_is_idempotent(): + async with httpx.AsyncClient(base_url="http://test") as client: + stream = AsyncThreadStream(client=client, thread_id="t-1", assistant_id="agent") + await stream.close() + await stream.close() # must not raise + assert stream._closed is True + + +async def test_threads_stream_returns_async_thread_stream_with_explicit_id(): + async with httpx.AsyncClient(base_url="http://test") as raw: + threads = ThreadsClient(HttpClient(raw)) + stream = threads.stream(thread_id="my-thread", assistant_id="agent") + assert stream.thread_id == "my-thread" + assert stream.assistant_id == "agent" + + +async def test_threads_stream_mints_uuid4_when_thread_id_none(): + async with httpx.AsyncClient(base_url="http://test") as raw: + threads = ThreadsClient(HttpClient(raw)) + stream = threads.stream(assistant_id="agent") + # uuid4 format: 8-4-4-4-12 hex + assert re.fullmatch( + r"[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}", + stream.thread_id, + ) + # And it's actually parseable as a v4 UUID. + assert uuid.UUID(stream.thread_id).version == 4 + + +async def test_threads_stream_requires_assistant_id(): + async with httpx.AsyncClient(base_url="http://test") as raw: + threads = ThreadsClient(HttpClient(raw)) + with pytest.raises(TypeError): + threads.stream(thread_id="t-1") # ty: ignore[missing-argument] + + +async def test_threads_stream_accepts_headers_kwarg(): + """`headers` is accepted as a kwarg even though it isn't forwarded yet.""" + from langgraph_sdk._async.http import HttpClient + from langgraph_sdk._async.threads import ThreadsClient + + async with httpx.AsyncClient(base_url="http://test") as raw: + threads = ThreadsClient(HttpClient(raw)) + stream = threads.stream( + thread_id="t-1", + assistant_id="agent", + headers={"X-Foo": "bar"}, + ) + assert stream.thread_id == "t-1" + + +async def test_aenter_constructs_transport_with_thread_id(): + fake = FakeServer() + transport = httpx.ASGITransport(app=fake.app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as raw: + from langgraph_sdk._async.http import HttpClient + from langgraph_sdk._async.threads import ThreadsClient + + threads = ThreadsClient(HttpClient(raw)) + stream = threads.stream(thread_id="t-1", assistant_id="agent") + async with stream: + assert stream._transport is not None + assert stream._transport.thread_id == "t-1" + + +async def test_aexit_closes_transport(): + fake = FakeServer() + transport = httpx.ASGITransport(app=fake.app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as raw: + from langgraph_sdk._async.http import HttpClient + from langgraph_sdk._async.threads import ThreadsClient + + threads = ThreadsClient(HttpClient(raw)) + stream = threads.stream(thread_id="t-1", assistant_id="agent") + async with stream: + inner_transport = stream._transport + assert inner_transport is not None + assert inner_transport._closed is True + + +async def test_run_start_sends_command_with_assistant_id(): + fake = FakeServer() + transport = httpx.ASGITransport(app=fake.app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as raw: + from langgraph_sdk._async.http import HttpClient + from langgraph_sdk._async.threads import ThreadsClient + + threads = ThreadsClient(HttpClient(raw)) + async with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + result = await thread.run.start(input={"x": 1}) + assert result == {"run_id": "run-1"} + command = fake.received_commands[0] + assert command["method"] == "run.start" + assert command["params"]["assistant_id"] == "agent" + assert command["params"]["input"] == {"x": 1} + assert command["id"] == 1 + + +async def test_command_ids_are_monotonic(): + fake = FakeServer() + transport = httpx.ASGITransport(app=fake.app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as raw: + from langgraph_sdk._async.http import HttpClient + from langgraph_sdk._async.threads import ThreadsClient + + threads = ThreadsClient(HttpClient(raw)) + async with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + await thread.run.start(input={"x": 1}) + await thread.run.start(input={"x": 2}) + assert [c["id"] for c in fake.received_commands] == [1, 2] + + +async def test_run_start_forwards_config_and_metadata(): + fake = FakeServer() + transport = httpx.ASGITransport(app=fake.app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as raw: + from langgraph_sdk._async.http import HttpClient + from langgraph_sdk._async.threads import ThreadsClient + + threads = ThreadsClient(HttpClient(raw)) + async with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + await thread.run.start( + input={"x": 1}, + config={"recursion_limit": 5}, + metadata={"trace": "abc"}, + ) + params = fake.received_commands[0]["params"] + assert params["config"] == {"recursion_limit": 5} + assert params["metadata"] == {"trace": "abc"} + + +async def test_run_start_raises_outside_context_manager(): + import pytest + + async with httpx.AsyncClient(base_url="http://test") as raw: + stream = AsyncThreadStream(client=raw, thread_id="t-1", assistant_id="agent") + with pytest.raises(RuntimeError, match="async with"): + await stream.run.start(input={"x": 1}) + + +async def test_run_start_raises_on_error_envelope(): + from starlette.applications import Starlette + from starlette.responses import JSONResponse + from starlette.routing import Route + + async def commands(_request): + return JSONResponse( + { + "type": "error", + "id": 1, + "error": "invalid_argument", + "message": "run.start requires an assistant_id.", + } + ) + + app = Starlette( + routes=[Route("/threads/{thread_id}/commands", commands, methods=["POST"])] + ) + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as raw: + import pytest + + from langgraph_sdk._async.http import HttpClient + from langgraph_sdk._async.threads import ThreadsClient + + threads = ThreadsClient(HttpClient(raw)) + async with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + with pytest.raises(RuntimeError, match="invalid_argument"): + await thread.run.start(input={"x": 1}) + + +async def test_events_yields_raw_events_after_run_start(): + fake = FakeServer() + fake.script( + [ + lifecycle_event(seq=0), + values_event(seq=1), + ] + ) + transport = httpx.ASGITransport(app=fake.app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as raw: + from langgraph_sdk._async.http import HttpClient + from langgraph_sdk._async.threads import ThreadsClient + + threads = ThreadsClient(HttpClient(raw)) + async with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + await thread.run.start(input={}) + received = [e async for e in thread.events] + methods = [e["method"] for e in received] + assert methods == ["lifecycle", "values"] + + +async def test_events_subscribes_to_all_channels(): + fake = FakeServer() + fake.script([]) + transport = httpx.ASGITransport(app=fake.app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as raw: + from langgraph_sdk._async.http import HttpClient + from langgraph_sdk._async.threads import ThreadsClient + + threads = ThreadsClient(HttpClient(raw)) + async with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + await thread.run.start(input={}) + _ = [e async for e in thread.events] + body = fake.stream_request_bodies[0] + assert set(body["channels"]) == { + "values", + "updates", + "messages", + "tools", + "lifecycle", + "input", + "checkpoints", + "tasks", + "custom", + } + + +async def test_events_terminates_on_aexit(): + import asyncio + + import pytest + + fake = FakeServer() + fake.script([lifecycle_event(seq=i) for i in range(5)]) + transport = httpx.ASGITransport(app=fake.app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as raw: + from langgraph_sdk._async.http import HttpClient + from langgraph_sdk._async.threads import ThreadsClient + + threads = ThreadsClient(HttpClient(raw)) + stream = threads.stream(thread_id="t-1", assistant_id="agent") + async with stream as thread: + await thread.run.start(input={}) + handle_events = thread.events + # After __aexit__, further iteration must terminate cleanly. + with pytest.raises(StopAsyncIteration): + await asyncio.wait_for(handle_events.__anext__(), timeout=1.0) + + +async def test_events_raises_outside_context_manager(): + async with httpx.AsyncClient(base_url="http://test") as raw: + stream = AsyncThreadStream(client=raw, thread_id="t-1", assistant_id="agent") + with pytest.raises(RuntimeError, match="async with"): + _ = stream.events + + +async def test_aexit_preserves_original_exception_if_close_raises(): + """If the body of `async with` raises, AND close() also raises, the + body's exception must propagate. close()'s error is suppressed (chained + as context on close_err, but does not replace the original).""" + async with httpx.AsyncClient(base_url="http://test") as raw: + thread = AsyncThreadStream(client=raw, thread_id="t-1", assistant_id="agent") + + async def failing_close(): + raise RuntimeError("close failed") + + thread.close = failing_close # ty:ignore[invalid-assignment] + + with pytest.raises(ValueError, match="original"): + async with thread: + raise ValueError("original") + + +async def test_events_property_returns_fresh_iterator_each_access(): + """Two separate accesses of `thread.events` must return independent + subscriptions — the second access should produce a fresh iterator, + even if both are accessed before either is drained.""" + fake = FakeServer() + fake.script([]) + transport = httpx.ASGITransport(app=fake.app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as raw: + from langgraph_sdk._async.http import HttpClient + from langgraph_sdk._async.threads import ThreadsClient + + threads = ThreadsClient(HttpClient(raw)) + async with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + first_iter = thread.events + second_iter = thread.events + # Each property access must return a distinct iterator object. + assert first_iter is not second_iter + + +async def test_fresh_thread_happy_path_end_to_end(): + """User passes no thread_id; SDK mints one and uses it in all URLs. + + Validates the thread-stream surface end-to-end: + - uuid4 minted at client.threads.stream() + - run.start posted to /threads//commands + - events SSE opened at /threads//stream/events + - scripted events delivered to the user iterator + """ + fake = FakeServer() + fake.script([lifecycle_event(seq=0), values_event(seq=1)]) + + posted_paths: list[str] = [] + + class _PathSpyTransport(httpx.ASGITransport): + async def handle_async_request(self, request): + posted_paths.append(str(request.url.path)) + return await super().handle_async_request(request) + + spy = _PathSpyTransport(app=fake.app) + async with httpx.AsyncClient(transport=spy, base_url="http://test") as raw: + threads = ThreadsClient(HttpClient(raw)) + async with threads.stream(assistant_id="agent") as thread: + assert uuid.UUID(thread.thread_id).version == 4 + result = await thread.run.start(input={"x": 1}) + assert result == {"run_id": "run-1"} + received = [e async for e in thread.events] + assert [e["method"] for e in received] == ["lifecycle", "values"] + # Both POSTs must include the minted thread_id in the path. + minted_id_paths = [p for p in posted_paths if thread.thread_id in p] + assert any(p.endswith("/commands") for p in minted_id_paths) + assert any(p.endswith("/stream/events") for p in minted_id_paths) diff --git a/libs/sdk-py/tests/streaming/test_transport_http.py b/libs/sdk-py/tests/streaming/test_transport_http.py index 370e81be6..5904f81bb 100644 --- a/libs/sdk-py/tests/streaming/test_transport_http.py +++ b/libs/sdk-py/tests/streaming/test_transport_http.py @@ -42,13 +42,13 @@ async def test_send_command_posts_json_and_returns_response(): sse = ProtocolSseTransport(client=client, thread_id="t-1") result = await sse.send_command( { - "command_id": 7, + "id": 7, "method": "run.start", "params": {"input": {"x": 1}}, } ) - assert result == {"command_id": 7, "result": {"run_id": "run-1"}} - assert fake.received_commands[0]["command_id"] == 7 + assert result == {"type": "success", "id": 7, "result": {"run_id": "run-1"}} + assert fake.received_commands[0]["id"] == 7 async def test_send_command_returns_none_on_202(): @@ -68,9 +68,7 @@ async def test_send_command_returns_none_on_202(): transport = httpx.ASGITransport(app=app) async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: sse = ProtocolSseTransport(client=client, thread_id="t-1") - result = await sse.send_command( - {"command_id": 1, "method": "noop", "params": {}} - ) + result = await sse.send_command({"id": 1, "method": "noop", "params": {}}) assert result is None assert len(received) == 1 @@ -84,7 +82,7 @@ async def test_send_command_raises_when_closed(): sse = ProtocolSseTransport(client=client, thread_id="t-1") await sse.close() with pytest.raises(RuntimeError, match="closed"): - await sse.send_command({"command_id": 1, "method": "noop", "params": {}}) + await sse.send_command({"id": 1, "method": "noop", "params": {}}) async def test_send_command_raises_http_error_on_4xx(): @@ -102,7 +100,7 @@ async def test_send_command_raises_http_error_on_4xx(): async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: sse = ProtocolSseTransport(client=client, thread_id="t-1") with pytest.raises(httpx.HTTPStatusError): - await sse.send_command({"command_id": 1, "method": "noop", "params": {}}) + await sse.send_command({"id": 1, "method": "noop", "params": {}}) async def test_open_event_stream_yields_scripted_events(): @@ -339,3 +337,100 @@ async def test_cancel_event_prevents_post_cancel_flush(): async for _ in handle.events: pytest.fail("event yielded after close()") assert len(received) == 1 + + +@pytest.mark.anyio +async def test_open_event_stream_ready_rejects_on_5xx(): + from starlette.applications import Starlette + from starlette.responses import JSONResponse + from starlette.routing import Route + + async def stream_events(_request): + return JSONResponse({"error": "boom"}, status_code=500) + + app = Starlette( + routes=[ + Route( + "/threads/{thread_id}/stream/events", + stream_events, + methods=["POST"], + ) + ] + ) + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + sse = ProtocolSseTransport(client=client, thread_id="t-1") + handle = sse.open_event_stream({"channels": ["lifecycle"]}) + with pytest.raises(httpx.HTTPStatusError): + await asyncio.wait_for(handle.ready, timeout=1.0) + # Iterator should terminate cleanly (no hang). + with pytest.raises(StopAsyncIteration): + await asyncio.wait_for(handle.events.__anext__(), timeout=1.0) + await handle.close() + + +def test_build_event_stream_body_minimal_channels_only(): + from langgraph_sdk.stream.transport.http import _build_event_stream_body + + body = _build_event_stream_body({"channels": ["values"]}) + assert body == {"channels": ["values"]} + + +def test_build_event_stream_body_includes_all_optional_fields(): + from langgraph_sdk.stream.transport.http import _build_event_stream_body + + body = _build_event_stream_body( + { + "channels": ["values", "messages"], + "namespaces": [["fetcher"]], + "depth": 2, + "since": 7, + } + ) + assert body == { + "channels": ["values", "messages"], + "namespaces": [["fetcher"]], + "depth": 2, + "since": 7, + } + + +def test_build_event_stream_body_omits_since_when_not_int(): + from langgraph_sdk.stream.transport.http import _build_event_stream_body + + body = _build_event_stream_body({"channels": ["values"], "since": None}) + assert "since" not in body + + +async def test_open_event_stream_raises_when_closed(): + from streaming._fake_server import FakeServer + + fake = FakeServer() + transport = httpx.ASGITransport(app=fake.app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + sse = ProtocolSseTransport(client=client, thread_id="t-1") + await sse.close() + with pytest.raises(RuntimeError, match="closed"): + sse.open_event_stream({"channels": ["lifecycle"]}) + + +async def test_transport_close_cancels_open_event_streams(): + from streaming._events import lifecycle_event + from streaming._fake_server import FakeServer + + fake = FakeServer() + fake.script([lifecycle_event(seq=i) for i in range(5)], delay=0.05) + transport = httpx.ASGITransport(app=fake.app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + sse = ProtocolSseTransport(client=client, thread_id="t-1") + handle = sse.open_event_stream({"channels": ["lifecycle"]}) + await asyncio.wait_for(handle.ready, timeout=1.0) + # Closing the transport must terminate the open stream within a bounded time. + await asyncio.wait_for(sse.close(), timeout=1.0) + + # Drain any already-queued events; the stream must end (not hang). + async def drain() -> None: + async for _ in handle.events: + pass + + await asyncio.wait_for(drain(), timeout=1.0) diff --git a/libs/sdk-py/tests/test_api_parity.py b/libs/sdk-py/tests/test_api_parity.py index 4158bd134..32e0ecb60 100644 --- a/libs/sdk-py/tests/test_api_parity.py +++ b/libs/sdk-py/tests/test_api_parity.py @@ -48,6 +48,12 @@ def _normalize_return_annotation(ann: object) -> str: return s +# Methods that exist only on the async client surface. +ASYNC_ONLY_METHODS: dict[str, set[str]] = { + "ThreadsClient": {"stream"}, +} + + @pytest.mark.parametrize( "async_cls,sync_cls", [ @@ -62,12 +68,16 @@ def test_sync_api_matches_async(async_cls, sync_cls): async_methods = _public_methods(async_cls) sync_methods = _public_methods(sync_cls) - # Method name parity - assert set(sync_methods.keys()) == set(async_methods.keys()), ( - f"Method sets differ: async-only={set(async_methods) - set(sync_methods)}, sync-only={set(sync_methods) - set(async_methods)}" + allowlist = ASYNC_ONLY_METHODS.get(async_cls.__name__, set()) + async_method_names = set(async_methods.keys()) - allowlist + + # Method name parity (modulo the async-only allowlist). + assert sync_methods.keys() == async_method_names, ( + f"Method sets differ: async-only={async_method_names - set(sync_methods)}, sync-only={set(sync_methods) - async_method_names}" ) - for name, async_fn in async_methods.items(): + for name in async_method_names: + async_fn = async_methods[name] sync_fn = sync_methods[name] # Use inspect.signature for parameter names (robust across versions)