mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-17 21:25:46 +02:00
feat(sdk-py): add sync thread stream core (#7826)
This commit is contained in:
Nick Hollon
@@ -1182,6 +1182,7 @@ class AsyncThreadStream:
|
||||
self._run_start_ready: asyncio.Future[None] | None = None
|
||||
self._run_seen: bool = False
|
||||
self._run_done: asyncio.Future[_RunTerminal] | None = None
|
||||
self._cursor: int | None = None
|
||||
self._active_message_streams: set[AsyncChatModelStream] = set()
|
||||
self._active_tool_calls: set[ToolCallHandle] = set()
|
||||
# Root-scope inbox: populated by `_SubgraphsProjection` when it consumes
|
||||
@@ -1320,6 +1321,16 @@ class AsyncThreadStream:
|
||||
handle._fail(err)
|
||||
self._active_tool_calls.clear()
|
||||
|
||||
def observe_applied_through_seq(self, seq: Any) -> None:
|
||||
"""Advance the reconnect cursor from a command response meta sequence."""
|
||||
if isinstance(seq, int) and (self._cursor is None or seq > self._cursor):
|
||||
self._cursor = seq
|
||||
|
||||
def _observe_event(self, event: Event) -> None:
|
||||
seq = event.get("seq")
|
||||
if isinstance(seq, int) and (self._cursor is None or seq > self._cursor):
|
||||
self._cursor = seq
|
||||
|
||||
def subscribe(
|
||||
self,
|
||||
channels: list[str],
|
||||
@@ -1381,6 +1392,7 @@ class AsyncThreadStream:
|
||||
async for event in self._dedup_iter(shared.events):
|
||||
if self._closed:
|
||||
break
|
||||
self._observe_event(event)
|
||||
for sub in list(self._subscriptions.values()):
|
||||
if matches_subscription(event, sub.params):
|
||||
sub.queue.put_nowait(event)
|
||||
@@ -1423,7 +1435,10 @@ class AsyncThreadStream:
|
||||
return # Existing stream is sufficient.
|
||||
|
||||
new_filter = self._compute_current_union(extra=candidate_filter)
|
||||
new_stream = self._transport.open_event_stream(new_filter)
|
||||
stream_params: dict[str, Any] = dict(new_filter)
|
||||
if self._cursor is not None:
|
||||
stream_params["since"] = self._cursor
|
||||
new_stream = self._transport.open_event_stream(stream_params)
|
||||
old_stream = self._shared_stream
|
||||
self._shared_stream = new_stream
|
||||
self._shared_stream_filter = new_filter
|
||||
@@ -1480,6 +1495,11 @@ class AsyncThreadStream:
|
||||
code = response.get("error", "unknown")
|
||||
message = response.get("message", "")
|
||||
raise RuntimeError(f"Protocol error [{code}]: {message}")
|
||||
meta = response.get("meta")
|
||||
if isinstance(meta, dict):
|
||||
applied_through_seq = meta.get("applied_through_seq")
|
||||
if self._controller is not None:
|
||||
self._controller.observe_applied_through_seq(applied_through_seq)
|
||||
return response.get("result", {})
|
||||
|
||||
async def _await_run_start_gate(self, *, timeout: float | None = None) -> None:
|
||||
|
||||
@@ -0,0 +1,503 @@
|
||||
"""Synchronous thread-centric streaming surface for the v3 protocol.
|
||||
|
||||
`SyncThreadStream` is a synchronous context manager that owns a
|
||||
`SyncProtocolSseTransport` for one thread, dispatches commands (`run.start`,
|
||||
`run.respond`), exposes subscriptions over a single shared SSE, and surfaces
|
||||
lifecycle state (`interrupted`, `interrupts`) via an always-on lifecycle watcher
|
||||
thread.
|
||||
|
||||
Sync mirror of `libs/sdk-py/langgraph_sdk/_async/stream.py`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import threading
|
||||
from collections.abc import Iterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Literal, TypedDict
|
||||
|
||||
from langchain_protocol import Event, SubscribeParams
|
||||
|
||||
from langgraph_sdk._sync.http import SyncHttpClient
|
||||
from langgraph_sdk.stream.sync_controller import SyncStreamController, _SyncSubscription
|
||||
from langgraph_sdk.stream.transport.sync_http import (
|
||||
SyncEventStreamHandle,
|
||||
SyncProtocolSseTransport,
|
||||
)
|
||||
|
||||
|
||||
class InterruptPayload(TypedDict):
|
||||
"""Payload surfaced when the server requests human input for a thread."""
|
||||
|
||||
interrupt_id: str
|
||||
value: Any
|
||||
namespace: list[str]
|
||||
|
||||
|
||||
@dataclass
|
||||
class _RunTerminal:
|
||||
"""Terminal state record resolved into `_run_done` on lifecycle completion."""
|
||||
|
||||
status: Literal["completed", "errored"]
|
||||
error: BaseException | None = None
|
||||
|
||||
|
||||
_ALL_CHANNELS: list[str] = [
|
||||
"values",
|
||||
"updates",
|
||||
"messages",
|
||||
"tools",
|
||||
"lifecycle",
|
||||
"input",
|
||||
"checkpoints",
|
||||
"tasks",
|
||||
"custom",
|
||||
]
|
||||
|
||||
|
||||
class _BlockingResult:
|
||||
def __init__(self) -> None:
|
||||
self._event = threading.Event()
|
||||
self._value: Any = None
|
||||
self._error: BaseException | None = None
|
||||
|
||||
def set_result(self, value: Any) -> None:
|
||||
if self._event.is_set():
|
||||
return
|
||||
self._value = value
|
||||
self._event.set()
|
||||
|
||||
def set_exception(self, error: BaseException) -> None:
|
||||
if self._event.is_set():
|
||||
return
|
||||
self._error = error
|
||||
self._event.set()
|
||||
|
||||
def result(self, timeout: float | None = None) -> Any:
|
||||
if not self._event.wait(timeout):
|
||||
raise TimeoutError("Result was not set before timeout.")
|
||||
if self._error is not None:
|
||||
raise self._error
|
||||
return self._value
|
||||
|
||||
def done(self) -> bool:
|
||||
return self._event.is_set()
|
||||
|
||||
|
||||
class SyncRunModule:
|
||||
"""Command dispatcher for `run.start`.
|
||||
|
||||
Bound to one `SyncThreadStream`; accesses its transport and id allocator.
|
||||
"""
|
||||
|
||||
def __init__(self, owner: SyncThreadStream) -> None:
|
||||
self._owner = owner
|
||||
|
||||
def start(
|
||||
self,
|
||||
*,
|
||||
input: Any = None,
|
||||
config: dict[str, Any] | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Send `run.start` to the server. Returns the result (`{"run_id": ...}`)."""
|
||||
params: dict[str, Any] = {"assistant_id": self._owner.assistant_id}
|
||||
if input is not None:
|
||||
params["input"] = input
|
||||
if config is not None:
|
||||
params["config"] = config
|
||||
if metadata is not None:
|
||||
params["metadata"] = metadata
|
||||
result = self._owner._send_command("run.start", params)
|
||||
self._owner._run_seen = True
|
||||
controller = self._owner._controller
|
||||
if controller is not None and controller._run_start_gate is not None:
|
||||
controller._run_start_gate.set()
|
||||
return result
|
||||
|
||||
def respond(
|
||||
self,
|
||||
response: Any,
|
||||
*,
|
||||
interrupt_id: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Reply to a server-side interrupt and resume the run.
|
||||
|
||||
Args:
|
||||
response: the response value forwarded as `params.response` on the wire.
|
||||
interrupt_id: optional explicit id. When omitted, requires exactly one
|
||||
outstanding interrupt.
|
||||
|
||||
Raises:
|
||||
RuntimeError: no outstanding interrupts; `interrupt_id` is None but
|
||||
multiple interrupts are outstanding; or the explicit `interrupt_id`
|
||||
doesn't match any outstanding interrupt.
|
||||
"""
|
||||
outstanding = self._owner.interrupts
|
||||
if interrupt_id is None:
|
||||
if len(outstanding) == 0:
|
||||
raise RuntimeError(
|
||||
"thread.run.respond: no outstanding interrupt. Provide an "
|
||||
"explicit `interrupt_id` or wait for `thread.interrupted`."
|
||||
)
|
||||
if len(outstanding) > 1:
|
||||
ids = [p["interrupt_id"] for p in outstanding]
|
||||
raise RuntimeError(
|
||||
f"thread.run.respond: ambiguous — {len(outstanding)} "
|
||||
f"outstanding interrupts ({ids!r}). Provide an explicit "
|
||||
"`interrupt_id`."
|
||||
)
|
||||
match = outstanding[0]
|
||||
else:
|
||||
match = next(
|
||||
(p for p in outstanding if p["interrupt_id"] == interrupt_id),
|
||||
None,
|
||||
)
|
||||
if match is None:
|
||||
raise RuntimeError(
|
||||
f"thread.run.respond: interrupt_id {interrupt_id!r} does not "
|
||||
"match any outstanding interrupt in `thread.interrupts`."
|
||||
)
|
||||
params = {
|
||||
"interrupt_id": match["interrupt_id"],
|
||||
"namespace": match["namespace"],
|
||||
"response": response,
|
||||
}
|
||||
return self._owner._send_command("input.respond", params)
|
||||
|
||||
|
||||
class _SyncValuesProjection:
|
||||
"""Typed projection for `thread.values`.
|
||||
|
||||
Supports `for snapshot in thread.values` (REST snapshot then live stream
|
||||
events) and `thread.values.get()` (delegates to `thread.output`).
|
||||
"""
|
||||
|
||||
def __init__(self, thread: SyncThreadStream) -> None:
|
||||
self._thread = thread
|
||||
|
||||
def __iter__(self) -> Iterator[Any]:
|
||||
"""Iterate over state snapshots: REST state first, then live values events."""
|
||||
return self._values_iter()
|
||||
|
||||
def get(self) -> Any:
|
||||
"""Return terminal state values; equivalent to `thread.output`."""
|
||||
return self._thread.output
|
||||
|
||||
def _values_iter(self) -> Iterator[Any]:
|
||||
if self._thread._transport is None:
|
||||
raise RuntimeError("SyncThreadStream not entered — use `with`.")
|
||||
params: SubscribeParams = {"channels": ["values"]}
|
||||
sub = self._thread._register_subscription(params)
|
||||
try:
|
||||
self._thread._reconcile_stream(params)
|
||||
self._thread._ensure_fanout_running()
|
||||
state = self._thread._fetch_state()
|
||||
yield state["values"]
|
||||
while True:
|
||||
item = sub.queue.get()
|
||||
if item is None:
|
||||
return
|
||||
params_field = item.get("params") or {}
|
||||
data = (
|
||||
params_field.get("data") if isinstance(params_field, dict) else None
|
||||
)
|
||||
if data is not None:
|
||||
yield data
|
||||
finally:
|
||||
self._thread._unregister_subscription(sub.id)
|
||||
|
||||
|
||||
class SyncThreadStream:
|
||||
"""Synchronous context manager for one thread's v3 streaming session.
|
||||
|
||||
Construct via `client.threads.stream(thread_id=None, *, assistant_id, ...)`
|
||||
rather than instantiating directly.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
http: SyncHttpClient,
|
||||
thread_id: str,
|
||||
assistant_id: str,
|
||||
headers: Mapping[str, str] | None = None,
|
||||
run_start_timeout: float | None = None,
|
||||
explicit_thread_id: bool = False,
|
||||
) -> None:
|
||||
self._http = http
|
||||
self._headers = dict(headers or {})
|
||||
self.thread_id = thread_id
|
||||
self.assistant_id = assistant_id
|
||||
self._run_start_timeout = run_start_timeout
|
||||
self._explicit_thread_id = explicit_thread_id
|
||||
self._closed = False
|
||||
self._transport: SyncProtocolSseTransport | None = None
|
||||
self._controller: SyncStreamController | None = None
|
||||
self._command_id_lock = threading.Lock()
|
||||
self._next_command_id = 1
|
||||
self.interrupted: bool = False
|
||||
self.interrupts: list[InterruptPayload] = []
|
||||
self._lifecycle_watcher_thread: threading.Thread | None = None
|
||||
self._lifecycle_watcher_handle: SyncEventStreamHandle | None = None
|
||||
self._run_seen: bool = False
|
||||
self._run_done: _BlockingResult | None = None
|
||||
self.run = SyncRunModule(self)
|
||||
self.values = _SyncValuesProjection(self)
|
||||
|
||||
def __enter__(self) -> SyncThreadStream:
|
||||
if self._closed:
|
||||
raise RuntimeError("SyncThreadStream is closed and cannot be re-entered.")
|
||||
self._transport = SyncProtocolSseTransport(
|
||||
client=self._http.client,
|
||||
thread_id=self.thread_id,
|
||||
headers=self._headers,
|
||||
)
|
||||
# Gate is unset; SyncRunModule.start (or an explicit set) clears it so
|
||||
# that subscriptions opening before run.start block until the server
|
||||
# has accepted the run command.
|
||||
run_start_gate = threading.Event()
|
||||
self._controller = SyncStreamController(
|
||||
self._transport, run_start_gate=run_start_gate
|
||||
)
|
||||
self._run_done = _BlockingResult()
|
||||
self._ensure_lifecycle_watcher_running()
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
|
||||
self.close()
|
||||
|
||||
@property
|
||||
def output(self) -> Any:
|
||||
"""Fetch terminal thread state (blocking); waits for lifecycle completion.
|
||||
|
||||
Raises:
|
||||
RuntimeError: stream not entered, or no run started and no explicit
|
||||
thread_id was provided.
|
||||
"""
|
||||
if self._can_return_existing_state_immediately():
|
||||
state = self._fetch_state()
|
||||
if self._state_is_terminal(state):
|
||||
return state["values"]
|
||||
terminal = self._wait_for_run_done()
|
||||
if terminal.error is not None:
|
||||
raise terminal.error
|
||||
state = self._fetch_state()
|
||||
return state["values"]
|
||||
|
||||
@property
|
||||
def events(self) -> Iterator[Event]:
|
||||
"""Raw iterator of every `Event` the server emits for this thread."""
|
||||
if self._transport is None:
|
||||
raise RuntimeError("SyncThreadStream not entered — use `with`.")
|
||||
return self.subscribe(_ALL_CHANNELS)
|
||||
|
||||
def close(self) -> None:
|
||||
"""Tear down the thread stream. Idempotent."""
|
||||
if self._closed:
|
||||
return
|
||||
self._closed = True
|
||||
run_done = self._run_done
|
||||
if run_done is not None and not run_done.done():
|
||||
run_done.set_exception(RuntimeError("SyncThreadStream closed"))
|
||||
# Close the controller first so active subscription iterators receive
|
||||
# their None sentinel immediately, before the lifecycle watcher join
|
||||
# (which may block for up to 1s).
|
||||
if self._controller is not None:
|
||||
self._controller.close()
|
||||
handle = self._lifecycle_watcher_handle
|
||||
if handle is not None:
|
||||
with contextlib.suppress(Exception):
|
||||
handle.close()
|
||||
thread = self._lifecycle_watcher_thread
|
||||
if thread is not None and thread.is_alive():
|
||||
with contextlib.suppress(RuntimeError):
|
||||
thread.join(timeout=1.0)
|
||||
if self._transport is not None:
|
||||
self._transport.close()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Delegation to SyncStreamController
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _register_subscription(self, params: SubscribeParams) -> _SyncSubscription:
|
||||
if self._controller is None:
|
||||
raise RuntimeError("SyncThreadStream not entered — use `with`.")
|
||||
return self._controller.register_subscription(params)
|
||||
|
||||
def _unregister_subscription(self, subscription_id: int) -> None:
|
||||
if self._controller is not None:
|
||||
self._controller.unregister_subscription(subscription_id)
|
||||
|
||||
def _ensure_fanout_running(self) -> None:
|
||||
if self._controller is not None:
|
||||
self._controller.ensure_fanout_running()
|
||||
|
||||
def _reconcile_stream(self, candidate_filter: SubscribeParams) -> None:
|
||||
if self._controller is None:
|
||||
raise RuntimeError("SyncThreadStream not entered — use `with`.")
|
||||
self._controller.reconcile_stream(candidate_filter)
|
||||
|
||||
def subscribe(
|
||||
self,
|
||||
channels: list[str],
|
||||
*,
|
||||
namespaces: list[list[str]] | None = None,
|
||||
depth: int | None = None,
|
||||
) -> Iterator[Event]:
|
||||
"""Open a typed subscription against the shared SSE.
|
||||
|
||||
Returns an 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("SyncThreadStream not entered — use `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)
|
||||
|
||||
def _subscription_iter(self, params: SubscribeParams) -> Iterator[Event]:
|
||||
sub = self._register_subscription(params)
|
||||
try:
|
||||
if self._closed:
|
||||
return
|
||||
self._reconcile_stream(params)
|
||||
self._ensure_fanout_running()
|
||||
while True:
|
||||
item = sub.queue.get()
|
||||
if item is None:
|
||||
return
|
||||
yield item
|
||||
finally:
|
||||
self._unregister_subscription(sub.id)
|
||||
|
||||
def _send_command(self, method: str, params: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Send a protocol command and return the `result` payload."""
|
||||
if self._transport is None:
|
||||
raise RuntimeError("SyncThreadStream not entered — use `with`.")
|
||||
with self._command_id_lock:
|
||||
command_id = self._next_command_id
|
||||
self._next_command_id += 1
|
||||
response = self._transport.send_command(
|
||||
{"id": command_id, "method": method, "params": params}
|
||||
)
|
||||
if response is None:
|
||||
return {}
|
||||
if response.get("type") == "error":
|
||||
code = response.get("error", "unknown")
|
||||
message = response.get("message", "")
|
||||
raise RuntimeError(f"Protocol error [{code}]: {message}")
|
||||
meta = response.get("meta")
|
||||
if isinstance(meta, dict):
|
||||
applied_through_seq = meta.get("applied_through_seq")
|
||||
if self._controller is not None:
|
||||
self._controller.observe_applied_through_seq(applied_through_seq)
|
||||
return response.get("result", {})
|
||||
|
||||
def _ensure_lifecycle_watcher_running(self) -> None:
|
||||
if self._lifecycle_watcher_thread is not None:
|
||||
return
|
||||
self._lifecycle_watcher_thread = threading.Thread(
|
||||
target=self._run_lifecycle_watcher,
|
||||
name="langgraph-sdk-sync-lifecycle",
|
||||
daemon=True,
|
||||
)
|
||||
self._lifecycle_watcher_thread.start()
|
||||
|
||||
def _run_lifecycle_watcher(self) -> None:
|
||||
"""Always-on thread consuming lifecycle + input channels."""
|
||||
if self._transport is None:
|
||||
return
|
||||
try:
|
||||
handle = self._transport.open_event_stream(
|
||||
{"channels": ["lifecycle", "input"]}
|
||||
)
|
||||
self._lifecycle_watcher_handle = handle
|
||||
for event in handle.events:
|
||||
if self._closed:
|
||||
return
|
||||
self._apply_lifecycle_event(event)
|
||||
except Exception as exc:
|
||||
run_done = self._run_done
|
||||
if run_done is not None and not run_done.done():
|
||||
run_done.set_result(
|
||||
_RunTerminal(
|
||||
status="errored",
|
||||
error=RuntimeError(f"Lifecycle transport failed: {exc}"),
|
||||
)
|
||||
)
|
||||
|
||||
def _fetch_state(self) -> dict[str, Any]:
|
||||
"""Fetch the current thread state from the REST endpoint."""
|
||||
return self._http.get(
|
||||
f"/threads/{self.thread_id}/state",
|
||||
headers=self._headers or None,
|
||||
)
|
||||
|
||||
def _state_is_terminal(self, state: dict[str, Any]) -> bool:
|
||||
"""Return `True` if the thread state has no pending tasks or next nodes."""
|
||||
return not state.get("next") and not state.get("tasks")
|
||||
|
||||
def _can_return_existing_state_immediately(self) -> bool:
|
||||
"""Return `True` if we can try the REST state before waiting on the lifecycle."""
|
||||
return self._explicit_thread_id and not self._run_seen
|
||||
|
||||
def _wait_for_run_done(self) -> _RunTerminal:
|
||||
"""Block until lifecycle completion.
|
||||
|
||||
Raises:
|
||||
RuntimeError: stream not entered, or no run started and no explicit
|
||||
thread_id was provided.
|
||||
"""
|
||||
if self._run_done is None:
|
||||
raise RuntimeError("SyncThreadStream not entered — use `with`.")
|
||||
if not self._run_seen and not self._explicit_thread_id:
|
||||
raise RuntimeError(
|
||||
"thread.output: no run has been started and no explicit thread_id "
|
||||
"was provided. Call thread.run.start() first."
|
||||
)
|
||||
return self._run_done.result()
|
||||
|
||||
def _apply_lifecycle_event(self, event: Event) -> None:
|
||||
"""Update `interrupted` / `interrupts` / `_run_done` 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 ("started", "running"):
|
||||
self._run_seen = True
|
||||
elif phase in ("completed", "errored"):
|
||||
self.interrupted = False
|
||||
self.interrupts = []
|
||||
run_done = self._run_done
|
||||
if run_done is not None and not run_done.done():
|
||||
if phase == "errored":
|
||||
error_msg = (
|
||||
data.get("error") if isinstance(data, dict) else None
|
||||
)
|
||||
error = RuntimeError(
|
||||
f"Run errored: {error_msg}" if error_msg else "Run errored"
|
||||
)
|
||||
run_done.set_result(_RunTerminal(status="errored", error=error))
|
||||
else:
|
||||
run_done.set_result(_RunTerminal(status="completed"))
|
||||
@@ -2,11 +2,13 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from typing import Any, Literal, overload
|
||||
|
||||
from langgraph_sdk._shared.utilities import _quote_path_param
|
||||
from langgraph_sdk._sync.http import SyncHttpClient
|
||||
from langgraph_sdk._sync.stream import SyncThreadStream
|
||||
from langgraph_sdk.schema import (
|
||||
Checkpoint,
|
||||
Json,
|
||||
@@ -722,6 +724,35 @@ class SyncThreadsClient:
|
||||
params=params,
|
||||
)
|
||||
|
||||
def stream(
|
||||
self,
|
||||
thread_id: str | None = None,
|
||||
*,
|
||||
assistant_id: str,
|
||||
headers: Mapping[str, str] | None = None,
|
||||
run_start_timeout: float | None = None,
|
||||
) -> SyncThreadStream:
|
||||
"""Open a v3 thread-centric streaming session.
|
||||
|
||||
Args:
|
||||
thread_id: optional explicit thread identifier. Defaults to a
|
||||
fresh UUIDv4.
|
||||
assistant_id: assistant the run will use. Required.
|
||||
headers: optional headers forwarded on every command and SSE
|
||||
request for this stream session.
|
||||
|
||||
Returns:
|
||||
A `SyncThreadStream` to use as a context manager.
|
||||
"""
|
||||
return SyncThreadStream(
|
||||
http=self.http,
|
||||
thread_id=thread_id if thread_id is not None else str(uuid.uuid4()),
|
||||
assistant_id=assistant_id,
|
||||
headers=headers,
|
||||
run_start_timeout=run_start_timeout,
|
||||
explicit_thread_id=thread_id is not None,
|
||||
)
|
||||
|
||||
def join_stream(
|
||||
self,
|
||||
thread_id: str,
|
||||
|
||||
@@ -0,0 +1,269 @@
|
||||
"""Synchronous shared-stream fan-out controller for v3 thread streaming."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import logging
|
||||
import random
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from queue import Queue as _Queue
|
||||
from typing import Any
|
||||
|
||||
from langchain_protocol import Event, SubscribeParams
|
||||
|
||||
from langgraph_sdk.stream.subscription import compute_union_filter, filter_covers
|
||||
from langgraph_sdk.stream.transport.sync_http import (
|
||||
SyncEventStreamHandle,
|
||||
SyncProtocolSseTransport,
|
||||
)
|
||||
|
||||
_logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _SyncSubscription:
|
||||
id: int
|
||||
params: SubscribeParams
|
||||
queue: _Queue[Event | None] = field(default_factory=_Queue)
|
||||
# Why: using `queue.Queue` in the annotation causes ty to resolve `queue`
|
||||
# as the field being defined (name shadowing), not the stdlib module.
|
||||
|
||||
|
||||
_DEFAULT_RUN_START_TIMEOUT: float = 30.0
|
||||
|
||||
|
||||
class SyncStreamController:
|
||||
"""Owns the sync shared SSE handle, subscription registry, and fan-out thread."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
transport: SyncProtocolSseTransport,
|
||||
*,
|
||||
run_start_gate: threading.Event | None = None,
|
||||
run_start_timeout: float = _DEFAULT_RUN_START_TIMEOUT,
|
||||
max_reconnect_attempts: int = 5,
|
||||
reconnect_backoff_base: float = 0.1,
|
||||
reconnect_backoff_cap: float = 10.0,
|
||||
) -> None:
|
||||
self._transport = transport
|
||||
self._next_subscription_id = 1
|
||||
self._subscriptions: dict[int, _SyncSubscription] = {}
|
||||
self._seen_event_ids: set[str] = set()
|
||||
self._shared_stream: SyncEventStreamHandle | None = None
|
||||
self._shared_stream_filter: dict[str, Any] | None = None
|
||||
self._fanout_thread: threading.Thread | None = None
|
||||
self._closed = False
|
||||
self._lock = threading.RLock()
|
||||
self._cursor: int | None = None
|
||||
# When None, no gate is applied and reconcile_stream proceeds immediately.
|
||||
# SyncThreadStream passes an un-set Event so subscriptions wait until
|
||||
# run.start completes.
|
||||
self._run_start_gate = run_start_gate
|
||||
self._run_start_timeout = run_start_timeout
|
||||
self._max_reconnect_attempts = max_reconnect_attempts
|
||||
self._reconnect_backoff_base = reconnect_backoff_base
|
||||
self._reconnect_backoff_cap = reconnect_backoff_cap
|
||||
self._drain_threads: set[threading.Thread] = set()
|
||||
|
||||
def register_subscription(self, params: SubscribeParams) -> _SyncSubscription:
|
||||
with self._lock:
|
||||
sub = _SyncSubscription(id=self._next_subscription_id, params=params)
|
||||
self._next_subscription_id += 1
|
||||
self._subscriptions[sub.id] = sub
|
||||
return sub
|
||||
|
||||
def unregister_subscription(self, subscription_id: int) -> None:
|
||||
with self._lock:
|
||||
self._subscriptions.pop(subscription_id, None)
|
||||
|
||||
def reconcile_stream(self, candidate_filter: SubscribeParams) -> None:
|
||||
if self._run_start_gate is not None and not self._run_start_gate.wait(
|
||||
timeout=self._run_start_timeout
|
||||
):
|
||||
raise TimeoutError("Sync run.start gate timeout.")
|
||||
with self._lock:
|
||||
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
|
||||
new_filter = self._compute_current_union(extra=candidate_filter)
|
||||
old_stream = self._shared_stream
|
||||
self._shared_stream = self._transport.open_event_stream(
|
||||
self._filter_with_since(new_filter)
|
||||
)
|
||||
self._shared_stream_filter = new_filter
|
||||
if old_stream is not None:
|
||||
drain_thread = threading.Thread(
|
||||
target=self._drain_and_close,
|
||||
args=(old_stream,),
|
||||
daemon=True,
|
||||
name="langgraph-sdk-sync-rotation-drain",
|
||||
)
|
||||
self._drain_threads.add(drain_thread)
|
||||
drain_thread.start()
|
||||
|
||||
def ensure_fanout_running(self) -> None:
|
||||
with self._lock:
|
||||
if self._fanout_thread is not None and self._fanout_thread.is_alive():
|
||||
return
|
||||
self._fanout_thread = threading.Thread(
|
||||
target=self._fanout,
|
||||
name="langgraph-sdk-sync-stream-fanout",
|
||||
daemon=True,
|
||||
)
|
||||
self._fanout_thread.start()
|
||||
|
||||
def _fanout(self) -> None:
|
||||
from langgraph_sdk.stream.subscription import matches_subscription
|
||||
|
||||
while True:
|
||||
with self._lock:
|
||||
if self._closed:
|
||||
return
|
||||
shared = self._shared_stream
|
||||
if shared is None:
|
||||
return
|
||||
try:
|
||||
for event in self._dedup_iter(shared.events):
|
||||
with self._lock:
|
||||
if self._closed:
|
||||
break
|
||||
subscriptions = list(self._subscriptions.values())
|
||||
for sub in subscriptions:
|
||||
if matches_subscription(event, sub.params):
|
||||
sub.queue.put(event)
|
||||
except Exception:
|
||||
with self._lock:
|
||||
for sub in self._subscriptions.values():
|
||||
sub.queue.put(None)
|
||||
raise
|
||||
with self._lock:
|
||||
if self._shared_stream is shared:
|
||||
break
|
||||
|
||||
with self._lock:
|
||||
for sub in self._subscriptions.values():
|
||||
sub.queue.put(None)
|
||||
|
||||
def _compute_current_union(
|
||||
self, extra: SubscribeParams | None = None
|
||||
) -> dict[str, Any]:
|
||||
filters = [dict(sub.params) for sub in self._subscriptions.values()]
|
||||
if extra is not None:
|
||||
filters.append(dict(extra))
|
||||
return compute_union_filter(filters)
|
||||
|
||||
def observe_applied_through_seq(self, seq: Any) -> None:
|
||||
"""Advance the reconnect cursor from a command response meta sequence."""
|
||||
with self._lock:
|
||||
self._observe_seq(seq)
|
||||
|
||||
def _observe_event(self, event: Event) -> None:
|
||||
with self._lock:
|
||||
self._observe_seq(event.get("seq"))
|
||||
|
||||
def _observe_seq(self, seq: Any) -> None:
|
||||
if isinstance(seq, int) and (self._cursor is None or seq > self._cursor):
|
||||
self._cursor = seq
|
||||
|
||||
def _filter_with_since(self, params: dict[str, Any]) -> dict[str, Any]:
|
||||
out = dict(params)
|
||||
if self._cursor is not None:
|
||||
out["since"] = self._cursor
|
||||
return out
|
||||
|
||||
def _dedup_iter(self, source: Any) -> Any:
|
||||
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)
|
||||
self._observe_event(event)
|
||||
yield event
|
||||
|
||||
def _drain_and_close(self, handle: SyncEventStreamHandle) -> None:
|
||||
"""Drain remaining events from an old handle before closing it.
|
||||
|
||||
Runs in a background thread spawned by `reconcile_stream` on rotation
|
||||
so buffered events are not lost when the shared stream is replaced.
|
||||
Events are dispatched to subscribers regardless of `_closed` so that
|
||||
already-buffered events reach consumers before the handle is closed.
|
||||
"""
|
||||
from langgraph_sdk.stream.subscription import matches_subscription
|
||||
|
||||
try:
|
||||
for event in self._dedup_iter(handle.events):
|
||||
with self._lock:
|
||||
subscriptions = list(self._subscriptions.values())
|
||||
for sub in subscriptions:
|
||||
if matches_subscription(event, sub.params):
|
||||
sub.queue.put(event)
|
||||
except Exception as err:
|
||||
_logger.debug("rotation drain exception: %r", err)
|
||||
finally:
|
||||
with contextlib.suppress(Exception):
|
||||
handle.close()
|
||||
with self._lock:
|
||||
self._drain_threads.discard(threading.current_thread())
|
||||
|
||||
def _reconnect_sleep(self, attempt: int) -> None:
|
||||
"""Sleep with exponential backoff + jitter before a reconnect attempt."""
|
||||
base = self._reconnect_backoff_base
|
||||
cap = self._reconnect_backoff_cap
|
||||
delay = min(cap, base * (2**attempt))
|
||||
jitter = random.uniform(0, delay * 0.25)
|
||||
time.sleep(delay + jitter)
|
||||
|
||||
def _reconnect_shared_stream(self) -> bool:
|
||||
"""Attempt to reopen the shared stream after a transport drop.
|
||||
|
||||
Returns True if a new stream was successfully opened, False if all
|
||||
reconnect attempts were exhausted or the controller was closed.
|
||||
"""
|
||||
base_filter = self._shared_stream_filter
|
||||
if base_filter is None:
|
||||
return False
|
||||
for attempt in range(self._max_reconnect_attempts):
|
||||
if self._closed:
|
||||
return False
|
||||
if attempt > 0:
|
||||
self._reconnect_sleep(attempt - 1)
|
||||
try:
|
||||
new_handle = self._transport.open_event_stream(
|
||||
self._filter_with_since(base_filter)
|
||||
)
|
||||
old = self._shared_stream
|
||||
self._shared_stream = new_handle
|
||||
if old is not None:
|
||||
with contextlib.suppress(Exception):
|
||||
old.close()
|
||||
return True
|
||||
except Exception as err:
|
||||
_logger.debug("sync reconnect attempt %d failed: %r", attempt, err)
|
||||
return False
|
||||
|
||||
def close(self) -> None:
|
||||
with self._lock:
|
||||
if self._closed:
|
||||
return
|
||||
self._closed = True
|
||||
shared = self._shared_stream
|
||||
for sub in self._subscriptions.values():
|
||||
sub.queue.put(None)
|
||||
if shared is not None:
|
||||
shared.close()
|
||||
thread = self._fanout_thread
|
||||
if thread is not None and thread.is_alive():
|
||||
with contextlib.suppress(RuntimeError):
|
||||
thread.join(timeout=1.0)
|
||||
with self._lock:
|
||||
drain_threads = set(self._drain_threads)
|
||||
for drain in drain_threads:
|
||||
if drain.is_alive():
|
||||
with contextlib.suppress(RuntimeError):
|
||||
drain.join(timeout=1.0)
|
||||
@@ -11,7 +11,7 @@ from langgraph_sdk._async.threads import ThreadsClient
|
||||
from langgraph_sdk.stream.controller import StreamController
|
||||
from langgraph_sdk.stream.transport.http import EventStreamHandle
|
||||
from streaming._events import lifecycle_event, values_event
|
||||
from streaming._fake_server import FakeServer
|
||||
from streaming._fake_server import FakeServer, _StreamScript
|
||||
|
||||
|
||||
async def test_shared_stream_serves_single_subscription():
|
||||
@@ -327,3 +327,28 @@ async def test_shared_stream_reconnect_dedupes_replayed_overlap():
|
||||
|
||||
assert [event["seq"] for event in received if event is not None] == [1, 2]
|
||||
assert received[-1] is None
|
||||
|
||||
|
||||
async def test_send_command_applied_through_seq_seeds_shared_stream_since():
|
||||
fake = FakeServer()
|
||||
fake.script_sequence([_StreamScript(events=[]), _StreamScript(events=[])])
|
||||
fake.script_command_response(
|
||||
{
|
||||
"type": "success",
|
||||
"id": None,
|
||||
"result": {"run_id": "run-1"},
|
||||
"meta": {"applied_through_seq": 17},
|
||||
}
|
||||
)
|
||||
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={})
|
||||
_ = [event async for event in thread.subscribe(["values"])]
|
||||
|
||||
values_requests = [
|
||||
b for b in fake.stream_request_bodies if b.get("channels") == ["values"]
|
||||
]
|
||||
assert len(values_requests) == 1
|
||||
assert values_requests[0]["since"] == 17
|
||||
|
||||
@@ -0,0 +1,25 @@
|
||||
"""Sync v3 streaming projection tests."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import httpx
|
||||
|
||||
from langgraph_sdk._sync.http import SyncHttpClient
|
||||
from langgraph_sdk._sync.threads import SyncThreadsClient
|
||||
from streaming._events import lifecycle_completed_event
|
||||
from streaming._sync_fake_server import SyncFakeServer
|
||||
|
||||
|
||||
def test_sync_values_first_yield_is_rest_state_and_output_returns_final_state():
|
||||
fake = SyncFakeServer()
|
||||
fake.script([lifecycle_completed_event(seq=1)])
|
||||
fake.set_state({"answer": 42})
|
||||
with httpx.Client(transport=fake.transport, base_url="http://test") as raw:
|
||||
threads = SyncThreadsClient(SyncHttpClient(raw))
|
||||
with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
|
||||
thread.run.start(input={})
|
||||
first = next(iter(thread.values))
|
||||
output = thread.output
|
||||
|
||||
assert first == {"answer": 42}
|
||||
assert output == {"answer": 42}
|
||||
@@ -0,0 +1,51 @@
|
||||
"""Sync shared stream controller tests."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import httpx
|
||||
|
||||
from langgraph_sdk._sync.http import SyncHttpClient
|
||||
from langgraph_sdk._sync.threads import SyncThreadsClient
|
||||
from langgraph_sdk.stream.sync_controller import SyncStreamController
|
||||
from langgraph_sdk.stream.transport.sync_http import SyncProtocolSseTransport
|
||||
from streaming._events import values_event
|
||||
from streaming._sync_fake_server import SyncFakeServer, SyncStreamScript
|
||||
|
||||
|
||||
def test_sync_controller_fans_out_to_subscription():
|
||||
fake = SyncFakeServer()
|
||||
fake.script([values_event(seq=1, counter=1)])
|
||||
with httpx.Client(transport=fake.transport, base_url="http://test") as raw:
|
||||
transport = SyncProtocolSseTransport(client=raw, thread_id="t-1")
|
||||
controller = SyncStreamController(transport)
|
||||
sub = controller.register_subscription({"channels": ["values"]})
|
||||
controller.reconcile_stream({"channels": ["values"]})
|
||||
controller.ensure_fanout_running()
|
||||
|
||||
assert sub.queue.get(timeout=1) == values_event(seq=1, counter=1)
|
||||
assert sub.queue.get(timeout=1) is None
|
||||
controller.close()
|
||||
|
||||
|
||||
def test_sync_send_command_applied_through_seq_seeds_shared_stream_since():
|
||||
fake = SyncFakeServer()
|
||||
fake.script_sequence([SyncStreamScript(events=[]), SyncStreamScript(events=[])])
|
||||
fake.script_command_response(
|
||||
{
|
||||
"type": "success",
|
||||
"id": None,
|
||||
"result": {"run_id": "run-1"},
|
||||
"meta": {"applied_through_seq": 17},
|
||||
}
|
||||
)
|
||||
with httpx.Client(transport=fake.transport, base_url="http://test") as raw:
|
||||
threads = SyncThreadsClient(SyncHttpClient(raw))
|
||||
with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
|
||||
thread.run.start(input={})
|
||||
assert list(thread.subscribe(["values"])) == []
|
||||
|
||||
values_requests = [
|
||||
b for b in fake.stream_request_bodies if b.get("channels") == ["values"]
|
||||
]
|
||||
assert len(values_requests) == 1
|
||||
assert values_requests[0]["since"] == 17
|
||||
@@ -0,0 +1,352 @@
|
||||
"""Tests for SyncThreadStream — Tasks 9.1 through 9.6."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Iterator
|
||||
|
||||
import httpx
|
||||
|
||||
from langgraph_sdk._sync.http import SyncHttpClient
|
||||
from langgraph_sdk._sync.threads import SyncThreadsClient
|
||||
from langgraph_sdk.stream.transport.sync_http import (
|
||||
SyncEventStreamHandle,
|
||||
SyncProtocolSseTransport,
|
||||
)
|
||||
from streaming._sync_fake_server import SyncFakeServer, SyncStreamScript
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Task 9.1 — run_start_gate
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_sync_subscribe_before_run_start_waits_on_gate():
|
||||
"""A subscribe issued before run.start completes must block until the
|
||||
gate is set, mirroring async behavior."""
|
||||
fake = SyncFakeServer()
|
||||
# Lifecycle + fanout streams: empty so threads terminate cleanly.
|
||||
fake.script_sequence(
|
||||
[
|
||||
SyncStreamScript(events=[]), # lifecycle watcher
|
||||
SyncStreamScript(events=[]), # first subscribe
|
||||
]
|
||||
)
|
||||
|
||||
with httpx.Client(transport=fake.transport, base_url="http://test") as raw:
|
||||
threads = SyncThreadsClient(SyncHttpClient(raw))
|
||||
with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
|
||||
controller = thread._controller
|
||||
assert controller is not None
|
||||
|
||||
started = threading.Event()
|
||||
|
||||
def slow_subscriber() -> None:
|
||||
started.set()
|
||||
list(thread.subscribe(["values"]))
|
||||
|
||||
t = threading.Thread(target=slow_subscriber)
|
||||
t.start()
|
||||
started.wait(timeout=0.5)
|
||||
|
||||
# Set the gate manually (simulating run.start completing)
|
||||
time.sleep(0.05)
|
||||
assert controller._run_start_gate is not None
|
||||
controller._run_start_gate.set()
|
||||
|
||||
t.join(timeout=2.0)
|
||||
|
||||
# The subscriber should have unblocked and terminated cleanly.
|
||||
assert not t.is_alive(), "subscriber thread should have terminated"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Task 9.2 — reconnect backoff + ready check
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_sync_reconnect_uses_backoff_between_attempts(monkeypatch):
|
||||
"""_reconnect_shared_stream sleeps between retry attempts with exp+jitter
|
||||
backoff, mirroring the async reconnect behavior."""
|
||||
import langgraph_sdk.stream.sync_controller as _ctrl_mod
|
||||
|
||||
sleeps: list[float] = []
|
||||
monkeypatch.setattr(_ctrl_mod.time, "sleep", lambda d: sleeps.append(d))
|
||||
|
||||
from langgraph_sdk.stream.sync_controller import SyncStreamController
|
||||
from langgraph_sdk.stream.transport.sync_http import SyncProtocolSseTransport
|
||||
|
||||
class _FailingTransport(SyncProtocolSseTransport):
|
||||
"""Transport that always raises on open_event_stream."""
|
||||
|
||||
def open_event_stream(self, params: dict) -> SyncEventStreamHandle: # noqa: ARG002
|
||||
raise RuntimeError("scripted transport failure")
|
||||
|
||||
with httpx.Client(base_url="http://test") as raw:
|
||||
transport = _FailingTransport(client=raw, thread_id="t-1")
|
||||
controller = SyncStreamController(transport, max_reconnect_attempts=5)
|
||||
controller._shared_stream_filter = {"channels": ["values"]}
|
||||
result = controller._reconnect_shared_stream()
|
||||
|
||||
assert result is False, "all attempts should have failed"
|
||||
# Attempts 0..4 → sleeps before attempts 1..4 → 4 sleeps
|
||||
assert len(sleeps) == 4, f"Expected 4 sleeps, got {sleeps}"
|
||||
# Backoff should grow (each delay is larger than previous, ignoring jitter)
|
||||
delays_without_jitter = [0.1 * (2**i) for i in range(4)]
|
||||
for i, (sleep, expected_base) in enumerate(
|
||||
zip(sleeps, delays_without_jitter, strict=False)
|
||||
):
|
||||
assert sleep >= expected_base, (
|
||||
f"sleep[{i}]={sleep} < expected base {expected_base}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Task 9.3 — rotation drains buffered events from old stream
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_sync_rotation_does_not_lose_buffered_events():
|
||||
"""When the shared stream rotates, old-stream events already in the queue
|
||||
are not dropped. _drain_and_close dispatches remaining events from the
|
||||
old handle to subscribers before closing it."""
|
||||
import queue
|
||||
from typing import Any
|
||||
|
||||
from langgraph_sdk.stream.sync_controller import SyncStreamController
|
||||
from langgraph_sdk.stream.transport.sync_http import (
|
||||
SyncEventStreamHandle,
|
||||
SyncProtocolSseTransport,
|
||||
)
|
||||
from streaming._events import values_event
|
||||
|
||||
event_a = values_event(seq=1, counter=1)
|
||||
|
||||
class _ScriptedTransport(SyncProtocolSseTransport):
|
||||
"""First call produces event_a; second call produces an empty stream."""
|
||||
|
||||
def open_event_stream(self, params: dict) -> SyncEventStreamHandle: # noqa: ARG002
|
||||
def _gen_a() -> Iterator[Any]:
|
||||
yield event_a
|
||||
|
||||
def _gen_empty() -> Iterator[Any]:
|
||||
return
|
||||
yield # pragma: no cover
|
||||
|
||||
# Alternate: first call → a, second → empty.
|
||||
if not hasattr(self, "_call_count"):
|
||||
self._call_count = 0
|
||||
self._call_count += 1
|
||||
events_gen: Iterator[Any] = (
|
||||
_gen_a() if self._call_count == 1 else _gen_empty()
|
||||
)
|
||||
return SyncEventStreamHandle(
|
||||
events=events_gen,
|
||||
error=lambda: None,
|
||||
close=lambda: None,
|
||||
)
|
||||
|
||||
with httpx.Client(base_url="http://test") as raw:
|
||||
transport = _ScriptedTransport(client=raw, thread_id="t-1")
|
||||
controller = SyncStreamController(transport)
|
||||
sub = controller.register_subscription({"channels": ["values"]})
|
||||
|
||||
# First reconcile — opens old stream (event_a available immediately).
|
||||
controller.reconcile_stream({"channels": ["values"]})
|
||||
# Do NOT start fanout; let reconcile_stream cause a rotation directly.
|
||||
|
||||
# Second reconcile: rotates to empty stream; drain thread handles old.
|
||||
controller.reconcile_stream({"channels": ["values", "updates"]})
|
||||
|
||||
# Start fanout AFTER rotation (picks up the new empty stream).
|
||||
controller.ensure_fanout_running()
|
||||
|
||||
# Allow drain thread to finish before collecting results.
|
||||
controller.close()
|
||||
|
||||
received = []
|
||||
while True:
|
||||
try:
|
||||
item = sub.queue.get_nowait()
|
||||
if item is None:
|
||||
continue
|
||||
received.append(item)
|
||||
except queue.Empty:
|
||||
break
|
||||
|
||||
seqs = [e.get("seq") for e in received]
|
||||
assert 1 in seqs, f"event_a (seq=1) not received via drain; got seqs={seqs}"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Task 9.4 — _next_command_id lock
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_sync_concurrent_commands_do_not_share_command_id():
|
||||
"""50 concurrent threads calling _send_command must each get a unique id."""
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import Any
|
||||
|
||||
captured_ids: list[int] = []
|
||||
ids_lock = threading.Lock()
|
||||
|
||||
class _CapturingTransport(SyncProtocolSseTransport):
|
||||
"""Captures command ids; always returns success."""
|
||||
|
||||
def send_command(self, command: dict) -> dict:
|
||||
with ids_lock:
|
||||
captured_ids.append(command["id"])
|
||||
return {"type": "success", "id": command["id"], "result": {}}
|
||||
|
||||
def open_event_stream(self, params: dict) -> SyncEventStreamHandle: # noqa: ARG002
|
||||
def _gen() -> Iterator[Any]:
|
||||
return
|
||||
yield
|
||||
|
||||
return SyncEventStreamHandle(
|
||||
events=_gen(), error=lambda: None, close=lambda: None
|
||||
)
|
||||
|
||||
fake = SyncFakeServer()
|
||||
fake.script_sequence([SyncStreamScript(events=[])])
|
||||
|
||||
with httpx.Client(transport=fake.transport, base_url="http://test") as raw:
|
||||
threads_client = SyncThreadsClient(SyncHttpClient(raw))
|
||||
with threads_client.stream(thread_id="t-cmd", assistant_id="agent") as stream:
|
||||
# Pre-set gate so _send_command doesn't wait.
|
||||
if stream._controller and stream._controller._run_start_gate:
|
||||
stream._controller._run_start_gate.set()
|
||||
# Replace transport with capturing transport.
|
||||
capture_transport = _CapturingTransport(client=raw, thread_id="t-cmd")
|
||||
stream._transport = capture_transport
|
||||
|
||||
with ThreadPoolExecutor(max_workers=50) as ex:
|
||||
futures = [
|
||||
ex.submit(stream._send_command, "noop", {}) for _ in range(50)
|
||||
]
|
||||
for f in futures:
|
||||
f.result()
|
||||
|
||||
assert len(set(captured_ids)) == 50, (
|
||||
f"Expected 50 unique command ids, got {len(set(captured_ids))} unique "
|
||||
f"out of {len(captured_ids)} total: {sorted(captured_ids)}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Task 9.5 — sync events returns fresh iterator per access
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_sync_events_returns_fresh_iterator_each_access():
|
||||
"""Two accesses of `thread.events` yield independent subscriptions,
|
||||
mirroring the async semantics where each access opens a new subscriber."""
|
||||
fake = SyncFakeServer()
|
||||
from streaming._events import values_event
|
||||
|
||||
event_1 = values_event(seq=1, counter=1)
|
||||
fake.script_sequence(
|
||||
[
|
||||
SyncStreamScript(events=[]), # lifecycle watcher
|
||||
SyncStreamScript(events=[event_1]), # first events access
|
||||
SyncStreamScript(events=[event_1]), # second events access
|
||||
]
|
||||
)
|
||||
|
||||
with httpx.Client(transport=fake.transport, base_url="http://test") as raw:
|
||||
threads_client = SyncThreadsClient(SyncHttpClient(raw))
|
||||
with threads_client.stream(thread_id="t-5", assistant_id="agent") as thread:
|
||||
# Pre-set gate.
|
||||
if thread._controller and thread._controller._run_start_gate:
|
||||
thread._controller._run_start_gate.set()
|
||||
|
||||
iter1 = thread.events
|
||||
iter2 = thread.events
|
||||
|
||||
# They must be independent objects (different subscription iterators).
|
||||
assert iter1 is not iter2
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Task 9.6 — close ordering: fail active streams before controller close
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_close_unblocks_active_subscription_before_lifecycle_join():
|
||||
"""close() must send None to active subscriptions BEFORE joining the
|
||||
lifecycle watcher thread, so callers wake quickly even if the watcher
|
||||
thread blocks for up to 1s."""
|
||||
import queue
|
||||
|
||||
# Gate that keeps the lifecycle watcher thread alive for 0.4s.
|
||||
lifecycle_block = threading.Event()
|
||||
unblock_times: list[float] = []
|
||||
close_times: list[float] = []
|
||||
|
||||
class _BlockingFakeServer(SyncFakeServer):
|
||||
"""Lifecycle stream blocks until gate set; subscribe stream is empty."""
|
||||
|
||||
def _handle(self, request: httpx.Request) -> httpx.Response:
|
||||
path = request.url.path
|
||||
if path.endswith("/stream/events"):
|
||||
import orjson
|
||||
|
||||
body = orjson.loads(request.content)
|
||||
channels = body.get("channels", [])
|
||||
if "lifecycle" in channels:
|
||||
# Block lifecycle watcher for 0.4s.
|
||||
lifecycle_block.wait(timeout=0.4)
|
||||
return super()._handle(request)
|
||||
|
||||
fake = _BlockingFakeServer()
|
||||
fake.script_sequence(
|
||||
[
|
||||
SyncStreamScript(events=[]), # lifecycle watcher
|
||||
SyncStreamScript(events=[]), # subscribe fanout stream
|
||||
]
|
||||
)
|
||||
|
||||
with httpx.Client(transport=fake.transport, base_url="http://test") as raw:
|
||||
threads_client = SyncThreadsClient(SyncHttpClient(raw))
|
||||
with threads_client.stream(thread_id="t-6", assistant_id="agent") as thread:
|
||||
if thread._controller and thread._controller._run_start_gate:
|
||||
thread._controller._run_start_gate.set()
|
||||
|
||||
assert thread._controller is not None
|
||||
sub = thread._controller.register_subscription({"channels": ["values"]})
|
||||
thread._controller.reconcile_stream({"channels": ["values"]})
|
||||
thread._controller.ensure_fanout_running()
|
||||
|
||||
consumer_ready = threading.Event()
|
||||
|
||||
def _consume() -> None:
|
||||
consumer_ready.set()
|
||||
while True:
|
||||
try:
|
||||
item = sub.queue.get(timeout=2.0)
|
||||
if item is None:
|
||||
unblock_times.append(time.monotonic())
|
||||
return
|
||||
except queue.Empty:
|
||||
return
|
||||
|
||||
t = threading.Thread(target=_consume)
|
||||
t.start()
|
||||
consumer_ready.wait(timeout=1.0)
|
||||
time.sleep(0.02)
|
||||
|
||||
close_times.append(time.monotonic())
|
||||
# __exit__ calls close() here.
|
||||
|
||||
lifecycle_block.set() # Unblock watcher so test can finish.
|
||||
t.join(timeout=2.0)
|
||||
assert not t.is_alive(), "consumer thread should have unblocked"
|
||||
assert unblock_times, "consumer never received sentinel"
|
||||
elapsed = unblock_times[0] - close_times[0]
|
||||
# With controller closed BEFORE lifecycle join, sentinel arrives fast.
|
||||
# Lifecycle watcher blocks for 0.4s but that should not delay the sentinel.
|
||||
assert elapsed < 0.3, (
|
||||
f"consumer woke {elapsed:.3f}s after close() — "
|
||||
"controller.close() should precede the lifecycle thread join"
|
||||
)
|
||||
@@ -45,13 +45,11 @@ def _normalize_return_annotation(ann: object) -> str:
|
||||
s = re.sub(r"Generator\[([^,\]]+)(?:,[^\]]*)?\]", r"Iterator[\1]", s)
|
||||
s = re.sub(r"AsyncIterator\[(.+)\]", r"Iterator[\1]", s)
|
||||
s = re.sub(r"AsyncIterable\[(.+)\]", r"Iterable[\1]", s)
|
||||
s = s.replace("AsyncThreadStream", "SyncThreadStream")
|
||||
return s
|
||||
|
||||
|
||||
# Methods that exist only on the async client surface.
|
||||
ASYNC_ONLY_METHODS: dict[str, set[str]] = {
|
||||
"ThreadsClient": {"stream"},
|
||||
}
|
||||
ASYNC_ONLY_METHODS: dict[str, set[str]] = {}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
||||
Reference in New Issue
Block a user