"""Tests for v2 streaming format (StreamPart TypedDicts). This file is checked by ty directly. Type-narrowing is validated via `assert_type` calls in `_check_type_narrowing`. """ from __future__ import annotations import operator import sys from dataclasses import dataclass from typing import Annotated, Any, TypeVar import pytest from langchain_core.messages import AIMessage, BaseMessage from langgraph.checkpoint.memory import InMemorySaver from pydantic import BaseModel, ValidationError from typing_extensions import TypedDict, assert_type from langgraph._internal._constants import INTERRUPT from langgraph.constants import END, START from langgraph.errors import GraphDrained from langgraph.func import entrypoint from langgraph.graph import StateGraph from langgraph.graph.message import MessagesState from langgraph.runtime import RunControl from langgraph.types import ( CheckpointPayload, CheckpointStreamPart, CustomStreamPart, DebugPayload, DebugStreamPart, GraphOutput, Interrupt, MessagesStreamPart, StreamPart, StreamWriter, TaskPayload, TaskResultPayload, TasksStreamPart, UpdatesStreamPart, ValuesStreamPart, interrupt, ) from tests.fake_chat import FakeChatModel NEEDS_CONTEXTVARS = pytest.mark.skipif( sys.version_info < (3, 11), reason="Python 3.11+ is required for async contextvars support", ) # --- state and graph builders --- class SimpleState(TypedDict): value: str items: Annotated[list[str], operator.add] _SIMPLE_INPUT: SimpleState = {"value": "x", "items": []} _MSG_INPUT: MessagesState = {"messages": "hi"} def _make_simple_graph() -> StateGraph[SimpleState, None, SimpleState, SimpleState]: def node_a(state: SimpleState) -> dict[str, Any]: return {"value": state["value"] + "_a", "items": ["a"]} def node_b(state: SimpleState) -> dict[str, Any]: return {"value": state["value"] + "_b", "items": ["b"]} builder = StateGraph(SimpleState, input_schema=SimpleState) builder.add_node("node_a", node_a) builder.add_node("node_b", node_b) builder.add_edge(START, "node_a") builder.add_edge("node_a", "node_b") builder.add_edge("node_b", END) return builder def _make_messages_graph() -> StateGraph[ MessagesState, None, MessagesState, MessagesState ]: model = FakeChatModel(messages=[AIMessage(content="hello world")]) def call_model(state: MessagesState) -> dict[str, Any]: return {"messages": model.invoke(state["messages"])} builder = StateGraph(MessagesState, input_schema=MessagesState) builder.add_node("call_model", call_model) builder.add_edge(START, "call_model") builder.add_edge("call_model", END) return builder def _make_custom_graph() -> Any: @entrypoint() def graph(inputs: Any, *, writer: StreamWriter) -> Any: writer("hello") writer(42) return inputs return graph def _make_subgraph() -> Any: inner = _make_simple_graph().compile() outer_builder = StateGraph(SimpleState, input_schema=SimpleState) outer_builder.add_node("inner", inner) outer_builder.add_edge(START, "inner") outer_builder.add_edge("inner", END) return outer_builder.compile() # --- shared assertion helpers --- _STREAM_PART_KEYS = {"type", "ns", "data"} def _assert_stream_part_shape(part: StreamPart[Any, Any]) -> None: """Assert a v2 stream part has the required keys and correct types.""" assert isinstance(part, dict), f"Expected dict, got {type(part)}" assert _STREAM_PART_KEYS <= part.keys(), ( f"Missing keys: {_STREAM_PART_KEYS - part.keys()}" ) assert isinstance(part["type"], str) assert isinstance(part["ns"], tuple) for elem in part["ns"]: assert isinstance(elem, str) if part["type"] == "values": assert "interrupts" in part, "values stream part missing 'interrupts' field" assert isinstance(part["interrupts"], tuple) # --- v1 backwards compatibility --- class TestV1BackwardsCompat: def test_stream_default_is_v1(self) -> None: graph = _make_simple_graph().compile() chunks = list(graph.stream(_SIMPLE_INPUT)) for chunk in chunks: assert isinstance(chunk, dict) def test_stream_v1_updates_mode(self) -> None: graph = _make_simple_graph().compile() chunks = list(graph.stream(_SIMPLE_INPUT, stream_mode="updates")) assert len(chunks) == 2 assert "node_a" in chunks[0] assert "node_b" in chunks[1] def test_stream_v1_list_mode(self) -> None: graph = _make_simple_graph().compile() chunks = list(graph.stream(_SIMPLE_INPUT, stream_mode=["values", "updates"])) for chunk in chunks: assert isinstance(chunk, tuple) and len(chunk) == 2 mode, _data = chunk assert mode in ("values", "updates") def test_stream_v1_subgraphs(self) -> None: graph = _make_simple_graph().compile() chunks = list( graph.stream(_SIMPLE_INPUT, stream_mode="updates", subgraphs=True) ) for chunk in chunks: assert isinstance(chunk, tuple) and len(chunk) == 2 ns, _data = chunk assert isinstance(ns, tuple) # --- v2 sync stream --- class TestV2Stream: def test_values(self) -> None: graph = _make_simple_graph().compile() chunks = list(graph.stream(_SIMPLE_INPUT, stream_mode="values", version="v2")) assert len(chunks) >= 1 for c in chunks: _assert_stream_part_shape(c) assert c["type"] == "values" assert c["ns"] == () assert isinstance(c["data"], dict) def test_updates(self) -> None: graph = _make_simple_graph().compile() chunks = list(graph.stream(_SIMPLE_INPUT, stream_mode="updates", version="v2")) assert len(chunks) == 2 for c in chunks: _assert_stream_part_shape(c) assert c["type"] == "updates" assert c["ns"] == () assert "node_a" in chunks[0]["data"] assert "node_b" in chunks[1]["data"] def test_messages(self) -> None: graph = _make_messages_graph().compile() chunks = list(graph.stream(_MSG_INPUT, stream_mode="messages", version="v2")) msg_chunks = [c for c in chunks if c["type"] == "messages"] assert len(msg_chunks) >= 1 for c in msg_chunks: _assert_stream_part_shape(c) assert c["ns"] == () data = c["data"] assert isinstance(data, tuple) and len(data) == 2 message, metadata = data assert isinstance(message, BaseMessage) assert isinstance(metadata, dict) assert "langgraph_node" in metadata def test_custom(self) -> None: graph = _make_custom_graph() chunks = list(graph.stream({"key": "val"}, stream_mode="custom", version="v2")) custom = [c for c in chunks if c["type"] == "custom"] assert len(custom) == 2 for c in custom: _assert_stream_part_shape(c) assert custom[0]["data"] == "hello" assert custom[1]["data"] == 42 def test_multiple_modes(self) -> None: graph = _make_simple_graph().compile() chunks = list( graph.stream( _SIMPLE_INPUT, stream_mode=["values", "updates"], version="v2", ) ) types_seen = {c["type"] for c in chunks} assert {"values", "updates"} <= types_seen for c in chunks: _assert_stream_part_shape(c) def test_stream_events_v3_accepts_control_for_drain(self) -> None: class DrainState(TypedDict, total=False): value: str skipped: str control = RunControl() def first_node(state: DrainState) -> dict[str, str]: control.request_drain("sigterm") return {"value": "done"} def second_node(state: DrainState) -> dict[str, str]: return {"skipped": "nope"} builder = StateGraph(DrainState) builder.add_node("first", first_node) builder.add_node("second", second_node) builder.add_edge(START, "first") builder.add_edge("first", "second") builder.add_edge("second", END) graph = builder.compile() run = graph.stream_events({}, control=control, version="v3") with pytest.raises(GraphDrained, match="sigterm"): list(run.values) def test_subgraphs_ns(self) -> None: outer = _make_subgraph() chunks = list( outer.stream( _SIMPLE_INPUT, stream_mode="updates", subgraphs=True, version="v2", ) ) for c in chunks: _assert_stream_part_shape(c) root = [c for c in chunks if c["ns"] == ()] sub = [c for c in chunks if c["ns"] != ()] assert len(root) >= 1 assert len(sub) >= 1 def test_checkpoints(self) -> None: graph = _make_simple_graph().compile(checkpointer=InMemorySaver()) config: Any = {"configurable": {"thread_id": "test-v2-ckpt"}} chunks = list( graph.stream( _SIMPLE_INPUT, config, stream_mode="checkpoints", version="v2", ) ) ckpt = [c for c in chunks if c["type"] == "checkpoints"] assert len(ckpt) >= 1 for c in ckpt: _assert_stream_part_shape(c) assert c["ns"] == () payload = c["data"] assert {"config", "metadata", "values", "next", "tasks"} <= payload.keys() def test_tasks(self) -> None: graph = _make_simple_graph().compile(checkpointer=InMemorySaver()) config: Any = {"configurable": {"thread_id": "test-v2-tasks"}} chunks = list( graph.stream( _SIMPLE_INPUT, config, stream_mode="tasks", version="v2", ) ) tasks = [c for c in chunks if c["type"] == "tasks"] assert len(tasks) >= 2 for c in tasks: _assert_stream_part_shape(c) assert c["ns"] == () assert "id" in c["data"] and "name" in c["data"] starts = [c for c in tasks if "triggers" in c["data"]] results = [c for c in tasks if "result" in c["data"]] assert len(starts) >= 2 assert len(results) >= 2 def test_debug(self) -> None: graph = _make_simple_graph().compile(checkpointer=InMemorySaver()) config: Any = {"configurable": {"thread_id": "test-v2-debug"}} chunks = list( graph.stream( _SIMPLE_INPUT, config, stream_mode="debug", version="v2", ) ) debug = [c for c in chunks if c["type"] == "debug"] assert len(debug) >= 1 for c in debug: _assert_stream_part_shape(c) assert c["ns"] == () envelope = c["data"] assert {"step", "timestamp", "type", "payload"} <= envelope.keys() assert envelope["type"] in ("checkpoint", "task", "task_result") def test_subgraphs_param_does_not_change_format(self) -> None: """In v2, subgraphs=True/False should not change the output format.""" graph = _make_simple_graph().compile() chunks_no_sub = list( graph.stream( _SIMPLE_INPUT, stream_mode="updates", subgraphs=False, version="v2", ) ) chunks_with_sub = list( graph.stream( _SIMPLE_INPUT, stream_mode="updates", subgraphs=True, version="v2", ) ) for c in chunks_no_sub + chunks_with_sub: _assert_stream_part_shape(c) # --- v2 sync invoke --- class TestV2Invoke: def test_values_default(self) -> None: graph = _make_simple_graph().compile() result = graph.invoke(_SIMPLE_INPUT, version="v2") assert isinstance(result, GraphOutput) assert result.value == {"value": "x_a_b", "items": ["a", "b"]} assert result.interrupts == () # backward compat dict access assert result["value"] == "x_a_b" assert result["items"] == ["a", "b"] def test_invoke_v2_graph_output_with_interrupts(self) -> None: def my_node(state: SimpleState) -> dict[str, Any]: answer = interrupt("what is your name?") return {"value": answer, "items": ["done"]} builder: StateGraph = StateGraph(SimpleState) builder.add_node("my_node", my_node) builder.add_edge(START, "my_node") builder.add_edge("my_node", END) graph = builder.compile(checkpointer=InMemorySaver()) config: Any = {"configurable": {"thread_id": "test-invoke-v2-interrupts"}} result = graph.invoke({"value": "x", "items": []}, config, version="v2") assert isinstance(result, GraphOutput) assert len(result.interrupts) > 0 for intr in result.interrupts: assert isinstance(intr, Interrupt) # value should still be the state (not None or empty) assert isinstance(result.value, dict) def test_invoke_v2_graph_output_interrupt_compat(self) -> None: """result['__interrupt__'] works via __getitem__.""" def my_node(state: SimpleState) -> dict[str, Any]: answer = interrupt("what is your name?") return {"value": answer, "items": ["done"]} builder: StateGraph = StateGraph(SimpleState) builder.add_node("my_node", my_node) builder.add_edge(START, "my_node") builder.add_edge("my_node", END) graph = builder.compile(checkpointer=InMemorySaver()) config: Any = {"configurable": {"thread_id": "test-invoke-v2-compat"}} result = graph.invoke({"value": "x", "items": []}, config, version="v2") assert isinstance(result, GraphOutput) assert INTERRUPT in result assert result[INTERRUPT] == result.interrupts assert len(result[INTERRUPT]) > 0 def test_invoke_v2_graph_output_no_interrupts(self) -> None: graph = _make_simple_graph().compile() result = graph.invoke(_SIMPLE_INPUT, version="v2") assert isinstance(result, GraphOutput) assert result.interrupts == () assert INTERRUPT not in result def test_invoke_v2_pydantic_state(self) -> None: """invoke with v2 and pydantic state returns GraphOutput with pydantic value.""" def node_a(state: PydanticState) -> dict[str, Any]: return {"value": state.value + "_a", "items": ["a"]} builder: StateGraph = StateGraph(PydanticState) builder.add_node("node_a", node_a) builder.add_edge(START, "node_a") builder.add_edge("node_a", END) graph = builder.compile() result = graph.invoke({"value": "x", "items": []}, version="v2") assert isinstance(result, GraphOutput) assert isinstance(result.value, PydanticState) assert result.value.value == "x_a" assert result.interrupts == () def test_invoke_v2_dataclass_state(self) -> None: """invoke with v2 and dataclass state returns GraphOutput with dataclass value.""" def node_a(state: DataclassState) -> dict[str, Any]: return {"value": state.value + "_a", "items": ["a"]} builder: StateGraph = StateGraph(DataclassState) builder.add_node("node_a", node_a) builder.add_edge(START, "node_a") builder.add_edge("node_a", END) graph = builder.compile() result = graph.invoke({"value": "x", "items": []}, version="v2") assert isinstance(result, GraphOutput) assert isinstance(result.value, DataclassState) assert result.value.value == "x_a" assert result.value.items == ["a"] assert result.interrupts == () def test_invoke_v2_non_values_mode_pydantic(self) -> None: """invoke with v2 + non-values mode + pydantic state returns list[StreamPart].""" def node_a(state: PydanticState) -> dict[str, Any]: return {"value": state.value + "_a", "items": ["a"]} builder: StateGraph = StateGraph(PydanticState) builder.add_node("node_a", node_a) builder.add_edge(START, "node_a") builder.add_edge("node_a", END) graph = builder.compile() result = graph.invoke( {"value": "x", "items": []}, stream_mode="updates", version="v2" ) assert isinstance(result, list) for chunk in result: _assert_stream_part_shape(chunk) assert chunk["type"] == "updates" # updates data should be plain dicts, not coerced to pydantic assert isinstance(chunk["data"], dict) def test_updates_mode(self) -> None: graph = _make_simple_graph().compile() result = graph.invoke(_SIMPLE_INPUT, stream_mode="updates", version="v2") assert isinstance(result, list) and len(result) == 2 for chunk in result: _assert_stream_part_shape(chunk) assert chunk["type"] == "updates" assert chunk["ns"] == () assert "node_a" in result[0]["data"] assert "node_b" in result[1]["data"] def test_multiple_modes(self) -> None: graph = _make_simple_graph().compile() modes: Any = ["values", "updates"] result = graph.invoke(_SIMPLE_INPUT, stream_mode=modes, version="v2") assert isinstance(result, list) types_seen = {c["type"] for c in result} assert {"values", "updates"} <= types_seen for c in result: _assert_stream_part_shape(c) def test_v1_default_unchanged(self) -> None: graph = _make_simple_graph().compile() result = graph.invoke(_SIMPLE_INPUT) assert isinstance(result, dict) assert result["value"] == "x_a_b" assert result["items"] == ["a", "b"] def test_v1_updates_unchanged(self) -> None: graph = _make_simple_graph().compile() result = graph.invoke(_SIMPLE_INPUT, stream_mode="updates") assert isinstance(result, list) for chunk in result: assert "node_a" in chunk or "node_b" in chunk # --- v2 async stream --- class TestV2StreamAsync: @pytest.mark.anyio async def test_values(self) -> None: graph = _make_simple_graph().compile() chunks = [ c async for c in graph.astream( _SIMPLE_INPUT, stream_mode="values", version="v2" ) ] assert len(chunks) >= 1 for c in chunks: _assert_stream_part_shape(c) assert c["type"] == "values" assert c["ns"] == () assert isinstance(c["data"], dict) @pytest.mark.anyio async def test_updates(self) -> None: graph = _make_simple_graph().compile() chunks = [ c async for c in graph.astream( _SIMPLE_INPUT, stream_mode="updates", version="v2" ) ] assert len(chunks) == 2 for c in chunks: _assert_stream_part_shape(c) assert c["type"] == "updates" assert c["ns"] == () assert "node_a" in chunks[0]["data"] assert "node_b" in chunks[1]["data"] @pytest.mark.anyio async def test_messages(self) -> None: graph = _make_messages_graph().compile() chunks = [ c async for c in graph.astream( _MSG_INPUT, stream_mode="messages", version="v2" ) ] msg_chunks = [c for c in chunks if c["type"] == "messages"] assert len(msg_chunks) >= 1 for c in msg_chunks: _assert_stream_part_shape(c) assert c["ns"] == () data = c["data"] assert isinstance(data, tuple) and len(data) == 2 message, metadata = data assert isinstance(message, BaseMessage) assert isinstance(metadata, dict) assert "langgraph_node" in metadata @NEEDS_CONTEXTVARS @pytest.mark.anyio async def test_custom(self) -> None: graph = _make_custom_graph() chunks = [ c async for c in graph.astream( {"key": "val"}, stream_mode="custom", version="v2" ) ] custom = [c for c in chunks if c["type"] == "custom"] assert len(custom) == 2 for c in custom: _assert_stream_part_shape(c) assert custom[0]["data"] == "hello" assert custom[1]["data"] == 42 @pytest.mark.anyio async def test_multiple_modes(self) -> None: graph = _make_simple_graph().compile() chunks = [ c async for c in graph.astream( _SIMPLE_INPUT, stream_mode=["values", "updates"], version="v2", ) ] types_seen = {c["type"] for c in chunks} assert {"values", "updates"} <= types_seen for c in chunks: _assert_stream_part_shape(c) @pytest.mark.anyio async def test_subgraphs_ns(self) -> None: outer = _make_subgraph() chunks = [ c async for c in outer.astream( _SIMPLE_INPUT, stream_mode="updates", subgraphs=True, version="v2", ) ] for c in chunks: _assert_stream_part_shape(c) root = [c for c in chunks if c["ns"] == ()] sub = [c for c in chunks if c["ns"] != ()] assert len(root) >= 1 assert len(sub) >= 1 @pytest.mark.anyio async def test_checkpoints(self) -> None: graph = _make_simple_graph().compile(checkpointer=InMemorySaver()) config: Any = {"configurable": {"thread_id": "test-v2-ckpt-async"}} chunks = [ c async for c in graph.astream( _SIMPLE_INPUT, config, stream_mode="checkpoints", version="v2", ) ] ckpt = [c for c in chunks if c["type"] == "checkpoints"] assert len(ckpt) >= 1 for c in ckpt: _assert_stream_part_shape(c) assert c["ns"] == () payload = c["data"] assert {"config", "metadata", "values", "next", "tasks"} <= payload.keys() @pytest.mark.anyio async def test_tasks(self) -> None: graph = _make_simple_graph().compile(checkpointer=InMemorySaver()) config: Any = {"configurable": {"thread_id": "test-v2-tasks-async"}} chunks = [ c async for c in graph.astream( _SIMPLE_INPUT, config, stream_mode="tasks", version="v2", ) ] tasks = [c for c in chunks if c["type"] == "tasks"] assert len(tasks) >= 2 for c in tasks: _assert_stream_part_shape(c) assert c["ns"] == () assert "id" in c["data"] and "name" in c["data"] starts = [c for c in tasks if "triggers" in c["data"]] results = [c for c in tasks if "result" in c["data"]] assert len(starts) >= 2 assert len(results) >= 2 @pytest.mark.anyio async def test_debug(self) -> None: graph = _make_simple_graph().compile(checkpointer=InMemorySaver()) config: Any = {"configurable": {"thread_id": "test-v2-debug-async"}} chunks = [ c async for c in graph.astream( _SIMPLE_INPUT, config, stream_mode="debug", version="v2", ) ] debug = [c for c in chunks if c["type"] == "debug"] assert len(debug) >= 1 for c in debug: _assert_stream_part_shape(c) assert c["ns"] == () envelope = c["data"] assert {"step", "timestamp", "type", "payload"} <= envelope.keys() assert envelope["type"] in ("checkpoint", "task", "task_result") # --- v2 async invoke --- class TestV2InvokeAsync: @pytest.mark.anyio async def test_values_default(self) -> None: graph = _make_simple_graph().compile() result = await graph.ainvoke(_SIMPLE_INPUT, version="v2") assert isinstance(result, GraphOutput) assert result.value == {"value": "x_a_b", "items": ["a", "b"]} assert result.interrupts == () # backward compat dict access assert result["value"] == "x_a_b" assert result["items"] == ["a", "b"] @NEEDS_CONTEXTVARS @pytest.mark.anyio async def test_ainvoke_v2_graph_output_with_interrupts(self) -> None: def my_node(state: SimpleState) -> dict[str, Any]: answer = interrupt("what is your name?") return {"value": answer, "items": ["done"]} builder: StateGraph = StateGraph(SimpleState) builder.add_node("my_node", my_node) builder.add_edge(START, "my_node") builder.add_edge("my_node", END) graph = builder.compile(checkpointer=InMemorySaver()) config: Any = {"configurable": {"thread_id": "test-ainvoke-v2-interrupts"}} result = await graph.ainvoke({"value": "x", "items": []}, config, version="v2") assert isinstance(result, GraphOutput) assert len(result.interrupts) > 0 for intr in result.interrupts: assert isinstance(intr, Interrupt) assert isinstance(result.value, dict) @pytest.mark.anyio async def test_ainvoke_v2_pydantic_state(self) -> None: """ainvoke with v2 and pydantic state returns GraphOutput with pydantic value.""" def node_a(state: PydanticState) -> dict[str, Any]: return {"value": state.value + "_a", "items": ["a"]} builder: StateGraph = StateGraph(PydanticState) builder.add_node("node_a", node_a) builder.add_edge(START, "node_a") builder.add_edge("node_a", END) graph = builder.compile() result = await graph.ainvoke({"value": "x", "items": []}, version="v2") assert isinstance(result, GraphOutput) assert isinstance(result.value, PydanticState) assert result.value.value == "x_a" assert result.value.items == ["a"] assert result.interrupts == () @pytest.mark.anyio async def test_ainvoke_v2_dataclass_state(self) -> None: """ainvoke with v2 and dataclass state returns GraphOutput with dataclass value.""" def node_a(state: DataclassState) -> dict[str, Any]: return {"value": state.value + "_a", "items": ["a"]} builder: StateGraph = StateGraph(DataclassState) builder.add_node("node_a", node_a) builder.add_edge(START, "node_a") builder.add_edge("node_a", END) graph = builder.compile() result = await graph.ainvoke({"value": "x", "items": []}, version="v2") assert isinstance(result, GraphOutput) assert isinstance(result.value, DataclassState) assert result.value.value == "x_a" assert result.value.items == ["a"] assert result.interrupts == () @pytest.mark.anyio async def test_ainvoke_v2_graph_output_no_interrupts(self) -> None: graph = _make_simple_graph().compile() result = await graph.ainvoke(_SIMPLE_INPUT, version="v2") assert isinstance(result, GraphOutput) assert result.interrupts == () assert INTERRUPT not in result @pytest.mark.anyio async def test_updates_mode(self) -> None: graph = _make_simple_graph().compile() result = await graph.ainvoke(_SIMPLE_INPUT, stream_mode="updates", version="v2") assert isinstance(result, list) and len(result) == 2 for chunk in result: _assert_stream_part_shape(chunk) assert chunk["type"] == "updates" assert chunk["ns"] == () assert "node_a" in result[0]["data"] assert "node_b" in result[1]["data"] @pytest.mark.anyio async def test_multiple_modes(self) -> None: graph = _make_simple_graph().compile() modes: Any = ["values", "updates"] result = await graph.ainvoke(_SIMPLE_INPUT, stream_mode=modes, version="v2") assert isinstance(result, list) types_seen = {c["type"] for c in result} assert {"values", "updates"} <= types_seen for c in result: _assert_stream_part_shape(c) # --- type-safe streaming: coercion + interrupt separation --- class PydanticState(BaseModel): value: str items: Annotated[list[str], operator.add] @dataclass class DataclassState: value: str items: Annotated[list[str], operator.add] class TestV2TypeSafeStreaming: """Test that v2 streaming coerces values to pydantic/dataclass instances and separates interrupts into a dedicated field.""" def test_values_pydantic_state(self) -> None: """v2 values + pydantic state -> data is pydantic model instance.""" def node_a(state: PydanticState) -> dict[str, Any]: return {"value": state.value + "_a", "items": ["a"]} builder: StateGraph = StateGraph(PydanticState) builder.add_node("node_a", node_a) builder.add_edge(START, "node_a") builder.add_edge("node_a", END) graph = builder.compile() chunks = list( graph.stream( {"value": "x", "items": []}, stream_mode="values", version="v2", ) ) assert len(chunks) >= 1 for c in chunks: _assert_stream_part_shape(c) assert c["type"] == "values" assert isinstance(c["data"], PydanticState), ( f"Expected PydanticState, got {type(c['data'])}" ) assert c["interrupts"] == () def test_values_dataclass_state(self) -> None: """v2 values + dataclass state -> data is dataclass instance.""" def node_a(state: DataclassState) -> dict[str, Any]: return {"value": state.value + "_a", "items": ["a"]} builder: StateGraph = StateGraph(DataclassState) builder.add_node("node_a", node_a) builder.add_edge(START, "node_a") builder.add_edge("node_a", END) graph = builder.compile() chunks = list( graph.stream( {"value": "x", "items": []}, stream_mode="values", version="v2", ) ) assert len(chunks) >= 1 for c in chunks: _assert_stream_part_shape(c) assert c["type"] == "values" assert isinstance(c["data"], DataclassState), ( f"Expected DataclassState, got {type(c['data'])}" ) def test_values_typeddict_state(self) -> None: """v2 values + TypedDict state -> data stays plain dict (no coercion).""" graph = _make_simple_graph().compile() chunks = list(graph.stream(_SIMPLE_INPUT, stream_mode="values", version="v2")) assert len(chunks) >= 1 for c in chunks: _assert_stream_part_shape(c) assert c["type"] == "values" # TypedDict state should remain a plain dict assert isinstance(c["data"], dict) assert type(c["data"]) is dict def test_values_interrupt_v2(self) -> None: """v2 values + interrupt -> interrupts in typed field, not in data.""" def my_node(state: SimpleState) -> dict[str, Any]: answer = interrupt("what is your name?") return {"value": answer, "items": ["done"]} builder: StateGraph = StateGraph(SimpleState) builder.add_node("my_node", my_node) builder.add_edge(START, "my_node") builder.add_edge("my_node", END) graph = builder.compile(checkpointer=InMemorySaver()) config: Any = {"configurable": {"thread_id": "test-v2-interrupt"}} chunks = list( graph.stream( {"value": "x", "items": []}, config, stream_mode="values", version="v2", ) ) # should have at least one values chunk with interrupts interrupt_chunks = [c for c in chunks if c.get("interrupts", ())] assert len(interrupt_chunks) >= 1, f"Expected interrupt chunks, got {chunks}" for c in interrupt_chunks: assert c["type"] == "values" assert isinstance(c["interrupts"], tuple) assert len(c["interrupts"]) > 0 for intr in c["interrupts"]: assert isinstance(intr, Interrupt) # __interrupt__ should NOT be in data if isinstance(c["data"], dict): assert INTERRUPT not in c["data"] def test_values_interrupt_v1_compat(self) -> None: """v1 values + interrupt -> __interrupt__ still in dict (v1 compat).""" def my_node(state: SimpleState) -> dict[str, Any]: answer = interrupt("what is your name?") return {"value": answer, "items": ["done"]} builder: StateGraph = StateGraph(SimpleState) builder.add_node("my_node", my_node) builder.add_edge(START, "my_node") builder.add_edge("my_node", END) graph = builder.compile(checkpointer=InMemorySaver()) config: Any = {"configurable": {"thread_id": "test-v1-interrupt-compat"}} chunks = list( graph.stream( {"value": "x", "items": []}, config, stream_mode="values", ) ) # v1 format: should have __interrupt__ in dict interrupt_chunks = [c for c in chunks if isinstance(c, dict) and INTERRUPT in c] assert len(interrupt_chunks) >= 1, ( f"Expected v1 interrupt chunks with {INTERRUPT}, got {chunks}" ) def test_checkpoints_pydantic_state(self) -> None: """v2 checkpoints + pydantic state -> values is pydantic model instance (at least for checkpoints emitted after all channels are populated).""" def node_a(state: PydanticState) -> dict[str, Any]: return {"value": state.value + "_a", "items": ["a"]} builder: StateGraph = StateGraph(PydanticState) builder.add_node("node_a", node_a) builder.add_edge(START, "node_a") builder.add_edge("node_a", END) graph = builder.compile(checkpointer=InMemorySaver()) config: Any = {"configurable": {"thread_id": "test-v2-ckpt-pydantic"}} chunks = list( graph.stream( {"value": "x", "items": []}, config, stream_mode="checkpoints", version="v2", ) ) ckpt_chunks = [c for c in chunks if c["type"] == "checkpoints"] assert len(ckpt_chunks) >= 1 # At least one checkpoint (after first node runs) should have coerced values coerced_ckpts = [ c for c in ckpt_chunks if isinstance(c["data"]["values"], PydanticState) ] assert len(coerced_ckpts) >= 1, ( f"Expected at least one checkpoint with PydanticState values, got types: " f"{[type(c['data']['values']) for c in ckpt_chunks]}" ) def test_debug_pydantic_state(self) -> None: """v2 debug + pydantic state -> inner checkpoint payload has coerced values (at least for checkpoints emitted after all channels are populated).""" def node_a(state: PydanticState) -> dict[str, Any]: return {"value": state.value + "_a", "items": ["a"]} builder: StateGraph = StateGraph(PydanticState) builder.add_node("node_a", node_a) builder.add_edge(START, "node_a") builder.add_edge("node_a", END) graph = builder.compile(checkpointer=InMemorySaver()) config: Any = {"configurable": {"thread_id": "test-v2-debug-pydantic"}} chunks = list( graph.stream( {"value": "x", "items": []}, config, stream_mode="debug", version="v2", ) ) debug_chunks = [c for c in chunks if c["type"] == "debug"] checkpoint_debug = [ c for c in debug_chunks if c["data"]["type"] == "checkpoint" ] assert len(checkpoint_debug) >= 1 # At least one debug checkpoint should have coerced values coerced_debug = [ c for c in checkpoint_debug if isinstance(c["data"]["payload"]["values"], PydanticState) ] assert len(coerced_debug) >= 1, ( f"Expected at least one debug checkpoint with PydanticState values, got types: " f"{[type(c['data']['payload']['values']) for c in checkpoint_debug]}" ) def test_values_pydantic_interrupt(self) -> None: """v2 values + pydantic state + interrupt -> data is model, interrupts separated.""" def my_node(state: PydanticState) -> dict[str, Any]: answer = interrupt("what is your name?") return {"value": answer, "items": ["done"]} builder: StateGraph = StateGraph(PydanticState) builder.add_node("my_node", my_node) builder.add_edge(START, "my_node") builder.add_edge("my_node", END) graph = builder.compile(checkpointer=InMemorySaver()) config: Any = {"configurable": {"thread_id": "test-v2-pydantic-interrupt"}} chunks = list( graph.stream( {"value": "x", "items": []}, config, stream_mode="values", version="v2", ) ) interrupt_chunks = [c for c in chunks if c.get("interrupts", ())] assert len(interrupt_chunks) >= 1 for c in interrupt_chunks: assert isinstance(c["data"], PydanticState), ( f"Expected PydanticState, got {type(c['data'])}" ) assert isinstance(c["interrupts"], tuple) assert len(c["interrupts"]) > 0 def test_subgraph_different_pydantic_schema(self) -> None: """Subgraph with different pydantic schema -> subgraph data coerced with subgraph's schema.""" class InnerState(BaseModel): value: str class OuterState(BaseModel): value: str def inner_node(state: InnerState) -> dict[str, Any]: return {"value": state.value + "_inner"} def outer_node(state: OuterState) -> dict[str, Any]: return {"value": state.value + "_outer"} inner_builder: StateGraph = StateGraph(InnerState) inner_builder.add_node("inner_node", inner_node) inner_builder.add_edge(START, "inner_node") inner_builder.add_edge("inner_node", END) inner_graph = inner_builder.compile() outer_builder: StateGraph = StateGraph(OuterState) outer_builder.add_node("outer_node", outer_node) outer_builder.add_node("inner", inner_graph) outer_builder.add_edge(START, "outer_node") outer_builder.add_edge("outer_node", "inner") outer_builder.add_edge("inner", END) outer = outer_builder.compile() chunks = list( outer.stream( {"value": "x"}, stream_mode="values", subgraphs=True, version="v2", ) ) # Root-level values should be OuterState instances root_values = [c for c in chunks if c["type"] == "values" and c["ns"] == ()] assert len(root_values) >= 1 for c in root_values: assert isinstance(c["data"], OuterState), ( f"Expected OuterState, got {type(c['data'])}" ) # Subgraph values are streamed from the subgraph's own stream() # which runs with default version="v1", so no coercion sub_values = [c for c in chunks if c["type"] == "values" and c["ns"] != ()] assert len(sub_values) >= 1 # --- v2 validation errors --- def _make_pydantic_graph() -> Any: """Build a simple graph with PydanticState for validation error tests.""" def node_a(state: PydanticState) -> dict[str, Any]: return {"value": state.value + "_a", "items": ["a"]} builder: StateGraph = StateGraph(PydanticState) builder.add_node("node_a", node_a) builder.add_edge(START, "node_a") builder.add_edge("node_a", END) return builder.compile() class TestV2ValidationErrors: """Validation errors propagate for pydantic state in both v1 and v2. Uses `value=[1, 2, 3]` which channels accept (LastValue stores anything) but pydantic rejects (list is not coercible to str even in lax mode). """ _INVALID_INPUT: dict[str, Any] = {"value": [1, 2, 3], "items": []} def test_stream_events_v3_pydantic_validation_error(self) -> None: """Invalid input to stream with v2 + pydantic state raises ValidationError.""" graph = _make_pydantic_graph() with pytest.raises(ValidationError): list( graph.stream( self._INVALID_INPUT, stream_mode="values", version="v2", ) ) def test_invoke_v2_pydantic_validation_error(self) -> None: """Invalid input to invoke with v2 + pydantic state raises ValidationError.""" graph = _make_pydantic_graph() with pytest.raises(ValidationError): graph.invoke(self._INVALID_INPUT, version="v2") def test_invoke_v1_pydantic_validation_error(self) -> None: """Regression: invalid input to invoke without version raises ValidationError.""" graph = _make_pydantic_graph() with pytest.raises(ValidationError): graph.invoke(self._INVALID_INPUT) # --- type narrowing compile-time checks --- # These assert_type calls verify that mypy narrows the union correctly. _OutputT = TypeVar("_OutputT") _StateT = TypeVar("_StateT") def _check_type_narrowing(part: StreamPart[_StateT, _OutputT]) -> None: """Compile-time type narrowing checks — never called at runtime.""" if part["type"] == "values": assert_type(part, ValuesStreamPart[_OutputT]) elif part["type"] == "updates": assert_type(part, UpdatesStreamPart) assert_type(part["data"], dict[str, Any]) elif part["type"] == "messages": assert_type(part, MessagesStreamPart) elif part["type"] == "custom": assert_type(part, CustomStreamPart) elif part["type"] == "checkpoints": assert_type(part, CheckpointStreamPart[_StateT]) assert_type(part["data"], CheckpointPayload[_StateT]) elif part["type"] == "tasks": assert_type(part, TasksStreamPart) assert_type(part["data"], TaskPayload | TaskResultPayload) elif part["type"] == "debug": assert_type(part, DebugStreamPart[_StateT]) assert_type(part["data"], DebugPayload[_StateT]) assert_type(part["ns"], tuple[str, ...])