chore: import upstream snapshot with attribution

libs/langgraph/tests/test_stream_data_transformers.py

@@ -0,0 +1,707 @@
+"""Tests for CustomTransformer, UpdatesTransformer, CheckpointsTransformer, DebugTransformer, TasksTransformer.
+
+These transformers capture raw protocol events for their respective stream
+modes and expose them as native projections on the run stream (run.custom,
+run.updates, run.checkpoints, run.debug, run.tasks). Tests dispatch synthetic
+protocol events through a StreamMux to isolate transformer logic; the final
+group exercises real graphs through stream_events(version="v3").
+"""
+
+from __future__ import annotations
+
+import operator
+import time
+from typing import Annotated, Any
+
+from typing_extensions import TypedDict
+
+from langgraph.constants import END, START
+from langgraph.graph import StateGraph
+from langgraph.stream._mux import StreamMux
+from langgraph.stream.stream_channel import StreamChannel
+from langgraph.stream.transformers import (
+    CheckpointsTransformer,
+    CustomTransformer,
+    DebugTransformer,
+    LifecycleTransformer,
+    TasksTransformer,
+    UpdatesTransformer,
+)
+
+TS = int(time.time() * 1000)
+
+
+def _custom_event(namespace: list[str], data: Any) -> dict[str, Any]:
+    return {
+        "type": "event",
+        "method": "custom",
+        "params": {"namespace": namespace, "timestamp": TS, "data": data},
+    }
+
+
+def _checkpoints_event(namespace: list[str], data: Any) -> dict[str, Any]:
+    return {
+        "type": "event",
+        "method": "checkpoints",
+        "params": {"namespace": namespace, "timestamp": TS, "data": data},
+    }
+
+
+def _debug_event(namespace: list[str], data: Any) -> dict[str, Any]:
+    return {
+        "type": "event",
+        "method": "debug",
+        "params": {"namespace": namespace, "timestamp": TS, "data": data},
+    }
+
+
+def _tasks_event(namespace: list[str], data: Any) -> dict[str, Any]:
+    return {
+        "type": "event",
+        "method": "tasks",
+        "params": {"namespace": namespace, "timestamp": TS, "data": data},
+    }
+
+
+def _updates_event(namespace: list[str], data: Any) -> dict[str, Any]:
+    return {
+        "type": "event",
+        "method": "updates",
+        "params": {"namespace": namespace, "timestamp": TS, "data": data},
+    }
+
+
+def _arm(mux: StreamMux, transformer: Any) -> None:
+    """Force projection logs to accept pushes (skip lazy-subscribe gate)."""
+    mux._events._subscribed = True
+    transformer._log._subscribed = True
+
+
+def _unstamped(items):
+    """Strip push stamps from a StreamChannel's internal buffer."""
+    return [item for _stamp, item in items]
+
+
+def _drain(transformer: Any) -> list[Any]:
+    return _unstamped(transformer._log._items)
+
+
+# ---------------------------------------------------------------------------
+# CustomTransformer
+# ---------------------------------------------------------------------------
+
+
+def test_custom_captures_root_scope_events() -> None:
+    t = CustomTransformer()
+    mux = StreamMux([t], is_async=False)
+    _arm(mux, t)
+
+    mux.push(_custom_event([], {"status": "processing"}))
+    mux.push(_custom_event([], {"status": "done"}))
+
+    items = _drain(t)
+    assert items == [{"status": "processing"}, {"status": "done"}]
+
+
+def test_custom_ignores_subgraph_scope_events() -> None:
+    t = CustomTransformer()
+    mux = StreamMux([t], is_async=False)
+    _arm(mux, t)
+
+    mux.push(_custom_event(["subgraph:abc"], {"from": "child"}))
+
+    assert _drain(t) == []
+
+
+def test_custom_scoped_transformer_captures_own_scope() -> None:
+    t = CustomTransformer(scope=("agent:abc",))
+    mux = StreamMux([t], is_async=False)
+    _arm(mux, t)
+
+    mux.push(_custom_event([], {"from": "root"}))
+    mux.push(_custom_event(["agent:abc"], {"from": "self"}))
+    mux.push(_custom_event(["agent:abc", "deep:def"], {"from": "child"}))
+
+    items = _drain(t)
+    assert items == [{"from": "self"}]
+
+
+def test_custom_preserves_any_payload_type() -> None:
+    t = CustomTransformer()
+    mux = StreamMux([t], is_async=False)
+    _arm(mux, t)
+
+    mux.push(_custom_event([], "string_payload"))
+    mux.push(_custom_event([], 42))
+    mux.push(_custom_event([], [1, 2, 3]))
+
+    assert _drain(t) == ["string_payload", 42, [1, 2, 3]]
+
+
+def test_custom_does_not_suppress_from_main_log() -> None:
+    t = CustomTransformer()
+    mux = StreamMux([t], is_async=False)
+    _arm(mux, t)
+
+    mux.push(_custom_event([], "data"))
+
+    methods = [evt["method"] for evt in _unstamped(mux._events._items)]
+    assert "custom" in methods
+
+
+def test_custom_ignores_other_methods() -> None:
+    t = CustomTransformer()
+    mux = StreamMux([t], is_async=False)
+    _arm(mux, t)
+
+    mux.push(
+        {
+            "type": "event",
+            "method": "values",
+            "params": {"namespace": [], "timestamp": TS, "data": {}},
+        }
+    )
+    assert _drain(t) == []
+
+
+def test_custom_required_stream_modes() -> None:
+    assert CustomTransformer.required_stream_modes == ("custom",)
+
+
+def test_custom_is_native() -> None:
+    assert getattr(CustomTransformer, "_native", False) is True
+
+
+def test_custom_init_returns_correct_key() -> None:
+    t = CustomTransformer()
+    projection = t.init()
+    assert "custom" in projection
+    assert isinstance(projection["custom"], StreamChannel)
+
+
+# ---------------------------------------------------------------------------
+# CheckpointsTransformer
+# ---------------------------------------------------------------------------
+
+
+def test_checkpoints_captures_root_scope_events() -> None:
+    t = CheckpointsTransformer()
+    mux = StreamMux([t], is_async=False)
+    _arm(mux, t)
+
+    checkpoint_data = {"values": {"x": 1}, "next": ["node_b"]}
+    mux.push(_checkpoints_event([], checkpoint_data))
+
+    items = _drain(t)
+    assert items == [checkpoint_data]
+
+
+def test_checkpoints_ignores_subgraph_events() -> None:
+    t = CheckpointsTransformer()
+    mux = StreamMux([t], is_async=False)
+    _arm(mux, t)
+
+    mux.push(_checkpoints_event(["child:abc"], {"values": {"x": 1}}))
+
+    assert _drain(t) == []
+
+
+def test_checkpoints_scoped_transformer() -> None:
+    t = CheckpointsTransformer(scope=("sub:abc",))
+    mux = StreamMux([t], is_async=False)
+    _arm(mux, t)
+
+    mux.push(_checkpoints_event([], {"from": "root"}))
+    mux.push(_checkpoints_event(["sub:abc"], {"from": "self"}))
+
+    assert _drain(t) == [{"from": "self"}]
+
+
+def test_checkpoints_does_not_suppress_from_main_log() -> None:
+    t = CheckpointsTransformer()
+    mux = StreamMux([t], is_async=False)
+    _arm(mux, t)
+
+    mux.push(_checkpoints_event([], {"values": {}}))
+
+    methods = [evt["method"] for evt in _unstamped(mux._events._items)]
+    assert "checkpoints" in methods
+
+
+def test_checkpoints_required_stream_modes() -> None:
+    assert CheckpointsTransformer.required_stream_modes == ("checkpoints",)
+
+
+def test_checkpoints_is_native() -> None:
+    assert getattr(CheckpointsTransformer, "_native", False) is True
+
+
+# ---------------------------------------------------------------------------
+# DebugTransformer
+# ---------------------------------------------------------------------------
+
+
+def test_debug_captures_root_scope_events() -> None:
+    t = DebugTransformer()
+    mux = StreamMux([t], is_async=False)
+    _arm(mux, t)
+
+    debug_data = {
+        "step": 0,
+        "type": "checkpoint",
+        "timestamp": "2026-01-01T00:00:00Z",
+        "payload": {"values": {"x": 1}},
+    }
+    mux.push(_debug_event([], debug_data))
+
+    items = _drain(t)
+    assert items == [debug_data]
+
+
+def test_debug_ignores_subgraph_events() -> None:
+    t = DebugTransformer()
+    mux = StreamMux([t], is_async=False)
+    _arm(mux, t)
+
+    mux.push(_debug_event(["child:abc"], {"step": 0, "type": "task"}))
+
+    assert _drain(t) == []
+
+
+def test_debug_captures_multiple_event_types() -> None:
+    t = DebugTransformer()
+    mux = StreamMux([t], is_async=False)
+    _arm(mux, t)
+
+    mux.push(_debug_event([], {"step": 0, "type": "checkpoint", "payload": {}}))
+    mux.push(_debug_event([], {"step": 1, "type": "task", "payload": {}}))
+    mux.push(_debug_event([], {"step": 1, "type": "task_result", "payload": {}}))
+
+    items = _drain(t)
+    assert len(items) == 3
+    assert [d["type"] for d in items] == ["checkpoint", "task", "task_result"]
+
+
+def test_debug_does_not_suppress_from_main_log() -> None:
+    t = DebugTransformer()
+    mux = StreamMux([t], is_async=False)
+    _arm(mux, t)
+
+    mux.push(_debug_event([], {"step": 0}))
+
+    methods = [evt["method"] for evt in _unstamped(mux._events._items)]
+    assert "debug" in methods
+
+
+def test_debug_required_stream_modes() -> None:
+    assert DebugTransformer.required_stream_modes == ("debug",)
+
+
+def test_debug_is_native() -> None:
+    assert getattr(DebugTransformer, "_native", False) is True
+
+
+# ---------------------------------------------------------------------------
+# TasksTransformer
+# ---------------------------------------------------------------------------
+
+
+def test_tasks_captures_root_scope_events() -> None:
+    t = TasksTransformer()
+    mux = StreamMux([t], is_async=False)
+    _arm(mux, t)
+
+    task_start = {"id": "t1", "name": "my_node", "input": None, "triggers": []}
+    mux.push(_tasks_event([], task_start))
+
+    items = _drain(t)
+    assert items == [task_start]
+
+
+def test_tasks_captures_start_and_result() -> None:
+    t = TasksTransformer()
+    mux = StreamMux([t], is_async=False)
+    _arm(mux, t)
+
+    start = {"id": "t1", "name": "a", "input": None, "triggers": []}
+    result = {"id": "t1", "name": "a", "result": {"output": 42}, "error": None}
+    mux.push(_tasks_event([], start))
+    mux.push(_tasks_event([], result))
+
+    items = _drain(t)
+    assert items == [start, result]
+
+
+def test_tasks_ignores_subgraph_events() -> None:
+    t = TasksTransformer()
+    mux = StreamMux([t], is_async=False)
+    _arm(mux, t)
+
+    mux.push(_tasks_event(["child:abc"], {"id": "t1", "name": "x"}))
+
+    assert _drain(t) == []
+
+
+def test_tasks_scoped_transformer() -> None:
+    t = TasksTransformer(scope=("agent:abc",))
+    mux = StreamMux([t], is_async=False)
+    _arm(mux, t)
+
+    mux.push(_tasks_event([], {"id": "t1"}))
+    mux.push(_tasks_event(["agent:abc"], {"id": "t2"}))
+    mux.push(_tasks_event(["agent:abc", "deep:def"], {"id": "t3"}))
+
+    assert _drain(t) == [{"id": "t2"}]
+
+
+def test_tasks_does_not_suppress_from_main_log() -> None:
+    """TasksTransformer returns True — it doesn't suppress tasks events.
+
+    (LifecycleTransformer suppresses them, but that's independent.)
+    """
+    t = TasksTransformer()
+    mux = StreamMux([t], is_async=False)
+    _arm(mux, t)
+
+    mux.push(_tasks_event([], {"id": "t1"}))
+
+    methods = [evt["method"] for evt in _unstamped(mux._events._items)]
+    assert "tasks" in methods
+
+
+def test_tasks_required_stream_modes() -> None:
+    assert TasksTransformer.required_stream_modes == ("tasks",)
+
+
+def test_tasks_is_native() -> None:
+    assert getattr(TasksTransformer, "_native", False) is True
+
+
+# ---------------------------------------------------------------------------
+# UpdatesTransformer
+# ---------------------------------------------------------------------------
+
+
+def test_updates_captures_root_scope_events() -> None:
+    t = UpdatesTransformer()
+    mux = StreamMux([t], is_async=False)
+    _arm(mux, t)
+
+    update = {"my_node": {"value": "hello!"}}
+    mux.push(_updates_event([], update))
+
+    items = _drain(t)
+    assert items == [update]
+
+
+def test_updates_captures_multiple_steps() -> None:
+    t = UpdatesTransformer()
+    mux = StreamMux([t], is_async=False)
+    _arm(mux, t)
+
+    mux.push(_updates_event([], {"node_a": {"x": 1}}))
+    mux.push(_updates_event([], {"node_b": {"x": 2}}))
+
+    items = _drain(t)
+    assert items == [{"node_a": {"x": 1}}, {"node_b": {"x": 2}}]
+
+
+def test_updates_ignores_subgraph_events() -> None:
+    t = UpdatesTransformer()
+    mux = StreamMux([t], is_async=False)
+    _arm(mux, t)
+
+    mux.push(_updates_event(["child:abc"], {"inner_node": {"v": 1}}))
+
+    assert _drain(t) == []
+
+
+def test_updates_scoped_transformer() -> None:
+    t = UpdatesTransformer(scope=("agent:abc",))
+    mux = StreamMux([t], is_async=False)
+    _arm(mux, t)
+
+    mux.push(_updates_event([], {"from": "root"}))
+    mux.push(_updates_event(["agent:abc"], {"from": "self"}))
+
+    assert _drain(t) == [{"from": "self"}]
+
+
+def test_updates_does_not_suppress_from_main_log() -> None:
+    t = UpdatesTransformer()
+    mux = StreamMux([t], is_async=False)
+    _arm(mux, t)
+
+    mux.push(_updates_event([], {"n": {}}))
+
+    methods = [evt["method"] for evt in _unstamped(mux._events._items)]
+    assert "updates" in methods
+
+
+def test_updates_required_stream_modes() -> None:
+    assert UpdatesTransformer.required_stream_modes == ("updates",)
+
+
+def test_updates_is_native() -> None:
+    assert getattr(UpdatesTransformer, "_native", False) is True
+
+
+# ---------------------------------------------------------------------------
+# Cross-transformer: unrelated events pass through
+# ---------------------------------------------------------------------------
+
+
+def test_unrelated_events_ignored_by_all() -> None:
+    """Non-matching method events don't land in any transformer's log."""
+    transformers = [
+        CustomTransformer(),
+        UpdatesTransformer(),
+        CheckpointsTransformer(),
+        DebugTransformer(),
+        TasksTransformer(),
+    ]
+    mux = StreamMux(transformers, is_async=False)
+    mux._events._subscribed = True
+    for t in transformers:
+        t._log._subscribed = True
+
+    mux.push(
+        {
+            "type": "event",
+            "method": "values",
+            "params": {"namespace": [], "timestamp": TS, "data": {"x": 1}},
+        }
+    )
+
+    for t in transformers:
+        assert _unstamped(t._log._items) == []
+
+
+# ---------------------------------------------------------------------------
+# End-to-end: real graphs through stream_events(version="v3")
+# ---------------------------------------------------------------------------
+
+
+class _State(TypedDict):
+    value: str
+    items: Annotated[list[str], operator.add]
+
+
+def _my_node(state: _State) -> dict[str, Any]:
+    from langgraph.config import get_stream_writer
+
+    writer = get_stream_writer()
+    writer({"status": "working", "node": "my_node"})
+    return {"value": state["value"] + "!", "items": ["done"]}
+
+
+def _make_simple_graph() -> Any:
+    builder = StateGraph(_State, input_schema=_State)
+    builder.add_node("my_node", _my_node)
+    builder.add_edge(START, "my_node")
+    builder.add_edge("my_node", END)
+    return builder.compile()
+
+
+def test_stream_events_v3_custom_projection_opt_in() -> None:
+    """run.custom surfaces get_stream_writer() payloads when opted in."""
+    graph = _make_simple_graph()
+    run = graph.stream_events(
+        {"value": "hello", "items": []}, version="v3", transformers=[CustomTransformer]
+    )
+
+    custom_events = list(run.custom)
+    assert len(custom_events) >= 1
+    assert any(e.get("status") == "working" for e in custom_events)
+
+
+def test_stream_events_v3_custom_and_values_coexist() -> None:
+    """Both run.custom and run.values work in the same run."""
+    graph = _make_simple_graph()
+    run = graph.stream_events(
+        {"value": "hello", "items": []}, version="v3", transformers=[CustomTransformer]
+    )
+
+    custom_events = list(run.custom)
+    assert run.output is not None
+    assert run.output["value"] == "hello!"
+    assert len(custom_events) >= 1
+
+
+def test_stream_events_v3_tasks_projection_opt_in() -> None:
+    """run.tasks surfaces raw task events when opted in via transformers=."""
+    graph = _make_simple_graph()
+    run = graph.stream_events(
+        {"value": "x", "items": []}, transformers=[TasksTransformer], version="v3"
+    )
+
+    tasks_events = list(run.tasks)
+    assert len(tasks_events) >= 1
+    names = [t.get("name") for t in tasks_events if "name" in t]
+    assert "my_node" in names
+
+
+def test_stream_events_v3_debug_projection_opt_in() -> None:
+    """run.debug surfaces debug events when opted in via transformers=."""
+    graph = _make_simple_graph()
+    run = graph.stream_events(
+        {"value": "x", "items": []}, transformers=[DebugTransformer], version="v3"
+    )
+
+    debug_events = list(run.debug)
+    assert len(debug_events) >= 1
+    types = {d.get("type") for d in debug_events}
+    assert types & {"checkpoint", "task", "task_result"}
+
+
+def test_stream_events_v3_updates_projection_opt_in() -> None:
+    """run.updates surfaces node output dicts when opted in via transformers=."""
+    graph = _make_simple_graph()
+    run = graph.stream_events(
+        {"value": "x", "items": []}, version="v3", transformers=[UpdatesTransformer]
+    )
+
+    updates = list(run.updates)
+    assert len(updates) >= 1
+    node_names = {k for u in updates for k in u if k != "__interrupt__"}
+    assert "my_node" in node_names
+
+
+def test_stream_events_v3_all_transformers_interleaved() -> None:
+    """All five transformers registered together, consumed via interleave."""
+    graph = _make_simple_graph()
+    run = graph.stream_events(
+        {"value": "x", "items": []},
+        version="v3",
+        transformers=[
+            CustomTransformer,
+            UpdatesTransformer,
+            CheckpointsTransformer,
+            DebugTransformer,
+            TasksTransformer,
+        ],
+    )
+
+    collected: dict[str, list[Any]] = {
+        "custom": [],
+        "updates": [],
+        "debug": [],
+        "tasks": [],
+    }
+    for name, item in run.interleave("custom", "updates", "debug", "tasks"):
+        collected[name].append(item)
+
+    assert len(collected["custom"]) >= 1
+    assert len(collected["updates"]) >= 1
+    assert len(collected["tasks"]) >= 1
+    assert len(collected["debug"]) >= 1
+    types = {d.get("type") for d in collected["debug"]}
+    assert types & {"checkpoint", "task", "task_result"}
+    node_names = {k for u in collected["updates"] for k in u if k != "__interrupt__"}
+    assert "my_node" in node_names
+
+    assert run.output is not None
+    assert run.output["value"] == "x!"
+
+
+def test_stream_events_v3_all_transformers_with_checkpointer() -> None:
+    """All transformers with a checkpointer — run.checkpoints populated."""
+    from langgraph.checkpoint.memory import InMemorySaver
+
+    builder = StateGraph(_State, input_schema=_State)
+    builder.add_node("my_node", _my_node)
+    builder.add_edge(START, "my_node")
+    builder.add_edge("my_node", END)
+    graph = builder.compile(checkpointer=InMemorySaver())
+
+    run = graph.stream_events(
+        {"value": "x", "items": []},
+        version="v3",
+        config={"configurable": {"thread_id": "test-all"}},
+        transformers=[
+            CustomTransformer,
+            UpdatesTransformer,
+            CheckpointsTransformer,
+            DebugTransformer,
+            TasksTransformer,
+        ],
+    )
+
+    collected: dict[str, list[Any]] = {
+        "custom": [],
+        "updates": [],
+        "checkpoints": [],
+        "debug": [],
+        "tasks": [],
+    }
+    for name, item in run.interleave(
+        "custom", "updates", "checkpoints", "debug", "tasks"
+    ):
+        collected[name].append(item)
+
+    assert len(collected["checkpoints"]) >= 1
+    assert len(collected["custom"]) >= 1
+
+
+def test_stream_events_v3_checkpoints_projection_opt_in() -> None:
+    """run.checkpoints surfaces checkpoint data when opted in with a checkpointer."""
+    from langgraph.checkpoint.memory import InMemorySaver
+
+    builder = StateGraph(_State, input_schema=_State)
+    builder.add_node("my_node", _my_node)
+    builder.add_edge(START, "my_node")
+    builder.add_edge("my_node", END)
+    graph = builder.compile(checkpointer=InMemorySaver())
+
+    run = graph.stream_events(
+        {"value": "x", "items": []},
+        version="v3",
+        config={"configurable": {"thread_id": "test-ckpt-standalone"}},
+        transformers=[CheckpointsTransformer],
+    )
+
+    checkpoints = list(run.checkpoints)
+    assert len(checkpoints) >= 1
+
+
+# ---------------------------------------------------------------------------
+# TasksTransformer + LifecycleTransformer co-registration
+# ---------------------------------------------------------------------------
+
+
+def test_tasks_and_lifecycle_coregistration() -> None:
+    """When both are in the same StreamMux, LifecycleTransformer suppresses
+    tasks events from the main log (returns False) while TasksTransformer
+    still captures them into its own log.
+    """
+    lifecycle = LifecycleTransformer()
+    tasks = TasksTransformer()
+    mux = StreamMux([lifecycle, tasks], is_async=False)
+    mux._events._subscribed = True
+    tasks._log._subscribed = True
+    lifecycle._channel._subscribed = True
+
+    task_data = {"id": "t1", "name": "my_node", "input": None, "triggers": []}
+    mux.push(_tasks_event([], task_data))
+
+    assert _drain(tasks) == [task_data]
+
+    methods = [evt["method"] for evt in _unstamped(mux._events._items)]
+    assert "tasks" not in methods
+
+
+def test_tasks_and_lifecycle_coregistration_e2e() -> None:
+    """E2e: TasksTransformer captures task events even when LifecycleTransformer
+    is present and suppressing them from the main log.
+    """
+    graph = _make_simple_graph()
+    run = graph.stream_events(
+        {"value": "x", "items": []},
+        version="v3",
+        transformers=[TasksTransformer],
+    )
+
+    tasks_events = list(run.tasks)
+    assert len(tasks_events) >= 1
+    names = [t.get("name") for t in tasks_events if "name" in t]
+    assert "my_node" in names