From ae3c82349913b0fb406eb3cb2c387bd8f1f79c44 Mon Sep 17 00:00:00 2001 From: Nick Hollon Date: Wed, 15 Apr 2026 19:14:18 -0400 Subject: [PATCH] Make sync streaming caller-driven, no background thread Replace the daemon thread pump with a pull-based model where the caller's iteration on any projection drives the graph forward. EventLog uses a _request_more callback instead of threading.Condition. Matches v1's model where the caller's for loop is the pump. Async path is unchanged (background task on the event loop). --- libs/langgraph/langgraph/stream/_event_log.py | 45 +++++++----- libs/langgraph/langgraph/stream/run_stream.py | 70 +++++++++++++++++-- .../langgraph/stream/streaming_handler.py | 38 +++++----- 3 files changed, 105 insertions(+), 48 deletions(-) diff --git a/libs/langgraph/langgraph/stream/_event_log.py b/libs/langgraph/langgraph/stream/_event_log.py index 9a6b10f30..6fa774067 100644 --- a/libs/langgraph/langgraph/stream/_event_log.py +++ b/libs/langgraph/langgraph/stream/_event_log.py @@ -2,7 +2,7 @@ from __future__ import annotations import asyncio import threading -from collections.abc import AsyncIterator, Iterator +from collections.abc import AsyncIterator, Callable, Iterator from typing import Generic, TypeVar T = TypeVar("T") @@ -52,22 +52,23 @@ class _EventLogBase(Generic[T]): class EventLog(_EventLogBase[T]): - """Sync event log with multi-cursor iteration. + """Sync event log with pull-based iteration. Each call to ``__iter__`` creates a new cursor starting from the - beginning. Cursors block via ``threading.Condition`` when they - catch up to the producer. + beginning. When a cursor catches up to the buffer and the log is + not yet closed, it calls ``_request_more`` to pull more data from + the producer (typically the graph iterator via the run stream). + + If no ``_request_more`` callback is set, the cursor returns + immediately when it reaches the end of the buffer — this is the + behavior used in unit tests where items are pushed before iteration. Use ``AsyncEventLog`` for async consumers. """ def __init__(self) -> None: super().__init__() - self._cond = threading.Condition(self._lock) - - def _notify(self) -> None: - with self._lock: - self._cond.notify_all() + self._request_more: Callable[[], bool] | None = None def __iter__(self) -> Iterator[T]: """Return a new independent sync cursor over the log.""" @@ -76,17 +77,23 @@ class EventLog(_EventLogBase[T]): def _sync_cursor(self) -> Iterator[T]: cursor = 0 while True: - with self._lock: - while cursor >= len(self._items) and not self._closed: - self._cond.wait() - if cursor < len(self._items): - item = self._items[cursor] - cursor += 1 - elif self._error is not None: + if cursor < len(self._items): + item = self._items[cursor] + cursor += 1 + yield item + elif self._closed: + if self._error is not None: raise self._error - else: - return - yield item + 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). + while cursor >= len(self._items) and not self._closed: + if not self._request_more(): + break + else: + # No producer callback and not closed — buffer is complete. + return class AsyncEventLog(_EventLogBase[T]): diff --git a/libs/langgraph/langgraph/stream/run_stream.py b/libs/langgraph/langgraph/stream/run_stream.py index 51058305b..53bb02e7f 100644 --- a/libs/langgraph/langgraph/stream/run_stream.py +++ b/libs/langgraph/langgraph/stream/run_stream.py @@ -1,17 +1,23 @@ from __future__ import annotations import asyncio -import threading from collections.abc import AsyncIterator, Iterator from typing import Any +from langgraph.stream._event_log import EventLog from langgraph.stream._mux import StreamMux from langgraph.stream._types import ProtocolEvent +from langgraph.stream.stream_channel import StreamChannel from langgraph.stream.transformers import ValuesTransformer class GraphRunStream: - """Sync run stream with transformer-driven projections. + """Sync run stream with caller-driven pumping. + + The caller's iteration on any projection (``values``, ``messages``, + raw events, or ``output``) drives the graph forward. No background + thread is used — this matches v1's model where the caller's ``for`` + loop is the pump. All transformer projections live in ``extensions``. Native transformer projections (those with ``_native = True``) are also set as direct @@ -23,20 +29,65 @@ class GraphRunStream: def __init__( self, + graph_iter: Iterator[Any], mux: StreamMux, extensions: dict[str, Any], values_transformer: ValuesTransformer, - pump_thread: threading.Thread, ) -> None: + self._graph_iter = graph_iter self._mux = mux self.extensions = extensions self._values_transformer = values_transformer - self._pump_thread = pump_thread + self._exhausted = False + # Wire pull-based iteration: every sync EventLog calls _pump_next + # when its cursor catches up to the buffer. + self._wire_request_more(mux, extensions) + + def _wire_request_more( + self, mux: StreamMux, extensions: dict[str, Any] + ) -> None: + """Set _request_more on all sync EventLogs so iteration drives the graph.""" + if isinstance(mux._events, EventLog): + mux._events._request_more = self._pump_next + for value in extensions.values(): + if isinstance(value, EventLog): + value._request_more = self._pump_next + elif isinstance(value, StreamChannel) and isinstance( + value._log, EventLog + ): + value._log._request_more = self._pump_next + + def _pump_next(self) -> bool: + """Pull one event from the graph and push through the mux. + + Returns True if an event was pulled, False if the graph is exhausted. + """ + if self._exhausted: + return False + try: + part = next(self._graph_iter) + except StopIteration: + self._mux.close() + self._exhausted = True + return False + except BaseException as e: + self._mux.fail(e) + self._exhausted = True + return False + from langgraph.stream._convert import convert_to_protocol_event + + self._mux.push(convert_to_protocol_event(part)) + return True + + def _pump_all(self) -> None: + """Drain the graph completely.""" + while self._pump_next(): + pass @property def output(self) -> dict[str, Any] | None: """Block until the run completes and return the final state.""" - self._pump_thread.join() + self._pump_all() if self._values_transformer._log._error is not None: raise self._values_transformer._log._error return self._values_transformer._latest @@ -44,13 +95,13 @@ class GraphRunStream: @property def interrupted(self) -> bool: """Block until the run completes, then return whether it was interrupted.""" - self._pump_thread.join() + self._pump_all() return self._values_transformer._interrupted @property def interrupts(self) -> list[Any]: """Block until the run completes, then return interrupt payloads.""" - self._pump_thread.join() + self._pump_all() return self._values_transformer._interrupts def __iter__(self) -> Iterator[ProtocolEvent]: @@ -61,6 +112,11 @@ class GraphRunStream: class AsyncGraphRunStream: """Async run stream with transformer-driven projections. + A background asyncio task pumps events from the graph into the + transformer pipeline. This is the standard async pattern — the task + runs on the same event loop and async consumers can iterate multiple + projections concurrently. + All transformer projections live in ``extensions``. Native transformer projections (those with ``_native = True``) are also set as direct attributes on this instance (e.g. ``run.values``, ``run.messages``). diff --git a/libs/langgraph/langgraph/stream/streaming_handler.py b/libs/langgraph/langgraph/stream/streaming_handler.py index 1ac2dc498..116ad0da3 100644 --- a/libs/langgraph/langgraph/stream/streaming_handler.py +++ b/libs/langgraph/langgraph/stream/streaming_handler.py @@ -1,7 +1,6 @@ from __future__ import annotations import asyncio -import threading from collections.abc import Sequence from typing import Any @@ -60,33 +59,28 @@ class StreamingHandler: ) -> GraphRunStream: """Start a sync streaming run. - Returns a `GraphRunStream` immediately. A background daemon thread - pumps events from the graph into the transformer pipeline. + 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. """ mux, extensions, native_keys, values_t = self._setup( transformers, is_async=False ) - def pump() -> None: - try: - for part in self._graph.stream( - input, - config, - stream_mode=STREAM_V2_MODES, - subgraphs=True, - version="v2", - interrupt_before=interrupt_before, - interrupt_after=interrupt_after, - ): - mux.push(convert_to_protocol_event(part)) - mux.close() - except BaseException as e: - mux.fail(e) + 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, + ) + ) - thread = threading.Thread(target=pump, daemon=True) - thread.start() - - run = GraphRunStream(mux, extensions, values_t, thread) + run = GraphRunStream(graph_iter, mux, extensions, values_t) for key in native_keys: setattr(run, key, extensions[key]) return run