From 5f24a0356a19401aab62aa6ff11f476cbfa2c4fd Mon Sep 17 00:00:00 2001 From: Nick Hollon Date: Thu, 16 Apr 2026 15:54:14 -0400 Subject: [PATCH] Tighten streaming run stream API and close review footguns - AsyncGraphRunStream.output/interrupted/interrupts are now methods (await run.output()), not properties returning coroutines. Forgetting `await` now fails at type-check time and at runtime on the common operations (bool/len/iter), instead of silently yielding a live coroutine that's truthy, lenless, and never awaited. - interrupted/interrupts re-raise the run's error on both lanes so a failed run doesn't silently return the last-known interrupt state. - Narrow the async pump catch from BaseException to Exception so CancelledError / KeyboardInterrupt propagate. - Wrap run.extensions with types.MappingProxyType so users can't add or remove projection keys behind the mux's back. - Add ValuesTransformer.error accessor; run stream stops reaching into _log._error. - Tighten StreamingHandler graph type from Any to Pregel and widen convert_to_protocol_event to accept StreamPart. - Projection-conflict ValueError now names the transformer that owns each colliding key, not just the incoming transformer. - Replace deprecated asyncio.get_event_loop() in the async iteration test with asyncio.create_task. - Document wall-clock semantics of ProtocolEvent.params.timestamp, the subgraph-namespace drop in MessagesTransformer, and the transformer-pipeline bypass for StreamChannel auto-forwarded events. - Add tests for the new error-raising behavior on interrupted / interrupts and for the read-only extensions contract. --- libs/langgraph/langgraph/stream/_convert.py | 20 ++-- libs/langgraph/langgraph/stream/_event_log.py | 5 +- libs/langgraph/langgraph/stream/_mux.py | 10 +- libs/langgraph/langgraph/stream/_types.py | 9 +- libs/langgraph/langgraph/stream/run_stream.py | 91 ++++++++++++------- .../langgraph/stream/stream_channel.py | 8 ++ .../langgraph/stream/streaming_handler.py | 8 +- .../langgraph/stream/transformers.py | 18 +++- .../langgraph/tests/test_streaming_handler.py | 68 ++++++++++++-- 9 files changed, 178 insertions(+), 59 deletions(-) diff --git a/libs/langgraph/langgraph/stream/_convert.py b/libs/langgraph/langgraph/stream/_convert.py index 1274fa7e7..c1a14b5b0 100644 --- a/libs/langgraph/langgraph/stream/_convert.py +++ b/libs/langgraph/langgraph/stream/_convert.py @@ -1,30 +1,32 @@ from __future__ import annotations import time -from typing import Any +from typing import Any, cast from langgraph.stream._types import ProtocolEvent, _ProtocolEventParams +from langgraph.types import StreamPart -def convert_to_protocol_event(part: dict[str, Any]) -> ProtocolEvent: - """Convert a v2 StreamPart dict to a ProtocolEvent. +def convert_to_protocol_event(part: StreamPart) -> ProtocolEvent: + """Convert a v2 StreamPart to a ProtocolEvent. Args: - part: A stream part dict with keys `type`, `ns`, `data`, and + part: A stream part with keys `type`, `ns`, `data`, and optionally `interrupts` (present on values events). Returns: The equivalent ProtocolEvent. """ + part_dict = cast(dict[str, Any], part) params: _ProtocolEventParams = { - "namespace": list(part["ns"]), + "namespace": list(part_dict["ns"]), "timestamp": int(time.time() * 1000), - "data": part["data"], + "data": part_dict["data"], } - if "interrupts" in part: - params["interrupts"] = part["interrupts"] + if "interrupts" in part_dict: + params["interrupts"] = part_dict["interrupts"] return { "type": "event", - "method": part["type"], + "method": part_dict["type"], "params": params, } diff --git a/libs/langgraph/langgraph/stream/_event_log.py b/libs/langgraph/langgraph/stream/_event_log.py index 01c22a442..f87c3b75f 100644 --- a/libs/langgraph/langgraph/stream/_event_log.py +++ b/libs/langgraph/langgraph/stream/_event_log.py @@ -198,7 +198,10 @@ class EventLog(Generic[T]): return elif self._request_more is not None: # Pull from the producer until this log gets a new item - # or the graph is exhausted (which closes the log). + # or the graph is exhausted (which closes the log). A push + # may evict the item this cursor was about to read; in that + # case the inner loop breaks and the outer `seq < _first_seq` + # check catches the overflow on the next iteration. while (seq - self._first_seq) >= len(self._items) and not self._closed: if not self._request_more(): break diff --git a/libs/langgraph/langgraph/stream/_mux.py b/libs/langgraph/langgraph/stream/_mux.py index 0481cd3eb..f3fab3965 100644 --- a/libs/langgraph/langgraph/stream/_mux.py +++ b/libs/langgraph/langgraph/stream/_mux.py @@ -78,6 +78,7 @@ class StreamMux: self.extensions: dict[str, Any] = {} self.native_keys: set[str] = set() + self._projection_owners: dict[str, str] = {} for transformer in transformers or (): self._register(transformer) @@ -103,14 +104,21 @@ class StreamMux: ) conflicts = set(projection) & set(self.extensions) if conflicts: + attributions = ", ".join( + f"{key!r} (owned by {self._projection_owners[key]})" + for key in sorted(conflicts) + ) raise ValueError( f"Transformer {type(transformer).__name__} returned " f"projection keys that conflict with already-registered " - f"keys: {conflicts}" + f"keys: {attributions}" ) self._transformers.append(transformer) self._bind_and_wire(projection) self.extensions.update(projection) + owner_name = type(transformer).__name__ + for key in projection: + self._projection_owners[key] = owner_name if getattr(transformer, "_native", False): self.native_keys.update(projection.keys()) diff --git a/libs/langgraph/langgraph/stream/_types.py b/libs/langgraph/langgraph/stream/_types.py index 60e3b5994..80fb630d4 100644 --- a/libs/langgraph/langgraph/stream/_types.py +++ b/libs/langgraph/langgraph/stream/_types.py @@ -12,7 +12,12 @@ _logger = logging.getLogger(__name__) class _ProtocolEventParams(TypedDict): - """Parameters for a protocol event.""" + """Parameters for a protocol event. + + `timestamp` is wall-clock milliseconds since the epoch and can go + backwards across NTP adjustments — use `ProtocolEvent.seq` for + ordering. + """ namespace: list[str] timestamp: int @@ -25,6 +30,8 @@ class ProtocolEvent(TypedDict): Wraps a raw stream part (values, messages, custom, etc.) in a uniform envelope with a monotonic sequence number assigned by the StreamMux. + Consumers that need a total order across events should use `seq`, not + `params.timestamp` (which is wall-clock and not monotonic). """ type: Literal["event"] diff --git a/libs/langgraph/langgraph/stream/run_stream.py b/libs/langgraph/langgraph/stream/run_stream.py index c17d5db76..dedfd9ab7 100644 --- a/libs/langgraph/langgraph/stream/run_stream.py +++ b/libs/langgraph/langgraph/stream/run_stream.py @@ -1,7 +1,8 @@ from __future__ import annotations import asyncio -from collections.abc import AsyncIterator, Iterator +from collections.abc import AsyncIterator, Iterator, Mapping +from types import MappingProxyType from typing import Any from langgraph.stream._convert import convert_to_protocol_event @@ -44,7 +45,7 @@ class GraphRunStream: """ self._graph_iter = graph_iter self._mux = mux - self.extensions = mux.extensions + self.extensions: Mapping[str, Any] = MappingProxyType(mux.extensions) self._values_transformer = values_transformer self._exhausted = False # Native-transformer projections also show up as direct attributes. @@ -98,20 +99,35 @@ class GraphRunStream: def output(self) -> dict[str, Any] | None: """Block until the run completes and return the final state.""" self._pump_all() - if self._values_transformer._log._error is not None: - raise self._values_transformer._log._error + err = self._values_transformer.error + if err is not None: + raise err return self._values_transformer._latest @property def interrupted(self) -> bool: - """Block until the run completes, then return whether it was interrupted.""" + """Block until the run completes, then return whether it was interrupted. + + Raises: + BaseException: If the run ended with an error. + """ self._pump_all() + err = self._values_transformer.error + if err is not None: + raise err return self._values_transformer._interrupted @property def interrupts(self) -> list[Any]: - """Block until the run completes, then return interrupt payloads.""" + """Block until the run completes, then return interrupt payloads. + + Raises: + BaseException: If the run ended with an error. + """ self._pump_all() + err = self._values_transformer.error + if err is not None: + raise err return self._values_transformer._interrupts def __iter__(self) -> Iterator[ProtocolEvent]: @@ -150,67 +166,76 @@ class AsyncGraphRunStream: pump_task: Background task pumping graph events into the mux. """ self._mux = mux - self.extensions = mux.extensions + self.extensions: Mapping[str, Any] = MappingProxyType(mux.extensions) self._values_transformer = values_transformer self._pump_task = pump_task # Native-transformer projections also show up as direct attributes. for key in mux.native_keys: setattr(self, key, mux.extensions[key]) - @property - def output(self) -> Any: - """Return an awaitable that resolves to the final state. + async def output(self) -> dict[str, Any] | None: + """Wait for the run to complete and return the final state. + + Methods (not properties) on the async lane so `run.output` without + `await` raises at type-check time instead of silently yielding a + coroutine object that's truthy, lenless, and never awaited. + + The pump routes any Exception into `mux.afail`, which surfaces on + `ValuesTransformer.error`. CancelledError / KeyboardInterrupt + propagate so cancellation isn't silently dropped. Example: ```python - output = await run.output + output = await run.output() ``` - """ - return self._get_output() - async def _get_output(self) -> dict[str, Any] | None: + Raises: + BaseException: If the run ended with an error. + """ try: await self._pump_task - except BaseException: + except Exception: pass - if self._values_transformer._log._error is not None: - raise self._values_transformer._log._error + if (err := self._values_transformer.error) is not None: + raise err return self._values_transformer._latest - @property - def interrupted(self) -> Any: - """Return an awaitable that resolves to whether the run was interrupted. + async def interrupted(self) -> bool: + """Wait for the run to complete and return whether it was interrupted. Example: ```python - interrupted = await run.interrupted + interrupted = await run.interrupted() ``` - """ - return self._get_interrupted() - async def _get_interrupted(self) -> bool: + Raises: + BaseException: If the run ended with an error. + """ try: await self._pump_task - except BaseException: + except Exception: pass + if (err := self._values_transformer.error) is not None: + raise err return self._values_transformer._interrupted - @property - def interrupts(self) -> Any: - """Return an awaitable that resolves to interrupt payloads. + async def interrupts(self) -> list[Any]: + """Wait for the run to complete and return interrupt payloads. Example: ```python - interrupts = await run.interrupts + interrupts = await run.interrupts() ``` - """ - return self._get_interrupts() - async def _get_interrupts(self) -> list[Any]: + Raises: + BaseException: If the run ended with an error. + """ try: await self._pump_task - except BaseException: + except Exception: pass + if (err := self._values_transformer.error) is not None: + raise err return self._values_transformer._interrupts def __aiter__(self) -> AsyncIterator[ProtocolEvent]: diff --git a/libs/langgraph/langgraph/stream/stream_channel.py b/libs/langgraph/langgraph/stream/stream_channel.py index a863c556f..7174002df 100644 --- a/libs/langgraph/langgraph/stream/stream_channel.py +++ b/libs/langgraph/langgraph/stream/stream_channel.py @@ -17,6 +17,14 @@ class StreamChannel(Generic[T]): `ProtocolEvent` into the main event stream using the channel's name as the method. + Auto-forwarded events bypass the transformer pipeline — other + transformers' `process()` / `aprocess()` methods do not see + `custom:` events produced by a channel push. This prevents a + transformer that pushes to its own channel during `process()` from + re-triggering itself, but it also means filter- or tap-style + transformers cannot observe channel output from peer transformers. + Consumers that need that should iterate the main event stream. + In-process consumers iterate the channel directly (`for item in ch` or `async for item in ch`). Remote SDK clients subscribe via `session.subscribe("custom:")`. diff --git a/libs/langgraph/langgraph/stream/streaming_handler.py b/libs/langgraph/langgraph/stream/streaming_handler.py index 373d74c72..9748c48b5 100644 --- a/libs/langgraph/langgraph/stream/streaming_handler.py +++ b/libs/langgraph/langgraph/stream/streaming_handler.py @@ -6,6 +6,7 @@ from typing import Any from langchain_core.runnables import RunnableConfig +from langgraph.pregel import Pregel from langgraph.stream._convert import convert_to_protocol_event from langgraph.stream._mux import StreamMux from langgraph.stream._types import StreamTransformer @@ -38,15 +39,16 @@ class StreamingHandler: print(state) output = run.output - # Async + # Async — terminal accessors are methods so a missing `await` + # fails loudly instead of silently yielding a coroutine. run = await handler.astream(input_data) async for state in run.values: print(state) - output = await run.output + output = await run.output() ``` """ - def __init__(self, graph: Any) -> None: + def __init__(self, graph: Pregel) -> None: """Initialize the handler. Args: diff --git a/libs/langgraph/langgraph/stream/transformers.py b/libs/langgraph/langgraph/stream/transformers.py index 1c6805a39..3fe7c5ea9 100644 --- a/libs/langgraph/langgraph/stream/transformers.py +++ b/libs/langgraph/langgraph/stream/transformers.py @@ -11,6 +11,9 @@ class ValuesTransformer(StreamTransformer): Native transformer — projection keys are exposed as direct attributes on the run stream (e.g. `run.values`). + + Only root-namespace values events are captured; subgraph state + snapshots are ignored. """ _native = True @@ -24,11 +27,18 @@ class ValuesTransformer(StreamTransformer): def init(self) -> dict[str, Any]: return {"values": self._log} + @property + def error(self) -> BaseException | None: + """The error that ended the run, or `None` if it succeeded. + + Set by the mux when it auto-fails the projection log. + """ + return self._log._error + def process(self, event: ProtocolEvent) -> bool: if event["method"] != "values": return True params = event["params"] - # Only capture root namespace events if params["namespace"]: return True self._latest = params["data"] @@ -47,6 +57,11 @@ class MessagesTransformer(StreamTransformer): A follow-on PR will replace this with a richer transformer that produces ChatModelStream objects using the protocol handler. + Only root-namespace messages events are captured; tokens emitted + from subgraphs are dropped from the `messages` projection. Consumers + that need subgraph tokens should iterate the raw event stream or + register a custom transformer. + Native transformer — projection keys are exposed as direct attributes on the run stream (e.g. `run.messages`). """ @@ -63,7 +78,6 @@ class MessagesTransformer(StreamTransformer): if event["method"] != "messages": return True params = event["params"] - # Only capture root namespace events if params["namespace"]: return True self._log.push(params["data"]) diff --git a/libs/langgraph/tests/test_streaming_handler.py b/libs/langgraph/tests/test_streaming_handler.py index 58020ea0f..d4797e4c1 100644 --- a/libs/langgraph/tests/test_streaming_handler.py +++ b/libs/langgraph/tests/test_streaming_handler.py @@ -176,8 +176,9 @@ class TestEventLog: log.push(i) log.close() - asyncio.get_event_loop().call_soon(lambda: asyncio.ensure_future(producer())) + producer_task = asyncio.create_task(producer()) items = [item async for item in log] + await producer_task assert items == [0, 1, 2] @pytest.mark.anyio @@ -420,6 +421,16 @@ class TestStreamingHandlerSync: assert run.values is run.extensions["values"] assert run.messages is run.extensions["messages"] + def test_extensions_is_read_only(self) -> None: + """`run.extensions` must reject mutations so users can't corrupt mux state.""" + graph = _build_simple_graph() + handler = StreamingHandler(graph) + run = handler.stream({"value": "x", "items": []}) + with pytest.raises(TypeError): + run.extensions["new_key"] = object() # type: ignore[index] + with pytest.raises(TypeError): + del run.extensions["values"] # type: ignore[attr-defined] + def test_custom_stream_events(self) -> None: graph = _build_custom_stream_graph() handler = StreamingHandler(graph) @@ -452,6 +463,22 @@ class TestStreamingHandlerSyncErrors: with pytest.raises(ValueError, match="boom"): list(run) + def test_error_propagation_interrupted(self) -> None: + """`run.interrupted` should raise on a failed run, not silently return False.""" + graph = _build_error_graph() + handler = StreamingHandler(graph) + run = handler.stream({"value": "x", "items": []}) + with pytest.raises(ValueError, match="boom"): + _ = run.interrupted + + def test_error_propagation_interrupts(self) -> None: + """`run.interrupts` should raise on a failed run.""" + graph = _build_error_graph() + handler = StreamingHandler(graph) + run = handler.stream({"value": "x", "items": []}) + with pytest.raises(ValueError, match="boom"): + _ = run.interrupts + class TestStreamingHandlerSyncInterrupt: def test_interrupted(self) -> None: @@ -490,7 +517,7 @@ class TestStreamingHandlerAsync: graph = _build_simple_graph() handler = StreamingHandler(graph) run = await handler.astream({"value": "x", "items": []}) - output = await run.output + output = await run.output() assert output is not None assert output["value"] == "xAB" assert output["items"] == ["a", "b"] @@ -512,7 +539,7 @@ class TestStreamingHandlerAsync: graph = _build_simple_graph() handler = StreamingHandler(graph) run = await handler.astream({"value": "x", "items": []}) - _ = await run.output + _ = await run.output() assert "values" in run.extensions assert "messages" in run.extensions assert run.values is run.extensions["values"] @@ -527,7 +554,7 @@ class TestStreamingHandlerAsyncErrors: handler = StreamingHandler(graph) run = await handler.astream({"value": "x", "items": []}) with pytest.raises(ValueError, match="boom"): - await run.output + await run.output() @pytest.mark.anyio @NEEDS_CONTEXTVARS @@ -549,6 +576,26 @@ class TestStreamingHandlerAsyncErrors: async for _ in run: pass + @pytest.mark.anyio + @NEEDS_CONTEXTVARS + async def test_error_propagation_interrupted(self) -> None: + """`await run.interrupted()` should raise on a failed async run.""" + graph = _build_error_graph() + handler = StreamingHandler(graph) + run = await handler.astream({"value": "x", "items": []}) + with pytest.raises(ValueError, match="boom"): + await run.interrupted() + + @pytest.mark.anyio + @NEEDS_CONTEXTVARS + async def test_error_propagation_interrupts(self) -> None: + """`await run.interrupts()` should raise on a failed async run.""" + graph = _build_error_graph() + handler = StreamingHandler(graph) + run = await handler.astream({"value": "x", "items": []}) + with pytest.raises(ValueError, match="boom"): + await run.interrupts() + class TestStreamingHandlerAsyncInterrupt: @pytest.mark.anyio @@ -560,9 +607,9 @@ class TestStreamingHandlerAsyncInterrupt: {"value": "x", "items": []}, {"configurable": {"thread_id": "t2"}}, ) - _ = await run.output - assert await run.interrupted is True - assert len(await run.interrupts) > 0 + _ = await run.output() + assert await run.interrupted() is True + assert len(await run.interrupts()) > 0 class TestStreamingHandlerAsyncCustom: @@ -1042,7 +1089,10 @@ class TestCustomTransformer: graph = _build_simple_graph() handler = StreamingHandler(graph) - with pytest.raises(ValueError, match="conflict.*{'values'}"): + with pytest.raises( + ValueError, + match=r"conflict.*'values'.*ValuesTransformer", + ): handler.stream( {"value": "x", "items": []}, transformers=[ConflictTransformer()], @@ -1442,7 +1492,7 @@ class TestAsyncTransformerLane: {"value": "x", "items": []}, transformers=[Scorer()], ) - _ = await run.output + _ = await run.output() scores = [x async for x in run.extensions["scores"]] assert scores and all(s == 42 for s in scores)