a7d6d88f6f
CI / changes (push) Has been cancelled
CI / cd libs/checkpoint (push) Has been cancelled
CI / cd libs/checkpoint-conformance (push) Has been cancelled
CI / cd libs/checkpoint-postgres (push) Has been cancelled
CI / cd libs/checkpoint-sqlite (push) Has been cancelled
CI / cd libs/cli (push) Has been cancelled
CI / cd libs/prebuilt (push) Has been cancelled
CI / cd libs/sdk-py (push) Has been cancelled
CI / cd libs/langgraph (push) Has been cancelled
CI / Check SDK methods matching (push) Has been cancelled
CI / Check CLI schema hasn't changed #3.13 (push) Has been cancelled
CI / CLI integration test (push) Has been cancelled
CI / sdk-py integration test (push) Has been cancelled
CI / CI Success (push) Has been cancelled
baseline / benchmark (push) Has been cancelled
Deploy Redirects to GitHub Pages / deploy (push) Has been cancelled
708 lines
21 KiB
Python
708 lines
21 KiB
Python
"""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
|