From 8215a9d02436a2969746f82ddcdb0aaabc341653 Mon Sep 17 00:00:00 2001 From: Nick Hollon Date: Thu, 21 May 2026 11:54:33 -0400 Subject: [PATCH] feat(langgraph): add `before_builtins` opt-in for stream transformers (#7882) --- libs/langgraph/langgraph/stream/_mux.py | 35 ++- libs/langgraph/langgraph/stream/_types.py | 17 + .../tests/test_stream_before_builtins.py | 296 ++++++++++++++++++ 3 files changed, 343 insertions(+), 5 deletions(-) create mode 100644 libs/langgraph/tests/test_stream_before_builtins.py diff --git a/libs/langgraph/langgraph/stream/_mux.py b/libs/langgraph/langgraph/stream/_mux.py index eb09154e4..061166ee1 100644 --- a/libs/langgraph/langgraph/stream/_mux.py +++ b/libs/langgraph/langgraph/stream/_mux.py @@ -62,6 +62,14 @@ class StreamMux: `extensions`, `_native` keys are recorded in `native_keys`, and any StreamChannel instances are bound and (if named) wired. + Transformers with `StreamTransformer.before_builtins = True` are + registered ahead of the rest, preserving relative order within + each lane. This lets content-mutating transformers (PII + redaction, content filters, etc.) run before built-ins like + `MessagesTransformer` that eagerly snapshot text fields into + their projections. See `StreamTransformer.before_builtins` for + the contract and foot-guns. + Args: transformers: Already-built transformer instances. Registered only on this mux — they are NOT cloned into child @@ -112,14 +120,31 @@ class StreamMux: self._pump_fn: Callable[[], bool] | None = None self._apump_fn: Callable[[], Awaitable[bool]] | None = None - # Factories run first (they propagate to child mini-muxes - # via `_make_child`), then any pre-built `transformers=` - # instances are registered as root-only — they aren't cloned - # for child scopes. + # Factories run first (they propagate to child mini-muxes via + # `_make_child`), then any pre-built `transformers=` instances + # are registered as root-only — they aren't cloned for child + # scopes. Within each group, transformers with + # `before_builtins = True` are registered ahead of the rest so + # they observe (and may mutate) events before built-ins like + # `MessagesTransformer`. The order *within* each lane matches + # the supplied sequence. + pre: list[StreamTransformer] = [] + rest: list[StreamTransformer] = [] if factories is not None: for factory in factories: - self._register(factory(scope)) + transformer = factory(scope) + ( + pre if getattr(transformer, "before_builtins", False) else rest + ).append(transformer) + for transformer in (*pre, *rest): + self._register(transformer) + pre.clear() + rest.clear() for transformer in transformers or (): + (pre if getattr(transformer, "before_builtins", False) else rest).append( + transformer + ) + for transformer in (*pre, *rest): self._register(transformer) def transformer_by_key(self, key: str) -> StreamTransformer | None: diff --git a/libs/langgraph/langgraph/stream/_types.py b/libs/langgraph/langgraph/stream/_types.py index 38a4e78ad..9f50884fd 100644 --- a/libs/langgraph/langgraph/stream/_types.py +++ b/libs/langgraph/langgraph/stream/_types.py @@ -91,11 +91,28 @@ class StreamTransformer(ABC): which modes a `stream_events(version="v3")` run requests from the graph. Empty tuple means the transformer consumes only synthetic events (or is purely passive). + before_builtins: Opt-in for transformers that must run *before* + built-in transformers like `MessagesTransformer` and + `ToolCallTransformer`. The mux partitions factories by this + flag at registration time: `before_builtins = True` + transformers are registered first, then everything else in + the order supplied. Within each lane, registration order is + preserved. This is the supported hook for content-mutating + transformers (PII redaction, profanity filters, etc.) whose + mutations must land before built-ins eagerly snapshot text + fields into their projections. **Foot-gun:** transformers + in this lane see `tasks` events before `LifecycleTransformer` + and `SubgraphTransformer` consume them — mutating + `event["params"]["namespace"]` or the data dict's + `id` / `result` / `error` / `interrupts` fields will desync + their bookkeeping. Observe freely; mutate only fields no + built-in reads (e.g. `delta.text` on `messages` events). """ requires_async: ClassVar[bool] = False supports_sync: ClassVar[bool] = False required_stream_modes: ClassVar[tuple[str, ...]] = () + before_builtins: ClassVar[bool] = False def __init__(self, scope: tuple[str, ...] = ()) -> None: """Initialize the transformer with its mux's scope. diff --git a/libs/langgraph/tests/test_stream_before_builtins.py b/libs/langgraph/tests/test_stream_before_builtins.py new file mode 100644 index 000000000..6e01dfb39 --- /dev/null +++ b/libs/langgraph/tests/test_stream_before_builtins.py @@ -0,0 +1,296 @@ +"""Tests for `StreamTransformer.before_builtins` lane ordering. + +`before_builtins = True` transformers are registered ahead of the +rest, preserving relative order within each lane. This lets +content-mutating transformers run before built-ins like +`MessagesTransformer` that eagerly snapshot text fields into their +projections. +""" + +from __future__ import annotations + +import time +from typing import Any, ClassVar + +from langgraph.stream._mux import StreamMux +from langgraph.stream._types import StreamTransformer +from langgraph.stream.stream_channel import StreamChannel +from langgraph.stream.transformers import ( + LifecycleTransformer, + MessagesTransformer, + TasksTransformer, +) + +TS = int(time.time() * 1000) + + +def _messages_event(namespace: list[str], data: Any) -> dict[str, Any]: + return { + "type": "event", + "method": "messages", + "params": {"namespace": namespace, "timestamp": TS, "data": data}, + } + + +class _Tap(StreamTransformer): + """Records the order it observed each event.""" + + required_stream_modes: ClassVar[tuple[str, ...]] = () + + def __init__(self, scope: tuple[str, ...] = (), *, label: str = "tap") -> None: + super().__init__(scope) + self.label = label + self.log: list[str] = [] + self._channel: StreamChannel[str] = StreamChannel() + + def init(self) -> dict[str, Any]: + return {f"tap_{self.label}": self._channel} + + def process(self, event: dict[str, Any]) -> bool: + self.log.append(self.label) + return True + + +class _PreTap(_Tap): + before_builtins: ClassVar[bool] = True + + +class _TextRedactor(StreamTransformer): + """Mutates `text-delta` events in place to a fixed redacted string.""" + + before_builtins: ClassVar[bool] = True + required_stream_modes: ClassVar[tuple[str, ...]] = ("messages",) + + def __init__(self, scope: tuple[str, ...] = ()) -> None: + super().__init__(scope) + self._channel: StreamChannel[str] = StreamChannel() + + def init(self) -> dict[str, Any]: + return {"redactor": self._channel} + + def process(self, event: dict[str, Any]) -> bool: + if event.get("method") != "messages": + return True + payload, _meta = event["params"]["data"] + if isinstance(payload, dict) and payload.get("event") == "content-block-delta": + delta = payload.get("delta") or {} + if delta.get("type") == "text-delta": + delta["text"] = "[REDACTED]" + return True + + +def test_before_builtins_factories_run_before_others() -> None: + """A `before_builtins=True` factory is registered ahead of the rest.""" + + seen: list[type[StreamTransformer]] = [] + + class _PostTap(_Tap): + def __init__(self, scope: tuple[str, ...] = ()) -> None: + super().__init__(scope, label="post") + + def process(self, event: dict[str, Any]) -> bool: + seen.append(_PostTap) + return True + + class _EagerTap(_Tap): + before_builtins: ClassVar[bool] = True + + def __init__(self, scope: tuple[str, ...] = ()) -> None: + super().__init__(scope, label="eager") + + def process(self, event: dict[str, Any]) -> bool: + seen.append(_EagerTap) + return True + + mux = StreamMux( + factories=[_PostTap, _EagerTap], + scope=(), + is_async=False, + ) + # `_EagerTap` was supplied second but should be registered first. + types_in_order = [type(t) for t in mux._transformers] + assert types_in_order.index(_EagerTap) < types_in_order.index(_PostTap) + + mux.push( + _messages_event([], ({"event": "message-start", "role": "ai", "id": "m1"}, {})) + ) + # And it ran first when the event was dispatched. + assert seen == [_EagerTap, _PostTap] + + +def test_within_lane_order_preserved() -> None: + """Within each lane, the supplied order is the registration order.""" + + class _A(_PreTap): + pass + + class _B(_PreTap): + pass + + class _C(_Tap): + pass + + class _D(_Tap): + pass + + mux = StreamMux( + factories=[ + lambda scope: _C(scope, label="c"), + lambda scope: _A(scope, label="a"), + lambda scope: _D(scope, label="d"), + lambda scope: _B(scope, label="b"), + ], + scope=(), + is_async=False, + ) + order = [t.label for t in mux._transformers] # type: ignore[attr-defined] + # Pre lane (a, b) ahead of default lane (c, d). Within each, supplied order kept. + assert order == ["a", "b", "c", "d"] + + +def test_redactor_runs_before_messages_transformer() -> None: + """Content mutated by a pre-lane transformer reaches `MessagesTransformer`.""" + + # Order supplied: built-ins first (as in pregel/main.py), then the + # opt-in pre-lane redactor. Partitioning should still register the + # redactor first. + mux = StreamMux( + factories=[MessagesTransformer, _TextRedactor], + scope=(), + is_async=False, + ) + types_in_order = [type(t) for t in mux._transformers] + assert types_in_order.index(_TextRedactor) < types_in_order.index( + MessagesTransformer + ) + + +def test_lifecycle_unaffected_by_pre_lane_observer() -> None: + """An observer-only pre-lane transformer doesn't break lifecycle bookkeeping.""" + + class _NoopPreObserver(StreamTransformer): + before_builtins: ClassVar[bool] = True + required_stream_modes: ClassVar[tuple[str, ...]] = ("tasks",) + + def __init__(self, scope: tuple[str, ...] = ()) -> None: + super().__init__(scope) + self._channel: StreamChannel[str] = StreamChannel() + self.seen: list[tuple[str, ...]] = [] + + def init(self) -> dict[str, Any]: + return {"noop_observer": self._channel} + + def process(self, event: dict[str, Any]) -> bool: + if event.get("method") == "tasks": + self.seen.append(tuple(event["params"]["namespace"])) + return True + + mux = StreamMux( + factories=[LifecycleTransformer, TasksTransformer, _NoopPreObserver], + scope=(), + is_async=False, + ) + types_in_order = [type(t) for t in mux._transformers] + assert types_in_order.index(_NoopPreObserver) < types_in_order.index( + LifecycleTransformer + ) + + observer = next(t for t in mux._transformers if isinstance(t, _NoopPreObserver)) + lifecycle = next( + t for t in mux._transformers if isinstance(t, LifecycleTransformer) + ) + + # Push a synthetic `tasks` event that lifecycle would normally track. + mux.push( + { + "type": "event", + "method": "tasks", + "params": { + "namespace": ["child:abc"], + "timestamp": TS, + "data": {"name": "child"}, + }, + } + ) + + # Pre-lane observer saw the event, AND lifecycle's bookkeeping still + # registered the new namespace (the observer didn't mutate anything). + assert observer.seen == [("child:abc",)] + assert ("child:abc",) in lifecycle._seen # type: ignore[attr-defined] + + +def test_default_is_false() -> None: + """`StreamTransformer.before_builtins` defaults to False.""" + + assert StreamTransformer.before_builtins is False + assert MessagesTransformer.before_builtins is False + assert LifecycleTransformer.before_builtins is False + + +def test_pre_lane_mutation_lands_in_messages_projection() -> None: + """End-to-end: text mutated by a pre-lane transformer is what + `MessagesTransformer` snapshots into its `ChatModelStream` projection. + + Without `before_builtins`, the redactor would run after + MessagesTransformer's eager extraction and the projection would + contain the raw, un-redacted text. + """ + + mux = StreamMux( + factories=[MessagesTransformer, _TextRedactor], + scope=(), + is_async=False, + ) + messages_transformer = next( + t for t in mux._transformers if isinstance(t, MessagesTransformer) + ) + # Unblock both the mux's main log and the messages projection log so + # synthetic pushes are accepted without a real consumer attached. + mux._events._subscribed = True + messages_transformer._log._subscribed = True + + meta = {"langgraph_node": "model", "run_id": "run-1"} + + # message-start → MessagesTransformer creates a ChatModelStream. + mux.push( + _messages_event( + [], + ({"event": "message-start", "role": "ai", "id": "msg-1"}, meta), + ) + ) + + # content-block-delta carrying the secret. The redactor (pre-lane) + # mutates `delta.text` BEFORE MessagesTransformer snapshots it. + mux.push( + _messages_event( + [], + ( + { + "event": "content-block-delta", + "index": 0, + "delta": {"type": "text-delta", "text": "secret@example.com"}, + }, + meta, + ), + ) + ) + + # Capture the still-open stream before message-finish removes it from + # MessagesTransformer's `_by_run` dict. + chat_stream = messages_transformer._by_run["run-1"] # type: ignore[attr-defined] + + # message-finish → closes the stream. + mux.push( + _messages_event( + [], + ({"event": "message-finish"}, meta), + ) + ) + + # The redactor mutated `delta.text` to "[REDACTED]" before + # MessagesTransformer snapshotted the string into the text + # accumulator. Without `before_builtins`, the accumulator would hold + # the raw "secret@example.com". + assert chat_stream._text_acc == "[REDACTED]", ( # type: ignore[attr-defined] + f"expected redacted text in projection, got {chat_stream._text_acc!r}" + )