mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-17 21:25:46 +02:00
feat(sdk-py): harden streaming reconnects (#7829)
This commit is contained in:
Nick Hollon
@@ -14,6 +14,7 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import random
|
||||
from collections.abc import AsyncGenerator, AsyncIterator, Generator, Mapping
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Literal, TypedDict
|
||||
@@ -1179,6 +1180,17 @@ class AsyncThreadStream:
|
||||
self._interrupts_lock = asyncio.Lock()
|
||||
self._lifecycle_watcher_task: asyncio.Task[None] | None = None
|
||||
self._lifecycle_watcher_handle: EventStreamHandle | None = None
|
||||
self._lifecycle_cursor: int | None = None
|
||||
self._lifecycle_max_reconnect_attempts = 5
|
||||
# Shared-stream reconnect knobs: applied by `_fanout` after a post-ready
|
||||
# transport drop so subscribers (messages/tools/tasks/values projections,
|
||||
# subgraph child handles) survive a brief SSE disconnect without losing
|
||||
# buffered events. Cursor (`_cursor`) is replayed as `since` so the
|
||||
# server resumes from where the prior stream left off; per-event
|
||||
# `event_id` dedup in `_dedup_iter` drops any overlap on the new stream.
|
||||
self._shared_max_reconnect_attempts = 5
|
||||
self._shared_reconnect_backoff_base = 0.1
|
||||
self._shared_reconnect_backoff_cap = 2.0
|
||||
self._run_start_ready: asyncio.Future[None] | None = None
|
||||
self._run_seen: bool = False
|
||||
self._run_done: asyncio.Future[_RunTerminal] | None = None
|
||||
@@ -1381,6 +1393,12 @@ class AsyncThreadStream:
|
||||
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.
|
||||
|
||||
On a post-ready transport drop (non-cancelled error in `shared.done`),
|
||||
attempts to reconnect up to `_shared_max_reconnect_attempts` times so
|
||||
scoped projections (subgraph child handles, message streams) survive
|
||||
without losing buffered events. The reconnect replays `since=<cursor>`
|
||||
and `_dedup_iter` drops any overlap.
|
||||
"""
|
||||
from langgraph_sdk.stream.subscription import matches_subscription
|
||||
|
||||
@@ -1397,13 +1415,23 @@ class AsyncThreadStream:
|
||||
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
|
||||
# Pump errored — fall through to error-handling/reconnect.
|
||||
pass
|
||||
if self._shared_stream is shared:
|
||||
# No rotation happened; stream genuinely ended.
|
||||
# No rotation happened; the stream genuinely ended. Check
|
||||
# `shared.done` for a post-ready drop and, if so, attempt to
|
||||
# reconnect with `since=<cursor>` so subscribers don't lose
|
||||
# buffered events on a transient transport failure.
|
||||
err = await shared.done
|
||||
if (
|
||||
err is not None
|
||||
and not isinstance(err, asyncio.CancelledError)
|
||||
and not self._closed
|
||||
):
|
||||
with contextlib.suppress(Exception):
|
||||
await shared.close()
|
||||
if await self._reconnect_shared_stream():
|
||||
continue
|
||||
break
|
||||
# Rotation: loop again to pick up the new _shared_stream.
|
||||
|
||||
@@ -1411,6 +1439,47 @@ class AsyncThreadStream:
|
||||
for sub in self._subscriptions.values():
|
||||
sub.queue.put_nowait(None)
|
||||
|
||||
async def _reconnect_sleep(self, attempt: int) -> None:
|
||||
"""Sleep with exponential backoff and jitter for reconnect attempt `attempt`."""
|
||||
base = self._shared_reconnect_backoff_base
|
||||
cap = self._shared_reconnect_backoff_cap
|
||||
delay = min(cap, base * (2**attempt))
|
||||
jitter = random.uniform(0, delay * 0.25)
|
||||
await asyncio.sleep(delay + jitter)
|
||||
|
||||
async def _reconnect_shared_stream(self) -> bool:
|
||||
"""Attempt to reopen the shared stream after a post-ready transport drop.
|
||||
|
||||
Returns:
|
||||
`True` if a new stream was opened (caller should resume fanout),
|
||||
`False` if all reconnect attempts were exhausted or the controller
|
||||
was closed in the meantime.
|
||||
"""
|
||||
if self._transport is None:
|
||||
return False
|
||||
# Use the current shared-stream filter (latest computed union); if
|
||||
# subscriptions changed during the drop, this picks up the new shape.
|
||||
base_filter = self._shared_stream_filter
|
||||
if base_filter is None:
|
||||
return False
|
||||
for attempt in range(self._shared_max_reconnect_attempts):
|
||||
if self._closed:
|
||||
return False
|
||||
stream_params: dict[str, Any] = dict(base_filter)
|
||||
if self._cursor is not None:
|
||||
stream_params["since"] = self._cursor
|
||||
try:
|
||||
new_stream = self._transport.open_event_stream(stream_params)
|
||||
await new_stream.ready
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception:
|
||||
await self._reconnect_sleep(attempt)
|
||||
continue
|
||||
self._shared_stream = new_stream
|
||||
return True
|
||||
return False
|
||||
|
||||
async def _reconcile_stream(self, candidate_filter: SubscribeParams) -> None:
|
||||
"""Ensure the shared SSE covers `candidate_filter`. Rotate if not.
|
||||
|
||||
@@ -1518,74 +1587,79 @@ class AsyncThreadStream:
|
||||
await asyncio.wait_for(asyncio.shield(gate), timeout=timeout)
|
||||
|
||||
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.
|
||||
def _observe_lifecycle_event(self, event: Event) -> None:
|
||||
seq = event.get("seq")
|
||||
if isinstance(seq, int) and (
|
||||
self._lifecycle_cursor is None or seq > self._lifecycle_cursor
|
||||
):
|
||||
self._lifecycle_cursor = seq
|
||||
|
||||
Independent of the union-filter shared stream so that interrupts
|
||||
surface even when no other subscription is active. Starts immediately
|
||||
on session entry (before any run.start) so reattach and thread.output
|
||||
work for existing runs.
|
||||
"""
|
||||
def _lifecycle_stream_params(self) -> dict[str, Any]:
|
||||
params: dict[str, Any] = {"channels": ["lifecycle", "input"]}
|
||||
if self._lifecycle_cursor is not None:
|
||||
params["since"] = self._lifecycle_cursor
|
||||
return params
|
||||
|
||||
async def _run_lifecycle_watcher(self) -> None:
|
||||
"""Always-on SSE 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
|
||||
await asyncio.wait_for(handle.ready, timeout=5.0)
|
||||
async for event in handle.events:
|
||||
if self._closed:
|
||||
reconnect_attempts = 0
|
||||
while not self._closed:
|
||||
try:
|
||||
handle = self._transport.open_event_stream(
|
||||
self._lifecycle_stream_params()
|
||||
)
|
||||
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._observe_lifecycle_event(event)
|
||||
await self._apply_lifecycle_event(event)
|
||||
err = await handle.done
|
||||
if err is None or isinstance(err, asyncio.CancelledError):
|
||||
# Clean EOF: stream ended without a terminal lifecycle
|
||||
# event. Resolve `_run_done` as errored so awaiters of
|
||||
# `thread.output` don't hang.
|
||||
if err is None:
|
||||
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(
|
||||
"lifecycle stream ended before terminal event"
|
||||
),
|
||||
)
|
||||
)
|
||||
return
|
||||
await self._apply_lifecycle_event(event)
|
||||
# Why: iterator exhausted without `_run_done` being resolved by a
|
||||
# terminal lifecycle event. Surface any transport error captured
|
||||
# on `handle.done`, otherwise treat the clean EOF as errored so
|
||||
# awaiters of `_run_done` (e.g. `thread.output`) don't hang.
|
||||
err = await handle.done
|
||||
run_done = self._run_done
|
||||
if run_done is not None and not run_done.done():
|
||||
if err is not None:
|
||||
run_done.set_result(
|
||||
_RunTerminal(
|
||||
status="errored",
|
||||
error=RuntimeError(f"Lifecycle transport failed: {err}"),
|
||||
)
|
||||
)
|
||||
else:
|
||||
run_done.set_result(
|
||||
_RunTerminal(
|
||||
status="errored",
|
||||
error=RuntimeError(
|
||||
"lifecycle stream ended before terminal event"
|
||||
),
|
||||
)
|
||||
)
|
||||
return
|
||||
except (Exception, asyncio.CancelledError) as exc:
|
||||
# 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.
|
||||
# Resolve _run_done with an error so thread.output doesn't wait
|
||||
# forever when the lifecycle transport fails.
|
||||
run_done = self._run_done
|
||||
if run_done is not None and not run_done.done():
|
||||
if not isinstance(exc, asyncio.CancelledError):
|
||||
reconnect_attempts += 1
|
||||
if reconnect_attempts > self._lifecycle_max_reconnect_attempts:
|
||||
raise err
|
||||
await asyncio.sleep(0.05)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
reconnect_attempts += 1
|
||||
if reconnect_attempts <= self._lifecycle_max_reconnect_attempts:
|
||||
await asyncio.sleep(0.05)
|
||||
continue
|
||||
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}"),
|
||||
)
|
||||
)
|
||||
return
|
||||
return
|
||||
|
||||
async def _fetch_state(self) -> dict[str, Any]:
|
||||
"""Fetch the current thread state from the REST endpoint."""
|
||||
|
||||
@@ -1162,6 +1162,8 @@ class SyncThreadStream:
|
||||
self.interrupts: list[InterruptPayload] = []
|
||||
self._lifecycle_watcher_thread: threading.Thread | None = None
|
||||
self._lifecycle_watcher_handle: SyncEventStreamHandle | None = None
|
||||
self._lifecycle_cursor: int | None = None
|
||||
self._lifecycle_max_reconnect_attempts = 5
|
||||
self._run_seen: bool = False
|
||||
self._run_done: _BlockingResult | None = None
|
||||
self._active_message_streams: set[ChatModelStream] = set()
|
||||
@@ -1366,28 +1368,54 @@ class SyncThreadStream:
|
||||
)
|
||||
self._lifecycle_watcher_thread.start()
|
||||
|
||||
def _observe_lifecycle_event(self, event: Event) -> None:
|
||||
seq = event.get("seq")
|
||||
if isinstance(seq, int) and (
|
||||
self._lifecycle_cursor is None or seq > self._lifecycle_cursor
|
||||
):
|
||||
self._lifecycle_cursor = seq
|
||||
|
||||
def _lifecycle_stream_params(self) -> dict[str, Any]:
|
||||
params: dict[str, Any] = {"channels": ["lifecycle", "input"]}
|
||||
if self._lifecycle_cursor is not None:
|
||||
params["since"] = self._lifecycle_cursor
|
||||
return params
|
||||
|
||||
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}"),
|
||||
)
|
||||
reconnect_attempts = 0
|
||||
while not self._closed:
|
||||
try:
|
||||
handle = self._transport.open_event_stream(
|
||||
self._lifecycle_stream_params()
|
||||
)
|
||||
self._lifecycle_watcher_handle = handle
|
||||
for event in handle.events:
|
||||
if self._closed:
|
||||
return
|
||||
self._observe_lifecycle_event(event)
|
||||
self._apply_lifecycle_event(event)
|
||||
err = handle.error()
|
||||
if err is None:
|
||||
return
|
||||
reconnect_attempts += 1
|
||||
if reconnect_attempts > self._lifecycle_max_reconnect_attempts:
|
||||
raise err
|
||||
except Exception as exc:
|
||||
reconnect_attempts += 1
|
||||
if reconnect_attempts <= self._lifecycle_max_reconnect_attempts:
|
||||
continue
|
||||
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}"),
|
||||
)
|
||||
)
|
||||
return
|
||||
|
||||
def _fetch_state(self) -> dict[str, Any]:
|
||||
"""Fetch the current thread state from the REST endpoint."""
|
||||
|
||||
@@ -137,18 +137,41 @@ class SyncStreamController:
|
||||
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
|
||||
pass # transport drop — attempt reconnect below
|
||||
|
||||
with self._lock:
|
||||
if self._shared_stream is shared:
|
||||
break
|
||||
if self._shared_stream is not shared:
|
||||
continue # rotation happened; pick up new stream
|
||||
# No rotation — check if this was a transport drop
|
||||
err = shared.error()
|
||||
if err is not None and not self._closed:
|
||||
if self._reconnect_shared_stream():
|
||||
continue
|
||||
break
|
||||
|
||||
with self._lock:
|
||||
for sub in self._subscriptions.values():
|
||||
sub.queue.put(None)
|
||||
|
||||
def _reconnect_shared_stream(self) -> bool:
|
||||
with self._lock:
|
||||
base_filter = self._shared_stream_filter
|
||||
if base_filter is None:
|
||||
return False
|
||||
for _ in range(self._max_reconnect_attempts):
|
||||
with self._lock:
|
||||
if self._closed:
|
||||
return False
|
||||
params = self._filter_with_since(base_filter)
|
||||
try:
|
||||
new_stream = self._transport.open_event_stream(params)
|
||||
except Exception:
|
||||
continue
|
||||
with self._lock:
|
||||
self._shared_stream = new_stream
|
||||
return True
|
||||
return False
|
||||
|
||||
def _compute_current_union(
|
||||
self, extra: SubscribeParams | None = None
|
||||
) -> dict[str, Any]:
|
||||
|
||||
@@ -15,6 +15,7 @@ from collections.abc import AsyncIterator
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
import orjson
|
||||
from starlette.applications import Starlette
|
||||
from starlette.requests import Request
|
||||
@@ -22,6 +23,38 @@ from starlette.responses import JSONResponse, Response, StreamingResponse
|
||||
from starlette.routing import Route
|
||||
|
||||
|
||||
class _AsyncSseByteStream(httpx.AsyncByteStream):
|
||||
"""Async SSE byte stream that supports mid-stream errors via `fail_after`."""
|
||||
|
||||
def __init__(self, script: _StreamScript) -> None:
|
||||
self._script = script
|
||||
|
||||
async def __aiter__(self) -> AsyncIterator[bytes]:
|
||||
for index, event in enumerate(self._script.events, start=1):
|
||||
if self._script.delay:
|
||||
await asyncio.sleep(self._script.delay)
|
||||
payload = orjson.dumps(event).decode()
|
||||
yield f"id: {event.get('event_id', '')}\n".encode()
|
||||
yield f"event: message\ndata: {payload}\n\n".encode()
|
||||
if self._script.fail_after is not None and index >= self._script.fail_after:
|
||||
raise httpx.ReadError("scripted async stream failure")
|
||||
|
||||
|
||||
class _CountedAsyncSseByteStream(httpx.AsyncByteStream):
|
||||
"""Wraps `_AsyncSseByteStream` and decrements the server's open stream counter."""
|
||||
|
||||
def __init__(self, script: _StreamScript, server: Any) -> None:
|
||||
self._inner = _AsyncSseByteStream(script)
|
||||
self._server = server
|
||||
|
||||
async def __aiter__(self) -> AsyncIterator[bytes]:
|
||||
try:
|
||||
async for chunk in self._inner:
|
||||
yield chunk
|
||||
finally:
|
||||
self._server.open_event_streams -= 1
|
||||
|
||||
|
||||
@dataclass
|
||||
class _StreamScript:
|
||||
events: list[dict[str, Any]]
|
||||
@@ -58,6 +91,7 @@ class FakeServer:
|
||||
self.state_request_headers: list[dict[str, str]] = []
|
||||
self._stream_scripts: list[_StreamScript] = []
|
||||
self._command_response: dict[str, Any] | None = None
|
||||
self.transport: httpx.MockTransport = httpx.MockTransport(self._handle_request)
|
||||
|
||||
def script(
|
||||
self,
|
||||
@@ -178,3 +212,50 @@ class FakeServer:
|
||||
@property
|
||||
def peak_open_event_streams(self) -> int:
|
||||
return self._open_event_streams_max
|
||||
|
||||
async def _handle_request(self, request: httpx.Request) -> httpx.Response:
|
||||
"""Async handler for `httpx.MockTransport` — supports proper streaming failures."""
|
||||
path = request.url.path
|
||||
if path.endswith("/commands"):
|
||||
body = orjson.loads(request.content)
|
||||
self.received_commands.append(body)
|
||||
self.command_request_headers.append(dict(request.headers))
|
||||
command_id = body.get("id")
|
||||
if self._command_response is not None:
|
||||
response = dict(self._command_response)
|
||||
response["id"] = command_id
|
||||
return httpx.Response(200, json=response)
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"type": "success",
|
||||
"id": command_id,
|
||||
"result": {"run_id": "run-1"},
|
||||
},
|
||||
)
|
||||
if path.endswith("/stream/events"):
|
||||
self.stream_request_bodies.append(orjson.loads(request.content))
|
||||
self.stream_request_headers_list.append(dict(request.headers))
|
||||
script = (
|
||||
self._stream_scripts.pop(0)
|
||||
if self._stream_scripts
|
||||
else _StreamScript(
|
||||
events=list(self.scripted_events), delay=self._stream_delay
|
||||
)
|
||||
)
|
||||
self.open_event_streams += 1
|
||||
self._open_event_streams_max = max(
|
||||
self._open_event_streams_max, self.open_event_streams
|
||||
)
|
||||
# Wrap the stream to decrement open_event_streams on exhaustion.
|
||||
stream = _CountedAsyncSseByteStream(script, self)
|
||||
return httpx.Response(
|
||||
200,
|
||||
headers={"content-type": "text/event-stream"},
|
||||
stream=stream,
|
||||
)
|
||||
if path.endswith("/state"):
|
||||
self.state_request_count += 1
|
||||
self.state_request_headers.append(dict(request.headers))
|
||||
return httpx.Response(200, json=self.state)
|
||||
return httpx.Response(404, json={"error": f"unexpected path: {path}"})
|
||||
|
||||
@@ -10,8 +10,12 @@ import httpx
|
||||
|
||||
from langgraph_sdk._async.http import HttpClient
|
||||
from langgraph_sdk._async.threads import ThreadsClient
|
||||
from streaming._events import input_requested_event, lifecycle_event
|
||||
from streaming._fake_server import FakeServer
|
||||
from streaming._events import (
|
||||
input_requested_event,
|
||||
lifecycle_completed_event,
|
||||
lifecycle_event,
|
||||
)
|
||||
from streaming._fake_server import FakeServer, _StreamScript
|
||||
|
||||
|
||||
async def test_interrupted_starts_false():
|
||||
@@ -218,3 +222,35 @@ async def test_lifecycle_mid_iteration_error_resolves_run_done_with_error(
|
||||
assert "simulated transport error" in str(terminal.error)
|
||||
# Quiet unused-import warnings under strict configs.
|
||||
_ = contextlib
|
||||
|
||||
|
||||
async def test_lifecycle_watcher_reconnects_with_since_after_transport_drop():
|
||||
fake = FakeServer()
|
||||
fake.set_state({"ok": True})
|
||||
fake.script_sequence(
|
||||
[
|
||||
_StreamScript(
|
||||
events=[lifecycle_event(seq=1, phase="running")],
|
||||
fail_after=1,
|
||||
),
|
||||
_StreamScript(events=[lifecycle_completed_event(seq=2)]),
|
||||
]
|
||||
)
|
||||
async with httpx.AsyncClient(
|
||||
transport=fake.transport, base_url="http://test"
|
||||
) as raw:
|
||||
threads = ThreadsClient(HttpClient(raw))
|
||||
async with threads.stream(thread_id="existing", assistant_id="agent") as thread:
|
||||
for _ in range(20):
|
||||
run_done = thread._run_done
|
||||
if run_done is not None and run_done.done():
|
||||
break
|
||||
await asyncio.sleep(0.05)
|
||||
assert thread._run_done is not None
|
||||
terminal = thread._run_done.result()
|
||||
|
||||
assert terminal.status == "completed"
|
||||
assert terminal.error is None
|
||||
assert fake.stream_request_bodies[0]["channels"] == ["lifecycle", "input"]
|
||||
assert fake.stream_request_bodies[1]["channels"] == ["lifecycle", "input"]
|
||||
assert fake.stream_request_bodies[1]["since"] == 1
|
||||
|
||||
@@ -19,7 +19,7 @@ from streaming._events import (
|
||||
tool_output_delta_event,
|
||||
tool_started_event,
|
||||
)
|
||||
from streaming._fake_server import FakeServer
|
||||
from streaming._fake_server import FakeServer, _StreamScript
|
||||
|
||||
|
||||
async def test_subgraphs_subscribes_to_tasks_channel():
|
||||
@@ -577,3 +577,62 @@ def test_close_inboxes_enqueues_sentinel_on_iterated_inboxes():
|
||||
assert handle._messages_inbox.get_nowait() is None
|
||||
assert handle._tools_inbox.qsize() == 0
|
||||
assert handle._tasks_inbox.qsize() == 0
|
||||
|
||||
|
||||
async def test_subgraph_scoped_messages_survive_shared_stream_reconnect():
|
||||
fake = FakeServer()
|
||||
# Connection order (async): shared stream opens first (via _reconcile_stream),
|
||||
# then the lifecycle watcher task runs. Three scripts are needed:
|
||||
# 1. shared stream initial (fail_after=2 to trigger reconnect)
|
||||
# 2. lifecycle watcher
|
||||
# 3. shared stream reconnect (carries the message events)
|
||||
fake.script_sequence(
|
||||
[
|
||||
_StreamScript(
|
||||
events=[
|
||||
lifecycle_started_event(seq=0),
|
||||
tasks_start_event(
|
||||
seq=1, namespace=["worker:abc"], task_id="t-child"
|
||||
),
|
||||
],
|
||||
fail_after=2,
|
||||
),
|
||||
_StreamScript(
|
||||
events=[lifecycle_completed_event(seq=7)],
|
||||
),
|
||||
_StreamScript(
|
||||
events=[
|
||||
message_start_event(
|
||||
seq=2,
|
||||
namespace=["worker:abc"],
|
||||
message_id="child-msg",
|
||||
run_id="child-run",
|
||||
),
|
||||
message_text_delta_event(
|
||||
seq=3, namespace=["worker:abc"], text="child"
|
||||
),
|
||||
message_text_finish_event(
|
||||
seq=4, namespace=["worker:abc"], text="child"
|
||||
),
|
||||
message_finish_event(seq=5, namespace=["worker:abc"]),
|
||||
tasks_result_event(
|
||||
seq=6, namespace=[], task_id="abc", name="worker"
|
||||
),
|
||||
lifecycle_completed_event(seq=7),
|
||||
]
|
||||
),
|
||||
]
|
||||
)
|
||||
async with httpx.AsyncClient(
|
||||
transport=fake.transport, 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={})
|
||||
handles = [handle async for handle in thread.subgraphs]
|
||||
child = handles[0]
|
||||
chunks = [chunk async for chunk in child.messages]
|
||||
|
||||
assert child.path == ("worker:abc",)
|
||||
assert [await chunk.text for chunk in chunks] == ["child"]
|
||||
assert fake.stream_request_bodies[-1]["since"] >= 1
|
||||
|
||||
@@ -27,7 +27,7 @@ from streaming._events import (
|
||||
tool_output_delta_event,
|
||||
tool_started_event,
|
||||
)
|
||||
from streaming._sync_fake_server import SyncFakeServer
|
||||
from streaming._sync_fake_server import SyncFakeServer, SyncStreamScript
|
||||
|
||||
|
||||
def test_sync_values_first_yield_is_rest_state_and_output_returns_final_state():
|
||||
@@ -227,6 +227,67 @@ def test_sync_subgraph_scoped_messages_and_tool_calls():
|
||||
assert child_calls[0].output == {"child": True}
|
||||
|
||||
|
||||
def test_sync_subgraph_scoped_messages_survive_shared_stream_reconnect():
|
||||
fake = SyncFakeServer()
|
||||
# Connection order (sync): lifecycle watcher thread opens first (it is started
|
||||
# in __enter__ and races ahead during run.start), then the main thread opens
|
||||
# the shared stream. Three scripts are needed:
|
||||
# 1. lifecycle watcher
|
||||
# 2. shared stream initial (fail_after=2 to trigger reconnect)
|
||||
# 3. shared stream reconnect (carries the message events)
|
||||
fake.script_sequence(
|
||||
[
|
||||
SyncStreamScript(
|
||||
events=[
|
||||
lifecycle_started_event(seq=0),
|
||||
lifecycle_completed_event(seq=7),
|
||||
]
|
||||
),
|
||||
SyncStreamScript(
|
||||
events=[
|
||||
lifecycle_started_event(seq=0),
|
||||
tasks_start_event(
|
||||
seq=1, namespace=["worker:abc"], task_id="t-child"
|
||||
),
|
||||
],
|
||||
fail_after=2,
|
||||
),
|
||||
SyncStreamScript(
|
||||
events=[
|
||||
message_start_event(
|
||||
seq=2,
|
||||
namespace=["worker:abc"],
|
||||
message_id="child-msg",
|
||||
run_id="child-run",
|
||||
),
|
||||
message_text_delta_event(
|
||||
seq=3, namespace=["worker:abc"], text="child"
|
||||
),
|
||||
message_text_finish_event(
|
||||
seq=4, namespace=["worker:abc"], text="child"
|
||||
),
|
||||
message_finish_event(seq=5, namespace=["worker:abc"]),
|
||||
tasks_result_event(
|
||||
seq=6, namespace=[], task_id="abc", name="worker"
|
||||
),
|
||||
lifecycle_completed_event(seq=7),
|
||||
]
|
||||
),
|
||||
]
|
||||
)
|
||||
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={})
|
||||
handles = list(thread.subgraphs)
|
||||
child = handles[0]
|
||||
chunks = list(child.messages)
|
||||
|
||||
assert child.path == ("worker:abc",)
|
||||
assert [str(chunk.text) for chunk in chunks] == ["child"]
|
||||
assert fake.stream_request_bodies[-1]["since"] >= 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Task 10.7 — comprehensive sync tool_calls projection tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -49,3 +49,35 @@ def test_sync_send_command_applied_through_seq_seeds_shared_stream_since():
|
||||
]
|
||||
assert len(values_requests) == 1
|
||||
assert values_requests[0]["since"] == 17
|
||||
|
||||
|
||||
def test_sync_shared_stream_reconnects_with_since_after_transport_drop():
|
||||
fake = SyncFakeServer()
|
||||
fake.script_sequence(
|
||||
[
|
||||
SyncStreamScript(
|
||||
events=[values_event(seq=1, values={"counter": 1})],
|
||||
fail_after=1,
|
||||
),
|
||||
SyncStreamScript(events=[values_event(seq=2, values={"counter": 2})]),
|
||||
]
|
||||
)
|
||||
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()
|
||||
|
||||
first = sub.queue.get(timeout=1.0)
|
||||
second = sub.queue.get(timeout=1.0)
|
||||
end = sub.queue.get(timeout=1.0)
|
||||
controller.close()
|
||||
transport.close()
|
||||
|
||||
assert first is not None
|
||||
assert second is not None
|
||||
assert first["seq"] == 1
|
||||
assert second["seq"] == 2
|
||||
assert end is None
|
||||
assert fake.stream_request_bodies[1]["since"] == 1
|
||||
|
||||
@@ -2,8 +2,10 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Iterator
|
||||
|
||||
import httpx
|
||||
@@ -350,3 +352,67 @@ def test_close_unblocks_active_subscription_before_lifecycle_join():
|
||||
f"consumer woke {elapsed:.3f}s after close() — "
|
||||
"controller.close() should precede the lifecycle thread join"
|
||||
)
|
||||
|
||||
|
||||
def test_sync_threads_stream_mints_uuid4_when_thread_id_none():
|
||||
with httpx.Client(base_url="http://test") as raw:
|
||||
threads = SyncThreadsClient(SyncHttpClient(raw))
|
||||
stream = threads.stream(assistant_id="agent")
|
||||
assert re.fullmatch(
|
||||
r"[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}",
|
||||
stream.thread_id,
|
||||
)
|
||||
assert uuid.UUID(stream.thread_id).version == 4
|
||||
|
||||
|
||||
def test_sync_run_start_sends_command():
|
||||
from streaming._events import lifecycle_completed_event
|
||||
|
||||
fake = SyncFakeServer()
|
||||
fake.script([lifecycle_completed_event(seq=1)])
|
||||
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:
|
||||
result = thread.run.start(input={"x": 1})
|
||||
|
||||
assert result == {"run_id": "run-1"}
|
||||
assert fake.received_commands[0]["method"] == "run.start"
|
||||
assert fake.received_commands[0]["params"]["assistant_id"] == "agent"
|
||||
|
||||
|
||||
def test_sync_events_iterates_raw_events():
|
||||
from streaming._events import values_event
|
||||
|
||||
fake = SyncFakeServer()
|
||||
fake.script([values_event(seq=1, counter=1)])
|
||||
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={})
|
||||
events = list(thread.subscribe(["values"]))
|
||||
|
||||
assert events == [values_event(seq=1, counter=1)]
|
||||
|
||||
|
||||
def test_sync_lifecycle_watcher_reconnects_with_since_after_transport_drop():
|
||||
from streaming._events import lifecycle_completed_event, lifecycle_event
|
||||
|
||||
fake = SyncFakeServer()
|
||||
fake.set_state({"ok": True})
|
||||
fake.script_sequence(
|
||||
[
|
||||
SyncStreamScript(
|
||||
events=[lifecycle_event(seq=1, phase="running")],
|
||||
fail_after=1,
|
||||
),
|
||||
SyncStreamScript(events=[lifecycle_completed_event(seq=2)]),
|
||||
]
|
||||
)
|
||||
with httpx.Client(transport=fake.transport, base_url="http://test") as raw:
|
||||
threads = SyncThreadsClient(SyncHttpClient(raw))
|
||||
with threads.stream(thread_id="existing", assistant_id="agent") as thread:
|
||||
terminal = thread._wait_for_run_done()
|
||||
|
||||
assert terminal.status == "completed"
|
||||
assert terminal.error is None
|
||||
assert fake.stream_request_bodies[1]["since"] == 1
|
||||
|
||||
Reference in New Issue
Block a user