mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-17 21:25:46 +02:00
feat(sdk-py): add shared stream subscriptions (#7820)
This commit is contained in:
Nick Hollon
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
@@ -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"
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user