mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-31 20:29:46 +02:00
Elior Nataf LackritzandGitHub
ea5f9cc9fb
Follow-up to #8540, which turned on `PLC0415` (import-outside-top-level) for checkpoint-postgres and checkpoint-sqlite. This does the remaining six packages: checkpoint, checkpoint-conformance, langgraph, prebuilt, cli, sdk-py. Scoped to tests, per @sydney-runkle's call on #8540: library code is exempted with `per-file-ignores`, since it still has deferred imports nobody has reviewed and mixing that in would make this hard to read. ## What changed Function-level imports across 56 test files moved to module level. Nine could not move and carry an explicit `# noqa: PLC0415` with a reason: | File | Why it stays local | |---|---| | `libs/langgraph/tests/test_deprecation.py` (4) | the import has to run inside `pytest.warns` for the warning to be observed | | `libs/langgraph/tests/test_serde_allowlist.py` | try/except guard, skips when langchain_core is absent | | `libs/langgraph/tests/test_delta_channel_benchmark.py` | optional psycopg probe | | `libs/checkpoint/tests/test_conformance_delta.py` (3) | protected by a module-level `pytest.importorskip`; hoisting past the guard turns a skip into a collection error | That last one is the trap: an import moved above `pytest.importorskip` silently defeats the guard. I hit it locally and it turned the skip into a `ModuleNotFoundError` at collection. Every file with an `importorskip` or `except ImportError` was checked by hand for this. ## Verification `make lint` and `make test` in each of the six: | Package | Tests | |---|---| | checkpoint | 156 passed, 17 skipped | | checkpoint-conformance | 1 passed | | langgraph | 1968 passed, 4 skipped | | prebuilt | 284 passed | | cli | 336 passed | | sdk-py | 493 passed | Also confirmed the rule actually fires: a throwaway test file with a function-level import is flagged in all six packages, and the source exemption holds.
540 lines
19 KiB
Python
540 lines
19 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
|
|
import httpx
|
|
import orjson
|
|
import pytest
|
|
from starlette.applications import Starlette
|
|
from starlette.responses import JSONResponse, Response
|
|
from starlette.routing import Route
|
|
|
|
from langgraph_sdk.stream.transport.http import (
|
|
EventStreamHandle,
|
|
ProtocolSseTransport,
|
|
_build_event_stream_body,
|
|
)
|
|
from streaming._events import lifecycle_event, values_event
|
|
from streaming._fake_server import FakeServer
|
|
|
|
|
|
async def test_event_stream_handle_constructs_with_open_state():
|
|
loop = asyncio.get_running_loop()
|
|
ready: asyncio.Future[None] = loop.create_future()
|
|
done: asyncio.Future[BaseException | None] = loop.create_future()
|
|
done.set_result(None)
|
|
closed = False
|
|
|
|
async def aiter_events():
|
|
if False:
|
|
yield # pragma: no cover
|
|
|
|
async def closer():
|
|
nonlocal closed
|
|
closed = True
|
|
|
|
handle = EventStreamHandle(
|
|
events=aiter_events(), ready=ready, done=done, close=closer
|
|
)
|
|
assert handle.ready is ready
|
|
assert handle.done is done
|
|
await handle.close()
|
|
assert closed is True
|
|
|
|
|
|
async def test_send_command_posts_json_and_returns_response():
|
|
|
|
fake = FakeServer()
|
|
transport = httpx.ASGITransport(app=fake.app)
|
|
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
|
sse = ProtocolSseTransport(client=client, thread_id="t-1")
|
|
result = await sse.send_command(
|
|
{
|
|
"id": 7,
|
|
"method": "run.start",
|
|
"params": {"input": {"x": 1}},
|
|
}
|
|
)
|
|
assert result == {"type": "success", "id": 7, "result": {"run_id": "run-1"}}
|
|
assert fake.received_commands[0]["id"] == 7
|
|
|
|
|
|
async def test_send_command_returns_none_on_202():
|
|
|
|
received: list[dict] = []
|
|
|
|
async def commands(request):
|
|
received.append(orjson.loads(await request.body()))
|
|
return Response(status_code=202)
|
|
|
|
app = Starlette(
|
|
routes=[Route("/threads/{thread_id}/commands", commands, methods=["POST"])]
|
|
)
|
|
transport = httpx.ASGITransport(app=app)
|
|
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
|
sse = ProtocolSseTransport(client=client, thread_id="t-1")
|
|
result = await sse.send_command({"id": 1, "method": "noop", "params": {}})
|
|
assert result is None
|
|
assert len(received) == 1
|
|
|
|
|
|
async def test_send_command_raises_when_closed():
|
|
|
|
fake = FakeServer()
|
|
transport = httpx.ASGITransport(app=fake.app)
|
|
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
|
sse = ProtocolSseTransport(client=client, thread_id="t-1")
|
|
await sse.close()
|
|
with pytest.raises(RuntimeError, match="closed"):
|
|
await sse.send_command({"id": 1, "method": "noop", "params": {}})
|
|
|
|
|
|
async def test_send_command_raises_http_error_on_4xx():
|
|
|
|
async def commands(_request):
|
|
return JSONResponse({"error": "bad request"}, status_code=400)
|
|
|
|
app = Starlette(
|
|
routes=[Route("/threads/{thread_id}/commands", commands, methods=["POST"])]
|
|
)
|
|
transport = httpx.ASGITransport(app=app)
|
|
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
|
sse = ProtocolSseTransport(client=client, thread_id="t-1")
|
|
with pytest.raises(httpx.HTTPStatusError):
|
|
await sse.send_command({"id": 1, "method": "noop", "params": {}})
|
|
|
|
|
|
async def test_open_event_stream_yields_scripted_events():
|
|
|
|
fake = FakeServer()
|
|
fake.script(
|
|
[
|
|
lifecycle_event(seq=0),
|
|
values_event(seq=1),
|
|
]
|
|
)
|
|
transport = httpx.ASGITransport(app=fake.app)
|
|
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
|
sse = ProtocolSseTransport(client=client, thread_id="t-1")
|
|
handle = sse.open_event_stream(
|
|
{"channels": ["lifecycle", "values"], "namespaces": [[]]}
|
|
)
|
|
await asyncio.wait_for(handle.ready, timeout=1.0)
|
|
received = [e async for e in handle.events]
|
|
await handle.close()
|
|
methods = [e["method"] for e in received]
|
|
assert methods == ["lifecycle", "values"]
|
|
|
|
|
|
async def test_open_event_stream_passes_since_in_body():
|
|
|
|
fake = FakeServer()
|
|
fake.script([])
|
|
transport = httpx.ASGITransport(app=fake.app)
|
|
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
|
sse = ProtocolSseTransport(client=client, thread_id="t-1")
|
|
handle = sse.open_event_stream({"channels": ["values"], "since": 42})
|
|
await asyncio.wait_for(handle.ready, timeout=1.0)
|
|
_ = [e async for e in handle.events]
|
|
await handle.close()
|
|
assert fake.stream_request_bodies[0]["since"] == 42
|
|
assert fake.stream_request_bodies[0]["channels"] == ["values"]
|
|
|
|
|
|
async def test_open_event_stream_close_cancels_in_flight_iteration():
|
|
|
|
fake = FakeServer()
|
|
fake.script(
|
|
[lifecycle_event(seq=i) for i in range(50)],
|
|
)
|
|
transport = httpx.ASGITransport(app=fake.app)
|
|
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
|
sse = ProtocolSseTransport(client=client, thread_id="t-1")
|
|
handle = sse.open_event_stream({"channels": ["lifecycle"]})
|
|
await asyncio.wait_for(handle.ready, timeout=1.0)
|
|
agen = handle.events
|
|
first = await agen.__anext__()
|
|
await handle.close()
|
|
# Further iteration should terminate cleanly, not hang.
|
|
with pytest.raises(StopAsyncIteration):
|
|
await asyncio.wait_for(agen.__anext__(), timeout=1.0)
|
|
assert first["method"] == "lifecycle"
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_transport_accepts_max_queue_size_kwarg():
|
|
transport = ProtocolSseTransport(
|
|
client=httpx.AsyncClient(),
|
|
thread_id="t1",
|
|
max_queue_size=42,
|
|
)
|
|
assert transport._max_queue_size == 42
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_transport_default_max_queue_size_is_1024():
|
|
transport = ProtocolSseTransport(
|
|
client=httpx.AsyncClient(),
|
|
thread_id="t1",
|
|
)
|
|
assert transport._max_queue_size == 1024
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_pump_backpressures_when_queue_full():
|
|
"""Slow consumer should not cause unbounded queue growth.
|
|
|
|
With maxsize=2 and a producer that emits 100 events before any consumption,
|
|
the pump must suspend on queue.put after the second enqueue rather than
|
|
buffer all 100. We verify by counting queue items observed at the suspension
|
|
point.
|
|
"""
|
|
q: asyncio.Queue[int] = asyncio.Queue(maxsize=2)
|
|
|
|
produced: list[int] = []
|
|
|
|
async def producer():
|
|
for i in range(100):
|
|
await q.put(i)
|
|
produced.append(i)
|
|
|
|
task = asyncio.create_task(producer())
|
|
await asyncio.sleep(0.05) # let producer fill and suspend
|
|
task.cancel()
|
|
with contextlib.suppress(asyncio.CancelledError):
|
|
await task
|
|
|
|
# Producer should have enqueued 2 items, then blocked waiting on a third
|
|
# put. The third item was attempted but never completed.
|
|
assert len(produced) == 2
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_mid_stream_error_after_ready_surfaces_on_done():
|
|
"""If the SSE response body iteration raises after headers/ready, the
|
|
error must be exposed on handle.done so callers can distinguish a clean
|
|
end from a transport failure."""
|
|
|
|
def handler(_request: httpx.Request) -> httpx.Response:
|
|
async def body():
|
|
yield b'event: message\ndata: {"jsonrpc": "2.0"}\n\n'
|
|
raise RuntimeError("simulated mid-stream drop")
|
|
|
|
return httpx.Response(200, content=body())
|
|
|
|
mock = httpx.MockTransport(handler)
|
|
async with httpx.AsyncClient(
|
|
transport=mock, base_url="http://example.com"
|
|
) as client:
|
|
transport = ProtocolSseTransport(
|
|
client=client,
|
|
thread_id="t1",
|
|
)
|
|
handle = transport.open_event_stream({"channels": ["values"]})
|
|
await handle.ready
|
|
async for _ in handle.events:
|
|
pass
|
|
err = await handle.done
|
|
assert isinstance(err, RuntimeError)
|
|
assert "mid-stream drop" in str(err)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_clean_stream_end_done_resolves_with_none():
|
|
"""A stream that ends without error must resolve `done` with None."""
|
|
|
|
def handler(_request: httpx.Request) -> httpx.Response:
|
|
async def body():
|
|
yield b'event: message\ndata: {"jsonrpc": "2.0"}\n\n'
|
|
|
|
return httpx.Response(200, content=body())
|
|
|
|
mock = httpx.MockTransport(handler)
|
|
async with httpx.AsyncClient(
|
|
transport=mock, base_url="http://example.com"
|
|
) as client:
|
|
transport = ProtocolSseTransport(
|
|
client=client,
|
|
thread_id="t1",
|
|
)
|
|
handle = transport.open_event_stream({"channels": ["values"]})
|
|
await handle.ready
|
|
async for _ in handle.events:
|
|
pass
|
|
err = await handle.done
|
|
assert err is None
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_send_command_empty_200_body_raises_runtime_error_not_decoder_error():
|
|
"""A 200 response with empty body must raise RuntimeError matching the
|
|
'did not return a valid response' contract, not orjson.JSONDecodeError."""
|
|
|
|
def handler(_request: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(200, content=b"")
|
|
|
|
mock = httpx.MockTransport(handler)
|
|
async with httpx.AsyncClient(
|
|
transport=mock, base_url="http://example.com"
|
|
) as client:
|
|
transport = ProtocolSseTransport(
|
|
client=client,
|
|
thread_id="t1",
|
|
)
|
|
with pytest.raises(RuntimeError, match="did not return a valid response"):
|
|
await transport.send_command(
|
|
{"command_id": 1, "method": "run.start", "params": {}}
|
|
)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_cancel_event_prevents_post_cancel_flush():
|
|
"""When the consumer cancels the handle mid-stream, the pump's decoder
|
|
flush MUST NOT emit additional events after the cancel point."""
|
|
|
|
received: list = []
|
|
|
|
def handler(_request: httpx.Request) -> httpx.Response:
|
|
async def body():
|
|
yield b'event: message\ndata: {"seq": 1}\n\n'
|
|
yield b'event: message\ndata: {"seq": 2}\n\n'
|
|
yield b'event: message\ndata: {"seq": 3}\n\n'
|
|
|
|
return httpx.Response(200, content=body())
|
|
|
|
mock = httpx.MockTransport(handler)
|
|
async with httpx.AsyncClient(
|
|
transport=mock, base_url="http://example.com"
|
|
) as client:
|
|
transport = ProtocolSseTransport(
|
|
client=client,
|
|
thread_id="t1",
|
|
)
|
|
handle = transport.open_event_stream({"channels": ["values"]})
|
|
await handle.ready
|
|
async for event in handle.events:
|
|
received.append(event)
|
|
if len(received) == 1:
|
|
await handle.close()
|
|
break
|
|
# After close, no further events should drain from the flush.
|
|
async for _ in handle.events:
|
|
pytest.fail("event yielded after close()")
|
|
assert len(received) == 1
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_open_event_stream_ready_rejects_on_5xx():
|
|
|
|
async def stream_events(_request):
|
|
return JSONResponse({"error": "boom"}, status_code=500)
|
|
|
|
app = Starlette(
|
|
routes=[
|
|
Route(
|
|
"/threads/{thread_id}/stream/events",
|
|
stream_events,
|
|
methods=["POST"],
|
|
)
|
|
]
|
|
)
|
|
transport = httpx.ASGITransport(app=app)
|
|
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
|
sse = ProtocolSseTransport(client=client, thread_id="t-1")
|
|
handle = sse.open_event_stream({"channels": ["lifecycle"]})
|
|
with pytest.raises(httpx.HTTPStatusError):
|
|
await asyncio.wait_for(handle.ready, timeout=1.0)
|
|
# Iterator should terminate cleanly (no hang).
|
|
with pytest.raises(StopAsyncIteration):
|
|
await asyncio.wait_for(handle.events.__anext__(), timeout=1.0)
|
|
await handle.close()
|
|
|
|
|
|
def test_build_event_stream_body_minimal_channels_only():
|
|
|
|
body = _build_event_stream_body({"channels": ["values"]})
|
|
assert body == {"channels": ["values"]}
|
|
|
|
|
|
def test_build_event_stream_body_includes_all_optional_fields():
|
|
|
|
body = _build_event_stream_body(
|
|
{
|
|
"channels": ["values", "messages"],
|
|
"namespaces": [["fetcher"]],
|
|
"depth": 2,
|
|
"since": 7,
|
|
}
|
|
)
|
|
assert body == {
|
|
"channels": ["values", "messages"],
|
|
"namespaces": [["fetcher"]],
|
|
"depth": 2,
|
|
"since": 7,
|
|
}
|
|
|
|
|
|
def test_build_event_stream_body_omits_since_when_not_int():
|
|
|
|
body = _build_event_stream_body({"channels": ["values"], "since": None})
|
|
assert "since" not in body
|
|
|
|
|
|
async def test_open_event_stream_raises_when_closed():
|
|
|
|
fake = FakeServer()
|
|
transport = httpx.ASGITransport(app=fake.app)
|
|
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
|
sse = ProtocolSseTransport(client=client, thread_id="t-1")
|
|
await sse.close()
|
|
with pytest.raises(RuntimeError, match="closed"):
|
|
sse.open_event_stream({"channels": ["lifecycle"]})
|
|
|
|
|
|
async def test_transport_close_cancels_open_event_streams():
|
|
|
|
fake = FakeServer()
|
|
fake.script([lifecycle_event(seq=i) for i in range(5)], delay=0.05)
|
|
transport = httpx.ASGITransport(app=fake.app)
|
|
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
|
sse = ProtocolSseTransport(client=client, thread_id="t-1")
|
|
handle = sse.open_event_stream({"channels": ["lifecycle"]})
|
|
await asyncio.wait_for(handle.ready, timeout=1.0)
|
|
# Closing the transport must terminate the open stream within a bounded time.
|
|
await asyncio.wait_for(sse.close(), timeout=1.0)
|
|
|
|
# Drain any already-queued events; the stream must end (not hang).
|
|
async def drain() -> None:
|
|
async for _ in handle.events:
|
|
pass
|
|
|
|
await asyncio.wait_for(drain(), timeout=1.0)
|
|
|
|
|
|
async def test_default_headers_forwarded_to_send_command():
|
|
"""Headers passed at construction are sent on every command request."""
|
|
|
|
fake = FakeServer()
|
|
transport = httpx.ASGITransport(app=fake.app)
|
|
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
|
sse = ProtocolSseTransport(
|
|
client=client,
|
|
thread_id="t-1",
|
|
headers={"X-Trace-Id": "abc123"},
|
|
)
|
|
await sse.send_command({"id": 1, "method": "run.start", "params": {}})
|
|
assert fake.command_request_headers[0].get("x-trace-id") == "abc123"
|
|
# content-type must not be clobbered by default headers
|
|
assert "application/json" in fake.command_request_headers[0].get("content-type", "")
|
|
|
|
|
|
async def test_default_headers_forwarded_to_open_event_stream():
|
|
"""Headers passed at construction are sent on every SSE stream request."""
|
|
|
|
fake = FakeServer()
|
|
fake.script([])
|
|
transport = httpx.ASGITransport(app=fake.app)
|
|
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
|
sse = ProtocolSseTransport(
|
|
client=client,
|
|
thread_id="t-1",
|
|
headers={"X-Trace-Id": "abc123"},
|
|
)
|
|
handle = sse.open_event_stream({"channels": ["lifecycle"]})
|
|
await asyncio.wait_for(handle.ready, timeout=1.0)
|
|
_ = [e async for e in handle.events]
|
|
await handle.close()
|
|
assert fake.stream_request_headers_list[0].get("x-trace-id") == "abc123"
|
|
# Fixed SSE headers must not be clobbered by default headers
|
|
assert "text/event-stream" in fake.stream_request_headers_list[0].get("accept", "")
|
|
|
|
|
|
async def test_default_headers_cannot_override_sse_fixed_headers():
|
|
"""Caller-supplied default headers must not override content-type or accept."""
|
|
|
|
fake = FakeServer()
|
|
fake.script([])
|
|
transport = httpx.ASGITransport(app=fake.app)
|
|
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
|
sse = ProtocolSseTransport(
|
|
client=client,
|
|
thread_id="t-1",
|
|
headers={
|
|
"content-type": "text/plain",
|
|
"accept": "application/json",
|
|
"cache-control": "max-age=3600",
|
|
},
|
|
)
|
|
handle = sse.open_event_stream({"channels": ["lifecycle"]})
|
|
await asyncio.wait_for(handle.ready, timeout=1.0)
|
|
_ = [e async for e in handle.events]
|
|
await handle.close()
|
|
hdrs = fake.stream_request_headers_list[0]
|
|
assert "application/json" in hdrs.get("content-type", "")
|
|
assert "text/event-stream" in hdrs.get("accept", "")
|
|
assert hdrs.get("cache-control") == "no-store"
|
|
|
|
|
|
async def test_fake_server_state_endpoint():
|
|
"""State endpoint returns the set state and increments the counter."""
|
|
|
|
fake = FakeServer()
|
|
fake.set_state({"foo": "bar"}, next=["node_a"])
|
|
transport = httpx.ASGITransport(app=fake.app)
|
|
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
|
resp = await client.get("/threads/t-1/state")
|
|
assert resp.status_code == 200
|
|
body = resp.json()
|
|
assert body["values"] == {"foo": "bar"}
|
|
assert body["next"] == ["node_a"]
|
|
assert body["tasks"] == []
|
|
assert body["metadata"] == {}
|
|
assert body["checkpoint"] is None
|
|
assert body["created_at"] is None
|
|
assert fake.state_request_count == 1
|
|
assert len(fake.state_request_headers) == 1
|
|
|
|
|
|
def test_values_event_builder_shape():
|
|
"""values_event produces the expected shape with params.data as the snapshot."""
|
|
|
|
evt = values_event(seq=1, values={"foo": 1})
|
|
assert evt["event_id"] == "evt-1"
|
|
assert evt["method"] == "values"
|
|
assert evt["params"]["data"] == {"values": {"foo": 1}}
|
|
assert evt["params"]["namespace"] == []
|
|
|
|
|
|
async def test_open_event_stream_done_records_post_ready_error():
|
|
|
|
event_data = values_event(seq=1)
|
|
|
|
class _FailAfterOneStream(httpx.AsyncByteStream):
|
|
async def __aiter__(self):
|
|
|
|
payload = orjson.dumps(event_data).decode()
|
|
yield f"id: {event_data.get('event_id', '')}\n".encode()
|
|
yield f"event: message\ndata: {payload}\n\n".encode()
|
|
raise RuntimeError("scripted async stream failure")
|
|
|
|
async def handler(_request: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(
|
|
200,
|
|
headers={"content-type": "text/event-stream"},
|
|
stream=_FailAfterOneStream(),
|
|
)
|
|
|
|
transport = httpx.MockTransport(handler)
|
|
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
|
sse = ProtocolSseTransport(client=client, thread_id="t-1")
|
|
handle = sse.open_event_stream({"channels": ["values"]})
|
|
await asyncio.wait_for(handle.ready, timeout=1.0)
|
|
received = [event async for event in handle.events]
|
|
err = await asyncio.wait_for(handle.done, timeout=1.0)
|
|
await handle.close()
|
|
|
|
assert received == [values_event(seq=1)]
|
|
assert isinstance(err, RuntimeError)
|
|
assert "scripted async stream failure" in str(err)
|