Add async lane to StreamTransformer; roll registration into StreamMux

- StreamTransformer: aprocess/afinalize/afail + schedule() helper with
  on_error="log"|"raise". requires_async flag (plus override detection)
  makes sync stream() raise at registration rather than at first event.
- StreamMux: apush/aclose/afail for the async dispatch path. aclose
  awaits all scheduled tasks across transformers before afinalize;
  afail cancels and awaits them before afail hooks.
- StreamMux now takes transformers in __init__ and owns extensions /
  native_keys aggregation and conflict detection — register() is gone.
- GraphRunStream / AsyncGraphRunStream read extensions and native keys
  off the mux directly; StreamingHandler._setup() inlined.

libs/langgraph/langgraph/stream/_mux.py

@@ -1,11 +1,16 @@
 from __future__ import annotations
 
+import asyncio
 import time
 from collections.abc import Callable
 from typing import Any
 
 from langgraph.stream._event_log import EventLog
-from langgraph.stream._types import ProtocolEvent, StreamTransformer
+from langgraph.stream._types import (
+    ProtocolEvent,
+    StreamTransformer,
+    transformer_requires_async,
+)
 from langgraph.stream.stream_channel import StreamChannel
 
 
@@ -23,7 +28,23 @@ class StreamMux:
     automatically bound to the matching mode.
     """
 
-    def __init__(self, *, is_async: bool = False) -> None:
+    def __init__(
+        self,
+        transformers: list[StreamTransformer] | None = None,
+        *,
+        is_async: bool = False,
+    ) -> None:
+        """Initialize the mux and register *transformers* in order.
+
+        Transformers are fixed at construction time — there is no
+        post-init ``register()``. Each transformer's ``init()`` is called,
+        projections are merged into ``self.extensions``, ``_native``
+        keys are recorded in ``self.native_keys``, and any ``EventLog``
+        / ``StreamChannel`` instances are bound/wired.
+
+        Raises ``RuntimeError`` if any transformer requires an async run
+        under sync mode, and ``ValueError`` on projection-key conflicts.
+        """
         self._is_async = is_async
         self._events: EventLog[ProtocolEvent] = EventLog()
         self._events._bind(is_async=is_async)
@@ -32,22 +53,48 @@ class StreamMux:
         self._logs: list[EventLog[Any]] = []
         self._seq = 0
 
-    def register(self, transformer: StreamTransformer) -> dict[str, Any]:
-        """Register a transformer and return its projection dict.
+        #: Merged projection dict across all registered transformers.
+        #: Treat as read-only — mutations won't be reflected back in
+        #: individual transformers' state.
+        self.extensions: dict[str, Any] = {}
+        #: Projection keys from transformers with ``_native = True``.
+        self.native_keys: set[str] = set()
+
+        for transformer in transformers or ():
+            self._register(transformer)
+
+    def _register(self, transformer: StreamTransformer) -> None:
+        """Register a transformer.
 
         Calls ``transformer.init()``, stores the transformer for event
         processing, binds any ``EventLog`` or ``StreamChannel`` instances
-        in the projection, and returns the projection.
+        in the projection, and merges the projection into
+        ``self.extensions``.
         """
+        if transformer_requires_async(transformer) and not self._is_async:
+            raise RuntimeError(
+                f"{type(transformer).__name__} requires an async run — "
+                "it overrides aprocess/afinalize/afail or sets "
+                "requires_async=True. Use astream(), not stream()."
+            )
         projection = transformer.init()
         if not isinstance(projection, dict):
             raise TypeError(
                 f"StreamTransformer.init() must return a dict, "
                 f"got {type(projection).__name__}"
             )
+        conflicts = set(projection) & set(self.extensions)
+        if conflicts:
+            raise ValueError(
+                f"Transformer {type(transformer).__name__} returned "
+                f"projection keys that conflict with already-registered "
+                f"keys: {conflicts}"
+            )
         self._transformers.append(transformer)
         self._bind_and_wire(projection)
-        return projection
+        self.extensions.update(projection)
+        if getattr(transformer, "_native", False):
+            self.native_keys.update(projection.keys())
 
     def push(self, event: ProtocolEvent) -> None:
         """Route *event* through all transformers, then append to the main log.
@@ -119,6 +166,107 @@ class StreamMux:
                 ch._fail(err)
         self._events.fail(err)
 
+    # ------------------------------------------------------------------
+    # Async dispatch
+    # ------------------------------------------------------------------
+
+    async def apush(self, event: ProtocolEvent) -> None:
+        """Async counterpart to ``push``. Awaits each transformer's
+        ``aprocess`` in registration order before appending to the main log.
+
+        A slow ``aprocess`` serializes the pipeline by design — that's the
+        guarantee that lets a later transformer (or a synchronous consumer)
+        see the result of the async work. For decoupled work, use
+        ``schedule()`` from inside ``process``/``aprocess`` instead.
+        """
+        keep = True
+        for transformer in self._transformers:
+            if not await transformer.aprocess(event):
+                keep = False
+        if keep:
+            self._seq += 1
+            event["seq"] = self._seq
+            self._events.push(event)
+
+    async def aclose(self) -> None:
+        """Async counterpart to ``close``.
+
+        Awaits every task started via ``StreamTransformer.schedule()``
+        across all transformers, then calls ``afinalize()`` on each,
+        then auto-closes logs/channels and the main event log.
+
+        If any scheduled task raised under ``on_error="raise"``, or any
+        transformer's ``afinalize`` raises, the exception propagates.
+        The caller (the pump) handles it by routing into ``afail``.
+        """
+        pending = self._collect_scheduled_tasks()
+        if pending:
+            results = await asyncio.gather(*pending, return_exceptions=True)
+            first_err = next(
+                (
+                    r
+                    for r in results
+                    if isinstance(r, BaseException)
+                    and not isinstance(r, asyncio.CancelledError)
+                ),
+                None,
+            )
+            if first_err is not None:
+                raise first_err
+
+        first_error: BaseException | None = None
+        for transformer in self._transformers:
+            try:
+                await transformer.afinalize()
+            except BaseException as e:
+                if first_error is None:
+                    first_error = e
+        for log in self._logs:
+            if not log._closed:
+                log.close()
+        for ch in self._channels:
+            if not ch._log._closed:
+                ch._close()
+        self._events.close()
+        if first_error is not None:
+            raise first_error
+
+    async def afail(self, err: BaseException) -> None:
+        """Async counterpart to ``fail``.
+
+        Cancels every scheduled task across all transformers, awaits
+        them to completion, then runs each transformer's ``afail``
+        hook and auto-fails logs/channels and the main event log.
+        """
+        pending = self._collect_scheduled_tasks()
+        for task in pending:
+            task.cancel()
+        if pending:
+            await asyncio.gather(*pending, return_exceptions=True)
+
+        for transformer in self._transformers:
+            try:
+                await transformer.afail(err)
+            except BaseException:
+                pass
+        for log in self._logs:
+            if not log._closed:
+                log.fail(err)
+        for ch in self._channels:
+            if not ch._log._closed:
+                ch._fail(err)
+        if not self._events._closed:
+            self._events.fail(err)
+
+    def _collect_scheduled_tasks(self) -> list[asyncio.Task[Any]]:
+        """Snapshot of all in-flight tasks scheduled via transformers."""
+        return [
+            task
+            for transformer in self._transformers
+            for task in getattr(transformer, "_stream_scheduled_tasks", ())
+            if not task.done()
+        ]
+
     # ------------------------------------------------------------------
     # Binding and StreamChannel auto-wiring
     # ------------------------------------------------------------------

libs/langgraph/langgraph/stream/_types.py

@@ -1,10 +1,15 @@
 from __future__ import annotations
 
+import asyncio
+import logging
 from abc import ABC, abstractmethod
-from typing import Any, Literal
+from collections.abc import Coroutine
+from typing import Any, ClassVar, Literal
 
 from typing_extensions import NotRequired, TypedDict
 
+_logger = logging.getLogger(__name__)
+
 
 class _ProtocolEventParams(TypedDict):
     """Parameters for a protocol event."""
@@ -39,13 +44,37 @@ class StreamTransformer(ABC):
     exposed as direct attributes on the run stream (in addition to
     appearing in `run.extensions`).
 
-    Subclasses must implement `init` and `process`. The `finalize` and
-    `fail` hooks are optional — the default implementations are no-ops.
+    Subclasses must implement `init` and override at least one of
+    `process` / `aprocess`. The `finalize`/`afinalize` and `fail`/`afail`
+    hooks are optional — the default implementations are no-ops.
     EventLog and StreamChannel instances in the projection dict are
     auto-closed/failed by the mux, so most transformers don't need
     ``finalize`` or ``fail`` at all.
+
+    Async lane
+    ----------
+    Transformers that need async work pick the async lane by:
+
+    1. Overriding ``aprocess`` (and optionally ``afinalize``/``afail``), or
+    2. Calling ``self.schedule(coro)`` from inside a sync ``process``, or
+    3. Setting ``requires_async = True`` explicitly.
+
+    The mux detects these cases at registration and raises if they're
+    used under sync ``stream()`` — they only work under ``astream()``.
+
+    Use ``aprocess`` when the pump must *wait* for async work before
+    the next transformer sees the event (e.g. PII redaction that
+    mutates ``event`` in place). Use ``schedule()`` for decoupled async
+    work whose result lands on an independent projection (e.g. async
+    moderation scoring, cost lookup, external tracing).
     """
 
+    #: Explicit opt-in for transformers that need a running event loop but
+    #: don't override any async method (for example, transformers that call
+    #: ``schedule()`` from a sync ``process``). The mux also auto-detects the
+    #: async lane when ``aprocess``/``afinalize``/``afail`` is overridden.
+    requires_async: ClassVar[bool] = False
+
     @abstractmethod
     def init(self) -> dict[str, Any]:
         """Return the projection dict.
@@ -59,25 +88,144 @@ class StreamTransformer(ABC):
         """
         ...
 
-    @abstractmethod
     def process(self, event: ProtocolEvent) -> bool:
-        """Process a protocol event.
+        """Sync event handler. Override for the sync lane.
 
         Called for every event before it is appended to the main event log.
         Return False to suppress the event from the main log.
+
+        Subclasses must override either ``process`` or ``aprocess``. The
+        default raises so a missing override fails loudly rather than
+        silently passing every event through.
         """
-        ...
+        raise NotImplementedError(
+            f"{type(self).__name__} must override process() or aprocess()"
+        )
+
+    async def aprocess(self, event: ProtocolEvent) -> bool:
+        """Async event handler. Override for the async lane.
+
+        The mux awaits this before dispatching to the next transformer,
+        so a slow ``aprocess`` serializes the pipeline. Use this only
+        when a later transformer — or a consumer reading the event
+        synchronously — must see the result of the async work (e.g.
+        PII redaction that mutates ``event`` in place).
+
+        The default delegates to ``process``, so purely-sync transformers
+        run unchanged under ``astream()``.
+        """
+        return self.process(event)
 
     def finalize(self) -> None:
-        """Called when the run ends normally.
+        """Called when the run ends normally (sync lane).
 
         Override to close EventLogs, resolve promises, or perform other
         teardown. StreamChannel instances are auto-closed by the mux.
         """
 
+    async def afinalize(self) -> None:
+        """Called when the run ends normally (async lane).
+
+        By the time this runs, the mux has already awaited every task
+        started via ``schedule()``, so EventLogs can be closed here
+        without a last-task-wins race.
+
+        The default delegates to ``finalize``.
+        """
+        self.finalize()
+
     def fail(self, err: BaseException) -> None:
-        """Called when the run ends with an error.
+        """Called when the run ends with an error (sync lane).
 
         Override to fail EventLogs, reject promises, or perform other
         teardown. StreamChannel instances are auto-failed by the mux.
         """
+
+    async def afail(self, err: BaseException) -> None:
+        """Called when the run ends with an error (async lane).
+
+        The mux cancels and awaits every task started via ``schedule()``
+        before calling this, so cleanup doesn't race with in-flight work.
+
+        The default delegates to ``fail``.
+        """
+        self.fail(err)
+
+    # ------------------------------------------------------------------
+    # Scheduled async work
+    # ------------------------------------------------------------------
+
+    def schedule(
+        self,
+        coro: Coroutine[Any, Any, Any],
+        *,
+        on_error: Literal["log", "raise"] = "log",
+    ) -> asyncio.Task[Any]:
+        """Schedule a coroutine tied to this transformer's lifecycle.
+
+        The mux holds the task reference, awaits all scheduled tasks
+        during ``aclose()`` before calling ``afinalize()``, and cancels
+        them on ``afail()``. Authors don't need to track tasks or
+        implement the last-task-closes-the-log dance.
+
+        ``on_error="log"`` (default): exceptions are caught and logged;
+        a single failed call doesn't tear down the run.
+
+        ``on_error="raise"``: exceptions propagate when the mux joins
+        pendings, converting the close path into the fail path.
+
+        Requires a running event loop — call only under ``astream()``.
+        Set ``requires_async = True`` on the class so registration under
+        sync ``stream()`` fails fast with a clear message.
+        """
+        try:
+            asyncio.get_running_loop()
+        except RuntimeError:
+            raise RuntimeError(
+                f"{type(self).__name__}.schedule() requires a running "
+                "event loop; this transformer must run under astream(), "
+                "not stream(). Set requires_async=True on the class so "
+                "this fails at registration rather than at first event."
+            ) from None
+
+        wrapped = self._wrap_scheduled(coro) if on_error == "log" else coro
+        task = asyncio.create_task(wrapped)
+        tasks = self._scheduled_task_set()
+        tasks.add(task)
+        task.add_done_callback(tasks.discard)
+        return task
+
+    @staticmethod
+    async def _wrap_scheduled(coro: Coroutine[Any, Any, Any]) -> Any:
+        try:
+            return await coro
+        except asyncio.CancelledError:
+            raise
+        except BaseException:
+            _logger.exception("Scheduled StreamTransformer task failed")
+
+    def _scheduled_task_set(self) -> set[asyncio.Task[Any]]:
+        """Lazily-allocated task set. Avoids requiring super().__init__()."""
+        tasks: set[asyncio.Task[Any]] | None = getattr(
+            self, "_stream_scheduled_tasks", None
+        )
+        if tasks is None:
+            tasks = set()
+            self._stream_scheduled_tasks = tasks
+        return tasks
+
+
+def transformer_requires_async(transformer: StreamTransformer) -> bool:
+    """True if the transformer needs a running event loop.
+
+    A transformer requires async if it explicitly opts in
+    (``requires_async = True``) or overrides any of the async-lane methods
+    (``aprocess``, ``afinalize``, ``afail``).
+    """
+    if transformer.requires_async:
+        return True
+    cls = type(transformer)
+    for name in ("aprocess", "afinalize", "afail"):
+        if getattr(cls, name) is not getattr(StreamTransformer, name):
+            return True
+    return False

libs/langgraph/langgraph/stream/run_stream.py

@@ -32,22 +32,24 @@ class GraphRunStream:
         self,
         graph_iter: Iterator[Any],
         mux: StreamMux,
-        extensions: dict[str, Any],
         values_transformer: ValuesTransformer,
     ) -> None:
         self._graph_iter = graph_iter
         self._mux = mux
-        self.extensions = extensions
+        self.extensions = mux.extensions
         self._values_transformer = values_transformer
         self._exhausted = False
+        # Native-transformer projections also show up as direct attributes.
+        for key in mux.native_keys:
+            setattr(self, key, mux.extensions[key])
         # Wire pull-based iteration: every sync EventLog calls _pump_next
         # when its cursor catches up to the buffer.
-        self._wire_request_more(mux, extensions)
+        self._wire_request_more(mux)
 
-    def _wire_request_more(self, mux: StreamMux, extensions: dict[str, Any]) -> None:
+    def _wire_request_more(self, mux: StreamMux) -> None:
         """Set _request_more on all sync EventLogs so iteration drives the graph."""
         mux._events._request_more = self._pump_next
-        for value in extensions.values():
+        for value in mux.extensions.values():
             if isinstance(value, EventLog):
                 value._request_more = self._pump_next
             elif isinstance(value, StreamChannel):
@@ -122,14 +124,16 @@ class AsyncGraphRunStream:
     def __init__(
         self,
         mux: StreamMux,
-        extensions: dict[str, Any],
         values_transformer: ValuesTransformer,
         pump_task: asyncio.Task[None],
     ) -> None:
         self._mux = mux
-        self.extensions = extensions
+        self.extensions = mux.extensions
         self._values_transformer = values_transformer
         self._pump_task = pump_task
+        # Native-transformer projections also show up as direct attributes.
+        for key in mux.native_keys:
+            setattr(self, key, mux.extensions[key])
 
     @property
     def output(self) -> Any:

libs/langgraph/langgraph/stream/streaming_handler.py

@@ -64,8 +64,10 @@ class StreamingHandler:
         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
+        values_t = ValuesTransformer()
+        mux = StreamMux(
+            [values_t, MessagesTransformer(), *(transformers or ())],
+            is_async=False,
         )
 
         graph_iter = iter(
@@ -80,10 +82,7 @@ class StreamingHandler:
             )
         )
 
-        run = GraphRunStream(graph_iter, mux, extensions, values_t)
-        for key in native_keys:
-            setattr(run, key, extensions[key])
-        return run
+        return GraphRunStream(graph_iter, mux, values_t)
 
     async def astream(
         self,
@@ -99,8 +98,10 @@ class StreamingHandler:
         Returns an `AsyncGraphRunStream` immediately. A background asyncio
         task pumps events from the graph into the transformer pipeline.
         """
-        mux, extensions, native_keys, values_t = self._setup(
-            transformers, is_async=True
+        values_t = ValuesTransformer()
+        mux = StreamMux(
+            [values_t, MessagesTransformer(), *(transformers or ())],
+            is_async=True,
         )
 
         async def pump() -> None:
@@ -114,55 +115,11 @@ class StreamingHandler:
                     interrupt_before=interrupt_before,
                     interrupt_after=interrupt_after,
                 ):
-                    mux.push(convert_to_protocol_event(part))
-                mux.close()
+                    await mux.apush(convert_to_protocol_event(part))
+                await mux.aclose()
             except BaseException as e:
-                mux.fail(e)
+                await mux.afail(e)
 
         task = asyncio.create_task(pump())
 
-        run = AsyncGraphRunStream(mux, extensions, values_t, task)
-        for key in native_keys:
-            setattr(run, key, extensions[key])
-        return run
-
-    # ------------------------------------------------------------------
-    # Internal helpers
-    # ------------------------------------------------------------------
-
-    @staticmethod
-    def _setup(
-        user_transformers: list[StreamTransformer] | None,
-        *,
-        is_async: bool = False,
-    ) -> tuple[StreamMux, dict[str, Any], set[str], ValuesTransformer]:
-        """Create the mux, register all transformers.
-
-        Returns (mux, extensions, native_keys, values_transformer).
-        """
-        mux = StreamMux(is_async=is_async)
-
-        values_t = ValuesTransformer()
-        messages_t = MessagesTransformer()
-
-        all_transformers: list[StreamTransformer] = [values_t, messages_t]
-        if user_transformers:
-            all_transformers.extend(user_transformers)
-
-        extensions: dict[str, Any] = {}
-        native_keys: set[str] = set()
-
-        for t in all_transformers:
-            projection = mux.register(t)
-            conflicts = set(projection) & set(extensions)
-            if conflicts:
-                name = type(t).__name__
-                raise ValueError(
-                    f"Transformer {name} returned projection keys that "
-                    f"conflict with already-registered keys: {conflicts}"
-                )
-            extensions.update(projection)
-            if getattr(t, "_native", False):
-                native_keys.update(projection.keys())
-
-        return mux, extensions, native_keys, values_t
+        return AsyncGraphRunStream(mux, values_t, task)

libs/langgraph/tests/test_streaming_handler.py

@@ -627,7 +627,7 @@ class TestConvertToProtocolEvent:
 
 class TestStreamMux:
     def test_register_non_dict_raises(self) -> None:
-        """init() returning a non-dict should raise TypeError."""
+        """init() returning a non-dict should raise TypeError at construction."""
 
         class BadTransformer(StreamTransformer):
             def init(self) -> Any:
@@ -636,9 +636,8 @@ class TestStreamMux:
             def process(self, event: ProtocolEvent) -> bool:
                 return True
 
-        mux = StreamMux()
         with pytest.raises(TypeError, match="must return a dict"):
-            mux.register(BadTransformer())
+            StreamMux([BadTransformer()])
 
     def test_event_suppression(self) -> None:
         """When process() returns False, the event should not appear in the main log."""
@@ -651,8 +650,7 @@ class TestStreamMux:
                 # Suppress "updates" events
                 return event["method"] != "updates"
 
-        mux = StreamMux()
-        mux.register(FilterTransformer())
+        mux = StreamMux([FilterTransformer()])
 
         mux.push(_event("values", {"a": 1}))
         mux.push(_event("updates", {"b": 2}))
@@ -685,9 +683,7 @@ class TestStreamMux:
                 seen_by_second.append(event["method"])
                 return False
 
-        mux = StreamMux()
-        mux.register(PassTransformer())
-        mux.register(RejectTransformer())
+        mux = StreamMux([PassTransformer(), RejectTransformer()])
 
         mux.push(_event("values"))
         mux.close()
@@ -836,10 +832,8 @@ class TestStreamMuxResilience:
             def finalize(self) -> None:
                 self.finalized = True
 
-        mux = StreamMux()
-        mux.register(BrokenFinalizer())
         good = GoodTransformer()
-        mux.register(good)
+        mux = StreamMux([BrokenFinalizer(), good])
 
         mux.push(_event("values"))
 
@@ -876,10 +870,8 @@ class TestStreamMuxResilience:
             def fail(self, err: BaseException) -> None:
                 self.failed_with = err
 
-        mux = StreamMux()
-        mux.register(BrokenFailer())
         good = GoodTransformer()
-        mux.register(good)
+        mux = StreamMux([BrokenFailer(), good])
 
         original_error = ValueError("original")
         mux.fail(original_error)
@@ -904,8 +896,7 @@ class TestStreamMuxResilience:
                 raise RuntimeError("finalize broke")
 
         t = BrokenWithChannel()
-        mux = StreamMux()
-        mux.register(t)
+        mux = StreamMux([t])
 
         with pytest.raises(RuntimeError, match="finalize broke"):
             mux.close()
@@ -1022,8 +1013,7 @@ class TestCustomTransformer:
                 self._channel.push(f"saw:{event['method']}")
                 return True
 
-        mux = StreamMux()
-        mux.register(ChannelPusher())
+        mux = StreamMux([ChannelPusher()])
 
         mux.push(_event("values"))
         mux.push(_event("updates"))
@@ -1074,8 +1064,7 @@ class TestEventLogAutoLifecycle:
                 self._log.push("saw_event")
                 return True
 
-        mux = StreamMux()
-        mux.register(SimpleTransformer())
+        mux = StreamMux([SimpleTransformer()])
 
         mux.push(_event("values"))
         mux.close()
@@ -1100,8 +1089,7 @@ class TestEventLogAutoLifecycle:
                 return True
 
         t = SimpleTransformer()
-        mux = StreamMux()
-        mux.register(t)
+        mux = StreamMux([t])
 
         mux.push(_event("values"))
         mux.fail(ValueError("boom"))
@@ -1127,8 +1115,7 @@ class TestEventLogAutoLifecycle:
             def finalize(self) -> None:
                 self._log.close()
 
-        mux = StreamMux()
-        mux.register(ManualCloseTransformer())
+        mux = StreamMux([ManualCloseTransformer()])
         # Should not raise even though the log is closed by both
         # the transformer and the mux.
         mux.close()
@@ -1156,3 +1143,304 @@ class TestEventLogAutoLifecycle:
         _ = run.output
         items = list(run.extensions["minimal"])
         assert len(items) > 0
+
+
+# ---------------------------------------------------------------------------
+# Async transformer lane — aprocess / afinalize / afail / schedule()
+# ---------------------------------------------------------------------------
+
+
+class TestAsyncTransformerLane:
+    @pytest.mark.anyio
+    async def test_aprocess_is_awaited_before_next_transformer(self) -> None:
+        """aprocess must complete before the next transformer sees the event.
+
+        This is the load-bearing guarantee for mutating transformers
+        like PII redaction: the downstream transformer reads the mutated
+        event synchronously.
+        """
+        order: list[str] = []
+
+        class RedactTransformer(StreamTransformer):
+            requires_async = True
+
+            def init(self) -> dict[str, Any]:
+                return {}
+
+            async def aprocess(self, event: ProtocolEvent) -> bool:
+                await asyncio.sleep(0.01)
+                order.append("redact")
+                event["params"]["data"]["redacted"] = True
+                return True
+
+        class ObserverTransformer(StreamTransformer):
+            def init(self) -> dict[str, Any]:
+                return {}
+
+            def process(self, event: ProtocolEvent) -> bool:
+                if event["method"] == "values":
+                    order.append(f"observe:{event['params']['data'].get('redacted')}")
+                return True
+
+        mux = StreamMux([RedactTransformer(), ObserverTransformer()], is_async=True)
+
+        await mux.apush(_event("values", {"secret": "x"}))
+        await mux.aclose()
+
+        assert order == ["redact", "observe:True"]
+
+    @pytest.mark.anyio
+    async def test_schedule_joins_tasks_before_afinalize(self) -> None:
+        """Every scheduled task must complete before afinalize runs."""
+        phase: list[str] = []
+
+        class SchedTransformer(StreamTransformer):
+            requires_async = True
+
+            def __init__(self) -> None:
+                self._log: EventLog[str] = EventLog()
+
+            def init(self) -> dict[str, Any]:
+                return {"out": self._log}
+
+            def process(self, event: ProtocolEvent) -> bool:
+                if event["method"] == "values":
+
+                    async def work() -> None:
+                        await asyncio.sleep(0.01)
+                        phase.append("task")
+                        self._log.push("done")
+
+                    self.schedule(work())
+                return True
+
+            async def afinalize(self) -> None:
+                phase.append("afinalize")
+                self._log.close()
+
+        t = SchedTransformer()
+        mux = StreamMux([t], is_async=True)
+
+        await mux.apush(_event("values", {}))
+        await mux.apush(_event("values", {}))
+        await mux.aclose()
+
+        # Both tasks ran before afinalize; the log holds both pushes.
+        assert phase.count("task") == 2
+        assert phase[-1] == "afinalize"
+
+    @pytest.mark.anyio
+    async def test_sync_stream_rejects_async_transformer(self) -> None:
+        """Registering a requires_async transformer on a sync mux raises."""
+
+        class NeedsAsync(StreamTransformer):
+            requires_async = True
+
+            def init(self) -> dict[str, Any]:
+                return {}
+
+            def process(self, event: ProtocolEvent) -> bool:
+                return True
+
+        with pytest.raises(RuntimeError, match="requires an async run"):
+            StreamMux([NeedsAsync()], is_async=False)
+
+    @pytest.mark.anyio
+    async def test_sync_stream_rejects_aprocess_override(self) -> None:
+        """Overriding aprocess also marks the transformer as async-required."""
+
+        class HasAprocess(StreamTransformer):
+            def init(self) -> dict[str, Any]:
+                return {}
+
+            async def aprocess(self, event: ProtocolEvent) -> bool:
+                return True
+
+        with pytest.raises(RuntimeError, match="requires an async run"):
+            StreamMux([HasAprocess()], is_async=False)
+
+    def test_schedule_without_running_loop_raises(self) -> None:
+        """schedule() called outside an event loop fails with a clear message."""
+
+        class Sched(StreamTransformer):
+            requires_async = True
+
+            def init(self) -> dict[str, Any]:
+                return {}
+
+            def process(self, event: ProtocolEvent) -> bool:
+                return True
+
+        t = Sched()
+
+        async def noop() -> None:
+            pass
+
+        coro = noop()
+        try:
+            with pytest.raises(RuntimeError, match="requires a running event loop"):
+                t.schedule(coro)
+        finally:
+            coro.close()
+
+    @pytest.mark.anyio
+    async def test_schedule_on_error_log_swallows_exceptions(self) -> None:
+        """on_error="log" (default) keeps the run alive when a task fails."""
+
+        class Bad(StreamTransformer):
+            requires_async = True
+
+            def __init__(self) -> None:
+                self._log: EventLog[str] = EventLog()
+
+            def init(self) -> dict[str, Any]:
+                return {"out": self._log}
+
+            def process(self, event: ProtocolEvent) -> bool:
+                async def work() -> None:
+                    raise ValueError("boom")
+
+                self.schedule(work())  # default on_error="log"
+                return True
+
+            async def afinalize(self) -> None:
+                self._log.close()
+
+        mux = StreamMux([Bad()], is_async=True)
+
+        await mux.apush(_event("values", {}))
+        # Should not raise; the scheduled task's exception is logged.
+        await mux.aclose()
+
+    @pytest.mark.anyio
+    async def test_schedule_on_error_raise_fails_the_run(self) -> None:
+        """on_error="raise" propagates the exception through aclose."""
+
+        class Strict(StreamTransformer):
+            requires_async = True
+
+            def init(self) -> dict[str, Any]:
+                return {}
+
+            def process(self, event: ProtocolEvent) -> bool:
+                async def work() -> None:
+                    raise ValueError("strict boom")
+
+                self.schedule(work(), on_error="raise")
+                return True
+
+        mux = StreamMux([Strict()], is_async=True)
+
+        await mux.apush(_event("values", {}))
+        with pytest.raises(ValueError, match="strict boom"):
+            await mux.aclose()
+
+    @pytest.mark.anyio
+    async def test_afail_cancels_pending_scheduled_tasks(self) -> None:
+        """When the run fails, outstanding scheduled tasks are cancelled."""
+        cancelled = asyncio.Event()
+
+        class Sched(StreamTransformer):
+            requires_async = True
+
+            def init(self) -> dict[str, Any]:
+                return {}
+
+            def process(self, event: ProtocolEvent) -> bool:
+                async def work() -> None:
+                    try:
+                        await asyncio.sleep(5)
+                    except asyncio.CancelledError:
+                        cancelled.set()
+                        raise
+
+                self.schedule(work())
+                return True
+
+        mux = StreamMux([Sched()], is_async=True)
+
+        await mux.apush(_event("values", {}))
+        # Yield so the scheduled task actually starts before we cancel it;
+        # otherwise it's cancelled before its first step and the `except`
+        # inside work() never runs.
+        await asyncio.sleep(0)
+        await mux.afail(RuntimeError("run died"))
+
+        assert cancelled.is_set()
+
+    @pytest.mark.anyio
+    async def test_mixed_sync_and_async_transformers(self) -> None:
+        """Sync and async transformers coexist under astream."""
+        seen_sync: list[str] = []
+
+        class SyncOne(StreamTransformer):
+            def init(self) -> dict[str, Any]:
+                return {}
+
+            def process(self, event: ProtocolEvent) -> bool:
+                seen_sync.append(event["method"])
+                return True
+
+        class AsyncOne(StreamTransformer):
+            requires_async = True
+
+            def __init__(self) -> None:
+                self._log: EventLog[str] = EventLog()
+
+            def init(self) -> dict[str, Any]:
+                return {"seen": self._log}
+
+            async def aprocess(self, event: ProtocolEvent) -> bool:
+                await asyncio.sleep(0)
+                self._log.push(event["method"])
+                return True
+
+            async def afinalize(self) -> None:
+                self._log.close()
+
+        async_t = AsyncOne()
+        mux = StreamMux([SyncOne(), async_t], is_async=True)
+
+        await mux.apush(_event("values", {}))
+        await mux.apush(_event("updates", {}))
+        await mux.aclose()
+
+        assert seen_sync == ["values", "updates"]
+        items = [x async for x in async_t._log]
+        assert items == ["values", "updates"]
+
+    @pytest.mark.anyio
+    async def test_handler_astream_with_scheduled_work(self) -> None:
+        """End-to-end: transformer schedules work during an astream run."""
+
+        class Scorer(StreamTransformer):
+            requires_async = True
+
+            def __init__(self) -> None:
+                self._log: EventLog[int] = EventLog()
+
+            def init(self) -> dict[str, Any]:
+                return {"scores": self._log}
+
+            def process(self, event: ProtocolEvent) -> bool:
+                if event["method"] == "values":
+
+                    async def work() -> None:
+                        await asyncio.sleep(0.01)
+                        self._log.push(42)
+
+                    self.schedule(work())
+                return True
+
+            async def afinalize(self) -> None:
+                self._log.close()
+
+        graph = _build_simple_graph()
+        handler = StreamingHandler(graph)
+        run = await handler.astream(
+            {"value": "x", "items": []},
+            transformers=[Scorer()],
+        )
+        _ = await run.output
+        scores = [x async for x in run.extensions["scores"]]
+        assert scores and all(s == 42 for s in scores)