mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-17 21:25:46 +02:00
353 lines
14 KiB
Python
353 lines
14 KiB
Python
"""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"
|
|
)
|