Unify EventLog — remove sync/async split from transformer API

Merge EventLog and AsyncEventLog into a single class with a _bind()
mechanism. EventLog starts unbound; the StreamMux calls _bind(is_async)
after transformer registration so only the correct iteration protocol
is available. This removes the is_async parameter from EventLog,
StreamChannel, and all transformer constructors — transformers just
create EventLog() and never need to know whether they run in sync or
async context.
This commit is contained in:
Nick Hollon
2026-04-16 09:32:47 -04:00
parent ca5d9a6bd7
commit 28cf5ed78d
8 changed files with 176 additions and 122 deletions
+1 -2
View File
@@ -4,14 +4,13 @@ Provides a ``StreamingHandler`` that wraps a compiled graph and exposes
ergonomic streaming projections through a transformer pipeline.
"""
from langgraph.stream._event_log import AsyncEventLog, EventLog
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__ = [
"AsyncEventLog",
"AsyncGraphRunStream",
"EventLog",
"GraphRunStream",
+81 -56
View File
@@ -8,16 +8,24 @@ from typing import Generic, TypeVar
T = TypeVar("T")
class _EventLogBase(Generic[T]):
"""Shared producer API for sync and async event logs.
class EventLog(Generic[T]):
"""Append-only buffer that supports multiple independent consumers.
Append-only buffer that supports multiple independent consumers.
Subclasses provide the iteration protocol (sync or async).
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``.
Producer API (thread-safe):
Producer API (thread-safe, works before and after binding):
push(item) — append an item, notify all waiting cursors
close() — mark the log as done
fail(err) — mark the log as errored
Sync iteration is pull-based: when a cursor catches up it calls
``_request_more`` to drive the graph forward.
Async iteration uses ``asyncio.Future`` objects — the producer
wakes cursors via ``loop.call_soon_threadsafe``.
"""
def __init__(self) -> None:
@@ -26,6 +34,35 @@ class _EventLogBase(Generic[T]):
self._error: BaseException | None = None
self._lock = threading.Lock()
# Binding state — None means unbound.
self._is_async: bool | None = None
# Sync pull callback (set by the run stream, not by bind).
self._request_more: Callable[[], bool] | None = None
# Async waiters (allocated on bind).
self._async_waiters: list[asyncio.Future[None]] | None = None
# ------------------------------------------------------------------
# Binding
# ------------------------------------------------------------------
def _bind(self, *, is_async: bool) -> None:
"""Bind this log to sync or async mode.
Called by the ``StreamMux`` after transformer registration.
Must be called exactly once before any iteration.
"""
if self._is_async is not None:
raise RuntimeError("EventLog is already bound")
self._is_async = is_async
if is_async:
self._async_waiters = []
# ------------------------------------------------------------------
# Producer API (thread-safe, mode-agnostic)
# ------------------------------------------------------------------
def push(self, item: T) -> None:
"""Append *item* and wake all waiting cursors."""
with self._lock:
@@ -47,31 +84,39 @@ class _EventLogBase(Generic[T]):
self._closed = True
self._notify()
# ------------------------------------------------------------------
# Notification
# ------------------------------------------------------------------
def _notify(self) -> None:
"""Wake waiting consumers. Overridden by subclasses."""
"""Wake async waiters if bound to async mode."""
waiters = self._async_waiters
if not waiters:
return
self._async_waiters = []
for fut in waiters:
if not fut.done():
try:
fut.get_loop().call_soon_threadsafe(fut.set_result, None)
except RuntimeError:
# Event loop already closed — nothing to notify.
pass
class EventLog(_EventLogBase[T]):
"""Sync event log with pull-based iteration.
Each call to ``__iter__`` creates a new cursor starting from the
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._request_more: Callable[[], bool] | None = None
# ------------------------------------------------------------------
# Sync iteration (pull-based)
# ------------------------------------------------------------------
def __iter__(self) -> Iterator[T]:
"""Return a new independent sync cursor over the log."""
if self._is_async is None:
raise TypeError(
"EventLog has not been bound yet. "
"Register the transformer with a StreamMux first."
)
if self._is_async:
raise TypeError(
"This EventLog is bound to async mode — use 'async for' instead."
)
return self._sync_cursor()
def _sync_cursor(self) -> Iterator[T]:
@@ -95,40 +140,19 @@ class EventLog(_EventLogBase[T]):
# No producer callback and not closed — buffer is complete.
return
class AsyncEventLog(_EventLogBase[T]):
"""Async event log with multi-cursor iteration.
Each call to ``__aiter__`` creates a new cursor starting from the
beginning. Cursors await ``asyncio.Future`` objects when they
catch up to the producer.
The producer (``push``/``close``/``fail``) is safe to call from
any thread — async waiters are notified via
``loop.call_soon_threadsafe``.
Use ``EventLog`` for sync consumers.
"""
def __init__(self) -> None:
super().__init__()
self._async_waiters: list[asyncio.Future[None]] = []
def _notify(self) -> None:
waiters = self._async_waiters
if not waiters:
return
self._async_waiters = []
for fut in waiters:
if not fut.done():
try:
fut.get_loop().call_soon_threadsafe(fut.set_result, None)
except RuntimeError:
# Event loop already closed — nothing to notify.
pass
# ------------------------------------------------------------------
# Async iteration
# ------------------------------------------------------------------
def __aiter__(self) -> AsyncIterator[T]:
"""Return a new independent async cursor over the log."""
if self._is_async is None:
raise TypeError(
"EventLog has not been bound yet. "
"Register the transformer with a StreamMux first."
)
if not self._is_async:
raise TypeError("This EventLog is bound to sync mode — use 'for' instead.")
return self._async_cursor()
async def _async_cursor(self) -> AsyncIterator[T]:
@@ -144,5 +168,6 @@ class AsyncEventLog(_EventLogBase[T]):
else:
loop = asyncio.get_running_loop()
fut: asyncio.Future[None] = loop.create_future()
assert self._async_waiters is not None
self._async_waiters.append(fut)
await fut
+15 -16
View File
@@ -4,7 +4,7 @@ import time
from collections.abc import Callable
from typing import Any
from langgraph.stream._event_log import AsyncEventLog, EventLog, _EventLogBase
from langgraph.stream._event_log import EventLog
from langgraph.stream._types import ProtocolEvent, StreamTransformer
from langgraph.stream.stream_channel import StreamChannel
@@ -18,15 +18,15 @@ class StreamMux:
into the main log.
Pass ``is_async=True`` when the mux will be consumed via async
iteration (``handler.astream()``). This creates ``AsyncEventLog``
instances instead of ``EventLog`` instances.
iteration (``handler.astream()``). All ``EventLog`` and
``StreamChannel`` instances discovered during ``register()`` are
automatically bound to the matching mode.
"""
def __init__(self, *, is_async: bool = False) -> None:
self._is_async = is_async
self._events: _EventLogBase[ProtocolEvent] = (
AsyncEventLog() if is_async else EventLog()
)
self._events: EventLog[ProtocolEvent] = EventLog()
self._events._bind(is_async=is_async)
self._transformers: list[StreamTransformer] = []
self._channels: list[StreamChannel[Any]] = []
self._seq = 0
@@ -35,8 +35,8 @@ class StreamMux:
"""Register a transformer and return its projection dict.
Calls ``transformer.init()``, stores the transformer for event
processing, and returns the projection. StreamChannels in the
projection are auto-wired.
processing, binds any ``EventLog`` or ``StreamChannel`` instances
in the projection, and returns the projection.
"""
projection = transformer.init()
if not isinstance(projection, dict):
@@ -45,7 +45,7 @@ class StreamMux:
f"got {type(projection).__name__}"
)
self._transformers.append(transformer)
self._wire_channels(projection)
self._bind_and_wire(projection)
return projection
def push(self, event: ProtocolEvent) -> None:
@@ -107,17 +107,14 @@ class StreamMux:
self._events.fail(err)
# ------------------------------------------------------------------
# StreamChannel auto-wiring
# Binding and StreamChannel auto-wiring
# ------------------------------------------------------------------
def _wire_channels(self, projection: dict[str, Any]) -> None:
"""Find StreamChannel instances in *projection* and wire them."""
def _bind_and_wire(self, projection: dict[str, Any]) -> None:
"""Bind and wire EventLog / StreamChannel instances in *projection*."""
for value in projection.values():
if isinstance(value, StreamChannel):
# Ensure the channel's log matches the mux's mode.
if value._is_async != self._is_async:
value._is_async = self._is_async
value._log = AsyncEventLog() if self._is_async else EventLog()
value._bind(is_async=self._is_async)
self._channels.append(value)
channel_name = value.name
@@ -128,6 +125,8 @@ class StreamMux:
return _forward
value._wire(_make_forward(channel_name))
elif isinstance(value, EventLog):
value._bind(is_async=self._is_async)
def _forward(self, channel_name: str, item: Any) -> None:
"""Inject a ProtocolEvent for a StreamChannel push.
@@ -5,7 +5,7 @@ from collections.abc import AsyncIterator, Iterator
from typing import Any
from langgraph.stream._convert import convert_to_protocol_event
from langgraph.stream._event_log import AsyncEventLog, EventLog
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
@@ -46,12 +46,11 @@ class GraphRunStream:
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
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):
elif isinstance(value, StreamChannel):
value._log._request_more = self._pump_next
def _pump_next(self) -> bool:
@@ -101,7 +100,6 @@ class GraphRunStream:
def __iter__(self) -> Iterator[ProtocolEvent]:
"""Iterate all protocol events from the mux's main event log."""
assert isinstance(self._mux._events, EventLog)
return iter(self._mux._events)
@@ -168,5 +166,4 @@ class AsyncGraphRunStream:
def __aiter__(self) -> AsyncIterator[ProtocolEvent]:
"""Iterate all protocol events from the mux's main event log."""
assert isinstance(self._mux._events, AsyncEventLog)
return self._mux._events.__aiter__()
@@ -3,7 +3,7 @@ from __future__ import annotations
from collections.abc import AsyncIterator, Callable, Iterator
from typing import Generic, TypeVar
from langgraph.stream._event_log import AsyncEventLog, EventLog, _EventLogBase
from langgraph.stream._event_log import EventLog
T = TypeVar("T")
@@ -12,25 +12,32 @@ class StreamChannel(Generic[T]):
"""A named projection channel with optional protocol auto-forwarding.
Wraps an event log and declares a protocol channel name. When the
`StreamMux` detects a `StreamChannel` in a transformer's ``init()``
``StreamMux`` detects a ``StreamChannel`` in a transformer's ``init()``
return value, it automatically wires every ``push()`` to inject a
`ProtocolEvent` into the main event stream using the channel's name
``ProtocolEvent`` into the main event stream using the channel's name
as the ``method``.
In-process consumers iterate the channel directly (``for item in ch``
or ``async for item in ch``). Remote SDK clients subscribe via
``session.subscribe("custom:<channelName>")``.
Like ``EventLog``, a ``StreamChannel`` starts unbound. The mux
calls ``_bind(is_async)`` during registration so the correct
iteration protocol is available by the time user code sees it.
Lifecycle (``_close`` / ``_fail``) is managed by the mux — transformers
using only StreamChannels don't need ``finalize`` / ``fail`` hooks.
"""
def __init__(self, name: str, *, is_async: bool = False) -> None:
def __init__(self, name: str) -> None:
self.name = name
self._is_async = is_async
self._log: _EventLogBase[T] = AsyncEventLog() if is_async else EventLog()
self._log: EventLog[T] = EventLog()
self._wire_fn: Callable[[T], None] | None = None
def _bind(self, *, is_async: bool) -> None:
"""Bind the underlying event log to sync or async mode."""
self._log._bind(is_async=is_async)
def push(self, item: T) -> None:
"""Append *item* to the log and auto-forward if wired."""
self._log.push(item)
@@ -58,16 +65,7 @@ class StreamChannel(Generic[T]):
# ------------------------------------------------------------------
def __iter__(self) -> Iterator[T]:
if not isinstance(self._log, EventLog):
raise RuntimeError(
"Cannot use sync iteration on an async StreamChannel. "
"Use 'async for' instead."
)
return iter(self._log)
def __aiter__(self) -> AsyncIterator[T]:
if not isinstance(self._log, AsyncEventLog):
raise RuntimeError(
"Cannot use async iteration on a sync StreamChannel. Use 'for' instead."
)
return self._log.__aiter__()
@@ -142,8 +142,8 @@ class StreamingHandler:
"""
mux = StreamMux(is_async=is_async)
values_t = ValuesTransformer(is_async=is_async)
messages_t = MessagesTransformer(is_async=is_async)
values_t = ValuesTransformer()
messages_t = MessagesTransformer()
all_transformers: list[StreamTransformer] = [values_t, messages_t]
if user_transformers:
@@ -2,7 +2,7 @@ from __future__ import annotations
from typing import Any
from langgraph.stream._event_log import AsyncEventLog, EventLog, _EventLogBase
from langgraph.stream._event_log import EventLog
from langgraph.stream._types import ProtocolEvent, StreamTransformer
@@ -15,10 +15,8 @@ class ValuesTransformer(StreamTransformer):
_native = True
def __init__(self, *, is_async: bool = False) -> None:
self._log: _EventLogBase[dict[str, Any]] = (
AsyncEventLog() if is_async else EventLog()
)
def __init__(self) -> None:
self._log: EventLog[dict[str, Any]] = EventLog()
self._latest: dict[str, Any] | None = None
self._interrupted = False
self._interrupts: list[Any] = []
@@ -61,10 +59,8 @@ class MessagesTransformer(StreamTransformer):
_native = True
def __init__(self, *, is_async: bool = False) -> None:
self._log: _EventLogBase[tuple[Any, dict[str, Any]]] = (
AsyncEventLog() if is_async else EventLog()
)
def __init__(self) -> None:
self._log: EventLog[tuple[Any, dict[str, Any]]] = EventLog()
def init(self) -> dict[str, Any]:
return {"messages": self._log}
+56 -16
View File
@@ -15,7 +15,6 @@ from typing_extensions import TypedDict
from langgraph.constants import END, START
from langgraph.graph import StateGraph
from langgraph.stream import (
AsyncEventLog,
EventLog,
StreamChannel,
StreamingHandler,
@@ -141,6 +140,7 @@ def _build_custom_stream_graph():
class TestEventLog:
def test_sync_iteration(self) -> None:
log: EventLog[int] = EventLog()
log._bind(is_async=False)
log.push(1)
log.push(2)
log.push(3)
@@ -149,6 +149,7 @@ class TestEventLog:
def test_multi_cursor(self) -> None:
log: EventLog[str] = EventLog()
log._bind(is_async=False)
log.push("a")
log.push("b")
log.close()
@@ -158,6 +159,7 @@ class TestEventLog:
def test_fail_propagation(self) -> None:
log: EventLog[int] = EventLog()
log._bind(is_async=False)
log.push(1)
log.fail(ValueError("test error"))
with pytest.raises(ValueError, match="test error"):
@@ -165,7 +167,8 @@ class TestEventLog:
@pytest.mark.anyio
async def test_async_iteration(self) -> None:
log: AsyncEventLog[int] = AsyncEventLog()
log: EventLog[int] = EventLog()
log._bind(is_async=True)
async def producer():
for i in range(3):
@@ -178,7 +181,8 @@ class TestEventLog:
@pytest.mark.anyio
async def test_async_multi_cursor(self) -> None:
log: AsyncEventLog[str] = AsyncEventLog()
log: EventLog[str] = EventLog()
log._bind(is_async=True)
log.push("x")
log.push("y")
log.close()
@@ -189,7 +193,8 @@ class TestEventLog:
@pytest.mark.anyio
async def test_async_fail(self) -> None:
log: AsyncEventLog[int] = AsyncEventLog()
log: EventLog[int] = EventLog()
log._bind(is_async=True)
log.push(1)
log.fail(RuntimeError("async error"))
with pytest.raises(RuntimeError, match="async error"):
@@ -199,6 +204,7 @@ class TestEventLog:
def test_sync_cursor_yields_items_before_error(self) -> None:
"""Sync cursor should yield all buffered items before raising."""
log: EventLog[int] = EventLog()
log._bind(is_async=False)
log.push(1)
log.push(2)
log.push(3)
@@ -212,7 +218,8 @@ class TestEventLog:
@pytest.mark.anyio
async def test_async_cursor_yields_items_before_error(self) -> None:
"""Async cursor should yield all buffered items before raising."""
log: AsyncEventLog[int] = AsyncEventLog()
log: EventLog[int] = EventLog()
log._bind(is_async=True)
log.push(1)
log.push(2)
log.push(3)
@@ -241,19 +248,22 @@ class TestEventLog:
def test_empty_log_sync(self) -> None:
"""Iterating a closed empty log should yield nothing."""
log: EventLog[int] = EventLog()
log._bind(is_async=False)
log.close()
assert list(log) == []
@pytest.mark.anyio
async def test_empty_log_async(self) -> None:
"""Async-iterating a closed empty log should yield nothing."""
log: AsyncEventLog[int] = AsyncEventLog()
log: EventLog[int] = EventLog()
log._bind(is_async=True)
log.close()
assert [item async for item in log] == []
def test_empty_log_fail_sync(self) -> None:
"""Failing an empty log should raise immediately with no items."""
log: EventLog[int] = EventLog()
log._bind(is_async=False)
log.fail(ValueError("empty fail"))
with pytest.raises(ValueError, match="empty fail"):
list(log)
@@ -261,24 +271,43 @@ class TestEventLog:
@pytest.mark.anyio
async def test_empty_log_fail_async(self) -> None:
"""Failing an empty log should raise immediately with no items (async)."""
log: AsyncEventLog[int] = AsyncEventLog()
log: EventLog[int] = EventLog()
log._bind(is_async=True)
log.fail(ValueError("empty fail"))
with pytest.raises(ValueError, match="empty fail"):
async for _ in log:
pass
def test_sync_has_no_aiter(self) -> None:
"""EventLog (sync) should not support async iteration."""
def test_unbound_iter_raises(self) -> None:
"""Iterating an unbound EventLog should raise TypeError."""
log: EventLog[int] = EventLog()
log.close()
assert not hasattr(log, "__aiter__")
with pytest.raises(TypeError, match="has not been bound"):
list(log)
def test_sync_bound_aiter_raises(self) -> None:
"""Sync-bound EventLog should reject async iteration."""
log: EventLog[int] = EventLog()
log._bind(is_async=False)
log.close()
with pytest.raises(TypeError, match="bound to sync mode"):
log.__aiter__()
@pytest.mark.anyio
async def test_async_has_no_iter(self) -> None:
"""AsyncEventLog should not support sync iteration."""
log: AsyncEventLog[int] = AsyncEventLog()
async def test_async_bound_iter_raises(self) -> None:
"""Async-bound EventLog should reject sync iteration."""
log: EventLog[int] = EventLog()
log._bind(is_async=True)
log.close()
assert not hasattr(log, "__iter__")
with pytest.raises(TypeError, match="bound to async mode"):
iter(log)
def test_double_bind_raises(self) -> None:
"""Binding an already-bound EventLog should raise."""
log: EventLog[int] = EventLog()
log._bind(is_async=False)
with pytest.raises(RuntimeError, match="already bound"):
log._bind(is_async=True)
# ---------------------------------------------------------------------------
@@ -289,6 +318,7 @@ class TestEventLog:
class TestStreamChannel:
def test_push_and_iterate(self) -> None:
ch: StreamChannel[str] = StreamChannel("test")
ch._bind(is_async=False)
ch.push("a")
ch.push("b")
ch._close()
@@ -297,6 +327,7 @@ class TestStreamChannel:
def test_wire_callback(self) -> None:
forwarded: list[str] = []
ch: StreamChannel[str] = StreamChannel("test")
ch._bind(is_async=False)
ch._wire(lambda item: forwarded.append(item))
ch.push("x")
ch.push("y")
@@ -307,6 +338,7 @@ class TestStreamChannel:
def test_fail_propagation(self) -> None:
"""_fail() should propagate the error through the underlying log."""
ch: StreamChannel[str] = StreamChannel("test")
ch._bind(is_async=False)
ch.push("a")
ch._fail(ValueError("channel error"))
items: list[str] = []
@@ -317,8 +349,9 @@ class TestStreamChannel:
@pytest.mark.anyio
async def test_async_iteration(self) -> None:
"""Async iteration should delegate to the inner AsyncEventLog."""
ch: StreamChannel[str] = StreamChannel("test", is_async=True)
"""Async iteration should delegate to the inner event log."""
ch: StreamChannel[str] = StreamChannel("test")
ch._bind(is_async=True)
ch.push("x")
ch.push("y")
ch._close()
@@ -328,6 +361,7 @@ class TestStreamChannel:
def test_push_without_wire(self) -> None:
"""Push without a wire callback should still append to the log."""
ch: StreamChannel[int] = StreamChannel("test")
ch._bind(is_async=False)
assert ch._wire_fn is None
ch.push(42)
ch._close()
@@ -690,6 +724,7 @@ class TestValuesTransformer:
"""Values events from subgraphs (non-empty namespace) should be ignored."""
t = ValuesTransformer()
t.init()
t._log._bind(is_async=False)
t.process(_event("values", {"val": "root"}))
t.process(_event("values", {"val": "sub"}, namespace=["sub"]))
@@ -703,6 +738,7 @@ class TestValuesTransformer:
"""Non-values events should be passed through but not captured."""
t = ValuesTransformer()
t.init()
t._log._bind(is_async=False)
result = t.process(_event("updates", {"x": 1}))
assert result is True # passed through
@@ -729,6 +765,7 @@ class TestMessagesTransformer:
def test_captures_root_messages(self) -> None:
t = MessagesTransformer()
t.init()
t._log._bind(is_async=False)
t.process(_event("messages", ("chunk", {"meta": True})))
t.finalize()
@@ -739,6 +776,7 @@ class TestMessagesTransformer:
def test_ignores_non_root_namespace(self) -> None:
t = MessagesTransformer()
t.init()
t._log._bind(is_async=False)
t.process(_event("messages", ("chunk", {}), namespace=["sub"]))
t.finalize()
@@ -747,6 +785,7 @@ class TestMessagesTransformer:
def test_ignores_non_messages_methods(self) -> None:
t = MessagesTransformer()
t.init()
t._log._bind(is_async=False)
result = t.process(_event("values", {"v": 1}))
assert result is True
@@ -756,6 +795,7 @@ class TestMessagesTransformer:
def test_fail_propagates(self) -> None:
t = MessagesTransformer()
t.init()
t._log._bind(is_async=False)
t.fail(ValueError("msg error"))
with pytest.raises(ValueError, match="msg error"):
list(t._log)