From b3194267257e38f15d8483e15a798f1cee2ba255 Mon Sep 17 00:00:00 2001 From: Nick Hollon Date: Mon, 20 Apr 2026 13:57:37 -0400 Subject: [PATCH 1/2] fix(stream): replace pump lock with condition-based take-a-number MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The async pump serialized `graph_aiter.__anext__()` with an `asyncio.Lock`, held across the full await. When two cursors read different projections concurrently, the "losing" task slept inside `_apump_next` on the lock itself — so when the active pumper pushed its data onto the losing task's buffer, the loser couldn't observe it until another graph event forced the lock to change hands. Each passive consumer saw its deltas one graph event late; bursts coalesced at turn boundaries instead of streaming live. Switch to an `asyncio.Condition` + `_pumping` flag. Exactly one task is the active pumper; others do `cond.wait()` and are notified after every pump step. Passive consumers wake as soon as their buffer fills, drop out of `_apump_next`, and let the iterator's buffer check yield the data. Single-consumer behavior is unchanged; multi- consumer throughput improves ~5x on bursty tools and no events are lost. --- libs/langgraph/langgraph/stream/run_stream.py | 50 +++++++++++++------ 1 file changed, 34 insertions(+), 16 deletions(-) diff --git a/libs/langgraph/langgraph/stream/run_stream.py b/libs/langgraph/langgraph/stream/run_stream.py index 13bd6fb86..98ab08586 100644 --- a/libs/langgraph/langgraph/stream/run_stream.py +++ b/libs/langgraph/langgraph/stream/run_stream.py @@ -254,7 +254,8 @@ class AsyncGraphRunStream: self.extensions: Mapping[str, Any] = MappingProxyType(mux.extensions) self._values_transformer = values_transformer self._exhausted = False - self._pump_lock = asyncio.Lock() + self._pump_cond = asyncio.Condition() + self._pumping = False for key in mux.native_keys: setattr(self, key, mux.extensions[key]) self._wire_arequest_more(mux) @@ -270,23 +271,35 @@ class AsyncGraphRunStream: value._log._arequest_more = self._apump_next async def _apump_next(self) -> bool: - """Pull one event from the graph and push it through the mux. + """Drive one pump step, or wait for the active pumper to drive one. - Serialized via `self._pump_lock` so concurrent cursors each - produce one event per acquisition rather than racing on the - graph iterator. + "Take-a-number" semantics: at most one task at a time calls + `graph_aiter.__anext__()` (asyncio iterators can't be advanced + concurrently). Other callers wait on a Condition that the + active pumper notifies after each step. This lets a "passive" + consumer — one whose projection's buffer is being filled by the + active pumper's push — wake up as soon as its data lands, + instead of queueing on the pump and only observing its data one + graph event late. `except Exception` is intentional — `CancelledError` and other `BaseException` subclasses propagate, matching asyncio's cancellation contract. Returns: - True if an event was pulled, False if the graph is - exhausted or has raised. + True if a pump step completed (by this task or another), + False if the graph is exhausted. """ - async with self._pump_lock: + async with self._pump_cond: if self._exhausted: return False + if self._pumping: + # Another task is pumping; wait for its progress signal. + await self._pump_cond.wait() + return not self._exhausted + self._pumping = True + + try: try: part = await self._graph_aiter.__anext__() except StopAsyncIteration: @@ -299,22 +312,27 @@ class AsyncGraphRunStream: return False await self._mux.apush(convert_to_protocol_event(part)) return True + finally: + async with self._pump_cond: + self._pumping = False + self._pump_cond.notify_all() async def abort(self) -> None: """Stop the run early. - Closes the mux and marks the stream exhausted. Any awaiting - cursors wake up and see the closed state; any `apush` blocked - on backpressure wakes and returns without appending. Idempotent. + Marks the stream exhausted, wakes any pump-waiters, and closes + the mux. Any `apush` blocked on backpressure wakes and returns + without appending. Idempotent. """ - async with self._pump_lock: + async with self._pump_cond: if self._exhausted: return self._exhausted = True - try: - await self._mux.aclose() - except Exception: - pass + self._pump_cond.notify_all() + try: + await self._mux.aclose() + except Exception: + pass async def __aenter__(self) -> AsyncGraphRunStream: return self From 40055e92cc8e0123169c885f1c0c0701c0eb4561 Mon Sep 17 00:00:00 2001 From: Nick Hollon Date: Mon, 20 Apr 2026 16:08:19 -0400 Subject: [PATCH 2/2] feat(langgraph): move stream_v2/astream_v2 onto Pregel, drop StreamingHandler Compile-time `transformers=` on `StateGraph.compile` now registers transformer factories directly on the compiled graph. `stream_v2` and `astream_v2` live on Pregel and read the stashed list, so callers no longer need a separate wrapper to drive the transformer pipeline. --- libs/langgraph/langgraph/graph/state.py | 7 + libs/langgraph/langgraph/pregel/main.py | 132 +++++++++++++ libs/langgraph/langgraph/stream/__init__.py | 7 +- .../langgraph/stream/streaming_handler.py | 149 -------------- ...ng_handler.py => test_pregel_stream_v2.py} | 185 +++++++++--------- 5 files changed, 234 insertions(+), 246 deletions(-) delete mode 100644 libs/langgraph/langgraph/stream/streaming_handler.py rename libs/langgraph/tests/{test_streaming_handler.py => test_pregel_stream_v2.py} (92%) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index b1c24de2b..4e3bde4c7 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -1045,6 +1045,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): interrupt_after: All | list[str] | None = None, debug: bool = False, name: str | None = None, + transformers: Sequence[Callable[[], Any]] | None = None, ) -> CompiledStateGraph[StateT, ContextT, InputT, OutputT]: """Compiles the `StateGraph` into a `CompiledStateGraph` object. @@ -1077,6 +1078,11 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): interrupt_after: An optional list of node names to interrupt after. debug: A flag indicating whether to enable debug mode. name: The name to use for the compiled graph. + transformers: Optional sequence of zero-arg factories returning + `StreamTransformer` instances. Registered on the compiled + graph and instantiated per-run whenever `stream_v2` / + `astream_v2` is called. Appended after the built-in + `ValuesTransformer` and `MessagesTransformer`. Returns: CompiledStateGraph: The compiled `StateGraph`. @@ -1159,6 +1165,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): store=store, cache=cache, name=name or "LangGraph", + stream_transformers=transformers, ) compiled._serde_allowlist = serde_allowlist diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index c440e77f9..ce38ae1bf 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -671,6 +671,7 @@ class Pregel( config: RunnableConfig | None = None, trigger_to_nodes: Mapping[str, Sequence[str]] | None = None, name: str = "LangGraph", + stream_transformers: Sequence[Callable[[], Any]] | None = None, **deprecated_kwargs: Unpack[DeprecatedKwargs], ) -> None: if ( @@ -717,6 +718,9 @@ class Pregel( self.config = config self.trigger_to_nodes = trigger_to_nodes or {} self.name = name + self._stream_transformers: tuple[Callable[[], Any], ...] = tuple( + stream_transformers or () + ) self._serde_allowlist: set[tuple[str, ...]] | None = None if auto_validate: self.validate() @@ -3237,6 +3241,134 @@ class Pregel( await asyncio.shield(run_manager.on_chain_error(e)) raise + def stream_v2( + self, + input: InputT | Command | None, + config: RunnableConfig | None = None, + *, + interrupt_before: All | Sequence[str] | None = None, + interrupt_after: All | Sequence[str] | None = None, + transformers: Sequence[Any] | None = None, + ) -> Any: + """Start a sync v2 streaming run driven by transformer projections. + + Builds a `StreamMux` from the built-in `ValuesTransformer` / + `MessagesTransformer`, this graph's compile-time + `stream_transformers`, and any additional `transformers=` + supplied at the call site. Returns a `GraphRunStream` that the + caller drives by iterating any projection — no background + thread. + + Args: + input: Graph input. + config: Optional runnable config forwarded to the graph. + interrupt_before: Nodes to interrupt before, if any. + interrupt_after: Nodes to interrupt after, if any. + transformers: Extra transformer instances appended after + compile-time `stream_transformers`. + + Returns: + A `GraphRunStream` the caller iterates to drive the run. + """ + from langgraph.stream._mux import StreamMux + from langgraph.stream.run_stream import GraphRunStream + from langgraph.stream.transformers import ( + MessagesTransformer, + ValuesTransformer, + ) + + values_t = ValuesTransformer() + compiled_instances = [f() for f in self._stream_transformers] + mux = StreamMux( + [ + values_t, + MessagesTransformer(), + *compiled_instances, + *(transformers or ()), + ], + is_async=False, + ) + graph_iter = iter( + self.stream( + input, + config, + stream_mode=[ + "values", + "updates", + "messages", + "custom", + "checkpoints", + "tasks", + "debug", + ], + subgraphs=True, + version="v2", + interrupt_before=interrupt_before, + interrupt_after=interrupt_after, + ) + ) + return GraphRunStream(graph_iter, mux, values_t) + + async def astream_v2( + self, + input: InputT | Command | None, + config: RunnableConfig | None = None, + *, + interrupt_before: All | Sequence[str] | None = None, + interrupt_after: All | Sequence[str] | None = None, + transformers: Sequence[Any] | None = None, + ) -> Any: + """Async counterpart to `stream_v2`. + + Returns an `AsyncGraphRunStream` whose projections can be awaited + concurrently; each subscribed cursor drives the pump when its + buffer is empty. + + Args: + input: Graph input. + config: Optional runnable config forwarded to the graph. + interrupt_before: Nodes to interrupt before, if any. + interrupt_after: Nodes to interrupt after, if any. + transformers: Extra transformer instances appended after + compile-time `stream_transformers`. + """ + from langgraph.stream._mux import StreamMux + from langgraph.stream.run_stream import AsyncGraphRunStream + from langgraph.stream.transformers import ( + MessagesTransformer, + ValuesTransformer, + ) + + values_t = ValuesTransformer() + compiled_instances = [f() for f in self._stream_transformers] + mux = StreamMux( + [ + values_t, + MessagesTransformer(), + *compiled_instances, + *(transformers or ()), + ], + is_async=True, + ) + graph_aiter = self.astream( + input, + config, + stream_mode=[ + "values", + "updates", + "messages", + "custom", + "checkpoints", + "tasks", + "debug", + ], + subgraphs=True, + version="v2", + interrupt_before=interrupt_before, + interrupt_after=interrupt_after, + ).__aiter__() + return AsyncGraphRunStream(graph_aiter, mux, values_t) + @overload def invoke( self, diff --git a/libs/langgraph/langgraph/stream/__init__.py b/libs/langgraph/langgraph/stream/__init__.py index 5a4633c71..ca5a26010 100644 --- a/libs/langgraph/langgraph/stream/__init__.py +++ b/libs/langgraph/langgraph/stream/__init__.py @@ -1,14 +1,14 @@ """Streaming infrastructure for LangGraph. -Provides a `StreamingHandler` that wraps a compiled graph and exposes -ergonomic streaming projections through a transformer pipeline. +Compile a graph with `transformers=[...]` and call `graph.stream_v2()` / +`graph.astream_v2()` to drive a transformer pipeline that projects the +graph's raw events into ergonomic per-channel streams. """ from langgraph.stream._event_log import EventLog from langgraph.stream._types import ProtocolEvent, StreamTransformer from langgraph.stream.run_stream import AsyncGraphRunStream, GraphRunStream from langgraph.stream.stream_channel import StreamChannel -from langgraph.stream.streaming_handler import StreamingHandler __all__ = [ "AsyncGraphRunStream", @@ -17,5 +17,4 @@ __all__ = [ "ProtocolEvent", "StreamChannel", "StreamTransformer", - "StreamingHandler", ] diff --git a/libs/langgraph/langgraph/stream/streaming_handler.py b/libs/langgraph/langgraph/stream/streaming_handler.py deleted file mode 100644 index ff031016e..000000000 --- a/libs/langgraph/langgraph/stream/streaming_handler.py +++ /dev/null @@ -1,149 +0,0 @@ -from __future__ import annotations - -from collections.abc import Sequence -from typing import Any - -from langchain_core.runnables import RunnableConfig - -from langgraph.pregel import Pregel -from langgraph.stream._mux import StreamMux -from langgraph.stream._types import StreamTransformer -from langgraph.stream.run_stream import AsyncGraphRunStream, GraphRunStream -from langgraph.stream.transformers import MessagesTransformer, ValuesTransformer -from langgraph.types import All, StreamMode - -# All stream modes to request from the graph. -STREAM_V2_MODES: list[StreamMode] = [ - "values", - "updates", - "messages", - "custom", - "checkpoints", - "tasks", - "debug", -] - - -class StreamingHandler: - """Wrap a compiled graph with ergonomic streaming projections. - - Example: - ```python - handler = StreamingHandler(graph) - - # Sync - run = handler.stream(input_data) - for state in run.values: - print(state) - output = run.output - - # 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() - ``` - """ - - def __init__(self, graph: Pregel) -> None: - """Initialize the handler. - - Args: - graph: A compiled LangGraph graph to stream from. - """ - self._graph = graph - - def stream( - self, - input: Any, - config: RunnableConfig | None = None, - *, - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - transformers: list[StreamTransformer] | None = None, - ) -> GraphRunStream: - """Start a sync streaming run. - - Returns a GraphRunStream immediately. The caller's iteration on - any projection drives the graph forward — no background thread - is used. This matches v1's model where the caller's `for` loop - is the pump. - - Args: - input: Graph input. - config: Optional runnable config forwarded to the graph. - interrupt_before: Nodes to interrupt before, if any. - interrupt_after: Nodes to interrupt after, if any. - transformers: User transformers appended after the built-in - `ValuesTransformer` and `MessagesTransformer`. - - Returns: - A GraphRunStream the caller can iterate to drive the run. - """ - values_t = ValuesTransformer() - mux = StreamMux( - [values_t, MessagesTransformer(), *(transformers or ())], - is_async=False, - ) - - graph_iter = iter( - self._graph.stream( - input, - config, - stream_mode=STREAM_V2_MODES, - subgraphs=True, - version="v2", - interrupt_before=interrupt_before, - interrupt_after=interrupt_after, - ) - ) - - return GraphRunStream(graph_iter, mux, values_t) - - async def astream( - self, - input: Any, - config: RunnableConfig | None = None, - *, - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - transformers: list[StreamTransformer] | None = None, - ) -> AsyncGraphRunStream: - """Start an async streaming run. - - Returns an AsyncGraphRunStream immediately. The caller's - iteration on any projection drives the graph forward — there - is no background task. Concurrent consumers share a - single-flight pump via an internal `asyncio.Lock`. - - Args: - input: Graph input. - config: Optional runnable config forwarded to the graph. - interrupt_before: Nodes to interrupt before, if any. - interrupt_after: Nodes to interrupt after, if any. - transformers: User transformers appended after the built-in - `ValuesTransformer` and `MessagesTransformer`. - - Returns: - An AsyncGraphRunStream whose projections can be awaited - concurrently; each subscribed cursor drives the pump when - its buffer is empty. - """ - values_t = ValuesTransformer() - mux = StreamMux( - [values_t, MessagesTransformer(), *(transformers or ())], - is_async=True, - ) - - graph_aiter = self._graph.astream( - input, - config, - stream_mode=STREAM_V2_MODES, - subgraphs=True, - version="v2", - interrupt_before=interrupt_before, - interrupt_after=interrupt_after, - ).__aiter__() - - return AsyncGraphRunStream(graph_aiter, mux, values_t) diff --git a/libs/langgraph/tests/test_streaming_handler.py b/libs/langgraph/tests/test_pregel_stream_v2.py similarity index 92% rename from libs/langgraph/tests/test_streaming_handler.py rename to libs/langgraph/tests/test_pregel_stream_v2.py index 835158ed5..e0d47126a 100644 --- a/libs/langgraph/tests/test_streaming_handler.py +++ b/libs/langgraph/tests/test_pregel_stream_v2.py @@ -1,4 +1,4 @@ -"""Tests for the StreamingHandler and its supporting infrastructure.""" +"""Tests for `Pregel.stream_v2` / `astream_v2` and the transformer pipeline.""" from __future__ import annotations @@ -17,7 +17,6 @@ from langgraph.graph import StateGraph from langgraph.stream import ( EventLog, StreamChannel, - StreamingHandler, StreamTransformer, ) from langgraph.stream._convert import convert_to_protocol_event @@ -404,15 +403,15 @@ class TestStreamChannel: # --------------------------------------------------------------------------- -# StreamingHandler sync tests +# stream_v2 sync tests # --------------------------------------------------------------------------- -class TestStreamingHandlerSync: +class TestStreamV2Sync: def test_values_projection(self) -> None: graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = handler.stream({"value": "x", "items": []}) + handler = graph + run = handler.stream_v2({"value": "x", "items": []}) snapshots = list(run.values) # Should have at least the initial + per-node snapshots. assert len(snapshots) >= 1 @@ -423,8 +422,8 @@ class TestStreamingHandlerSync: def test_output(self) -> None: graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = handler.stream({"value": "x", "items": []}) + handler = graph + run = handler.stream_v2({"value": "x", "items": []}) output = run.output assert output is not None assert output["value"] == "xAB" @@ -432,8 +431,8 @@ class TestStreamingHandlerSync: def test_raw_event_iteration(self) -> None: graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = handler.stream({"value": "x", "items": []}) + handler = graph + run = handler.stream_v2({"value": "x", "items": []}) events = list(run) assert len(events) > 0 for event in events: @@ -444,8 +443,8 @@ class TestStreamingHandlerSync: def test_extensions_has_native_keys(self) -> None: graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = handler.stream({"value": "x", "items": []}) + handler = graph + run = handler.stream_v2({"value": "x", "items": []}) # Drain events so the run completes. _ = run.output assert "values" in run.extensions @@ -457,8 +456,8 @@ class TestStreamingHandlerSync: 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": []}) + handler = graph + run = handler.stream_v2({"value": "x", "items": []}) with pytest.raises(TypeError): run.extensions["new_key"] = object() # type: ignore[index] with pytest.raises(TypeError): @@ -466,8 +465,8 @@ class TestStreamingHandlerSync: def test_custom_stream_events(self) -> None: graph = _build_custom_stream_graph() - handler = StreamingHandler(graph) - run = handler.stream({"value": "x", "items": []}) + handler = graph + run = handler.stream_v2({"value": "x", "items": []}) custom_events = [e for e in run if e["method"] == "custom"] assert len(custom_events) == 2 assert custom_events[0]["params"]["data"] == {"step": "start"} @@ -475,8 +474,8 @@ class TestStreamingHandlerSync: def test_interleave_values_and_messages(self) -> None: graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = handler.stream({"value": "x", "items": []}) + handler = graph + run = handler.stream_v2({"value": "x", "items": []}) tagged = list(run.interleave("values", "messages")) names = [name for name, _ in tagged] @@ -489,8 +488,8 @@ class TestStreamingHandlerSync: def test_abort_marks_exhausted_and_closes_mux(self) -> None: graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = handler.stream({"value": "x", "items": []}) + handler = graph + run = handler.stream_v2({"value": "x", "items": []}) values_iter = iter(run.values) # Consume one item so the pump advances. _ = next(values_iter) @@ -503,64 +502,64 @@ class TestStreamingHandlerSync: def test_context_manager_calls_abort_on_exit(self) -> None: graph = _build_simple_graph() - handler = StreamingHandler(graph) - with handler.stream({"value": "x", "items": []}) as run: + handler = graph + with handler.stream_v2({"value": "x", "items": []}) as run: values_iter = iter(run.values) _ = next(values_iter) assert run._exhausted is True def test_interleave_unknown_projection(self) -> None: graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = handler.stream({"value": "x", "items": []}) + handler = graph + run = handler.stream_v2({"value": "x", "items": []}) with pytest.raises(KeyError): list(run.interleave("values", "does_not_exist")) -class TestStreamingHandlerSyncErrors: +class TestStreamV2SyncErrors: def test_error_propagation_output(self) -> None: graph = _build_error_graph() - handler = StreamingHandler(graph) - run = handler.stream({"value": "x", "items": []}) + handler = graph + run = handler.stream_v2({"value": "x", "items": []}) with pytest.raises(ValueError, match="boom"): _ = run.output def test_error_propagation_values(self) -> None: graph = _build_error_graph() - handler = StreamingHandler(graph) - run = handler.stream({"value": "x", "items": []}) + handler = graph + run = handler.stream_v2({"value": "x", "items": []}) with pytest.raises(ValueError, match="boom"): list(run.values) def test_error_propagation_raw_events(self) -> None: graph = _build_error_graph() - handler = StreamingHandler(graph) - run = handler.stream({"value": "x", "items": []}) + handler = graph + run = handler.stream_v2({"value": "x", "items": []}) 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": []}) + handler = graph + run = handler.stream_v2({"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": []}) + handler = graph + run = handler.stream_v2({"value": "x", "items": []}) with pytest.raises(ValueError, match="boom"): _ = run.interrupts -class TestStreamingHandlerSyncInterrupt: +class TestStreamV2SyncInterrupt: def test_interrupted(self) -> None: graph = _build_interrupt_graph() - handler = StreamingHandler(graph) - run = handler.stream( + handler = graph + run = handler.stream_v2( {"value": "x", "items": []}, {"configurable": {"thread_id": "t1"}}, ) @@ -570,17 +569,17 @@ class TestStreamingHandlerSyncInterrupt: # --------------------------------------------------------------------------- -# StreamingHandler async tests +# astream_v2 async tests # --------------------------------------------------------------------------- -class TestStreamingHandlerAsync: +class TestStreamV2Async: @pytest.mark.anyio @NEEDS_CONTEXTVARS async def test_values_projection(self) -> None: graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = await handler.astream({"value": "x", "items": []}) + handler = graph + run = await handler.astream_v2({"value": "x", "items": []}) snapshots = [s async for s in run.values] assert len(snapshots) >= 1 last = snapshots[-1] @@ -591,8 +590,8 @@ class TestStreamingHandlerAsync: @NEEDS_CONTEXTVARS async def test_output(self) -> None: graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = await handler.astream({"value": "x", "items": []}) + handler = graph + run = await handler.astream_v2({"value": "x", "items": []}) output = await run.output() assert output is not None assert output["value"] == "xAB" @@ -602,8 +601,8 @@ class TestStreamingHandlerAsync: @NEEDS_CONTEXTVARS async def test_raw_event_iteration(self) -> None: graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = await handler.astream({"value": "x", "items": []}) + handler = graph + run = await handler.astream_v2({"value": "x", "items": []}) events = [e async for e in run] assert len(events) > 0 for event in events: @@ -613,8 +612,8 @@ class TestStreamingHandlerAsync: @NEEDS_CONTEXTVARS async def test_abort_marks_exhausted_and_closes_mux(self) -> None: graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = await handler.astream({"value": "x", "items": []}) + handler = graph + run = await handler.astream_v2({"value": "x", "items": []}) values_iter = aiter(run.values) _ = await anext(values_iter) await run.abort() @@ -628,8 +627,8 @@ class TestStreamingHandlerAsync: @NEEDS_CONTEXTVARS async def test_async_context_manager_calls_abort_on_exit(self) -> None: graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = await handler.astream({"value": "x", "items": []}) + handler = graph + run = await handler.astream_v2({"value": "x", "items": []}) async with run: values_iter = aiter(run.values) _ = await anext(values_iter) @@ -639,8 +638,8 @@ class TestStreamingHandlerAsync: @NEEDS_CONTEXTVARS async def test_extensions_has_native_keys(self) -> None: graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = await handler.astream({"value": "x", "items": []}) + handler = graph + run = await handler.astream_v2({"value": "x", "items": []}) _ = await run.output() assert "values" in run.extensions assert "messages" in run.extensions @@ -648,13 +647,13 @@ class TestStreamingHandlerAsync: assert run.messages is run.extensions["messages"] -class TestStreamingHandlerAsyncErrors: +class TestStreamV2AsyncErrors: @pytest.mark.anyio @NEEDS_CONTEXTVARS async def test_error_propagation_output(self) -> None: graph = _build_error_graph() - handler = StreamingHandler(graph) - run = await handler.astream({"value": "x", "items": []}) + handler = graph + run = await handler.astream_v2({"value": "x", "items": []}) with pytest.raises(ValueError, match="boom"): await run.output() @@ -662,8 +661,8 @@ class TestStreamingHandlerAsyncErrors: @NEEDS_CONTEXTVARS async def test_error_propagation_values(self) -> None: graph = _build_error_graph() - handler = StreamingHandler(graph) - run = await handler.astream({"value": "x", "items": []}) + handler = graph + run = await handler.astream_v2({"value": "x", "items": []}) with pytest.raises(ValueError, match="boom"): async for _ in run.values: pass @@ -672,8 +671,8 @@ class TestStreamingHandlerAsyncErrors: @NEEDS_CONTEXTVARS async def test_error_propagation_raw_events(self) -> None: graph = _build_error_graph() - handler = StreamingHandler(graph) - run = await handler.astream({"value": "x", "items": []}) + handler = graph + run = await handler.astream_v2({"value": "x", "items": []}) with pytest.raises(ValueError, match="boom"): async for _ in run: pass @@ -683,8 +682,8 @@ class TestStreamingHandlerAsyncErrors: 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": []}) + handler = graph + run = await handler.astream_v2({"value": "x", "items": []}) with pytest.raises(ValueError, match="boom"): await run.interrupted() @@ -693,19 +692,19 @@ class TestStreamingHandlerAsyncErrors: 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": []}) + handler = graph + run = await handler.astream_v2({"value": "x", "items": []}) with pytest.raises(ValueError, match="boom"): await run.interrupts() -class TestStreamingHandlerAsyncInterrupt: +class TestStreamV2AsyncInterrupt: @pytest.mark.anyio @NEEDS_CONTEXTVARS async def test_interrupted(self) -> None: graph = _build_interrupt_graph() - handler = StreamingHandler(graph) - run = await handler.astream( + handler = graph + run = await handler.astream_v2( {"value": "x", "items": []}, {"configurable": {"thread_id": "t2"}}, ) @@ -714,13 +713,13 @@ class TestStreamingHandlerAsyncInterrupt: assert len(await run.interrupts()) > 0 -class TestStreamingHandlerAsyncCustom: +class TestStreamV2AsyncCustom: @pytest.mark.anyio @NEEDS_CONTEXTVARS async def test_custom_stream_events(self) -> None: graph = _build_custom_stream_graph() - handler = StreamingHandler(graph) - run = await handler.astream({"value": "x", "items": []}) + handler = graph + run = await handler.astream_v2({"value": "x", "items": []}) events = [e async for e in run] custom_events = [e for e in events if e["method"] == "custom"] assert len(custom_events) == 2 @@ -1082,9 +1081,9 @@ class TestCustomTransformer: return True graph = _build_simple_graph() - handler = StreamingHandler(graph) + handler = graph counter_t = CounterTransformer() - run = handler.stream({"value": "x", "items": []}, transformers=[counter_t]) + run = handler.stream_v2({"value": "x", "items": []}, transformers=[counter_t]) assert "counter" in run.extensions # Subscribe before driving the run so channel pushes are retained. counter_iter = iter(run.extensions["counter"]) @@ -1113,9 +1112,9 @@ class TestCustomTransformer: return True graph = _build_simple_graph() - handler = StreamingHandler(graph) + handler = graph foo_t = FooTransformer() - run = handler.stream({"value": "x", "items": []}, transformers=[foo_t]) + run = handler.stream_v2({"value": "x", "items": []}, transformers=[foo_t]) # Subscribe before driving the run. foo_iter = iter(run.foo) _ = run.output @@ -1143,8 +1142,8 @@ class TestCustomTransformer: return True graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = handler.stream( + handler = graph + run = handler.stream_v2( {"value": "x", "items": []}, transformers=[EmitterTransformer()] ) events = list(run) @@ -1202,12 +1201,12 @@ class TestCustomTransformer: return True graph = _build_simple_graph() - handler = StreamingHandler(graph) + handler = graph with pytest.raises( ValueError, match=r"conflict.*'values'.*ValuesTransformer", ): - handler.stream( + handler.stream_v2( {"value": "x", "items": []}, transformers=[ConflictTransformer()], ) @@ -1304,9 +1303,9 @@ class TestEventLogAutoLifecycle: return True graph = _build_simple_graph() - handler = StreamingHandler(graph) + handler = graph t = MinimalTransformer() - run = handler.stream({"value": "x", "items": []}, transformers=[t]) + run = handler.stream_v2({"value": "x", "items": []}, transformers=[t]) minimal_iter = iter(run.extensions["minimal"]) _ = run.output items = list(minimal_iter) @@ -1605,8 +1604,8 @@ class TestAsyncTransformerLane: self._log.close() graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = await handler.astream( + handler = graph + run = await handler.astream_v2( {"value": "x", "items": []}, transformers=[Scorer()], ) @@ -1631,8 +1630,8 @@ class TestMemoryBounds: """With a single sync consumer, the pump produces exactly one event per cursor advance, so the buffer never holds more than one.""" graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = handler.stream({"value": "x", "items": []}) + handler = graph + run = handler.stream_v2({"value": "x", "items": []}) events_iter = iter(run) max_buffered = 0 count = 0 @@ -1649,8 +1648,8 @@ class TestMemoryBounds: """Projections without a subscriber drop pushes silently — their buffers stay empty regardless of run length.""" graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = handler.stream({"value": "x", "items": []}) + handler = graph + run = handler.stream_v2({"value": "x", "items": []}) # Subscribe to main events only; leave values and messages unsubscribed. list(run) values_log = run.extensions["values"] @@ -1665,8 +1664,8 @@ class TestMemoryBounds: process() without populating the log, so the values log buffer stays empty even across a full run.""" graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = handler.stream({"value": "x", "items": []}) + handler = graph + run = handler.stream_v2({"value": "x", "items": []}) _ = run.output values_log = run.extensions["values"] assert len(values_log._items) == 0 @@ -1675,8 +1674,8 @@ class TestMemoryBounds: def test_drained_subscriber_buffer_returns_to_empty(self) -> None: """After fully draining a subscribed log, the internal deque is empty.""" graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = handler.stream({"value": "x", "items": []}) + handler = graph + run = handler.stream_v2({"value": "x", "items": []}) values_log = run.extensions["values"] list(run.values) assert len(values_log._items) == 0 @@ -1686,8 +1685,8 @@ class TestMemoryBounds: async def test_async_single_consumer_buffer_stays_at_most_one(self) -> None: """Same drain-on-consume guarantee for the async lane.""" graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = await handler.astream({"value": "x", "items": []}) + handler = graph + run = await handler.astream_v2({"value": "x", "items": []}) max_buffered = 0 count = 0 async for _ in run: @@ -1701,8 +1700,8 @@ class TestMemoryBounds: async def test_async_unsubscribed_projections_never_accumulate(self) -> None: """Projections with no async subscriber stay empty under astream.""" graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = await handler.astream({"value": "x", "items": []}) + handler = graph + run = await handler.astream_v2({"value": "x", "items": []}) _ = await run.output() values_log = run.extensions["values"] messages_log = run.extensions["messages"]