WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-10-19 14:31:47 +02:00
Code Issues Packages Projects Releases Wiki Activity
Files
1ffa5bc7406cb80fffa7590438e8c55d44aa4b7d
langgraph/libs/sdk-py/tests/streaming/test_shared_stream.py
T
Elior Nataf LackritzandGitHub 9e6fc45c6f fix(sdk-py): surface thread stream errors and stop retrying a rejected stream (#9317) 
## Summary

Thread streams in `langgraph_sdk` now report why they ended. An `id:
null` error frame becomes the error `thread.output` raises, the sync
client no longer blocks forever when the server closes the stream, and a
stream the server rejects with a 4xx ends at once instead of being
retried.

## Problem

- A server sends `{"type": "error", "id": null, ...}` just before it
closes a thread stream whose read failed. The SDK dropped the frame, so
`thread.output` raised a generic "lifecycle stream ended before terminal
event" (async) or blocked forever (sync: the lifecycle watcher returned
without settling the run).
- A 403 or 404 on the stream request was retried like a network error,
then reported as `RuntimeError("Lifecycle transport failed: ...")` with
no status. A server that rejects a WebSocket after accepting it closes
with 4000 plus the HTTP status (4403, 4404, 4422), since a browser can't
read a refused upgrade's status. The shared stream counted each such
reconnect as a success and reconnected forever.

## Fix

- The lifecycle watcher fails the run with the frame's error on a clean
close (both clients). The shared stream logs the frame instead of
dropping it: projections still end when the stream closes, as before,
and `thread.output` carries the error.
- A 4xx other than 408 and 429, as a status or as a 4000+status close
code, ends the stream without retries. `thread.output` now raises that
`httpx.HTTPStatusError` or `websockets.exceptions.ConnectionClosedError`
instead of the generic `RuntimeError`.
- For a minted `thread_id`, the lifecycle watcher starts after
`run.start` creates the thread. Without this, the 4xx stop would turn a
stream opened under custom auth for a not-yet-created thread from
retried into failed. The sync client starts it under a lock, since
`run.start` can be called from several threads.
- Drops the first of two `SyncStreamController._reconnect_shared_stream`
definitions. It was shadowed by the second and never ran.

Not in this PR: over WebSocket the server closes with 1011 after an
error frame, which is still retried, while over SSE the close is clean
and fails the stream. Reconnecting after a clean close, as
`langchain-ai/langgraphjs#2812` does for the JS SDK, is a separate
change.

## Test plan

- [x] The 10 new behavior tests (async and sync: error frame, clean
close, 404, 4404 close, reconnect on 404, minted `thread_id`) fail on
`main` and pass here; a race test checks that two concurrent `run.start`
calls start one watcher.
- [x] `make test` in `libs/sdk-py`: 541 passed. `make lint` clean,
`check_sdk_methods.py` passes.
- [x] `RemoteGraph` v3 tests in `libs/langgraph`
(`test_remote_graph_v3.py`,
`test_stream_events_v3_kwarg_forwarding.py`): 94 passed against this
SDK.
- [x] Against a local server with custom auth: a thread the caller can't
read fails at once with 404 (SSE) or 4404 (WebSocket), and a new thread
works with a 1 s pause before `run.start`, which fails on `main`.
2026-10-16 15:29:34 -04:00

416 lines
16 KiB
Python
Raw Blame History

from __future__ import annotations
import asyncio
from collections.abc import AsyncGenerator, AsyncIterator
from typing import Any, cast
from unittest.mock import MagicMock
import httpx
from websockets.exceptions import ConnectionClosedError
from websockets.frames import Close
from langgraph_sdk._async.http import HttpClient
from langgraph_sdk._async.threads import ThreadsClient
from langgraph_sdk.stream.controller import StreamController
from langgraph_sdk.stream.transport.http import EventStreamHandle, ProtocolSseTransport
from streaming._events import lifecycle_completed_event, lifecycle_event, values_event
from streaming._fake_server import FakeServer, _StreamScript
async def test_shared_stream_serves_single_subscription():
fake = FakeServer()
fake.script([lifecycle_event(seq=0), values_event(seq=1)])
asgi = httpx.ASGITransport(app=fake.app)
async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
threads = ThreadsClient(HttpClient(raw))
async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
await thread.run.start(input={})
await thread._reconcile_stream({"channels": ["lifecycle", "values"]})
assert thread._shared_stream is not None
received = [
e async for e in thread._dedup_iter(thread._shared_stream.events)
]
methods = [e["method"] for e in received]
assert methods == ["lifecycle", "values"]
assert fake.peak_open_event_streams == 1
async def test_seen_event_ids_dedupes_replayed_events():
fake = FakeServer()
fake.script(
[
lifecycle_event(seq=0),
lifecycle_event(seq=0), # duplicate event_id
values_event(seq=1),
]
)
asgi = httpx.ASGITransport(app=fake.app)
async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
threads = ThreadsClient(HttpClient(raw))
async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
await thread.run.start(input={})
await thread._reconcile_stream({"channels": ["lifecycle", "values"]})
assert thread._shared_stream is not None
received = [
e async for e in thread._dedup_iter(thread._shared_stream.events)
]
seqs = [e["seq"] for e in received]
assert seqs == [0, 1] # the duplicate seq=0 was deduped via event_id
async def test_rotation_when_new_subscription_widens_filter():
fake = FakeServer()
fake.script([lifecycle_event(seq=0), values_event(seq=1)])
asgi = httpx.ASGITransport(app=fake.app)
async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
threads = ThreadsClient(HttpClient(raw))
async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
await thread.run.start(input={})
# First subscription: lifecycle only.
await thread._reconcile_stream({"channels": ["lifecycle"]})
# Second subscription widens to lifecycle + values.
await thread._reconcile_stream({"channels": ["lifecycle", "values"]})
# Rotation: two separate SSE requests were opened (old + new).
assert len(fake.stream_request_bodies) >= 2
async def test_no_rotation_when_existing_filter_covers_new_subscription():
fake = FakeServer()
fake.script([])
asgi = httpx.ASGITransport(app=fake.app)
async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
threads = ThreadsClient(HttpClient(raw))
async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
await thread.run.start(input={})
await thread._reconcile_stream({"channels": ["lifecycle", "values"]})
# New subscription is a subset — existing filter covers it.
await thread._reconcile_stream({"channels": ["values"]})
# No rotation in the shared stream (1 shared SSE) plus 1 lifecycle watcher SSE = 2.
assert len(fake.stream_request_bodies) == 2
async def test_subscribe_yields_only_matching_events():
fake = FakeServer()
fake.script(
[
lifecycle_event(seq=0),
values_event(seq=1),
lifecycle_event(seq=2),
]
)
asgi = httpx.ASGITransport(app=fake.app)
async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
threads = ThreadsClient(HttpClient(raw))
async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
await thread.run.start(input={})
async def drain(channels):
return [e async for e in thread.subscribe(channels)]
lifecycle_events, values_events = await asyncio.gather(
drain(["lifecycle"]),
drain(["values"]),
)
assert [e["seq"] for e in lifecycle_events] == [0, 2]
assert [e["seq"] for e in values_events] == [1]
async def test_two_concurrent_subscribes_share_one_stream():
fake = FakeServer()
fake.script([lifecycle_event(seq=i) for i in range(5)])
asgi = httpx.ASGITransport(app=fake.app)
async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
threads = ThreadsClient(HttpClient(raw))
async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
await thread.run.start(input={})
async def drain(channels):
return [e async for e in thread.subscribe(channels)]
results = await asyncio.gather(
drain(["lifecycle"]),
drain(["lifecycle"]),
)
assert len(results[0]) == 5
assert len(results[1]) == 5
# Both subscriptions share one SSE (no rotation) plus 1 lifecycle watcher SSE = 2.
assert len(fake.stream_request_bodies) == 2
async def test_subscribe_does_not_leak_when_iterator_unconsumed():
"""Subscriptions register lazily on first __anext__, not at subscribe() call time.
Why: registering eagerly would leak the subscription if the caller
constructs the iterator but never iterates it. The lazy pattern ties
registration to the generator's lifecycle, which is bounded by aclose()
/ exhaustion / cancellation.
"""
fake = FakeServer()
fake.script([])
asgi = httpx.ASGITransport(app=fake.app)
async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
threads = ThreadsClient(HttpClient(raw))
async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
await thread.run.start(input={})
_ = thread.subscribe(["lifecycle"]) # construct but never iterate
# Subscription is not registered yet — the generator body hasn't run.
assert len(thread._subscriptions) == 0
async def test_values_projection_registers_via_delegation_not_controller_directly():
"""Values projection must register subscriptions through AsyncThreadStream delegation wrappers.
Verifies that subscriptions created by `thread.values` are visible via
`thread._subscriptions` (the delegated property) and also accessible in
`thread._controller._subscriptions`, confirming both views are consistent.
The test also confirms that the `_ValuesProjection` never bypasses
`AsyncThreadStream._register_subscription` to write to the controller
directly — the subscription count seen through the thread wrapper equals
the count inside the controller at the moment the subscription is live.
"""
fake = FakeServer()
fake.script([lifecycle_completed_event(seq=0)])
fake.set_state({"ok": True})
asgi = httpx.ASGITransport(app=fake.app)
counts_during: list[tuple[int, int]] = []
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={})
raw_controller = thread._controller
assert raw_controller is not None
# `_controller` returns `self` (AsyncThreadStream) as a duck-typed
# controller surface; StreamController is parallel groundwork
# used elsewhere.
controller: StreamController = raw_controller # ty: ignore[invalid-assignment]
assert len(thread._subscriptions) == 0
# Collect counts while the subscription is live (first snapshot only).
aiter: AsyncGenerator[Any, None] = cast(
"AsyncGenerator[Any, None]", thread.values.__aiter__()
)
# Advance to first item — subscription must be registered by now.
await aiter.__anext__()
counts_during.append(
(len(thread._subscriptions), len(controller._subscriptions))
)
# Close the iterator explicitly so the finally block runs immediately.
await aiter.aclose()
# Both views must have agreed — no bypass of the delegation wrapper.
assert len(counts_during) == 1
thread_count, ctrl_count = counts_during[0]
assert thread_count == ctrl_count
assert thread_count >= 1
def _make_handle(
events: list[dict[str, Any]],
err: BaseException | None = None,
) -> tuple[EventStreamHandle, asyncio.Queue]:
"""Build a synthetic EventStreamHandle for reconnect tests.
Returns the handle and the underlying queue so callers can inject events
or the end sentinel directly from test code.
"""
loop = asyncio.get_running_loop()
ready: asyncio.Future[None] = loop.create_future()
done: asyncio.Future[BaseException | None] = loop.create_future()
queue: asyncio.Queue = asyncio.Queue()
async def _pump() -> None:
ready.set_result(None)
for event in events:
await queue.put(event)
done.set_result(err)
await queue.put(None) # sentinel
asyncio.create_task(_pump()) # noqa: RUF006
async def _aiter():
while True:
item = await queue.get()
if item is None:
return
yield item
async def _close() -> None:
if not done.done():
done.set_result(None)
await queue.put(None)
return EventStreamHandle(
events=_aiter(), ready=ready, done=done, close=_close
), queue
async def test_shared_stream_reconnects_with_since_after_transport_drop():
"""StreamController reopens the stream with `since` after a post-ready drop."""
opened_params: list[dict[str, Any]] = []
handle1, _ = _make_handle(
[values_event(seq=1, values={"counter": 1})],
err=RuntimeError("scripted async stream failure"),
)
handle2, _ = _make_handle([values_event(seq=2, values={"counter": 2})])
handles = [handle1, handle2]
transport = MagicMock(spec=ProtocolSseTransport)
def _open(params: dict[str, Any]) -> EventStreamHandle:
opened_params.append(dict(params))
return handles.pop(0)
transport.open_event_stream.side_effect = _open
async def gate() -> None:
return None
controller = StreamController(transport=transport, run_start_gate=gate)
sub = controller.register_subscription({"channels": ["values"]})
await controller.reconcile_stream({"channels": ["values"]})
controller.ensure_fanout_running()
first = await asyncio.wait_for(sub.queue.get(), timeout=1.0)
second = await asyncio.wait_for(sub.queue.get(), timeout=1.0)
end = await asyncio.wait_for(sub.queue.get(), timeout=1.0)
await controller.close()
assert first["seq"] == 1
assert second["seq"] == 2
assert end is None
assert opened_params[0]["channels"] == ["values"]
assert "since" not in opened_params[0]
assert opened_params[1]["channels"] == ["values"]
assert opened_params[1]["since"] == 1
async def test_shared_stream_reconnect_dedupes_replayed_overlap():
"""StreamController deduplicates events replayed on reconnect."""
handle1, _ = _make_handle(
[values_event(seq=1, values={"counter": 1})],
err=RuntimeError("scripted async stream failure"),
)
handle2, _ = _make_handle(
[
values_event(seq=1, values={"counter": 1}), # replayed overlap
values_event(seq=2, values={"counter": 2}),
]
)
handles = [handle1, handle2]
transport = MagicMock(spec=ProtocolSseTransport)
transport.open_event_stream.side_effect = lambda _params: handles.pop(0)
async def gate() -> None:
return None
controller = StreamController(transport=transport, run_start_gate=gate)
sub = controller.register_subscription({"channels": ["values"]})
await controller.reconcile_stream({"channels": ["values"]})
controller.ensure_fanout_running()
received = [
await asyncio.wait_for(sub.queue.get(), timeout=1.0),
await asyncio.wait_for(sub.queue.get(), timeout=1.0),
await asyncio.wait_for(sub.queue.get(), timeout=1.0),
]
await controller.close()
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"])]
# The shared SSE filter is a union of subscription params plus
# ``lifecycle`` (added by ``_compute_current_union`` so the fanout
# consumer can detect root-terminal events for projection-iterator
# termination). Match any request whose channels include ``values``.
values_requests = [
b for b in fake.stream_request_bodies if "values" in (b.get("channels") or [])
]
assert len(values_requests) == 1
assert values_requests[0]["since"] == 17
async def _drain(events: AsyncIterator[Any]) -> list[Any]:
return [event async for event in events]
async def test_shared_stream_stops_reconnecting_on_a_404():
fake = FakeServer()
fake.script_sequence(
[
_StreamScript(events=[values_event(seq=1)], fail_after=1),
_StreamScript(events=[], status=404),
]
)
async with httpx.AsyncClient(
transport=fake.transport, base_url="http://test"
) as raw:
threads = ThreadsClient(HttpClient(raw))
async with threads.stream(assistant_id="agent") as thread:
received = await asyncio.wait_for(
_drain(thread.subscribe(["values"])), timeout=5.0
)
assert [event["seq"] for event in received] == [1]
assert len(fake.stream_request_bodies) == 2
async def test_shared_stream_does_not_reconnect_after_a_4404_close(
monkeypatch: Any,
) -> None:
opened: list[dict[str, Any]] = []
def open_then_close_with_4404(
_self: ProtocolSseTransport, params: dict[str, Any]
) -> EventStreamHandle:
opened.append(params)
loop = asyncio.get_running_loop()
ready: asyncio.Future[None] = loop.create_future()
ready.set_result(None)
done: asyncio.Future[BaseException | None] = loop.create_future()
done.set_result(ConnectionClosedError(Close(4404, "not found"), None))
async def events() -> AsyncGenerator[Any, None]:
await asyncio.sleep(0)
yield values_event(seq=1)
async def close() -> None:
return None
return EventStreamHandle(events=events(), ready=ready, done=done, close=close)
monkeypatch.setattr(
ProtocolSseTransport, "open_event_stream", open_then_close_with_4404
)
async with httpx.AsyncClient(base_url="http://test") as raw:
threads = ThreadsClient(HttpClient(raw))
async with threads.stream(assistant_id="agent") as thread:
received = await asyncio.wait_for(
_drain(thread.subscribe(["values"])), timeout=5.0
)
assert [event["seq"] for event in received] == [1]
assert len(opened) == 1
Reference in New Issue View Git Blame Copy Permalink