from __future__ import annotations import asyncio from collections import deque from collections.abc import AsyncIterator, Awaitable, Callable, Iterator from typing import TYPE_CHECKING, Generic, TypeVar if TYPE_CHECKING: from langgraph.stream._mux import StreamMux T = TypeVar("T") class StreamChannel(Generic[T]): """Single-consumer drainable queue for streaming events, with optional protocol auto-forwarding. When constructed with a `name`, the StreamMux auto-wires every `push()` to also inject a `ProtocolEvent` into the main event stream using the channel's name as the method. When constructed without a name, the channel is local-only — items are only visible to in-process consumers that iterate the channel directly. Items are popped off the front as the consumer advances — there is no retention beyond what's currently queued. A channel accepts exactly one subscriber; a second `__iter__` / `__aiter__` call raises. Use `tee(n)` / `atee(n)` for fan-out. Starts unbound — neither `__iter__` nor `__aiter__` is available until the StreamMux calls `_bind(is_async)`. After binding, only the matching iteration protocol works; the other raises `TypeError`. Pump wiring (set by the run stream, not by `_bind`): - `_request_more`: sync pump callable, returns True if a new event was produced. - `_arequest_more`: async pump coroutine factory, same contract. Memory is bounded by caller pace: both sync and async use caller- driven pumps, so each cursor advance produces at most one event. Lazy-subscribe: `push` appends to the local buffer only when a subscriber has registered. Auto-forward via `_wire_fn` always fires regardless of subscription state. Lifecycle (`close` / `fail`) is managed by the mux — transformers don't need to close their channels manually. """ def __init__(self, name: str | None = None, *, maxlen: int | None = None) -> None: """Initialize the channel. Args: name: Optional protocol channel name. When set, the StreamMux wires every `push()` to also inject a `ProtocolEvent` into the main event stream. Surfaced on the wire as `custom:` for user-defined transformers, or as `` for channels owned by a native transformer (`_native = True`). When `None`, the channel is local-only. maxlen: Accepted for forward compatibility; currently unused. The caller-driven pump bounds memory naturally for single-consumer use. Raises: ValueError: If `maxlen` is not a positive integer or `None`. """ if maxlen is not None and maxlen <= 0: raise ValueError("StreamChannel maxlen must be a positive int or None") self.name = name self._items: deque[tuple[int, T]] = deque() self._maxlen: int | None = maxlen self._closed = False self._error: BaseException | None = None self._is_async: bool | None = None self._subscribed = False self._request_more: Callable[[], bool] | None = None self._arequest_more: Callable[[], Awaitable[bool]] | None = None self._wire_fn: Callable[[T], None] | None = None self._mux: StreamMux | None = None # ------------------------------------------------------------------ # Binding # ------------------------------------------------------------------ def _bind_mux(self, mux: StreamMux) -> None: self._mux = mux def _bind(self, *, is_async: bool) -> None: """Bind this channel to sync or async mode. Called by the StreamMux after transformer registration. Must be called exactly once before any iteration. Args: is_async: True to enable async iteration, False for sync. Raises: RuntimeError: If the channel has already been bound. """ if self._is_async is not None: raise RuntimeError("StreamChannel is already bound") self._is_async = is_async # ------------------------------------------------------------------ # Mux wiring (not called by transformers directly) # ------------------------------------------------------------------ def _wire(self, fn: Callable[[T], None]) -> None: """Install the auto-forward callback (called by StreamMux).""" self._wire_fn = fn # ------------------------------------------------------------------ # Producer API # ------------------------------------------------------------------ def push(self, item: T) -> None: """Append an item. Auto-forwards if wired. The local buffer append is a no-op when no subscriber is registered, but auto-forwarding always fires so wired events reach the main event log regardless of subscription state. Items are stored as `(stamp, item)` tuples where stamp is a monotonic counter from the owning mux. Stamps are stripped by the default cursors; raw stamped tuples are visible on `_items`. Raises: RuntimeError: If the channel is closed (and subscribed). """ if self._subscribed: if self._closed: raise RuntimeError("Cannot push to a closed StreamChannel") stamp = self._mux._next_push_seq() if self._mux is not None else 0 self._items.append((stamp, item)) if self._wire_fn is not None: self._wire_fn(item) def close(self) -> None: """Mark the channel as complete.""" self._closed = True def fail(self, err: BaseException) -> None: """Mark the channel as errored. Args: err: The exception to surface to the subscriber. """ self._error = err self._closed = True # ------------------------------------------------------------------ # Sync iteration (caller-driven pump) # ------------------------------------------------------------------ def __iter__(self) -> Iterator[T]: """Subscribe and return a sync cursor. Can be called only once. Raises: TypeError: If the channel is unbound or bound to async mode. RuntimeError: If the channel already has a subscriber. """ if self._is_async is None: raise TypeError( "StreamChannel has not been bound yet. " "Register the transformer with a StreamMux first." ) if self._is_async: raise TypeError( "This StreamChannel is bound to async mode — use 'async for' instead." ) if self._subscribed: raise RuntimeError( "StreamChannel already has a subscriber; use .tee(n) for fan-out." ) self._subscribed = True return self._sync_cursor() def _sync_cursor(self) -> Iterator[T]: while True: if self._items: _stamp, item = self._items.popleft() yield item elif self._closed: if self._error is not None: raise self._error return elif self._request_more is not None: if not self._request_more(): if not self._items and not self._closed: return else: return # ------------------------------------------------------------------ # Async iteration (caller-driven pump) # ------------------------------------------------------------------ def __aiter__(self) -> AsyncIterator[T]: """Subscribe and return an async cursor. Can be called only once. Raises: TypeError: If the channel is unbound or bound to sync mode. RuntimeError: If the channel already has a subscriber. """ if self._is_async is None: raise TypeError( "StreamChannel has not been bound yet. " "Register the transformer with a StreamMux first." ) if not self._is_async: raise TypeError( "This StreamChannel is bound to sync mode — use 'for' instead." ) if self._subscribed: raise RuntimeError( "StreamChannel already has a subscriber; use .atee(n) for fan-out." ) self._subscribed = True return self._async_cursor() async def _async_cursor(self) -> AsyncIterator[T]: while True: if self._items: _stamp, item = self._items.popleft() yield item elif self._closed: if self._error is not None: raise self._error return elif self._arequest_more is not None: if not await self._arequest_more(): if not self._items and not self._closed: return else: return # ------------------------------------------------------------------ # Fan-out via tee # ------------------------------------------------------------------ def tee(self, n: int = 2) -> tuple[Iterator[T], ...]: """Subscribe and return `n` independent sync iterators. Each branch has its own buffer; items pulled from the underlying cursor are copied into every branch. Branches are naturally bounded by caller pace since the sync pump is caller-driven. Args: n: Number of branches to create. Must be >= 1. Returns: A tuple of `n` iterators over the same underlying stream. Raises: TypeError: If the channel is unbound or bound to async mode. RuntimeError: If the channel already has a subscriber. ValueError: If `n` < 1. """ if n < 1: raise ValueError("tee() requires n >= 1") source = self.__iter__() buffers: list[deque[T]] = [deque() for _ in range(n)] exhausted = [False] def branch(i: int) -> Iterator[T]: buf = buffers[i] while True: if buf: yield buf.popleft() elif exhausted[0]: return else: try: item = next(source) except StopIteration: exhausted[0] = True return for b in buffers: b.append(item) return tuple(branch(i) for i in range(n)) def atee(self, n: int = 2) -> tuple[AsyncIterator[T], ...]: """Subscribe and return `n` independent async iterators. Caller-driven fan-out: each branch's `__anext__` either pops from its own buffer or, under a shared `asyncio.Lock`, pulls one item from the underlying cursor and distributes it to every branch's buffer. Args: n: Number of branches to create. Must be >= 1. Returns: A tuple of `n` async iterators over the same underlying stream. Raises: TypeError: If the channel is unbound or bound to sync mode. RuntimeError: If the channel already has a subscriber. ValueError: If `n` < 1. """ if n < 1: raise ValueError("atee() requires n >= 1") source = self.__aiter__() buffers: list[deque[T]] = [deque() for _ in range(n)] exhausted = [False] error: list[BaseException | None] = [None] lock = asyncio.Lock() async def branch(i: int) -> AsyncIterator[T]: buf = buffers[i] while True: if buf: yield buf.popleft() continue if exhausted[0]: if error[0] is not None: raise error[0] return async with lock: if buf or exhausted[0]: continue try: item = await source.__anext__() except StopAsyncIteration: exhausted[0] = True continue except Exception as e: error[0] = e exhausted[0] = True continue for b in buffers: b.append(item) return tuple(branch(i) for i in range(n))