diff --git a/libs/langgraph/langgraph/stream/__init__.py b/libs/langgraph/langgraph/stream/__init__.py index d882191b5..bb800b7c4 100644 --- a/libs/langgraph/langgraph/stream/__init__.py +++ b/libs/langgraph/langgraph/stream/__init__.py @@ -2,7 +2,7 @@ from langgraph.stream._convert import STREAM_V2_MODES, convert_to_protocol_event from langgraph.stream._event_log import EventLog -from langgraph.stream._mux import StreamMux +from langgraph.stream._mux import AsyncStreamMux, StreamMux from langgraph.stream._types import ( InterruptPayload, ProtocolEvent, @@ -13,6 +13,7 @@ from langgraph.stream.run_stream import ( AsyncGraphRunStream, AsyncSubgraphRunStream, GraphRunStream, + SubgraphRunStream, create_async_graph_run_stream, create_graph_run_stream, ) @@ -25,6 +26,7 @@ from langgraph.stream.transformers import ( __all__ = [ "STREAM_V2_MODES", + "AsyncStreamMux", "AsyncChatModelStream", "AsyncGraphRunStream", "AsyncSubgraphRunStream", @@ -36,6 +38,7 @@ __all__ = [ "ProtocolEvent", "StreamChannel", "StreamMux", + "SubgraphRunStream", "StreamTransformer", "StreamingHandler", "ValuesTransformer", diff --git a/libs/langgraph/langgraph/stream/_event_log.py b/libs/langgraph/langgraph/stream/_event_log.py index 660a05531..c4faebaf6 100644 --- a/libs/langgraph/langgraph/stream/_event_log.py +++ b/libs/langgraph/langgraph/stream/_event_log.py @@ -15,6 +15,17 @@ from typing import Generic, TypeVar T = TypeVar("T") +def _resolve_future(fut: asyncio.Future[None]) -> None: + """Set a future's result if it hasn't already completed or been cancelled. + + Runs on the event loop thread (scheduled via ``call_soon_threadsafe``) + so that the ``done()`` check and ``set_result`` are atomic with + respect to cancellation. + """ + if not fut.done(): + fut.set_result(None) + + class EventLog(Generic[T]): """Append-only event buffer with cursor-based async iteration. @@ -56,14 +67,9 @@ class EventLog(Generic[T]): # -- Consumer API ------------------------------------------------------- - def subscribe(self, offset: int = 0) -> _Cursor[T]: - """Create a new cursor starting at *offset*. - - If *offset* is 0 the cursor replays the entire log. If - *offset* equals ``len(self)`` the cursor starts from the - current tip and only sees future events. - """ - return _Cursor(self, offset) + def __aiter__(self) -> _Cursor[T]: + """Return a fresh cursor from the beginning of the log.""" + return _Cursor(self) # -- Inspection --------------------------------------------------------- @@ -81,12 +87,11 @@ class EventLog(Generic[T]): def _wake_all(self) -> None: for fut in self._waiters: - if not fut.done(): - try: - fut.get_loop().call_soon_threadsafe(fut.set_result, None) - except RuntimeError: - # Loop already closed — ignore. - pass + try: + fut.get_loop().call_soon_threadsafe(_resolve_future, fut) + except RuntimeError: + # Loop already closed — ignore. + pass self._waiters.clear() @@ -95,9 +100,9 @@ class _Cursor(Generic[T]): __slots__ = ("_log", "_offset") - def __init__(self, log: EventLog[T], offset: int) -> None: + def __init__(self, log: EventLog[T]) -> None: self._log = log - self._offset = offset + self._offset = 0 def __aiter__(self) -> _Cursor[T]: return self @@ -114,18 +119,17 @@ class _Cursor(Generic[T]): if self._log._closed: raise StopAsyncIteration # Nothing available yet — register a waiter - loop = asyncio.get_running_loop() - fut: asyncio.Future[None] = loop.create_future() + fut: asyncio.Future[None] = asyncio.get_running_loop().create_future() self._log._waiters.append(fut) # Wait outside the lock try: await fut except asyncio.CancelledError: - # Remove our future so it doesn't accumulate in the list. - try: - self._log._waiters.remove(fut) - except ValueError: - pass # Already removed by _wake_all + with self._log._lock: + try: + self._log._waiters.remove(fut) + except ValueError: + pass # Already removed by _wake_all raise diff --git a/libs/langgraph/langgraph/stream/_mux.py b/libs/langgraph/langgraph/stream/_mux.py index 83082d143..7b02b5bb2 100644 --- a/libs/langgraph/langgraph/stream/_mux.py +++ b/libs/langgraph/langgraph/stream/_mux.py @@ -1,9 +1,12 @@ """Central event dispatcher with transformer pipeline for StreamingHandler. -``StreamMux`` is the core coordination point: it holds the main +``StreamMux`` is the sync-safe core: it holds the main :class:`EventLog`, tracks discovered namespaces for subgraph stream creation, and pipes every event through the registered :class:`StreamTransformer` pipeline before appending it to the log. + +``AsyncStreamMux`` extends the base with async subscription endpoints +(output futures, namespace waiters, filtered event iteration). """ from __future__ import annotations @@ -19,11 +22,13 @@ from langgraph.stream.stream_channel import StreamChannel, is_stream_channel class StreamMux: - """Central event dispatcher for the StreamingHandler infrastructure. + """Sync-safe event dispatcher for the StreamingHandler infrastructure. The mux owns the main event log, applies the transformer pipeline to - every incoming event, and provides subscription endpoints for - filtered event iteration and subgraph discovery. + every incoming event, and tracks namespace discovery and latest values. + + For async subscription endpoints (output futures, namespace waiters, + filtered event iteration), use :class:`AsyncStreamMux`. """ def __init__(self, transformers: list[StreamTransformer] | None = None) -> None: @@ -35,8 +40,6 @@ class StreamMux: # Namespace discovery: maps top-level ns segment → True self._discovered_ns: dict[str, bool] = {} - # Waiters for new namespace discovery - self._ns_waiters: list[asyncio.Future[None]] = [] # Latest values per namespace (list-of-strings key) self._latest_values: dict[str, Any] = {} @@ -45,9 +48,6 @@ class StreamMux: self._interrupts: list[InterruptPayload] = [] self._interrupted = False - # Output promise tracking - self._output_futures: dict[str, asyncio.Future[Any]] = {} - # Closed state self._closed = False self._error: BaseException | None = None @@ -76,7 +76,7 @@ class StreamMux: top_segment = ns[0] if top_segment not in self._discovered_ns: self._discovered_ns[top_segment] = True - self._wake_ns_waiters() + self._on_ns_discovered(top_segment) # Track values if event["method"] == "values": @@ -113,7 +113,7 @@ class StreamMux: self._event_log.append(event) def close(self, output: Any = None) -> None: - """Close the mux, resolving all output futures.""" + """Close the mux, finalizing transformers and the event log.""" if self._closed: return self._closed = True @@ -130,20 +130,8 @@ class StreamMux: # Close the event log self._event_log.close() - # Resolve output futures - for ns_key, fut in self._output_futures.items(): - if not fut.done(): - value = self._latest_values.get(ns_key) - try: - fut.get_loop().call_soon_threadsafe(fut.set_result, value) - except RuntimeError: - pass - - # Wake namespace waiters - self._wake_ns_waiters() - def fail(self, error: BaseException) -> None: - """Fail the mux, rejecting all output futures.""" + """Fail the mux, propagating the error to transformers and channels.""" if self._closed: return self._closed = True @@ -161,83 +149,6 @@ class StreamMux: # Fail the event log self._event_log.fail(error) - # Reject output futures - for fut in self._output_futures.values(): - if not fut.done(): - try: - fut.get_loop().call_soon_threadsafe(fut.set_exception, error) - except RuntimeError: - pass - - # Wake namespace waiters - self._wake_ns_waiters() - - # -- Consumer API ------------------------------------------------------- - - def subscribe_events( - self, path: list[str] | None = None, offset: int = 0 - ) -> AsyncIterator[ProtocolEvent]: - """Return an async iterator over events matching *path*. - - If *path* is ``None`` or empty, all events are yielded. - Otherwise, only events whose namespace starts with *path* - are yielded. - """ - cursor = self._event_log.subscribe(offset) - if not path: - return cursor - return _FilteredEventIterator(cursor, path) - - async def subscribe_subgraphs( - self, path: list[str] | None = None, offset: int = 0 - ) -> AsyncIterator[str]: - """Yield top-level namespace segments as they are discovered. - - Each yielded value is the first namespace segment of a newly - discovered subgraph (e.g. ``"agent:0"``). - """ - yielded: set[str] = set() - while True: - # Yield any newly discovered namespaces - for ns_segment in list(self._discovered_ns): - if ns_segment not in yielded: - # Filter by path prefix if specified - if path: - if not ns_segment.startswith(path[0]): - continue - yielded.add(ns_segment) - yield ns_segment - - if self._closed: - return - - # Wait for new namespaces - loop = asyncio.get_running_loop() - fut: asyncio.Future[None] = loop.create_future() - self._ns_waiters.append(fut) - await fut - - def get_output_future(self, ns: list[str] | None = None) -> asyncio.Future[Any]: - """Get or create an output future for a namespace. - - The future resolves to the latest ``values`` event data when - the mux is closed. - """ - ns_key = _ns_key(ns or []) - if ns_key not in self._output_futures: - loop = asyncio.get_running_loop() - self._output_futures[ns_key] = loop.create_future() - - # If already closed, resolve immediately - if self._closed: - value = self._latest_values.get(ns_key) - if self._error is not None: - self._output_futures[ns_key].set_exception(self._error) - else: - self._output_futures[ns_key].set_result(value) - - return self._output_futures[ns_key] - # -- Inspection --------------------------------------------------------- @property @@ -258,6 +169,13 @@ class StreamMux: # -- Internal ----------------------------------------------------------- + def _on_ns_discovered(self, segment: str) -> None: + """Hook called when a new top-level namespace segment is discovered. + + The base implementation is a no-op. :class:`AsyncStreamMux` + overrides this to wake namespace waiters. + """ + def register_transformer(self, transformer: StreamTransformer) -> None: """Register a new transformer and replay all buffered events through it. @@ -333,6 +251,135 @@ class StreamMux: channel._wire(_make_forwarder(channel)) + +class AsyncStreamMux(StreamMux): + """Async extension of :class:`StreamMux`. + + Adds output futures, namespace waiters, and async subscription + endpoints (``subscribe_events``, ``subscribe_subgraphs``, + ``get_output_future``). + """ + + def __init__(self, transformers: list[StreamTransformer] | None = None) -> None: + super().__init__(transformers) + # Waiters for new namespace discovery + self._ns_waiters: list[asyncio.Future[None]] = [] + # Output promise tracking + self._output_futures: dict[str, asyncio.Future[Any]] = {} + + # -- Producer API overrides --------------------------------------------- + + def close(self, output: Any = None) -> None: + """Close the mux, resolving all output futures.""" + if self._closed: + return + + # Let the base class finalize transformers, channels, and event log + super().close(output) + + # Resolve output futures + for ns_key, fut in self._output_futures.items(): + if not fut.done(): + value = self._latest_values.get(ns_key) + try: + fut.get_loop().call_soon_threadsafe(fut.set_result, value) + except RuntimeError: + pass + + # Wake namespace waiters + self._wake_ns_waiters() + + def fail(self, error: BaseException) -> None: + """Fail the mux, rejecting all output futures.""" + if self._closed: + return + + # Let the base class fail transformers, channels, and event log + super().fail(error) + + # Reject output futures + for fut in self._output_futures.values(): + if not fut.done(): + try: + fut.get_loop().call_soon_threadsafe(fut.set_exception, error) + except RuntimeError: + pass + + # Wake namespace waiters + self._wake_ns_waiters() + + # -- Consumer API ------------------------------------------------------- + + def subscribe_events( + self, path: list[str] | None = None + ) -> AsyncIterator[ProtocolEvent]: + """Return an async iterator over events matching *path*. + + If *path* is ``None`` or empty, all events are yielded. + Otherwise, only events whose namespace starts with *path* + are yielded. + """ + cursor = aiter(self._event_log) + if not path: + return cursor + return _FilteredEventIterator(cursor, path) + + async def subscribe_subgraphs( + self, path: list[str] | None = None, offset: int = 0 + ) -> AsyncIterator[str]: + """Yield top-level namespace segments as they are discovered. + + Each yielded value is the first namespace segment of a newly + discovered subgraph (e.g. ``"agent:0"``). + """ + yielded: set[str] = set() + while True: + # Yield any newly discovered namespaces + for ns_segment in list(self._discovered_ns): + if ns_segment not in yielded: + # Filter by path prefix if specified + if path: + if not ns_segment.startswith(path[0]): + continue + yielded.add(ns_segment) + yield ns_segment + + if self._closed: + return + + # Wait for new namespaces + loop = asyncio.get_running_loop() + fut: asyncio.Future[None] = loop.create_future() + self._ns_waiters.append(fut) + await fut + + def get_output_future(self, ns: list[str] | None = None) -> asyncio.Future[Any]: + """Get or create an output future for a namespace. + + The future resolves to the latest ``values`` event data when + the mux is closed. + """ + ns_key = _ns_key(ns or []) + if ns_key not in self._output_futures: + loop = asyncio.get_running_loop() + self._output_futures[ns_key] = loop.create_future() + + # If already closed, resolve immediately + if self._closed: + value = self._latest_values.get(ns_key) + if self._error is not None: + self._output_futures[ns_key].set_exception(self._error) + else: + self._output_futures[ns_key].set_result(value) + + return self._output_futures[ns_key] + + # -- Internal ----------------------------------------------------------- + + def _on_ns_discovered(self, segment: str) -> None: + """Wake namespace waiters when a new namespace is discovered.""" + self._wake_ns_waiters() + def _wake_ns_waiters(self) -> None: for fut in self._ns_waiters: if not fut.done(): @@ -375,4 +422,4 @@ def _ns_starts_with(ns: list[str], prefix: list[str]) -> bool: return ns[: len(prefix)] == prefix -__all__ = ["StreamMux"] +__all__ = ["AsyncStreamMux", "StreamMux"] diff --git a/libs/langgraph/langgraph/stream/chat_model_stream.py b/libs/langgraph/langgraph/stream/chat_model_stream.py index 1905696bc..2b6bc6ecc 100644 --- a/libs/langgraph/langgraph/stream/chat_model_stream.py +++ b/libs/langgraph/langgraph/stream/chat_model_stream.py @@ -13,11 +13,69 @@ or ``full = await msg.text``). from __future__ import annotations import asyncio -from collections.abc import Generator +from collections.abc import Callable, Generator, Iterator from typing import Any from langgraph.stream._types import UsageInfo + +# --------------------------------------------------------------------------- +# Sync dual projection — iterable of deltas, str() for accumulated text +# --------------------------------------------------------------------------- + + +class _SyncDualProjection: + """Pump-driven sync iterable of string deltas. + + Iterating yields incremental text fragments as the pump delivers + new ``content-block-delta`` events. Calling ``str()`` drains the + pump and returns the full accumulated string. + + This is the sync counterpart of :class:`_DualProjection` (the async + variant used by ``AsyncChatModelStream``). + """ + + __slots__ = ("_stream", "_attr", "_pump_one") + + def __init__( + self, + stream: ChatModelStream, + attr: str, + pump_one: Callable[[], bool], + ) -> None: + self._stream = stream + self._attr = attr + self._pump_one = pump_one + + def __iter__(self) -> Iterator[str]: + prev_len = 0 + while True: + cur = getattr(self._stream, self._attr) + if len(cur) > prev_len: + yield cur[prev_len:] + prev_len = len(cur) + if self._stream._done: + return + if not self._pump_one(): + # Source exhausted — yield any remaining + cur = getattr(self._stream, self._attr) + if len(cur) > prev_len: + yield cur[prev_len:] + return + + def __str__(self) -> str: + while not self._stream._done: + if not self._pump_one(): + break + return getattr(self._stream, self._attr) + + def __repr__(self) -> str: + return repr(getattr(self._stream, self._attr)) + + def __bool__(self) -> bool: + return bool(getattr(self._stream, self._attr)) + + # --------------------------------------------------------------------------- # Sync variant # --------------------------------------------------------------------------- @@ -56,21 +114,52 @@ class ChatModelStream: self._usage_value: UsageInfo | None = None self._done = False + # Optional pump for sync streaming (set via _bind_pump) + self._pump_one: Callable[[], bool] | None = None + + # -- Pump binding (called by GraphRunStream) --------------------------- + + def _bind_pump(self, pump_one: Callable[[], bool]) -> None: + """Bind a pump function for sync token-by-token streaming. + + When bound, ``.text`` and ``.reasoning`` return + :class:`_SyncDualProjection` instances that drive the pump and + yield deltas as the LLM produces tokens. + """ + self._pump_one = pump_one + # -- Public projections ------------------------------------------------ @property - def text(self) -> str: - """Accumulated text content.""" + def text(self) -> str | _SyncDualProjection: + """Text content. + + When a pump is bound (sync streaming), returns a + :class:`_SyncDualProjection` — iterable of deltas, + ``str()`` for the full accumulated text. Otherwise returns + the accumulated text string directly. + """ + if self._pump_one is not None and not self._done: + return _SyncDualProjection(self, "_text_acc", self._pump_one) return self._text_acc @property - def reasoning(self) -> str: - """Accumulated reasoning content.""" + def reasoning(self) -> str | _SyncDualProjection: + """Reasoning content. + + Same dual behavior as :attr:`text`. + """ + if self._pump_one is not None and not self._done: + return _SyncDualProjection(self, "_reasoning_acc", self._pump_one) return self._reasoning_acc @property def usage(self) -> UsageInfo | None: """Usage info, available after the message finishes.""" + if self._pump_one is not None and not self._done: + while not self._done: + if not self._pump_one(): + break return self._usage_value @property @@ -310,4 +399,4 @@ class AsyncChatModelStream(ChatModelStream): self._usage_proj._fail(error) -__all__ = ["AsyncChatModelStream", "ChatModelStream"] +__all__ = ["AsyncChatModelStream", "ChatModelStream", "_SyncDualProjection"] diff --git a/libs/langgraph/langgraph/stream/run_stream.py b/libs/langgraph/langgraph/stream/run_stream.py index 8044f1238..9cbb0d3db 100644 --- a/libs/langgraph/langgraph/stream/run_stream.py +++ b/libs/langgraph/langgraph/stream/run_stream.py @@ -15,7 +15,7 @@ from typing import Any from langgraph.stream._convert import convert_to_protocol_event from langgraph.stream._event_log import EventLog -from langgraph.stream._mux import StreamMux +from langgraph.stream._mux import AsyncStreamMux, StreamMux from langgraph.stream._types import InterruptPayload, ProtocolEvent, StreamTransformer from langgraph.stream.chat_model_stream import AsyncChatModelStream, ChatModelStream from langgraph.stream.transformers import MessagesTransformer, ValuesTransformer @@ -30,7 +30,7 @@ class _ValuesProjection: def __init__( self, - mux: StreamMux, + mux: AsyncStreamMux, values_transformer: ValuesTransformer, ns: list[str], mapper: Callable[[Any], Any] | None = None, @@ -62,7 +62,7 @@ class _ValuesIterator: ns: list[str], mapper: Callable[[Any], Any] | None = None, ) -> None: - self._cursor = transformer.values_log.subscribe(0) + self._cursor = aiter(transformer.values_log) self._ns = ns self._mapper = mapper @@ -92,7 +92,7 @@ class _MessagesProjection: self._transformer = messages_transformer def __aiter__(self) -> AsyncIterator[AsyncChatModelStream]: - return self._transformer.messages_log.subscribe(0) + return aiter(self._transformer.messages_log) # --------------------------------------------------------------------------- @@ -103,7 +103,7 @@ class _MessagesProjection: class _SubgraphsProjection: """Async iterable yielding :class:`AsyncSubgraphRunStream` for each discovered subgraph.""" - def __init__(self, mux: StreamMux, ns: list[str]) -> None: + def __init__(self, mux: AsyncStreamMux, ns: list[str]) -> None: self._mux = mux self._ns = ns @@ -143,7 +143,7 @@ class AsyncGraphRunStream: def __init__( self, *, - mux: StreamMux, + mux: AsyncStreamMux, namespace: list[str] | None = None, transformers: list[StreamTransformer], abort_event: asyncio.Event | None = None, @@ -308,7 +308,7 @@ async def create_async_graph_run_stream( if projection is not None: projections.append(projection) - mux = StreamMux(transformers=all_transformers) + mux = AsyncStreamMux(transformers=all_transformers) # Wire any StreamChannel instances found in transformer projections for projection in projections: @@ -492,20 +492,91 @@ class GraphRunStream: def messages(self) -> Iterator[ChatModelStream]: """Sync iterable of :class:`ChatModelStream` instances. - Each yielded ``ChatModelStream`` is fully populated (``done=True``) - so that sync consumers can read ``.text``, ``.reasoning``, and - ``.usage`` immediately. + Each ``ChatModelStream`` is yielded as soon as the LLM begins + responding (on ``message-start``). Its ``.text`` and + ``.reasoning`` properties are pump-driven + :class:`~langgraph.stream.chat_model_stream._SyncDualProjection` + instances that yield deltas as tokens arrive:: + + for msg in run.messages: + for delta in msg.text: + print(delta, end="", flush=True) + + If you don't need streaming, ``str(msg.text)`` pumps until + the message completes and returns the full text. + + After each message is consumed, the pump advances through + non-message events (tool completions, values, etc.) so that + other transformer state is up-to-date before the next message + is yielded. This means you can check + ``run.extensions["tools"]`` between messages and see inline + results. """ t = self._find_transformer("messages") if t is None: return - for msg in _PumpDrivenLog(t.value, self._pump_one): - # Pump until this message is complete so sync consumers - # get a fully populated ChatModelStream. - while not msg.done: + log = t.value + for msg in _PumpDrivenLog(log, self._pump_one): + msg._bind_pump(self._pump_one) + yield msg + # Advance the pump past non-message events so other + # transformers have up-to-date state before the next + # message is yielded. + prev_count = len(log) + while len(log) == prev_count: if not self._pump_one(): break - yield msg + + # -- Subgraphs ---------------------------------------------------------- + + @property + def subgraphs(self) -> Iterator[SubgraphRunStream]: + """Sync iterable of :class:`SubgraphRunStream` for child graphs. + + Namespaces are discovered lazily as events are pumped from the + source. Each yielded stream has its own ``values``, ``messages``, + and ``output`` projections scoped to the child namespace. + + After yielding a subgraph, the caller may consume its projections + (e.g. ``sub.values``), which pumps more events and can discover + new namespaces. The loop re-checks for newly discovered + namespaces after each yield before attempting another pump. + """ + yielded: set[str] = set() + + while True: + # Yield any newly discovered namespaces. Re-check after + # each yield because consuming a subgraph's projections + # can pump events that discover further namespaces. + found_new = False + for ns_segment in list(self._mux._discovered_ns): + if ns_segment in yielded: + continue + found_new = True + yielded.add(ns_segment) + child_ns = self._ns + [ns_segment] + child_transformers: list[StreamTransformer] = [ + ValuesTransformer(), + MessagesTransformer(namespace=child_ns), + ] + for t in child_transformers: + t.init() + self._mux.register_transformer(t) + + yield SubgraphRunStream( + mux=self._mux, + namespace=child_ns, + transformers=child_transformers, + pump_one=self._pump_one, + output_mapper=self._output_mapper, + ) + + if found_new: + continue # re-check before pumping + + # No new namespaces — pump one event + if not self._pump_one(): + break # -- State -------------------------------------------------------------- @@ -534,6 +605,133 @@ class GraphRunStream: return result +# --------------------------------------------------------------------------- +# SubgraphRunStream — sync child stream +# --------------------------------------------------------------------------- + + +class SubgraphRunStream: + """Synchronous run stream for a child subgraph. + + Shares the parent's :class:`StreamMux` and pump function. Has its + own transformer set registered on the shared mux so that projections + (``values``, ``messages``, ``output``) are scoped to the child + namespace. + + Adds ``.name`` and ``.index`` parsed from the last namespace segment + (e.g. ``"researcher:2"`` → ``name="researcher"``, ``index=2``). + """ + + def __init__( + self, + *, + mux: StreamMux, + namespace: list[str], + transformers: list[StreamTransformer], + pump_one: Callable[[], bool], + output_mapper: Callable[[Any], Any] | None = None, + ) -> None: + self._mux = mux + self._ns = namespace + self._transformers = transformers + self._pump_one = pump_one + self._output_mapper = output_mapper + + # -- Identity ----------------------------------------------------------- + + @property + def name(self) -> str: + if self._ns: + segment = self._ns[-1] + return segment.split(":")[0] if ":" in segment else segment + return "" + + @property + def index(self) -> int: + if self._ns: + segment = self._ns[-1] + if ":" in segment: + try: + return int(segment.split(":")[-1]) + except ValueError: + pass + return 0 + + # -- Transformer lookup ------------------------------------------------- + + def _find_transformer(self, name: str) -> StreamTransformer | None: + for t in self._transformers: + if getattr(t, "name", None) == name: + return t + return None + + # -- Helpers ------------------------------------------------------------ + + def _map(self, value: Any) -> Any: + if value is not None and self._output_mapper is not None: + return self._output_mapper(value) + return value + + def _pump_all(self) -> None: + while self._pump_one(): + pass + + # -- Raw event iteration (sync) ----------------------------------------- + + def __iter__(self) -> Iterator[ProtocolEvent]: + for event in _PumpDrivenLog(self._mux.event_log, self._pump_one): + ns = event["params"].get("namespace", []) + if ns[: len(self._ns)] == self._ns: + yield event + + # -- Named projections (sync) ------------------------------------------- + + @property + def output(self) -> Any: + """The final output state (blocking). Drains the source.""" + self._pump_all() + return self._map(self._mux.get_latest_values(self._ns)) + + @property + def values(self) -> Iterator[Any]: + """Sync iterable of intermediate state snapshots.""" + t = self._find_transformer("values") + if t is None: + return + for item in _PumpDrivenLog(t.value, self._pump_one): + if item.get("namespace", []) == self._ns: + yield self._map(item["data"]) + + @property + def messages(self) -> Iterator[ChatModelStream]: + """Sync iterable of :class:`ChatModelStream` instances. + + Each ``ChatModelStream`` is yielded as soon as the LLM begins + responding. See :attr:`GraphRunStream.messages` for usage. + """ + t = self._find_transformer("messages") + if t is None: + return + log = t.value + for msg in _PumpDrivenLog(log, self._pump_one): + msg._bind_pump(self._pump_one) + yield msg + prev_count = len(log) + while len(log) == prev_count: + if not self._pump_one(): + break + + # -- State -------------------------------------------------------------- + + @property + def interrupted(self) -> bool: + return self._mux.interrupted + + @property + def interrupts(self) -> list[InterruptPayload]: + return self._mux.interrupts + + def create_graph_run_stream( source: Iterator[tuple[tuple[str, ...], str, Any]], *, @@ -579,6 +777,7 @@ __all__ = [ "AsyncGraphRunStream", "AsyncSubgraphRunStream", "GraphRunStream", + "SubgraphRunStream", "create_async_graph_run_stream", "create_graph_run_stream", ] diff --git a/libs/langgraph/langgraph/stream/stream_channel.py b/libs/langgraph/langgraph/stream/stream_channel.py index fade0c2a1..5e0e6b8ff 100644 --- a/libs/langgraph/langgraph/stream/stream_channel.py +++ b/libs/langgraph/langgraph/stream/stream_channel.py @@ -51,7 +51,7 @@ class StreamChannel(Generic[T]): # -- Async iteration (in-process consumption) --------------------------- def __aiter__(self) -> AsyncIterator[T]: - return self._log.subscribe(0) + return aiter(self._log) # -- Internal (called by the mux) --------------------------------------- diff --git a/libs/langgraph/tests/test_stream_event_log.py b/libs/langgraph/tests/test_stream_event_log.py index 7bec10156..90b8b5baa 100644 --- a/libs/langgraph/tests/test_stream_event_log.py +++ b/libs/langgraph/tests/test_stream_event_log.py @@ -12,7 +12,7 @@ async def test_push_and_iterate_in_order(): log.append("b") log.append("c") log.close() - items = [item async for item in log.subscribe(0)] + items = [item async for item in aiter(log)] assert items == ["a", "b", "c"] @@ -22,8 +22,8 @@ async def test_multiple_independent_cursors(): log.append("x") log.append("y") log.close() - items1 = [item async for item in log.subscribe(0)] - items2 = [item async for item in log.subscribe(0)] + items1 = [item async for item in aiter(log)] + items2 = [item async for item in aiter(log)] assert items1 == ["x", "y"] assert items2 == ["x", "y"] @@ -32,7 +32,7 @@ async def test_multiple_independent_cursors(): async def test_close_ends_iteration(): log = EventLog() log.close() - items = [item async for item in log.subscribe(0)] + items = [item async for item in aiter(log)] assert items == [] @@ -41,28 +41,17 @@ async def test_fail_raises_error(): log = EventLog() log.fail(RuntimeError("boom")) with pytest.raises(RuntimeError, match="boom"): - async for _ in log.subscribe(0): + async for _ in aiter(log): pass -@pytest.mark.anyio -async def test_nonzero_offset(): - log = EventLog() - log.append("x") - log.append("y") - log.append("z") - log.close() - items = [item async for item in log.subscribe(2)] - assert items == ["z"] - - @pytest.mark.anyio async def test_concurrent_push_and_iterate(): log = EventLog() received = [] async def consumer(): - async for item in log.subscribe(0): + async for item in aiter(log): received.append(item) async def producer(): @@ -80,7 +69,7 @@ async def test_items_before_cursor_visible(): log = EventLog() log.append("a") log.append("b") - cursor = log.subscribe(0) + cursor = aiter(log) log.append("c") log.close() items = [item async for item in cursor] @@ -91,7 +80,7 @@ async def test_items_before_cursor_visible(): async def test_empty_log_closed_yields_nothing(): log = EventLog() log.close() - items = [item async for item in log.subscribe(0)] + items = [item async for item in aiter(log)] assert items == [] @@ -102,7 +91,7 @@ async def test_fail_mid_iteration(): received = [] async def consumer(): - async for item in log.subscribe(0): + async for item in aiter(log): received.append(item) async def producer(): @@ -129,7 +118,7 @@ async def test_abandoned_cursor_cleans_up_waiters(): log: EventLog[str] = EventLog() for _ in range(10): - cursor = log.subscribe(0) + cursor = aiter(log) task = asyncio.ensure_future(cursor.__anext__()) await asyncio.sleep(0) # let task register its waiter task.cancel() diff --git a/libs/langgraph/tests/test_stream_mux.py b/libs/langgraph/tests/test_stream_mux.py index d2160e5f1..8b76c9333 100644 --- a/libs/langgraph/tests/test_stream_mux.py +++ b/libs/langgraph/tests/test_stream_mux.py @@ -3,7 +3,7 @@ from typing import Any import pytest from langgraph.stream._convert import convert_to_protocol_event -from langgraph.stream._mux import StreamMux +from langgraph.stream._mux import AsyncStreamMux, StreamMux from langgraph.stream._types import ProtocolEvent from langgraph.stream.stream_channel import StreamChannel @@ -70,7 +70,7 @@ async def test_top_level_ns_only(): @pytest.mark.anyio async def test_subscribe_events_filter(): - mux = StreamMux() + mux = AsyncStreamMux() mux.push(_event("values", {"a": 1}, ns=["child:0"])) mux.push(_event("values", {"b": 2}, ns=["other:1"])) mux.push(_event("values", {"c": 3}, ns=["child:0"])) @@ -86,7 +86,7 @@ async def test_subscribe_events_filter(): @pytest.mark.anyio async def test_close_resolves_output(): - mux = StreamMux() + mux = AsyncStreamMux() fut = mux.get_output_future() mux.push(_event("values", {"v": 1})) mux.push(_event("values", {"v": 2})) @@ -97,7 +97,7 @@ async def test_close_resolves_output(): @pytest.mark.anyio async def test_fail_rejects_output(): - mux = StreamMux() + mux = AsyncStreamMux() fut = mux.get_output_future() mux.fail(ValueError("boom")) with pytest.raises(ValueError, match="boom"): @@ -155,7 +155,7 @@ async def test_push_after_close_ignored(): @pytest.mark.anyio async def test_fail_rejects_all_futures(): - mux = StreamMux() + mux = AsyncStreamMux() fut1 = mux.get_output_future([]) fut2 = mux.get_output_future(["child:0"]) mux.fail(ValueError("boom")) @@ -172,7 +172,7 @@ async def test_channel_events_bypass_transformer_pipeline(): the JS implementation and avoids re-entrancy bugs. """ mock = _MockTransformer() - mux = StreamMux(transformers=[mock]) + mux = AsyncStreamMux(transformers=[mock]) channel: StreamChannel[str] = StreamChannel("my_channel") mux.wire_channels({"ch": channel}) @@ -208,7 +208,7 @@ async def test_event_log_has_monotonic_seq_numbers(): while channel-emitted events use a separate counter (``_next_emit_seq``). When interleaved, seq numbers can duplicate. """ - mux = StreamMux() + mux = AsyncStreamMux() channel: StreamChannel[str] = StreamChannel("test_ch") mux.wire_channels({"ch": channel}) @@ -266,7 +266,7 @@ async def test_channel_push_during_process_preserves_namespace(): t1 = _ChannelTransformer("first") t2 = _ChannelTransformer("second") - mux = StreamMux(transformers=[t1, t2]) + mux = AsyncStreamMux(transformers=[t1, t2]) mux.wire_channels({"first": t1.channel}) mux.wire_channels({"second": t2.channel}) diff --git a/libs/langgraph/tests/test_stream_reducers.py b/libs/langgraph/tests/test_stream_reducers.py index c0d119168..20362ec75 100644 --- a/libs/langgraph/tests/test_stream_reducers.py +++ b/libs/langgraph/tests/test_stream_reducers.py @@ -31,7 +31,7 @@ async def test_values_captures_values_events(): reducer.finalize() collected = [] - async for item in reducer.values_log.subscribe(0): + async for item in reducer.values_log: collected.append(item) assert len(collected) == 2 assert collected[0]["data"] == {"a": 1} @@ -47,7 +47,7 @@ async def test_values_ignores_other_modes(): reducer.finalize() collected = [] - async for item in reducer.values_log.subscribe(0): + async for item in reducer.values_log: collected.append(item) assert len(collected) == 0 @@ -113,7 +113,7 @@ async def test_messages_groups_lifecycle(): reducer.finalize() collected = [] - async for stream in reducer.messages_log.subscribe(0): + async for stream in reducer.messages_log: collected.append(stream) assert len(collected) == 1 assert isinstance(collected[0], ChatModelStream) @@ -131,7 +131,7 @@ async def test_messages_multiple_sequential(): reducer.finalize() collected = [] - async for stream in reducer.messages_log.subscribe(0): + async for stream in reducer.messages_log: collected.append(stream) assert len(collected) == 2 @@ -147,7 +147,7 @@ async def test_messages_namespace_filter(): reducer.finalize() collected = [] - async for stream in reducer.messages_log.subscribe(0): + async for stream in reducer.messages_log: collected.append(stream) assert len(collected) == 1 @@ -163,7 +163,7 @@ async def test_messages_node_filter(): reducer.finalize() collected = [] - async for stream in reducer.messages_log.subscribe(0): + async for stream in reducer.messages_log: collected.append(stream) assert len(collected) == 1 @@ -181,7 +181,7 @@ async def test_messages_error_event(): reducer.finalize() collected: list[ChatModelStream] = [] - async for stream in reducer.messages_log.subscribe(0): + async for stream in reducer.messages_log: collected.append(stream) assert len(collected) == 1 assert collected[0].done @@ -198,7 +198,7 @@ async def test_messages_fail_propagates_to_active(): # The messages log should be failed too with pytest.raises(RuntimeError, match="graph failed"): - async for _ in reducer.messages_log.subscribe(0): + async for _ in reducer.messages_log: pass @@ -210,5 +210,5 @@ async def test_values_fail_propagates(): reducer.fail(RuntimeError("graph failed")) with pytest.raises(RuntimeError, match="graph failed"): - async for _ in reducer.values_log.subscribe(0): + async for _ in reducer.values_log: pass diff --git a/libs/langgraph/tests/test_stream_run_stream.py b/libs/langgraph/tests/test_stream_run_stream.py index 09cb5b67b..91c892d5f 100644 --- a/libs/langgraph/tests/test_stream_run_stream.py +++ b/libs/langgraph/tests/test_stream_run_stream.py @@ -4,12 +4,13 @@ from typing import Any import pytest -from langgraph.stream._mux import StreamMux +from langgraph.stream._mux import AsyncStreamMux from langgraph.stream._types import ProtocolEvent from langgraph.stream.chat_model_stream import ChatModelStream from langgraph.stream.run_stream import ( AsyncGraphRunStream, AsyncSubgraphRunStream, + SubgraphRunStream, create_async_graph_run_stream, create_graph_run_stream, ) @@ -44,7 +45,7 @@ async def test_aiter_yields_all_events(): @pytest.mark.anyio async def test_subgraph_name_and_index(): vr, mr = ValuesTransformer(), MessagesTransformer() - mux = StreamMux(transformers=[vr, mr]) + mux = AsyncStreamMux(transformers=[vr, mr]) sub = AsyncSubgraphRunStream( mux=mux, namespace=["researcher:2"], @@ -57,7 +58,7 @@ async def test_subgraph_name_and_index(): @pytest.mark.anyio async def test_subgraph_name_no_index(): vr, mr = ValuesTransformer(), MessagesTransformer() - mux = StreamMux(transformers=[vr, mr]) + mux = AsyncStreamMux(transformers=[vr, mr]) sub = AsyncSubgraphRunStream( mux=mux, namespace=["agent"], transformers=[vr, mr] ) @@ -132,7 +133,7 @@ async def test_messages_yields_streams(): @pytest.mark.anyio async def test_interrupted_false_by_default(): vr, mr = ValuesTransformer(), MessagesTransformer() - mux = StreamMux(transformers=[vr, mr]) + mux = AsyncStreamMux(transformers=[vr, mr]) run = AsyncGraphRunStream(mux=mux, transformers=[vr, mr]) assert run.interrupted is False @@ -140,7 +141,7 @@ async def test_interrupted_false_by_default(): @pytest.mark.anyio async def test_abort_sets_signal(): vr, mr = ValuesTransformer(), MessagesTransformer() - mux = StreamMux(transformers=[vr, mr]) + mux = AsyncStreamMux(transformers=[vr, mr]) run = AsyncGraphRunStream(mux=mux, transformers=[vr, mr]) assert not run.signal.is_set() run.abort() @@ -290,10 +291,8 @@ def test_sync_messages(): assert collected[0].done -def test_sync_messages_text_accessible(): - """Sync consumers should be able to read ChatModelStream text content - without an async event loop. - """ +def test_sync_messages_text_streaming(): + """Sync consumers can iterate msg.text for deltas.""" chunks = [ ((), "messages", {"event": "message-start", "message_id": "m1"}), ( @@ -314,18 +313,29 @@ def test_sync_messages_text_accessible(): ), ((), "messages", {"event": "message-finish", "reason": "stop"}), ] - run = create_graph_run_stream(_sync_source(chunks)) + # Iterate deltas + run = create_graph_run_stream(_sync_source(chunks)) for msg in run.messages: - assert isinstance(msg.text, str) - assert msg.text == "Hello world" + deltas = list(msg.text) + assert deltas == ["Hello", " world"] assert msg.done + # str() returns full text + run = create_graph_run_stream(_sync_source(chunks)) + for msg in run.messages: + assert str(msg.text) == "Hello world" -def test_sync_messages_content_populated_when_yielded(): - """When sync run.messages yields a ChatModelStream, its content should - be fully populated (done=True) with all text accumulated. - """ + # After message is done, .text returns plain str + run = create_graph_run_stream(_sync_source(chunks)) + for msg in run.messages: + list(msg.text) # exhaust deltas + assert isinstance(msg.text, str) + assert msg.text == "Hello world" + + +def test_sync_messages_multiple(): + """Multiple sync messages each stream their own deltas.""" chunks = [ ((), "messages", {"event": "message-start", "message_id": "m1"}), ( @@ -350,12 +360,10 @@ def test_sync_messages_content_populated_when_yielded(): ] run = create_graph_run_stream(_sync_source(chunks)) - messages = list(run.messages) - assert len(messages) == 2 - assert messages[0].done - assert messages[0].text == "answer" - assert messages[1].done - assert messages[1].text == "second" + all_deltas = [] + for msg in run.messages: + all_deltas.append(list(msg.text)) + assert all_deltas == [["answer"], ["second"]] def test_sync_output_mapper(): @@ -498,11 +506,14 @@ def test_sync_lazy_interleaved_projections(): assert next(vit) == {"v": 1} assert consumed == 1 - # Pull first message — pumps until message-finish (item 3) so the - # ChatModelStream is fully populated before yielding. + # Pull first message — yielded on message-start (item 2). + # Consuming str(msg.text) drives the pump to message-finish (item 3). mit = iter(run.messages) msg = next(mit) assert isinstance(msg, ChatModelStream) + assert consumed == 2 + assert not msg.done + str(msg.text) # pump until message completes assert msg.done assert consumed == 3 @@ -608,3 +619,362 @@ async def test_subgraph_child_values_receive_post_discovery_events(): f"Expected 2 child value snapshots but got {len(values)}: {values}. " "Child transformer missed post-discovery events." ) + + +# --------------------------------------------------------------------------- +# SubgraphRunStream — sync subgraph tests +# --------------------------------------------------------------------------- + + +def test_sync_subgraphs_discovery(): + """Iterating .subgraphs should discover child namespaces and yield + SubgraphRunStream instances with correct name and index. + """ + chunks = [ + (("agent:0",), "values", {"v": 1}), + (("agent:1",), "values", {"v": 2}), + ((), "values", {"done": True}), + ] + run = create_graph_run_stream(_sync_source(chunks)) + subs = list(run.subgraphs) + assert len(subs) == 2 + assert all(isinstance(s, SubgraphRunStream) for s in subs) + assert subs[0].name == "agent" + assert subs[0].index == 0 + assert subs[1].name == "agent" + assert subs[1].index == 1 + + +def test_sync_subgraph_name_no_index(): + """Subgraph without a colon-delimited index should have index=0.""" + chunks = [ + (("planner",), "values", {"v": 1}), + ((), "values", {"done": True}), + ] + run = create_graph_run_stream(_sync_source(chunks)) + subs = list(run.subgraphs) + assert len(subs) == 1 + assert subs[0].name == "planner" + assert subs[0].index == 0 + + +def test_sync_subgraph_no_subgraphs(): + """When all events are root-level, .subgraphs should yield nothing.""" + chunks = [ + ((), "values", {"v": 1}), + ((), "values", {"v": 2}), + ] + run = create_graph_run_stream(_sync_source(chunks)) + subs = list(run.subgraphs) + assert subs == [] + + +def test_sync_subgraph_values(): + """SubgraphRunStream.values should yield only values from the child namespace.""" + chunks = [ + (("child:0",), "values", {"v": 1}), + ((), "values", {"root": True}), + (("child:0",), "values", {"v": 2}), + ((), "values", {"done": True}), + ] + run = create_graph_run_stream(_sync_source(chunks)) + for sub in run.subgraphs: + vals = list(sub.values) + assert vals == [{"v": 1}, {"v": 2}] + + +def test_sync_subgraph_values_multiple_children(): + """Each child stream should only see its own values.""" + chunks = [ + (("a:0",), "values", {"who": "a0"}), + (("b:0",), "values", {"who": "b0"}), + (("a:0",), "values", {"who": "a0-2"}), + ((), "values", {"done": True}), + ] + run = create_graph_run_stream(_sync_source(chunks)) + children: dict[str, list[Any]] = {} + for sub in run.subgraphs: + children[f"{sub.name}:{sub.index}"] = list(sub.values) + + assert children["a:0"] == [{"who": "a0"}, {"who": "a0-2"}] + assert children["b:0"] == [{"who": "b0"}] + + +def test_sync_subgraph_output(): + """SubgraphRunStream.output should return the last values for the child.""" + chunks = [ + (("child:0",), "values", {"v": 1}), + (("child:0",), "values", {"v": 2}), + ((), "values", {"done": True}), + ] + run = create_graph_run_stream(_sync_source(chunks)) + for sub in run.subgraphs: + assert sub.output == {"v": 2} + + +def test_sync_subgraph_output_with_mapper(): + """Output mapper should apply to subgraph output.""" + chunks = [ + (("child:0",), "values", {"v": 42}), + ((), "values", {"done": True}), + ] + run = create_graph_run_stream( + _sync_source(chunks), output_mapper=lambda x: {"mapped": x.get("v")} + ) + for sub in run.subgraphs: + assert sub.output == {"mapped": 42} + + +def test_sync_subgraph_values_with_mapper(): + """Output mapper should apply to each yielded value snapshot.""" + chunks = [ + (("child:0",), "values", {"v": 1}), + (("child:0",), "values", {"v": 2}), + ((), "values", {"done": True}), + ] + run = create_graph_run_stream( + _sync_source(chunks), output_mapper=lambda x: {"m": x.get("v")} + ) + for sub in run.subgraphs: + vals = list(sub.values) + assert vals == [{"m": 1}, {"m": 2}] + + +def test_sync_subgraph_messages(): + """SubgraphRunStream.messages should yield fully populated ChatModelStream instances.""" + chunks = [ + ( + ("agent:0",), + "messages", + {"event": "message-start", "message_id": "m1", "__node__": "agent"}, + ), + ( + ("agent:0",), + "messages", + { + "event": "content-block-delta", + "content_block": {"type": "text", "text": "hello"}, + "__node__": "agent", + }, + ), + ( + ("agent:0",), + "messages", + {"event": "message-finish", "reason": "stop", "__node__": "agent"}, + ), + ((), "values", {"done": True}), + ] + run = create_graph_run_stream(_sync_source(chunks)) + for sub in run.subgraphs: + msgs = list(sub.messages) + assert len(msgs) == 1 + assert isinstance(msgs[0], ChatModelStream) + assert msgs[0].done + assert msgs[0].text == "hello" + + +def test_sync_subgraph_messages_isolated(): + """Messages from different subgraphs should not leak between children.""" + chunks = [ + ( + ("a:0",), + "messages", + {"event": "message-start", "message_id": "m-a", "__node__": "a"}, + ), + ( + ("a:0",), + "messages", + { + "event": "content-block-delta", + "content_block": {"type": "text", "text": "from-a"}, + "__node__": "a", + }, + ), + ( + ("a:0",), + "messages", + {"event": "message-finish", "reason": "stop", "__node__": "a"}, + ), + ( + ("b:0",), + "messages", + {"event": "message-start", "message_id": "m-b", "__node__": "b"}, + ), + ( + ("b:0",), + "messages", + { + "event": "content-block-delta", + "content_block": {"type": "text", "text": "from-b"}, + "__node__": "b", + }, + ), + ( + ("b:0",), + "messages", + {"event": "message-finish", "reason": "stop", "__node__": "b"}, + ), + ((), "values", {"done": True}), + ] + run = create_graph_run_stream(_sync_source(chunks)) + msg_texts: dict[str, list[str]] = {} + for sub in run.subgraphs: + msg_texts[sub.name] = [str(m.text) for m in sub.messages] + + assert msg_texts["a"] == ["from-a"] + assert msg_texts["b"] == ["from-b"] + + +def test_sync_subgraph_raw_iter(): + """Iterating a SubgraphRunStream directly should yield events scoped + to the child namespace. + """ + chunks = [ + (("child:0",), "values", {"v": 1}), + ((), "values", {"root": True}), + (("child:0",), "updates", {"node": "x"}), + ((), "values", {"done": True}), + ] + run = create_graph_run_stream(_sync_source(chunks)) + for sub in run.subgraphs: + events = list(sub) + methods = [e["method"] for e in events] + assert "values" in methods + assert "updates" in methods + # Root events should not appear + for e in events: + assert e["params"]["namespace"] == ["child:0"] + + +def test_sync_subgraph_events_after_discovery(): + """Events arriving after a namespace is first discovered should still + be visible in the child's values iteration. + """ + chunks = [ + (("child:0",), "values", {"v": 1}), # triggers discovery + ((), "values", {"root": 1}), + (("child:0",), "values", {"v": 2}), # after discovery + (("child:0",), "values", {"v": 3}), # after discovery + ((), "values", {"done": True}), + ] + run = create_graph_run_stream(_sync_source(chunks)) + for sub in run.subgraphs: + vals = list(sub.values) + assert vals == [{"v": 1}, {"v": 2}, {"v": 3}] + + +def test_sync_subgraph_lazy_pump(): + """Subgraph iteration should pump the source lazily.""" + consumed = 0 + + def counting_source(): + nonlocal consumed + for chunk in [ + (("child:0",), "values", {"v": 1}), + (("child:0",), "values", {"v": 2}), + (("child:0",), "values", {"v": 3}), + ((), "values", {"done": True}), + ]: + consumed += 1 + yield chunk + + run = create_graph_run_stream(counting_source()) + assert consumed == 0 + + for sub in run.subgraphs: + # Discovery pumped the first event + it = iter(sub.values) + v = next(it) + assert v == {"v": 1} + # Should not have consumed everything yet + assert consumed < 4 + break # don't exhaust subgraphs + + +def test_sync_subgraph_interleave_parent_values(): + """Parent values and subgraph values should both be accessible + when interleaving iteration. + """ + chunks = [ + ((), "values", {"root": 1}), + (("child:0",), "values", {"child": 1}), + ((), "values", {"root": 2}), + (("child:0",), "values", {"child": 2}), + ((), "values", {"root": 3}), + ] + run = create_graph_run_stream(_sync_source(chunks)) + + # First drain parent values + root_vals = list(run.values) + assert root_vals == [{"root": 1}, {"root": 2}, {"root": 3}] + + # Source is exhausted, but subgraph transformers were registered + # via replay — subgraph iteration should still see buffered events + # Note: subgraphs must be iterated while source is being pumped + # to discover namespaces. Since we drained via values, namespace + # "child:0" was already discovered. But subgraphs iteration also + # needs to pump — and the source is exhausted. Let's verify it + # yields the discovered child. + subs = list(run.subgraphs) + assert len(subs) == 1 + assert subs[0].name == "child" + # The child transformer was registered via replay, so it saw the events + vals = list(subs[0].values) + assert vals == [{"child": 1}, {"child": 2}] + + +def test_sync_subgraph_interrupted(): + """Subgraph .interrupted should reflect the mux's interrupt state.""" + + class _FakeInterrupt: + def __init__(self, id: str): + self.id = id + + chunks = [ + (("child:0",), "values", {"__interrupt__": [_FakeInterrupt("i1")]}), + ((), "values", {"done": True}), + ] + run = create_graph_run_stream(_sync_source(chunks)) + for sub in run.subgraphs: + # Pump to process the interrupt + _ = sub.output + assert sub.interrupted is True + assert len(sub.interrupts) == 1 + + +def test_sync_subgraph_source_error(): + """If the source raises mid-stream, subgraphs that were already + discovered should still have their buffered data. + """ + + def bad_source(): + yield (("child:0",), "values", {"v": 1}) + yield (("child:0",), "values", {"v": 2}) + raise ValueError("boom") + + run = create_graph_run_stream(bad_source()) + for sub in run.subgraphs: + vals = list(sub.values) + assert vals == [{"v": 1}, {"v": 2}] + assert run._mux._error is not None + + +def test_sync_subgraph_output_drains_source(): + """Accessing subgraph .output should drain the full source.""" + consumed = 0 + + def counting_source(): + nonlocal consumed + for chunk in [ + (("child:0",), "values", {"v": 1}), + (("child:0",), "values", {"v": 2}), + ((), "values", {"done": True}), + ]: + consumed += 1 + yield chunk + + run = create_graph_run_stream(counting_source()) + for sub in run.subgraphs: + result = sub.output + assert result == {"v": 2} + assert consumed == 3