From d03310abbb67056cd4175ae96bba231b08b080f3 Mon Sep 17 00:00:00 2001 From: Nick Hollon Date: Wed, 27 May 2026 10:53:26 -0400 Subject: [PATCH] feat(sdk-py): add shared stream subscriptions (#7820) --- libs/sdk-py/langgraph_sdk/_async/stream.py | 271 +++++++++++++++- libs/sdk-py/langgraph_sdk/_async/threads.py | 6 +- .../sdk-py/langgraph_sdk/stream/controller.py | 298 ++++++++++++++++++ .../langgraph_sdk/stream/subscription.py | 106 +++++++ libs/sdk-py/tests/streaming/_fake_server.py | 30 +- .../sdk-py/tests/streaming/test_controller.py | 154 +++++++++ .../tests/streaming/test_lifecycle_watcher.py | 38 +++ .../tests/streaming/test_shared_stream.py | 150 +++++++++ .../tests/streaming/test_subscription.py | 111 +++++++ .../tests/streaming/test_thread_stream.py | 33 ++ 10 files changed, 1181 insertions(+), 16 deletions(-) create mode 100644 libs/sdk-py/langgraph_sdk/stream/controller.py create mode 100644 libs/sdk-py/tests/streaming/test_controller.py create mode 100644 libs/sdk-py/tests/streaming/test_lifecycle_watcher.py create mode 100644 libs/sdk-py/tests/streaming/test_shared_stream.py diff --git a/libs/sdk-py/langgraph_sdk/_async/stream.py b/libs/sdk-py/langgraph_sdk/_async/stream.py index 7cb6b3a6d..3a41e4d36 100644 --- a/libs/sdk-py/langgraph_sdk/_async/stream.py +++ b/libs/sdk-py/langgraph_sdk/_async/stream.py @@ -9,14 +9,36 @@ Direct port of `libs/sdk/src/client/stream/index.ts`. from __future__ import annotations -from collections.abc import AsyncIterator -from typing import Any +import asyncio +from collections.abc import AsyncGenerator, AsyncIterator +from dataclasses import dataclass, field +from typing import Any, TypedDict import httpx -from langchain_protocol import Event +from langchain_protocol import Event, SubscribeParams from langgraph_sdk.stream.transport import EventStreamHandle, ProtocolSseTransport + +class InterruptPayload(TypedDict): + """Payload surfaced when the server requests human input for a thread.""" + + interrupt_id: str + value: Any + namespace: list[str] + + +@dataclass +class _Subscription: + """Internal record for one active subscription on an `AsyncThreadStream`.""" + + id: int + params: SubscribeParams + queue: asyncio.Queue = field(default_factory=asyncio.Queue) + # Why: asyncio.Queue[Event | None] as a subscript in the field annotation + # causes a type error with ty; bare asyncio.Queue is accepted. + + # All public protocol channels used by the raw `events` surface. _ALL_CHANNELS: list[str] = [ "values", @@ -55,9 +77,20 @@ class RunModule: params["config"] = config if metadata is not None: params["metadata"] = metadata + self._owner._ensure_lifecycle_watcher_running() return await self._owner._send_command("run.start", params) +async def _close_after(handle: EventStreamHandle, *, delay: float = 0.0) -> None: + """Close a handle, optionally after a brief delay. Used to detach + closing the old stream from the synchronous rotation step so the new + stream can absorb server-side replayed events first. + """ + if delay: + await asyncio.sleep(delay) + await handle.close() + + class AsyncThreadStream: """Async context manager for one thread's v3 streaming session. @@ -71,31 +104,46 @@ class AsyncThreadStream: client: httpx.AsyncClient, thread_id: str, assistant_id: str, + max_queue_size: int = 1024, ) -> None: self._http_client = client self.thread_id = thread_id self.assistant_id = assistant_id + self._max_queue_size = max_queue_size self._closed = False self._transport: ProtocolSseTransport | None = None self._open_handles: list[EventStreamHandle] = [] self._next_command_id = 1 + self._next_subscription_id = 1 + self._subscriptions: dict[int, _Subscription] = {} + self._seen_event_ids: set[str] = set() + self._shared_stream: EventStreamHandle | None = None + self._shared_stream_filter: dict[str, Any] | None = None + self._fanout_task: asyncio.Task[None] | None = None + self.interrupted: bool = False + self.interrupts: list[InterruptPayload] = [] + self._lifecycle_watcher_task: asyncio.Task[None] | None = None + self._lifecycle_watcher_handle: EventStreamHandle | None = None self.run = RunModule(self) async def __aenter__(self) -> AsyncThreadStream: + if self._closed: + raise RuntimeError("AsyncThreadStream is closed and cannot be re-entered.") self._transport = ProtocolSseTransport( client=self._http_client, thread_id=self.thread_id, + max_queue_size=self._max_queue_size, ) return self async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None: try: await self.close() - except BaseException: + except BaseException as close_err: 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. + # Original exception takes precedence; chain close error as context. + close_err.__context__ = exc @property def events(self) -> AsyncIterator[Event]: @@ -122,6 +170,157 @@ class AsyncThreadStream: if self._transport is not None: await self._transport.close() + def _register_subscription(self, params: SubscribeParams) -> _Subscription: + """Allocate a subscription id and add it to the registry.""" + sub = _Subscription( + id=self._next_subscription_id, + params=params, + queue=asyncio.Queue(maxsize=self._max_queue_size), + ) + self._next_subscription_id += 1 + self._subscriptions[sub.id] = sub + return sub + + def _unregister_subscription(self, subscription_id: int) -> None: + """Remove a subscription from the registry. No-op if already absent.""" + self._subscriptions.pop(subscription_id, None) + + def subscribe( + self, + channels: list[str], + *, + namespaces: list[list[str]] | None = None, + depth: int | None = None, + ) -> AsyncIterator[Event]: + """Open a typed subscription against the shared SSE. + + Returns an async iterator that yields raw `Event` dicts matching the + given filter. Multiple concurrent subscribes share one HTTP connection + whose union expands or rotates as subscriptions come and go. + """ + if self._transport is None: + raise RuntimeError("AsyncThreadStream not entered — use `async with`.") + params: SubscribeParams = {"channels": list(channels)} + if namespaces is not None: + params["namespaces"] = namespaces + if depth is not None: + params["depth"] = depth + return self._subscription_iter(params) + + async def _subscription_iter( + self, params: SubscribeParams + ) -> AsyncGenerator[Event, None]: + sub = self._register_subscription(params) + try: + if self._closed: + return + await self._reconcile_stream(params) + self._ensure_fanout_running() + while True: + item = await sub.queue.get() + if item is None: + return + yield item + finally: + self._unregister_subscription(sub.id) + + def _ensure_fanout_running(self) -> None: + if self._fanout_task is None or self._fanout_task.done(): + self._fanout_task = asyncio.create_task(self._fanout()) + + async def _fanout(self) -> None: + """Single consumer of the shared SSE; routes events to subscriptions. + + Why: rotation in `_reconcile_stream` replaces `_shared_stream` mid-loop. + Re-read `self._shared_stream` on each outer iteration so we always + consume from the current handle. The old handle's iterator exhausts + naturally after `_close_after` closes it. + """ + from langgraph_sdk.stream.subscription import matches_subscription + + while not self._closed: + shared = self._shared_stream + if shared is None: + return + try: + async for event in self._dedup_iter(shared.events): + if self._closed: + break + for sub in list(self._subscriptions.values()): + if matches_subscription(event, sub.params): + sub.queue.put_nowait(event) + except Exception: + # Pump errored — close all subscription queues so consumers + # don't hang. + for sub in self._subscriptions.values(): + sub.queue.put_nowait(None) + raise + if self._shared_stream is shared: + # No rotation happened; stream genuinely ended. + break + # Rotation: loop again to pick up the new _shared_stream. + + # Terminate consumers cleanly on shutdown / stream-end. + for sub in self._subscriptions.values(): + sub.queue.put_nowait(None) + + async def _reconcile_stream(self, candidate_filter: SubscribeParams) -> None: + """Ensure the shared SSE covers `candidate_filter`. Rotate if not. + + Open-new-before-close-old: any events buffered server-side between + the two opens are replayed on the new SSE, and the per-thread + `_seen_event_ids` set dedupes the overlap. Awaits `new_stream.ready` + so the HTTP connection is established before returning, guaranteeing + that both old and new streams are simultaneously connected during + rotation (enabling correct peak-count tracking and dedup correctness). + """ + from langgraph_sdk.stream.subscription import filter_covers + + if self._transport is None: + raise RuntimeError("AsyncThreadStream not entered — use `async with`.") + + if ( + self._shared_stream is not None + and self._shared_stream_filter is not None + and filter_covers(self._shared_stream_filter, dict(candidate_filter)) + ): + return # Existing stream is sufficient. + + new_filter = self._compute_current_union(extra=candidate_filter) + new_stream = self._transport.open_event_stream(new_filter) + old_stream = self._shared_stream + self._shared_stream = new_stream + self._shared_stream_filter = new_filter + # Await the new stream's ready future so the HTTP connection is + # established before we schedule the old stream's close. This ensures + # old and new are simultaneously open during the rotation window. + await new_stream.ready + if old_stream is not None: + # Schedule the old stream's close as a separate task so the + # caller doesn't pay close() latency in the rotation hot path. + asyncio.create_task(_close_after(old_stream)) # noqa: RUF006 + + def _compute_current_union( + self, extra: SubscribeParams | None = None + ) -> dict[str, Any]: + from langgraph_sdk.stream.subscription import compute_union_filter + + filters: list[dict[str, Any]] = [ + dict(sub.params) for sub in self._subscriptions.values() + ] + if extra is not None: + filters.append(dict(extra)) + return compute_union_filter(filters) + + async def _dedup_iter(self, source: AsyncIterator[Event]) -> AsyncIterator[Event]: + async for event in source: + event_id = event.get("event_id") + if event_id is not None: + if event_id in self._seen_event_ids: + continue + self._seen_event_ids.add(event_id) + yield event + async def _send_command( self, method: str, params: dict[str, Any] ) -> dict[str, Any]: @@ -146,3 +345,63 @@ class AsyncThreadStream: message = response.get("message", "") raise RuntimeError(f"Protocol error [{code}]: {message}") return response.get("result", {}) + + def _ensure_lifecycle_watcher_running(self) -> None: + # Why: this watcher is intentionally one-shot. If it crashes, it stays + # dead until the AsyncThreadStream is closed. + if self._lifecycle_watcher_task is not None: + return + self._lifecycle_watcher_task = asyncio.create_task( + self._run_lifecycle_watcher() + ) + + async def _run_lifecycle_watcher(self) -> None: + """Always-on SSE consuming lifecycle + input channels. + + Independent of the union-filter shared stream so that interrupts + surface even when no other subscription is active. + + The watcher waits for the run-start gate before opening so it does not + race server-side thread creation. + """ + if self._transport is None: + return + try: + handle = self._transport.open_event_stream( + {"channels": ["lifecycle", "input"]} + ) + self._lifecycle_watcher_handle = handle + await asyncio.wait_for(handle.ready, timeout=5.0) + async for event in handle.events: + if self._closed: + return + self._apply_lifecycle_event(event) + except (Exception, asyncio.CancelledError): + # Why: advisory-only watcher. Any error (HTTP failure, malformed + # event in `_apply_lifecycle_event`, cancellation on close) must + # not crash the caller; the watcher is one-shot best-effort. + return + + def _apply_lifecycle_event(self, event: Event) -> None: + """Update `interrupted` / `interrupts` state from a lifecycle or input event.""" + method = event.get("method") + if method == "input.requested": + params = event.get("params") or {} + data = params.get("data") if isinstance(params, dict) else None + interrupt_id = data.get("interrupt_id") if isinstance(data, dict) else None + if isinstance(interrupt_id, str): + payload: InterruptPayload = { + "interrupt_id": interrupt_id, + "value": data.get("value") if isinstance(data, dict) else None, + "namespace": params.get("namespace") or [] + if isinstance(params, dict) + else [], + } + self.interrupts.append(payload) + self.interrupted = True + elif method == "lifecycle": + params = event.get("params") or {} + data = params.get("data") if isinstance(params, dict) else None + phase = data.get("phase") if isinstance(data, dict) else None + if phase in ("completed", "errored"): + self.interrupted = False diff --git a/libs/sdk-py/langgraph_sdk/_async/threads.py b/libs/sdk-py/langgraph_sdk/_async/threads.py index 96efb4bdc..c56f77cb1 100644 --- a/libs/sdk-py/langgraph_sdk/_async/threads.py +++ b/libs/sdk-py/langgraph_sdk/_async/threads.py @@ -748,8 +748,10 @@ class ThreadsClient: 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`. + `run.start` (internal server detail — the SDK does not send any + `if_not_exists` flag). The v3 protocol response carries only + `run_id`, never `thread_id` — that's why the SDK mints the id + client-side. Args: thread_id: optional explicit thread identifier. Defaults to a diff --git a/libs/sdk-py/langgraph_sdk/stream/controller.py b/libs/sdk-py/langgraph_sdk/stream/controller.py new file mode 100644 index 000000000..186d4c6e8 --- /dev/null +++ b/libs/sdk-py/langgraph_sdk/stream/controller.py @@ -0,0 +1,298 @@ +"""Stream controller: subscription registry and fan-out for AsyncThreadStream. + +`StreamController` manages the set of active subscriptions against one shared +SSE connection, routing events from the shared stream to per-subscription +queues. It is the centralised place for: + +- subscription registration / teardown +- shared-stream lifecycle (open, rotate, close) +- dedup of replayed events across rotations +- fan-out from the shared stream to subscriber queues +""" + +from __future__ import annotations + +import asyncio +import contextlib +from collections import OrderedDict +from collections.abc import AsyncGenerator, AsyncIterator +from dataclasses import dataclass, field +from typing import Any + +from langchain_protocol import Event, SubscribeParams + +from langgraph_sdk.stream.transport import EventStreamHandle + +# --------------------------------------------------------------------------- +# Bounded LRU set for event-id dedup +# --------------------------------------------------------------------------- + + +class _SeenEventIds: + """LRU set of event ids with bounded memory.""" + + __slots__ = ("_data", "_maxsize") + + def __init__(self, maxsize: int = 10_000) -> None: + self._data: OrderedDict[str, None] = OrderedDict() + self._maxsize = maxsize + + def add(self, event_id: str) -> None: + if event_id in self._data: + self._data.move_to_end(event_id) + return + self._data[event_id] = None + if len(self._data) > self._maxsize: + self._data.popitem(last=False) + + def __contains__(self, event_id: object) -> bool: + return event_id in self._data + + def __iter__(self): + return iter(self._data) + + +# --------------------------------------------------------------------------- +# Per-subscription record +# --------------------------------------------------------------------------- + + +@dataclass +class _Subscription: + """Internal record for one active subscription on a `StreamController`.""" + + id: int + params: SubscribeParams + queue: asyncio.Queue = field(default_factory=asyncio.Queue) + # Why: asyncio.Queue[Event | None] as a subscript in the field annotation + # causes a type error with ty; bare asyncio.Queue is accepted. + + +# --------------------------------------------------------------------------- +# Rotation close helper +# --------------------------------------------------------------------------- + + +async def _close_after(handle: EventStreamHandle, *, delay: float = 0.0) -> None: + """Close a handle, optionally after a brief delay. + + Used to detach closing the old stream from the synchronous rotation step + so the new stream can absorb server-side replayed events first. + """ + if delay: + await asyncio.sleep(delay) + await handle.close() + + +# --------------------------------------------------------------------------- +# StreamController +# --------------------------------------------------------------------------- + + +class StreamController: + """Manages subscriptions and fan-out against one shared SSE connection. + + Responsibilities: + - subscription registry (register / unregister) + - shared-stream lifecycle (open on first subscribe, rotate on filter widen) + - dedup of replayed events via a bounded LRU `_SeenEventIds` + - fan-out from the shared stream to per-subscription queues + + Args: + transport: the transport used to open event streams. + max_queue_size: per-subscription queue bound (default 1024). + seen_event_ids_max: LRU cap for the dedup set (default 10_000). + """ + + def __init__( + self, + *, + transport: Any, + max_queue_size: int = 1024, + seen_event_ids_max: int = 10_000, + ) -> None: + self._transport = transport + self._max_queue_size = max_queue_size + self._seen_event_ids = _SeenEventIds(maxsize=seen_event_ids_max) + self._next_subscription_id = 1 + self._subscriptions: dict[int, _Subscription] = {} + self._shared_stream: EventStreamHandle | None = None + self._shared_stream_filter: dict[str, Any] | None = None + self._fanout_task: asyncio.Task[None] | None = None + self._rotation_close_tasks: set[asyncio.Task[None]] = set() + self._closed = False + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def subscribe( + self, + channels: list[str], + *, + namespaces: list[list[str]] | None = None, + depth: int | None = None, + ) -> AsyncIterator[Event]: + """Open a typed subscription against the shared SSE. + + Returns an async iterator that yields raw `Event` dicts matching the + given filter. Multiple concurrent subscribes share one HTTP connection + whose union expands or rotates as subscriptions come and go. + """ + params: SubscribeParams = {"channels": list(channels)} + if namespaces is not None: + params["namespaces"] = namespaces + if depth is not None: + params["depth"] = depth + return self._subscription_iter(params) + + async def close(self) -> None: + """Tear down the controller, awaiting any pending rotation closes.""" + if self._closed: + return + self._closed = True + if self._fanout_task is not None: + self._fanout_task.cancel() + with contextlib.suppress(Exception, asyncio.CancelledError): + await self._fanout_task + if self._shared_stream is not None: + await self._shared_stream.close() + if self._rotation_close_tasks: + await asyncio.gather(*self._rotation_close_tasks, return_exceptions=True) + + # ------------------------------------------------------------------ + # Subscription internals + # ------------------------------------------------------------------ + + def _register_subscription(self, params: SubscribeParams) -> _Subscription: + """Allocate a subscription id, create a bounded queue, add to registry.""" + sub = _Subscription( + id=self._next_subscription_id, + params=params, + queue=asyncio.Queue(maxsize=self._max_queue_size), + ) + self._next_subscription_id += 1 + self._subscriptions[sub.id] = sub + return sub + + def _unregister_subscription(self, subscription_id: int) -> None: + """Remove a subscription from the registry. No-op if already absent.""" + self._subscriptions.pop(subscription_id, None) + + async def _subscription_iter( + self, params: SubscribeParams + ) -> AsyncGenerator[Event, None]: + sub = self._register_subscription(params) + try: + if self._closed: + return + await self._reconcile_stream(params) + self._ensure_fanout_running() + while True: + item = await sub.queue.get() + if item is None: + return + yield item + finally: + self._unregister_subscription(sub.id) + + # ------------------------------------------------------------------ + # Fan-out + # ------------------------------------------------------------------ + + def _ensure_fanout_running(self) -> None: + if self._fanout_task is None or self._fanout_task.done(): + self._fanout_task = asyncio.create_task(self._fanout()) + + async def _fanout(self) -> None: + """Single consumer of the shared SSE; routes events to subscriptions. + + Why: rotation in `_reconcile_stream` replaces `_shared_stream` mid-loop. + Re-read `self._shared_stream` on each outer iteration so we always + consume from the current handle. The old handle's iterator exhausts + naturally after `_close_after` closes it. + """ + from langgraph_sdk.stream.subscription import matches_subscription + + while not self._closed: + shared = self._shared_stream + if shared is None: + return + try: + async for event in self._dedup_iter(shared.events): + if self._closed: + break + for sub in list(self._subscriptions.values()): + if matches_subscription(event, sub.params): + sub.queue.put_nowait(event) + except Exception: + # Pump errored — close all subscription queues so consumers + # don't hang. + for sub in self._subscriptions.values(): + sub.queue.put_nowait(None) + raise + if self._shared_stream is shared: + # No rotation happened; stream genuinely ended. + break + # Rotation: loop again to pick up the new _shared_stream. + + # Terminate consumers cleanly on shutdown / stream-end. + for sub in self._subscriptions.values(): + sub.queue.put_nowait(None) + + # ------------------------------------------------------------------ + # Stream rotation + # ------------------------------------------------------------------ + + async def _reconcile_stream(self, candidate_filter: SubscribeParams) -> None: + """Ensure the shared SSE covers `candidate_filter`. Rotate if not. + + Open-new-before-close-old: any events buffered server-side between + the two opens are replayed on the new SSE, and `_seen_event_ids` + dedupes the overlap. Awaits `new_stream.ready` so the HTTP connection + is established before returning. + """ + from langgraph_sdk.stream.subscription import filter_covers + + if ( + self._shared_stream is not None + and self._shared_stream_filter is not None + and filter_covers(self._shared_stream_filter, dict(candidate_filter)) + ): + return # Existing stream is sufficient. + + new_filter = self._compute_current_union(extra=candidate_filter) + new_stream = self._transport.open_event_stream(new_filter) + old_stream = self._shared_stream + self._shared_stream = new_stream + self._shared_stream_filter = new_filter + await new_stream.ready + if old_stream is not None: + task = asyncio.create_task(_close_after(old_stream)) + self._rotation_close_tasks.add(task) + task.add_done_callback(self._rotation_close_tasks.discard) + + def _compute_current_union( + self, extra: SubscribeParams | None = None + ) -> dict[str, Any]: + from langgraph_sdk.stream.subscription import compute_union_filter + + filters: list[dict[str, Any]] = [ + dict(sub.params) for sub in self._subscriptions.values() + ] + if extra is not None: + filters.append(dict(extra)) + return compute_union_filter(filters) + + # ------------------------------------------------------------------ + # Dedup + # ------------------------------------------------------------------ + + async def _dedup_iter(self, source: AsyncIterator[Event]) -> AsyncIterator[Event]: + async for event in source: + event_id = event.get("event_id") + if event_id is not None: + if event_id in self._seen_event_ids: + continue + self._seen_event_ids.add(event_id) + yield event diff --git a/libs/sdk-py/langgraph_sdk/stream/subscription.py b/libs/sdk-py/langgraph_sdk/stream/subscription.py index d4bbbf98a..aa8298e90 100644 --- a/libs/sdk-py/langgraph_sdk/stream/subscription.py +++ b/libs/sdk-py/langgraph_sdk/stream/subscription.py @@ -5,6 +5,8 @@ Direct port of `libs/sdk/src/client/stream/subscription.ts` from the JS SDK. from __future__ import annotations +from typing import Any + from langchain_protocol import Channel, Event, Namespace, SubscribeParams @@ -100,3 +102,107 @@ def matches_subscription(event: Event, definition: SubscribeParams) -> bool: definition.get("namespaces"), definition.get("depth"), ) + + +def compute_union_filter( + subscriptions: list[dict[str, Any]], +) -> dict[str, Any]: + """Aggregate a set of subscription filters into one covering filter. + + Direct port of `client/stream/index.ts:#computeUnionFilter`. + + - Channels are unioned. + - Namespaces: if any subscription omits `namespaces` (wildcard), the union + is unscoped (omits the key). Otherwise, deduplicated union. + - Depth: if any subscription omits `depth` (unbounded), the union is + unbounded (omits the key). Otherwise, take the max. `depth=0` is a + valid bounded value — never omit when all subscriptions provide it. + + Args: + subscriptions: list of `SubscribeParams`-shaped dicts. + + Returns: + A `SubscribeParams`-shaped dict covering every input. + """ + if not subscriptions: + return {"channels": []} + + channels: set[str] = set() + wildcard_namespaces = False + namespace_map: dict[tuple[str, ...], list[str]] = {} + unbounded_depth = False + max_depth = 0 + + for sub in subscriptions: + for ch in sub.get("channels", []): + channels.add(ch) + + sub_namespaces = sub.get("namespaces") + if sub_namespaces is None: + wildcard_namespaces = True + elif not wildcard_namespaces: + for ns in sub_namespaces: + namespace_map[tuple(ns)] = ns + + sub_depth = sub.get("depth") + if sub_depth is None: + unbounded_depth = True + elif not unbounded_depth and sub_depth > max_depth: + max_depth = sub_depth + + result: dict[str, Any] = {"channels": sorted(channels)} + if not wildcard_namespaces and namespace_map: + result["namespaces"] = list(namespace_map.values()) + if not unbounded_depth: + result["depth"] = max_depth + return result + + +def filter_covers(coverer: dict[str, Any], target: dict[str, Any]) -> bool: + """Whether `coverer` is a superset of `target`. + + Direct port of `client/stream/index.ts:filterCovers`. Depth coverage + accounts for namespace-prefix offset: a scoped coverer needs enough depth + to absorb the extra levels of any deeper target namespace prefix. + """ + coverer_channels = set(coverer.get("channels", [])) + for ch in target.get("channels", []): + if ch not in coverer_channels: + return False + + coverer_depth = coverer.get("depth") + target_depth = target.get("depth") + coverer_namespaces = coverer.get("namespaces") + target_namespaces = target.get("namespaces") + + # Unscoped coverer covers any namespace; depth is a simple scalar check. + if coverer_namespaces is None: + if coverer_depth is None: + return True + if target_depth is None: + return False + return target_depth <= coverer_depth + + # Scoped coverer cannot cover an unscoped target. + if target_namespaces is None: + return False + + # Each target namespace must be covered by SOME coverer namespace, + # AND the depth-with-offset must fit. + for tp in target_namespaces: + covered = False + for cp in coverer_namespaces: + if not is_prefix_match(tp, cp): + continue + if coverer_depth is None: + covered = True + break + if target_depth is None: + # target wants unbounded depth — coverer bounded can't cover. + continue + if len(tp) - len(cp) + target_depth <= coverer_depth: + covered = True + break + if not covered: + return False + return True diff --git a/libs/sdk-py/tests/streaming/_fake_server.py b/libs/sdk-py/tests/streaming/_fake_server.py index 7c47604bc..b9dce4531 100644 --- a/libs/sdk-py/tests/streaming/_fake_server.py +++ b/libs/sdk-py/tests/streaming/_fake_server.py @@ -35,6 +35,8 @@ class FakeServer: self.stream_request_bodies: list[dict[str, Any]] = [] self._stream_delay: float = 0.0 self._app: Starlette | None = None + self.open_event_streams = 0 + self._open_event_streams_max = 0 def script(self, events: list[dict[str, Any]], *, delay: float = 0.0) -> None: """Set the events the next /stream/events call will replay.""" @@ -79,11 +81,23 @@ class FakeServer: ) async def _sse_body(self) -> AsyncIterator[bytes]: - # Why: script() rebinds scripted_events; in-flight iterators retain a - # reference to the prior list and are unaffected by later script() calls. - for event in self.scripted_events: - if self._stream_delay: - await asyncio.sleep(self._stream_delay) - payload = orjson.dumps(event).decode() - yield f"id: {event.get('event_id', '')}\n".encode() - yield f"event: message\ndata: {payload}\n\n".encode() + self.open_event_streams += 1 + self._open_event_streams_max = max( + self._open_event_streams_max, self.open_event_streams + ) + try: + # Why: script() rebinds scripted_events; in-flight iterators retain + # a reference to the prior list and are unaffected by later + # script() calls. + for event in self.scripted_events: + if self._stream_delay: + await asyncio.sleep(self._stream_delay) + payload = orjson.dumps(event).decode() + yield f"id: {event.get('event_id', '')}\n".encode() + yield f"event: message\ndata: {payload}\n\n".encode() + finally: + self.open_event_streams -= 1 + + @property + def peak_open_event_streams(self) -> int: + return self._open_event_streams_max diff --git a/libs/sdk-py/tests/streaming/test_controller.py b/libs/sdk-py/tests/streaming/test_controller.py new file mode 100644 index 000000000..35e8c3c84 --- /dev/null +++ b/libs/sdk-py/tests/streaming/test_controller.py @@ -0,0 +1,154 @@ +"""Tests for StreamController, _SeenEventIds, and related stream/controller.py types.""" + +from __future__ import annotations + +import pytest + +from langgraph_sdk.stream.controller import StreamController, _SeenEventIds + +# --------------------------------------------------------------------------- +# Task 3.1: bounded subscription queues +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_subscription_queue_bounded_by_max_queue_size(): + """`StreamController` must create per-subscription queues bounded by `max_queue_size`.""" + import httpx + + from langgraph_sdk.stream.transport.http import ProtocolSseTransport + + transport = ProtocolSseTransport( + client=httpx.AsyncClient(base_url="http://test"), + thread_id="t-1", + ) + controller = StreamController(transport=transport, max_queue_size=4) + sub = controller._register_subscription({"channels": ["values"]}) + assert sub.queue.maxsize == 4 + + +@pytest.mark.asyncio +async def test_subscription_queue_default_max_queue_size_is_1024(): + """`StreamController` default `max_queue_size` is 1024.""" + import httpx + + from langgraph_sdk.stream.transport.http import ProtocolSseTransport + + transport = ProtocolSseTransport( + client=httpx.AsyncClient(base_url="http://test"), + thread_id="t-1", + ) + controller = StreamController(transport=transport) + sub = controller._register_subscription({"channels": ["values"]}) + assert sub.queue.maxsize == 1024 + + +# --------------------------------------------------------------------------- +# Task 3.2: bounded LRU seen-event-ids +# --------------------------------------------------------------------------- + + +def test_seen_event_ids_is_bounded_lru(): + """`_SeenEventIds` must evict oldest entries when capacity is exceeded. + + Default cap is 10_000; explicit kwarg overrides. + """ + seen = _SeenEventIds(maxsize=3) + seen.add("a") + seen.add("b") + seen.add("c") + assert "a" in seen + seen.add("d") + assert "a" not in seen # evicted + assert {"b", "c", "d"} <= set(seen) + + +def test_seen_event_ids_move_to_end_on_re_add(): + """`_SeenEventIds.add` of an existing key must promote it (LRU move-to-end).""" + seen = _SeenEventIds(maxsize=3) + seen.add("a") + seen.add("b") + seen.add("c") + # Re-adding "a" should promote it so "b" is evicted next. + seen.add("a") + seen.add("d") + assert "b" not in seen # "b" was the oldest, "a" was promoted + assert "a" in seen + + +def test_seen_event_ids_default_maxsize_is_10000(): + """Default `_SeenEventIds` max is 10_000.""" + seen = _SeenEventIds() + # Add 10_000 + 1 entries. + for i in range(10_001): + seen.add(str(i)) + # "0" (the first added) should have been evicted. + assert "0" not in seen + assert "10000" in seen + + +def test_seen_event_ids_contains_false_for_missing(): + seen = _SeenEventIds(maxsize=10) + assert "missing" not in seen + + +def test_seen_event_ids_iter_returns_keys(): + seen = _SeenEventIds(maxsize=10) + seen.add("x") + seen.add("y") + assert set(seen) == {"x", "y"} + + +# --------------------------------------------------------------------------- +# Task 3.3: close() awaits pending rotation closes +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_close_awaits_pending_rotation_closes(): + """When a rotation is mid-flight, controller.close() must await the old + stream close before returning.""" + import asyncio as _asyncio + + import httpx + + from langgraph_sdk.stream.controller import _close_after + from langgraph_sdk.stream.transport.http import ProtocolSseTransport + + rotation_close_done = _asyncio.Event() + + class _SlowHandle: + """A fake EventStreamHandle whose close() takes a moment.""" + + def __init__(self): + self.events = self._empty() + loop = _asyncio.get_running_loop() + self.ready: _asyncio.Future[None] = loop.create_future() + self.ready.set_result(None) + self.done: _asyncio.Future[None] = loop.create_future() + + async def _empty(self): + if False: + yield # pragma: no cover + + async def close(self): + await _asyncio.sleep(0.05) + rotation_close_done.set() + + transport = ProtocolSseTransport( + client=httpx.AsyncClient(base_url="http://test"), + thread_id="t-1", + ) + controller = StreamController(transport=transport) + + # Simulate a mid-flight rotation close by directly injecting a task. + slow_handle = _SlowHandle() + task = _asyncio.create_task( + _close_after(slow_handle) # ty: ignore[invalid-argument-type] + ) + controller._rotation_close_tasks.add(task) + task.add_done_callback(controller._rotation_close_tasks.discard) + + # close() must block until the rotation close completes. + await controller.close() + assert rotation_close_done.is_set() diff --git a/libs/sdk-py/tests/streaming/test_lifecycle_watcher.py b/libs/sdk-py/tests/streaming/test_lifecycle_watcher.py new file mode 100644 index 000000000..0422e27c6 --- /dev/null +++ b/libs/sdk-py/tests/streaming/test_lifecycle_watcher.py @@ -0,0 +1,38 @@ +"""Tests for the lifecycle watcher: `interrupted` / `interrupts` state.""" + +from __future__ import annotations + +import asyncio + +import httpx + +from langgraph_sdk._async.http import HttpClient +from langgraph_sdk._async.threads import ThreadsClient +from streaming._events import input_requested_event +from streaming._fake_server import FakeServer + + +async def test_interrupted_starts_false(): + async with httpx.AsyncClient(base_url="http://test") as raw: + threads = ThreadsClient(HttpClient(raw)) + async with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + assert thread.interrupted is False + assert thread.interrupts == [] + + +async def test_interrupts_populated_from_input_requested_event(): + fake = FakeServer() + fake.script([input_requested_event(seq=0)]) + asgi = httpx.ASGITransport(app=fake.app) + async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw: + threads = ThreadsClient(HttpClient(raw)) + async with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + await thread.run.start(input={}) + # Lifecycle watcher consumes asynchronously — poll briefly. + for _ in range(20): + if thread.interrupted: + break + await asyncio.sleep(0.05) + assert thread.interrupted is True + assert len(thread.interrupts) == 1 + assert thread.interrupts[0]["interrupt_id"] == "i-1" diff --git a/libs/sdk-py/tests/streaming/test_shared_stream.py b/libs/sdk-py/tests/streaming/test_shared_stream.py new file mode 100644 index 000000000..f23fe4eb6 --- /dev/null +++ b/libs/sdk-py/tests/streaming/test_shared_stream.py @@ -0,0 +1,150 @@ +from __future__ import annotations + +import asyncio + +import httpx + +from langgraph_sdk._async.http import HttpClient +from langgraph_sdk._async.threads import ThreadsClient +from streaming._events import lifecycle_event, values_event +from streaming._fake_server import FakeServer + + +async def test_shared_stream_serves_single_subscription(): + fake = FakeServer() + fake.script([lifecycle_event(seq=0), values_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={}) + await thread._reconcile_stream({"channels": ["lifecycle", "values"]}) + assert thread._shared_stream is not None + received = [ + e async for e in thread._dedup_iter(thread._shared_stream.events) + ] + methods = [e["method"] for e in received] + assert methods == ["lifecycle", "values"] + assert fake.peak_open_event_streams == 1 + + +async def test_seen_event_ids_dedupes_replayed_events(): + fake = FakeServer() + fake.script( + [ + lifecycle_event(seq=0), + lifecycle_event(seq=0), # duplicate event_id + values_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={}) + await thread._reconcile_stream({"channels": ["lifecycle", "values"]}) + assert thread._shared_stream is not None + received = [ + e async for e in thread._dedup_iter(thread._shared_stream.events) + ] + seqs = [e["seq"] for e in received] + assert seqs == [0, 1] # the duplicate seq=0 was deduped via event_id + + +async def test_rotation_when_new_subscription_widens_filter(): + fake = FakeServer() + fake.script([lifecycle_event(seq=0), values_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={}) + # First subscription: lifecycle only. + await thread._reconcile_stream({"channels": ["lifecycle"]}) + # Second subscription widens to lifecycle + values. + await thread._reconcile_stream({"channels": ["lifecycle", "values"]}) + # Rotation: two separate SSE requests were opened (old + new). + assert len(fake.stream_request_bodies) >= 2 + + +async def test_no_rotation_when_existing_filter_covers_new_subscription(): + fake = FakeServer() + fake.script([]) + 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={}) + await thread._reconcile_stream({"channels": ["lifecycle", "values"]}) + # New subscription is a subset — existing filter covers it. + await thread._reconcile_stream({"channels": ["values"]}) + # No rotation in the shared stream (1 shared SSE) plus 1 lifecycle watcher SSE = 2. + assert len(fake.stream_request_bodies) == 2 + + +async def test_subscribe_yields_only_matching_events(): + fake = FakeServer() + fake.script( + [ + lifecycle_event(seq=0), + values_event(seq=1), + lifecycle_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={}) + + async def drain(channels): + return [e async for e in thread.subscribe(channels)] + + lifecycle_events, values_events = await asyncio.gather( + drain(["lifecycle"]), + drain(["values"]), + ) + assert [e["seq"] for e in lifecycle_events] == [0, 2] + assert [e["seq"] for e in values_events] == [1] + + +async def test_two_concurrent_subscribes_share_one_stream(): + fake = FakeServer() + fake.script([lifecycle_event(seq=i) for i in range(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={}) + + async def drain(channels): + return [e async for e in thread.subscribe(channels)] + + results = await asyncio.gather( + drain(["lifecycle"]), + drain(["lifecycle"]), + ) + assert len(results[0]) == 5 + assert len(results[1]) == 5 + # Both subscriptions share one SSE (no rotation) plus 1 lifecycle watcher SSE = 2. + assert len(fake.stream_request_bodies) == 2 + + +async def test_subscribe_does_not_leak_when_iterator_unconsumed(): + """Subscriptions register lazily on first __anext__, not at subscribe() call time. + + Why: registering eagerly would leak the subscription if the caller + constructs the iterator but never iterates it. The lazy pattern ties + registration to the generator's lifecycle, which is bounded by aclose() + / exhaustion / cancellation. + """ + fake = FakeServer() + fake.script([]) + 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={}) + _ = thread.subscribe(["lifecycle"]) # construct but never iterate + # Subscription is not registered yet — the generator body hasn't run. + assert len(thread._subscriptions) == 0 diff --git a/libs/sdk-py/tests/streaming/test_subscription.py b/libs/sdk-py/tests/streaming/test_subscription.py index 066348fea..9e8bc92ae 100644 --- a/libs/sdk-py/tests/streaming/test_subscription.py +++ b/libs/sdk-py/tests/streaming/test_subscription.py @@ -3,6 +3,8 @@ from __future__ import annotations import pytest from langgraph_sdk.stream.subscription import ( + compute_union_filter, + filter_covers, infer_channel, is_prefix_match, matches_subscription, @@ -125,3 +127,112 @@ def test_matches_subscription_namespace_filter_applied(): def test_matches_subscription_bare_custom_event_matches_bare_custom_filter(): sub = {"channels": ["custom"]} assert matches_subscription(custom_event(name=""), sub) is True # ty: ignore[invalid-argument-type] + + +def test_compute_union_filter_merges_channels(): + a = {"channels": ["values"]} + b = {"channels": ["messages", "lifecycle"]} + result = compute_union_filter([a, b]) + assert set(result["channels"]) == {"values", "messages", "lifecycle"} + + +def test_compute_union_filter_drops_namespaces_when_any_subscription_is_unscoped(): + a = {"channels": ["values"], "namespaces": [["fetcher"]]} + b = {"channels": ["messages"]} # no namespaces == wildcard + result = compute_union_filter([a, b]) + # If any subscription is unscoped, the union must be unscoped. + assert "namespaces" not in result or result.get("namespaces") is None + + +def test_compute_union_filter_unions_namespaces_when_all_scoped(): + a = {"channels": ["values"], "namespaces": [["fetcher"]]} + b = {"channels": ["messages"], "namespaces": [["scorer"]]} + result = compute_union_filter([a, b]) + assert sorted(result["namespaces"]) == [["fetcher"], ["scorer"]] + + +def test_compute_union_filter_takes_max_depth(): + a = {"channels": ["values"], "depth": 1} + b = {"channels": ["messages"], "depth": 3} + result = compute_union_filter([a, b]) + assert result["depth"] == 3 + + +def test_compute_union_filter_omits_depth_when_any_subscription_omits(): + a = {"channels": ["values"], "depth": 2} + b = {"channels": ["messages"]} # no depth == unbounded + result = compute_union_filter([a, b]) + assert "depth" not in result or result.get("depth") is None + + +def test_compute_union_filter_empty_input_returns_empty_channel_filter(): + result = compute_union_filter([]) + assert result == {"channels": []} + + +def test_filter_covers_same_filter(): + f = {"channels": ["values", "messages"]} + assert filter_covers(f, f) is True + + +def test_filter_covers_superset_channels(): + coverer = {"channels": ["values", "messages", "lifecycle"]} + target = {"channels": ["values", "messages"]} + assert filter_covers(coverer, target) is True + + +def test_filter_covers_missing_channel(): + coverer = {"channels": ["values"]} + target = {"channels": ["values", "messages"]} + assert filter_covers(coverer, target) is False + + +def test_filter_covers_unscoped_covers_scoped(): + coverer = {"channels": ["values"]} # wildcard namespaces + target = {"channels": ["values"], "namespaces": [["fetcher"]]} + assert filter_covers(coverer, target) is True + + +def test_filter_covers_scoped_does_not_cover_unscoped(): + coverer = {"channels": ["values"], "namespaces": [["fetcher"]]} + target = {"channels": ["values"]} # wildcard + assert filter_covers(coverer, target) is False + + +def test_filter_covers_depth_with_namespace_offset(): + # Coverer at depth 1 from ["agent"] reaches ["agent", X]. + # Target needs depth 1 from ["agent", "tool"] — i.e., ["agent", "tool", X]. + # That's 2 levels past coverer's prefix, but coverer only allows 1. + coverer = {"channels": ["values"], "namespaces": [["agent"]], "depth": 1} + target = { + "channels": ["values"], + "namespaces": [["agent", "tool"]], + "depth": 1, + } + assert filter_covers(coverer, target) is False + + +def test_filter_covers_depth_with_offset_enough_depth(): + # Same setup but coverer depth=2 absorbs the offset. + coverer = {"channels": ["values"], "namespaces": [["agent"]], "depth": 2} + target = { + "channels": ["values"], + "namespaces": [["agent", "tool"]], + "depth": 1, + } + assert filter_covers(coverer, target) is True + + +def test_filter_covers_unscoped_coverer_with_bounded_depth(): + # Coverer is unscoped; depth comparison is the simple scalar form. + coverer = {"channels": ["values"], "depth": 2} + target = {"channels": ["values"], "depth": 1} + assert filter_covers(coverer, target) is True + target_too_deep = {"channels": ["values"], "depth": 3} + assert filter_covers(coverer, target_too_deep) is False + + +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 diff --git a/libs/sdk-py/tests/streaming/test_thread_stream.py b/libs/sdk-py/tests/streaming/test_thread_stream.py index 16eca70e2..ff3009a55 100644 --- a/libs/sdk-py/tests/streaming/test_thread_stream.py +++ b/libs/sdk-py/tests/streaming/test_thread_stream.py @@ -354,3 +354,36 @@ async def test_fresh_thread_happy_path_end_to_end(): 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) + + +async def test_aenter_raises_after_close(): + async with httpx.AsyncClient(base_url="http://test") as raw: + stream = AsyncThreadStream(client=raw, thread_id="t-1", assistant_id="agent") + async with stream: + pass + # After exit, the stream is closed; re-entering must raise rather than + # silently constructing a new transport that would leak on the next exit. + with pytest.raises(RuntimeError, match="closed and cannot be re-entered"): + async with stream: + pass + + +async def test_register_subscription_assigns_monotonic_ids(): + async with httpx.AsyncClient(base_url="http://test") as raw: + stream = AsyncThreadStream(client=raw, thread_id="t-1", assistant_id="agent") + async with stream: + sub_a = stream._register_subscription({"channels": ["values"]}) + sub_b = stream._register_subscription({"channels": ["messages"]}) + assert sub_a.id == 1 + assert sub_b.id == 2 + assert stream._subscriptions[sub_a.id] is sub_a + assert stream._subscriptions[sub_b.id] is sub_b + + +async def test_unregister_subscription_removes_from_registry(): + async with httpx.AsyncClient(base_url="http://test") as raw: + stream = AsyncThreadStream(client=raw, thread_id="t-1", assistant_id="agent") + async with stream: + sub = stream._register_subscription({"channels": ["values"]}) + stream._unregister_subscription(sub.id) + assert sub.id not in stream._subscriptions