feat(langgraph): run the cancel handler only for RunCancelled cancels

Plain cancels, like asyncio timeouts, task groups stopping siblings or a
process shutting down, skip the handler so the run can be resumed as is.

libs/langgraph/langgraph/errors.py

@@ -26,6 +26,7 @@ __all__ = (
     "NodeInterrupt",
     "NodeTimeoutError",
     "ParentCommand",
+    "RunCancelled",
     "EmptyInputError",
     "TaskNotFound",
     "UnknownInterruptError",
@@ -207,6 +208,25 @@ class NodeError:
     """Exception raised by the failed node."""
 
 
+class RunCancelled(Exception):
+    """Cancel message that asks the graph to run its cancel handler.
+
+    Pass an instance to `task.cancel()` on the task running the graph. A cancel
+    without one, such as `asyncio.timeout`, a task group stopping siblings, or a
+    process shutting down, skips the handler, so the run can be resumed as is.
+
+    Example:
+        ```python
+        run = asyncio.create_task(graph.ainvoke(inputs, config))
+        run.cancel(RunCancelled("user stop"))
+        ```
+    """
+
+    run_cancel_handler: bool = True
+    """Set to `False` in a subclass whose cancels should skip the handler, such as
+    a rollback that deletes the run."""
+
+
 class NodeCancelledError(Exception):
     """Raised when a node body raises ``asyncio.CancelledError`` itself.
 

libs/langgraph/langgraph/graph/state.py

@@ -77,6 +77,7 @@ from langgraph.pregel._write import (
     ChannelWriteEntry,
     ChannelWriteTupleEntry,
 )
+from langgraph.runtime import Runtime
 from langgraph.types import (
     All,
     CachePolicy,
@@ -264,6 +265,8 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
         self.context_schema = context_schema
 
         self._node_defaults: _NodeDefaults = _NodeDefaults()
+        self._cancel_handler: Callable[[Any, Runtime[ContextT]], Any] | None = None
+        self._cancel_timeout = 5.0
 
         self._add_schema(self.state_schema)
         self._add_schema(self.input_schema, allow_managed=False)
@@ -334,6 +337,41 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
             defaults.timeout = coerce_timeout_policy(timeout)
         return self
 
+    def set_cancel_handler(
+        self,
+        handler: Callable[[StateT, Runtime[ContextT]], Any],
+        *,
+        timeout: float = 5.0,
+    ) -> Self:
+        """Set a handler that runs once when an async run of this graph is cancelled.
+
+        A cancelled node never reaches its `error_handler`, so state can be left
+        mid-step, for example an AI message whose tool call never got a result.
+        After the cancelled run finishes saving, the handler gets the latest state
+        and the run's `Runtime`. A returned update is saved with `update_state`,
+        then the cancellation propagates.
+
+        The handler runs only when the cancel message is a `RunCancelled`, as in
+        `task.cancel(RunCancelled("user stop"))`. Plain cancels, like `asyncio.timeout`
+        or a process shutting down, skip it so the run can be resumed as is.
+
+        Only top-level async runs with a checkpointer call the handler. It gets
+        `timeout` seconds in total; if it runs over or raises, the error is logged
+        and the update is skipped.
+
+        Args:
+            handler: Sync or async callable taking `(state, runtime)` and returning
+                a state update or `None`.
+            timeout: Seconds the handler may take, including the wait for the
+                cancelled run to finish saving.
+
+        Returns:
+            Self: The builder instance, for chaining.
+        """
+        self._cancel_handler = handler
+        self._cancel_timeout = timeout
+        return self
+
     @property
     def _all_edges(self) -> set[tuple[str, str]]:
         return self.edges | {
@@ -1365,6 +1403,8 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
             store=store,
             cache=cache,
             node_error_handler_map=node_error_handler_map,
+            cancel_handler=self._cancel_handler,
+            cancel_timeout=self._cancel_timeout,
             name=name or "LangGraph",
             stream_transformers=transformers,
         )

libs/langgraph/langgraph/pregel/main.py

@@ -4,6 +4,7 @@ import asyncio
 import concurrent
 import concurrent.futures
 import contextlib
+import logging
 import queue
 import warnings
 import weakref
@@ -19,7 +20,7 @@ from collections.abc import (
 from dataclasses import is_dataclass, replace
 from datetime import timedelta
 from functools import partial
-from inspect import isclass
+from inspect import isawaitable, isclass
 from typing import (
     Any,
     Generic,
@@ -118,6 +119,7 @@ from langgraph.errors import (
     GraphDrained,
     GraphRecursionError,
     InvalidUpdateError,
+    RunCancelled,
     create_error_message,
 )
 from langgraph.managed.base import ManagedValueSpec
@@ -206,6 +208,8 @@ from langgraph.warnings import LangGraphDeprecatedSinceV10
 
 __all__ = ("NodeBuilder", "Pregel")
 
+logger = logging.getLogger(__name__)
+
 _WriteValue = Callable[[Input], Output] | Any
 
 
@@ -761,6 +765,15 @@ class Pregel(
     trigger_to_nodes: Mapping[str, Sequence[str]]
     node_error_handler_map: Mapping[str, str]
 
+    cancel_handler: Callable[[Any, Runtime], Any] | None = None
+    """Called with the latest state and the run's `Runtime` when a top-level async
+    run with a checkpointer is cancelled with a `RunCancelled` message. A non-`None`
+    return is saved as a state update before the cancellation propagates."""
+
+    cancel_timeout: float = 5.0
+    """Seconds the cancel handler may take, including the wait for the cancelled
+    run to finish saving. A handler that runs over or raises is logged and skipped."""
+
     def __init__(
         self,
         *,
@@ -785,6 +798,8 @@ class Pregel(
         config: RunnableConfig | None = None,
         trigger_to_nodes: Mapping[str, Sequence[str]] | None = None,
         node_error_handler_map: Mapping[str, str] | None = None,
+        cancel_handler: Callable[[Any, Runtime], Any] | None = None,
+        cancel_timeout: float = 5.0,
         name: str = "LangGraph",
         stream_transformers: Sequence[Callable[[tuple[str, ...]], Any]] | None = None,
         **deprecated_kwargs: Unpack[DeprecatedKwargs],
@@ -833,6 +848,8 @@ class Pregel(
         self.config = config
         self.trigger_to_nodes = trigger_to_nodes or {}
         self.node_error_handler_map = node_error_handler_map or {}
+        self.cancel_handler = cancel_handler
+        self.cancel_timeout = cancel_timeout
         self.name = name
         self.stream_transformers: tuple[Callable[[tuple[str, ...]], Any], ...] = tuple(
             stream_transformers or ()
@@ -884,20 +901,38 @@ class Pregel(
         )
 
     def _subgraph_for_namespace(
-        self, config: RunnableConfig, checkpointer: BaseCheckpointSaver
+        self,
+        config: RunnableConfig,
+        checkpointer: BaseCheckpointSaver,
+        *,
+        for_update: bool = False,
     ) -> tuple[PregelProtocol, RunnableConfig] | None:
         """The subgraph a state method's namespaced config addresses, and the
         config to call it with, lending it `checkpointer`. `None` when the
-        config is for this graph."""
+        config is for this graph, or names no subgraph this graph declares or
+        detects (such as one a tool invokes), so a read uses this graph's
+        channels.
+
+        Raises:
+            ValueError: If `for_update` and the namespace names no such
+                subgraph: this graph's nodes and channels would write the wrong
+                state into it.
+        """
         checkpoint_ns = config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "")
         if not checkpoint_ns or CONFIG_KEY_CHECKPOINTER in config[CONF]:
             return None
         recast = recast_checkpoint_ns(checkpoint_ns)
-        for _, pregel in self.get_subgraphs(namespace=recast, recurse=True):
-            return pregel, patch_configurable(
-                config, {CONFIG_KEY_CHECKPOINTER: checkpointer}
-            )
-        raise ValueError(f"Subgraph {recast} not found")
+        # Walk every subgraph rather than filter by `namespace=`: that filter
+        # matches by prefix, so for `research_review` it can yield
+        # `research|review` and skip a declared `research_review`.
+        for name, pregel in self.get_subgraphs(recurse=True):
+            if name == recast:
+                return pregel, patch_configurable(
+                    config, {CONFIG_KEY_CHECKPOINTER: checkpointer}
+                )
+        if for_update:
+            raise ValueError(f"Subgraph {recast} not found")
+        return None
 
     def _apply_checkpointer_allowlist(
         self, checkpointer: BaseCheckpointSaver | None
@@ -1641,7 +1676,9 @@ class Pregel(
         if any(len(u) == 0 for u in supersteps):
             raise ValueError("No updates provided")
 
-        if subgraph := self._subgraph_for_namespace(config, checkpointer):
+        if subgraph := self._subgraph_for_namespace(
+            config, checkpointer, for_update=True
+        ):
             pregel, subgraph_config = subgraph
             return pregel.bulk_update_state(subgraph_config, supersteps)
 
@@ -2142,7 +2179,9 @@ class Pregel(
         if any(len(u) == 0 for u in supersteps):
             raise ValueError("No updates provided")
 
-        if subgraph := self._subgraph_for_namespace(config, checkpointer):
+        if subgraph := self._subgraph_for_namespace(
+            config, checkpointer, for_update=True
+        ):
             pregel, subgraph_config = subgraph
             return await pregel.abulk_update_state(subgraph_config, supersteps)
 
@@ -3255,6 +3294,7 @@ class Pregel(
 
         config = ensure_config(self.config, config)
         run_manager = None
+        checkpointer = None
         try:
             # assign defaults
             (
@@ -3587,8 +3627,68 @@ class Pregel(
         except BaseException as e:
             if run_manager is not None:
                 await asyncio.shield(run_manager.on_chain_error(e))
+            if isinstance(e, asyncio.CancelledError):
+                await self._arun_cancel_handler(e, config, checkpointer)
             raise
 
+    async def _arun_cancel_handler(
+        self,
+        exc: asyncio.CancelledError,
+        config: RunnableConfig,
+        checkpointer: BaseCheckpointSaver | None,
+    ) -> None:
+        """Give the graph one bounded state update after its run is cancelled."""
+        handler = self.cancel_handler
+        reason = exc.args[0] if exc.args else None
+        # A nested graph's parent owns the cancel; without a checkpointer nothing was saved.
+        if (
+            handler is None
+            or checkpointer is None
+            or CONFIG_KEY_TASK_ID in config[CONF]
+            or not isinstance(reason, RunCancelled)
+            or not reason.run_cancel_handler
+        ):
+            return
+        runtime = config[CONF].get(CONFIG_KEY_RUNTIME, DEFAULT_RUNTIME)
+        latest = {
+            **config,
+            CONF: {
+                k: v for k, v in config[CONF].items() if k != CONFIG_KEY_CHECKPOINT_ID
+            },
+        }
+
+        async def run() -> None:
+            # The cancel can cut short the loop's exit, which saves the last checkpoint.
+            for arg in exc.args:
+                if isinstance(arg, asyncio.Task):
+                    with contextlib.suppress(BaseException):
+                        await arg
+            state = await self.aget_state(latest)
+            update = handler(state.values, runtime)
+            if isawaitable(update):
+                update = await update
+            if update is not None:
+                # Save the update as the cancelled task's result.
+                unfinished = [t for t in state.tasks if t.result is None]
+                if len({t.name for t in unfinished}) == 1:
+                    await self.aupdate_state(
+                        state.config, update, unfinished[0].name, unfinished[0].id
+                    )
+                else:
+                    await self.aupdate_state(state.config, update)
+
+        task = asyncio.create_task(run())
+        try:
+            done, _ = await asyncio.wait({task}, timeout=self.cancel_timeout)
+        except asyncio.CancelledError:
+            task.cancel()
+            raise
+        if not done:
+            task.cancel()
+            logger.warning("Cancel handler timed out after %ss", self.cancel_timeout)
+        elif (err := task.exception()) is not None:
+            logger.error("Cancel handler failed", exc_info=err)
+
     @beta(message="The v3 streaming protocol on Pregel is experimental.")
     def _pregel_stream_v3(
         self,

libs/langgraph/tests/test_cancel_handler.py

@@ -0,0 +1,246 @@
+import asyncio
+import logging
+import operator
+import sys
+from typing import Annotated, Any
+
+import pytest
+from langgraph.checkpoint.memory import InMemorySaver
+from typing_extensions import TypedDict
+
+from langgraph.errors import RunCancelled
+from langgraph.graph import START, StateGraph
+from langgraph.runtime import Runtime
+
+pytestmark = [
+    pytest.mark.anyio,
+    pytest.mark.skipif(
+        sys.version_info < (3, 11), reason="cancel messages propagate from Python 3.11"
+    ),
+]
+
+
+class State(TypedDict):
+    log: Annotated[list[str], operator.add]
+
+
+async def _call(state: State) -> dict[str, Any]:
+    return {"log": ["call"]}
+
+
+async def _slow_tool(state: State) -> dict[str, Any]:
+    await asyncio.sleep(10)
+    return {"log": ["result"]}
+
+
+def _builder() -> StateGraph:
+    return (
+        StateGraph(State)
+        .add_node("call", _call)
+        .add_node("tool", _slow_tool)
+        .add_edge(START, "call")
+        .add_edge("call", "tool")
+    )
+
+
+class Rollback(RunCancelled):
+    run_cancel_handler = False
+
+
+async def _cancel_mid_tool(
+    graph: Any,
+    config: dict,
+    value: list[str],
+    durability: Any = None,
+    msg: Any = RunCancelled("user stop"),
+) -> None:
+    run = asyncio.create_task(
+        graph.ainvoke({"log": value}, config, durability=durability)
+    )
+    await asyncio.sleep(0.1)
+    run.cancel(msg)
+    with pytest.raises(asyncio.CancelledError):
+        await run
+
+
+def _appends_cancelled(state: State, runtime: Runtime) -> dict[str, Any]:
+    return {"log": ["cancelled"]}
+
+
+@pytest.mark.parametrize("msg", [None, "stop", Rollback("rollback")])
+async def test_cancel_handler_skipped_without_run_cancelled(msg: Any) -> None:
+    graph = (
+        _builder()
+        .set_cancel_handler(_appends_cancelled)
+        .compile(checkpointer=InMemorySaver())
+    )
+    config = {"configurable": {"thread_id": "1"}}
+
+    await _cancel_mid_tool(graph, config, ["h0"], msg=msg)
+
+    state = await graph.aget_state(config)
+    assert state.values["log"] == ["h0", "call"]
+
+
+async def test_cancel_handler_skipped_on_wait_for_timeout() -> None:
+    graph = (
+        _builder()
+        .set_cancel_handler(_appends_cancelled)
+        .compile(checkpointer=InMemorySaver())
+    )
+    config = {"configurable": {"thread_id": "1"}}
+
+    with pytest.raises(asyncio.TimeoutError):
+        await asyncio.wait_for(graph.ainvoke({"log": ["h0"]}, config), timeout=0.1)
+
+    state = await graph.aget_state(config)
+    assert state.values["log"] == ["h0", "call"]
+
+
+@pytest.mark.parametrize("durability", ["sync", "async", "exit"])
+async def test_cancel_handler_saves_update_after_cancel(durability: str) -> None:
+    seen: list[Any] = []
+
+    def on_cancel(state: State, runtime: Runtime) -> dict[str, Any]:
+        seen.append((list(state["log"]), runtime.context))
+        return {"log": ["cancelled"]}
+
+    graph = (
+        _builder().set_cancel_handler(on_cancel).compile(checkpointer=InMemorySaver())
+    )
+    config = {"configurable": {"thread_id": "1"}}
+
+    await _cancel_mid_tool(graph, config, ["h0"], durability)
+
+    assert seen == [(["h0", "call"], None)]
+    state = await graph.aget_state(config)
+    assert state.values["log"] == ["h0", "call", "cancelled"]
+
+
+async def test_async_cancel_handler_and_next_run() -> None:
+    async def on_cancel(state: State, runtime: Runtime) -> dict[str, Any]:
+        await asyncio.sleep(0)
+        return {"log": ["cancelled"]}
+
+    graph = (
+        _builder().set_cancel_handler(on_cancel).compile(checkpointer=InMemorySaver())
+    )
+    config = {"configurable": {"thread_id": "1"}}
+
+    await _cancel_mid_tool(graph, config, ["h0"])
+    await _cancel_mid_tool(graph, config, ["h1"])
+
+    state = await graph.aget_state(config)
+    assert state.values["log"] == [
+        "h0",
+        "call",
+        "cancelled",
+        "h1",
+        "call",
+        "cancelled",
+    ]
+
+
+async def test_cancel_handler_returning_none_saves_nothing() -> None:
+    graph = (
+        _builder()
+        .set_cancel_handler(lambda state, runtime: None)
+        .compile(checkpointer=InMemorySaver())
+    )
+    config = {"configurable": {"thread_id": "1"}}
+
+    await _cancel_mid_tool(graph, config, ["h0"])
+
+    state = await graph.aget_state(config)
+    assert state.values["log"] == ["h0", "call"]
+
+
+async def test_cancel_handler_skipped_without_checkpointer() -> None:
+    calls: list[State] = []
+    graph = _builder().set_cancel_handler(lambda s, r: calls.append(s)).compile()
+
+    await _cancel_mid_tool(graph, {}, ["h0"])
+
+    assert calls == []
+
+
+async def test_cancel_handler_timeout_is_logged_and_skipped(
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    async def on_cancel(state: State, runtime: Runtime) -> dict[str, Any]:
+        await asyncio.sleep(10)
+        return {"log": ["cancelled"]}
+
+    graph = (
+        _builder()
+        .set_cancel_handler(on_cancel, timeout=0.1)
+        .compile(checkpointer=InMemorySaver())
+    )
+    config = {"configurable": {"thread_id": "1"}}
+
+    with caplog.at_level(logging.WARNING):
+        await asyncio.wait_for(_cancel_mid_tool(graph, config, ["h0"]), timeout=2)
+
+    assert "Cancel handler timed out" in caplog.text
+    state = await graph.aget_state(config)
+    assert state.values["log"] == ["h0", "call"]
+
+
+async def test_cancel_handler_error_is_logged_and_cancel_propagates(
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    def on_cancel(state: State, runtime: Runtime) -> dict[str, Any]:
+        raise ValueError("boom")
+
+    graph = (
+        _builder().set_cancel_handler(on_cancel).compile(checkpointer=InMemorySaver())
+    )
+    config = {"configurable": {"thread_id": "1"}}
+
+    with caplog.at_level(logging.ERROR):
+        await _cancel_mid_tool(graph, config, ["h0"])
+
+    assert "Cancel handler failed" in caplog.text
+    state = await graph.aget_state(config)
+    assert state.values["log"] == ["h0", "call"]
+
+
+async def test_cancel_handler_not_called_on_success_or_error() -> None:
+    calls: list[State] = []
+
+    async def fails(state: State) -> dict[str, Any]:
+        raise ValueError("boom")
+
+    graph = (
+        StateGraph(State)
+        .add_node("call", _call)
+        .add_node("fails", fails)
+        .add_edge(START, "call")
+        .add_edge("call", "fails")
+        .set_cancel_handler(lambda s, r: calls.append(s))
+        .compile(checkpointer=InMemorySaver())
+    )
+
+    with pytest.raises(ValueError, match="boom"):
+        await graph.ainvoke({"log": []}, {"configurable": {"thread_id": "1"}})
+
+    assert calls == []
+
+
+async def test_subgraph_cancel_handler_defers_to_parent() -> None:
+    calls: list[str] = []
+
+    subgraph = (
+        _builder().set_cancel_handler(lambda s, r: calls.append("child")).compile()
+    )
+    graph = (
+        StateGraph(State)
+        .add_node("child", subgraph)
+        .add_edge(START, "child")
+        .set_cancel_handler(lambda s, r: calls.append("parent"))
+        .compile(checkpointer=InMemorySaver())
+    )
+
+    await _cancel_mid_tool(graph, {"configurable": {"thread_id": "1"}}, ["h0"])
+
+    assert calls == ["parent"]

libs/langgraph/tests/test_subgraph_persistence.py

@@ -6,16 +6,23 @@ Covers three checkpointer settings for subgraph state:
   interrupt support, but state resets each invocation. This is the common case
   when an agent is invoked from inside a tool used by another agent.
 - checkpointer=True: "stateful" — state accumulates across invocations on the same thread id
+
+Also covers reading a subgraph's namespace from the parent graph, including
+one where no subgraph is declared, such as a subagent a tool invokes.
 """
 
+from typing import Annotated, Any
 from uuid import uuid4
 
-from langchain_core.messages import AIMessage, HumanMessage
+import pytest
+from langchain_core.messages import AIMessage, AnyMessage, HumanMessage
+from langchain_core.runnables import RunnableConfig
 from langgraph.checkpoint.base import BaseCheckpointSaver
 from typing_extensions import TypedDict
 
+from langgraph.channels.delta import DeltaChannel
 from langgraph.graph import START, StateGraph
-from langgraph.graph.message import MessagesState
+from langgraph.graph.message import MessagesState, _messages_delta_reducer
 from langgraph.types import Command, Interrupt, interrupt
 from tests.any_str import AnyStr
 
@@ -685,3 +692,227 @@ def test_child_with_own_thread_id_keeps_namespace(
     parent.invoke({"result": ""}, config=parent_config)
     state2 = child.get_state(child_config)
     assert state2.values.get("count") == 2
+
+
+# -- reading a subgraph's namespace from the parent graph --
+
+
+class NotesState(MessagesState):
+    notes: str  # a key only the subgraphs have
+
+
+class DeltaMessagesState(TypedDict):
+    # Stores `messages` as deltas between snapshots, as deepagents' parent does.
+    messages: Annotated[list[AnyMessage], DeltaChannel(_messages_delta_reducer)]
+
+
+class DeltaNotesState(DeltaMessagesState):
+    notes: str
+
+
+def _notes_graph(label: str, state_schema: type = NotesState) -> Any:
+    # Typed `dict` so it fits either schema (see `_run_tool_subgraph`).
+    def work(state: dict) -> dict:
+        return {"messages": [AIMessage(content=f"{label} done")], "notes": label}
+
+    return (
+        StateGraph(state_schema)
+        .add_node("work", work)
+        .add_edge(START, "work")
+        .compile()
+    )
+
+
+def _namespace_configs(parent: Any) -> dict[str, RunnableConfig]:
+    """Run `parent` on a new thread. Returns a config for each subgraph
+    namespace it wrote, keyed by the node the namespace starts with."""
+    thread_id = str(uuid4())
+    return {
+        ns[0].split(":")[0]: {
+            "configurable": {"thread_id": thread_id, "checkpoint_ns": ns[0]}
+        }
+        for ns, _ in parent.stream(
+            {"messages": [HumanMessage(content="go")]},
+            {"configurable": {"thread_id": thread_id}},
+            subgraphs=True,
+        )
+        if ns
+    }
+
+
+def _run_tool_subgraph(
+    checkpointer: BaseCheckpointSaver,
+    *,
+    parent_schema: type = MessagesState,
+    subagent_schema: type = NotesState,
+) -> tuple[Any, RunnableConfig]:
+    """Run a parent whose `tools` node invokes a subgraph it doesn't declare,
+    as deepagents runs subagents. Returns the parent and a config for the
+    subgraph's `tools:<task_id>` namespace."""
+    # Subgraph detection doesn't follow a subscript, so `tools` declares none.
+    subagents = {"researcher": _notes_graph("researcher", subagent_schema)}
+
+    # Typed `dict`: a schema here would add its `messages` channel to the
+    # parent's, and a plain one clashes with a `DeltaChannel` one.
+    def tools(state: dict) -> dict:
+        subagents["researcher"].invoke(
+            {"messages": [HumanMessage(content="research apples")]}
+        )
+        return {"messages": [AIMessage(content="tool done")]}
+
+    parent = (
+        StateGraph(parent_schema)
+        .add_node("tools", tools)
+        .add_edge(START, "tools")
+        .compile(checkpointer=checkpointer)
+    )
+    return parent, _namespace_configs(parent)["tools"]
+
+
+def test_undeclared_subgraph_get_state(sync_checkpointer: BaseCheckpointSaver) -> None:
+    """A namespace with no declared subgraph is read with the parent graph."""
+    parent, ns_config = _run_tool_subgraph(sync_checkpointer)
+
+    state = parent.get_state(ns_config)
+
+    assert [m.text for m in state.values["messages"]] == [
+        "research apples",
+        "researcher done",
+    ]
+    assert state.config["configurable"]["checkpoint_ns"] == AnyStr("tools:")
+
+
+def test_undeclared_subgraph_get_state_history(
+    sync_checkpointer: BaseCheckpointSaver,
+) -> None:
+    """History of a namespace with no declared subgraph is read with the parent
+    graph, and only holds that namespace's checkpoints."""
+    parent, ns_config = _run_tool_subgraph(sync_checkpointer)
+
+    history = list(parent.get_state_history(ns_config))
+
+    assert [[m.text for m in s.values.get("messages", [])] for s in history] == [
+        ["research apples", "researcher done"],
+        ["research apples"],
+        [],
+    ]
+    assert {s.config["configurable"]["checkpoint_ns"] for s in history} == {
+        ns_config["configurable"]["checkpoint_ns"]
+    }
+
+
+def test_undeclared_subgraph_update_state_raises(
+    sync_checkpointer: BaseCheckpointSaver,
+) -> None:
+    """Unlike a read, an update can't fall back to the parent graph: its nodes
+    and channels would write the wrong state into the subgraph's namespace."""
+    parent, ns_config = _run_tool_subgraph(sync_checkpointer)
+
+    with pytest.raises(ValueError, match="Subgraph tools not found"):
+        parent.update_state(ns_config, {"messages": [AIMessage(content="edit")]})
+
+    assert len(list(parent.get_state_history(ns_config))) == 3
+
+
+@pytest.mark.parametrize(
+    ("parent_schema", "subagent_schema"),
+    [(DeltaMessagesState, NotesState), (MessagesState, DeltaNotesState)],
+    ids=["delta_parent", "delta_subagent"],
+)
+def test_undeclared_subgraph_read_across_messages_channel_types(
+    sync_checkpointer: BaseCheckpointSaver,
+    parent_schema: type,
+    subagent_schema: type,
+) -> None:
+    """deepagents' parent stores `messages` in a `DeltaChannel`, and its
+    subagents in a plain `add_messages` channel. Reading a subagent's namespace
+    with the parent's channel restores its messages, either way round."""
+    parent, ns_config = _run_tool_subgraph(
+        sync_checkpointer,
+        parent_schema=parent_schema,
+        subagent_schema=subagent_schema,
+    )
+
+    state = parent.get_state(ns_config)
+    history = list(parent.get_state_history(ns_config))
+
+    assert [m.text for m in state.values["messages"]] == [
+        "research apples",
+        "researcher done",
+    ]
+    assert [[m.text for m in s.values.get("messages", [])] for s in history] == [
+        ["research apples", "researcher done"],
+        ["research apples"],
+        [],
+    ]
+
+
+@pytest.mark.parametrize("review_declared", [False, True])
+def test_prefix_named_namespaces_resolve_to_their_own_graph(
+    sync_checkpointer: BaseCheckpointSaver, review_declared: bool
+) -> None:
+    """`research_review` starts with `research`, whose subgraph has a `review`
+    node with its own subgraph. Each namespace still reads with its own graph:
+    a declared subgraph's channels, or the parent's when none is declared."""
+    research = (
+        StateGraph(NotesState)
+        .add_node("review", _notes_graph("research"))
+        .add_edge(START, "review")
+        .compile()
+    )
+
+    class ReviewState(NotesState):
+        verdict: str  # a key `research|review`'s graph doesn't have
+
+    def review(state: ReviewState) -> dict:
+        return {
+            "messages": [AIMessage(content="review done")],
+            "notes": "review",
+            "verdict": "approved",
+        }
+
+    # Subgraph detection doesn't follow a subscript, so `call_review` declares
+    # none.
+    subgraphs = {
+        "review": StateGraph(ReviewState)
+        .add_node("work", review)
+        .add_edge(START, "work")
+        .compile()
+    }
+
+    def call_review(state: MessagesState) -> dict:
+        subgraphs["review"].invoke({"messages": [HumanMessage(content="check it")]})
+        return {}
+
+    parent = (
+        StateGraph(MessagesState)
+        .add_node("research", research)
+        .add_node(
+            "research_review",
+            subgraphs["review"] if review_declared else call_review,
+        )
+        .add_edge(START, "research")
+        .add_edge("research", "research_review")
+        .compile(checkpointer=sync_checkpointer)
+    )
+    configs = _namespace_configs(parent)
+
+    def values(name: str) -> dict:
+        return {
+            k: [m.text for m in v] if k == "messages" else v
+            for k, v in parent.get_state(configs[name]).values.items()
+        }
+
+    assert values("research") == {
+        "messages": ["go", "research done"],
+        "notes": "research",
+    }
+    if review_declared:
+        assert values("research_review") == {
+            "messages": ["go", "research done", "review done"],
+            "notes": "review",
+            "verdict": "approved",
+        }
+    else:
+        # The parent's channels have no `notes` or `verdict`.
+        assert values("research_review") == {"messages": ["check it", "review done"]}

libs/langgraph/tests/test_subgraph_persistence_async.py

@@ -6,18 +6,24 @@ Covers three checkpointer settings for subgraph state:
   interrupt support, but state resets each invocation. This is the common case
   when an agent is invoked from inside a tool used by another agent.
 - checkpointer=True: "stateful" — state accumulates across invocations on the same thread id
+
+Also covers reading a subgraph's namespace from the parent graph, including
+one where no subgraph is declared, such as a subagent a tool invokes.
 """
 
 import sys
+from typing import Annotated, Any
 from uuid import uuid4
 
 import pytest
-from langchain_core.messages import AIMessage, HumanMessage
+from langchain_core.messages import AIMessage, AnyMessage, HumanMessage
+from langchain_core.runnables import RunnableConfig
 from langgraph.checkpoint.base import BaseCheckpointSaver
 from typing_extensions import TypedDict
 
+from langgraph.channels.delta import DeltaChannel
 from langgraph.graph import START, StateGraph
-from langgraph.graph.message import MessagesState
+from langgraph.graph.message import MessagesState, _messages_delta_reducer
 from langgraph.types import Command, Interrupt, interrupt
 from tests.any_str import AnyStr
 
@@ -707,3 +713,239 @@ async def test_child_with_own_thread_id_keeps_namespace_async(
     await parent.ainvoke({"result": ""}, config=parent_config)
     state2 = await child.aget_state(child_config)
     assert state2.values.get("count") == 2
+
+
+# -- reading a subgraph's namespace from the parent graph --
+
+
+class NotesState(MessagesState):
+    notes: str  # a key only the subgraphs have
+
+
+class DeltaMessagesState(TypedDict):
+    # Stores `messages` as deltas between snapshots, as deepagents' parent does.
+    messages: Annotated[list[AnyMessage], DeltaChannel(_messages_delta_reducer)]
+
+
+class DeltaNotesState(DeltaMessagesState):
+    notes: str
+
+
+def _notes_graph(label: str, state_schema: type = NotesState) -> Any:
+    # Typed `dict` so it fits either schema (see `_run_tool_subgraph`).
+    def work(state: dict) -> dict:
+        return {"messages": [AIMessage(content=f"{label} done")], "notes": label}
+
+    return (
+        StateGraph(state_schema)
+        .add_node("work", work)
+        .add_edge(START, "work")
+        .compile()
+    )
+
+
+async def _namespace_configs(parent: Any) -> dict[str, RunnableConfig]:
+    """Run `parent` on a new thread. Returns a config for each subgraph
+    namespace it wrote, keyed by the node the namespace starts with."""
+    thread_id = str(uuid4())
+    return {
+        ns[0].split(":")[0]: {
+            "configurable": {"thread_id": thread_id, "checkpoint_ns": ns[0]}
+        }
+        async for ns, _ in parent.astream(
+            {"messages": [HumanMessage(content="go")]},
+            {"configurable": {"thread_id": thread_id}},
+            subgraphs=True,
+        )
+        if ns
+    }
+
+
+async def _run_tool_subgraph(
+    checkpointer: BaseCheckpointSaver,
+    *,
+    parent_schema: type = MessagesState,
+    subagent_schema: type = NotesState,
+) -> tuple[Any, RunnableConfig]:
+    """Run a parent whose `tools` node invokes a subgraph it doesn't declare,
+    as deepagents runs subagents. Returns the parent and a config for the
+    subgraph's `tools:<task_id>` namespace."""
+    # Subgraph detection doesn't follow a subscript, so `tools` declares none.
+    subagents = {"researcher": _notes_graph("researcher", subagent_schema)}
+
+    # Typed `dict`: a schema here would add its `messages` channel to the
+    # parent's, and a plain one clashes with a `DeltaChannel` one.
+    async def tools(state: dict) -> dict:
+        await subagents["researcher"].ainvoke(
+            {"messages": [HumanMessage(content="research apples")]}
+        )
+        return {"messages": [AIMessage(content="tool done")]}
+
+    parent = (
+        StateGraph(parent_schema)
+        .add_node("tools", tools)
+        .add_edge(START, "tools")
+        .compile(checkpointer=checkpointer)
+    )
+    return parent, (await _namespace_configs(parent))["tools"]
+
+
+@NEEDS_CONTEXTVARS
+async def test_undeclared_subgraph_get_state_async(
+    async_checkpointer: BaseCheckpointSaver,
+) -> None:
+    """A namespace with no declared subgraph is read with the parent graph."""
+    parent, ns_config = await _run_tool_subgraph(async_checkpointer)
+
+    state = await parent.aget_state(ns_config)
+
+    assert [m.text for m in state.values["messages"]] == [
+        "research apples",
+        "researcher done",
+    ]
+    assert state.config["configurable"]["checkpoint_ns"] == AnyStr("tools:")
+
+
+@NEEDS_CONTEXTVARS
+async def test_undeclared_subgraph_get_state_history_async(
+    async_checkpointer: BaseCheckpointSaver,
+) -> None:
+    """History of a namespace with no declared subgraph is read with the parent
+    graph, and only holds that namespace's checkpoints."""
+    parent, ns_config = await _run_tool_subgraph(async_checkpointer)
+
+    history = [s async for s in parent.aget_state_history(ns_config)]
+
+    assert [[m.text for m in s.values.get("messages", [])] for s in history] == [
+        ["research apples", "researcher done"],
+        ["research apples"],
+        [],
+    ]
+    assert {s.config["configurable"]["checkpoint_ns"] for s in history} == {
+        ns_config["configurable"]["checkpoint_ns"]
+    }
+
+
+@NEEDS_CONTEXTVARS
+async def test_undeclared_subgraph_update_state_raises_async(
+    async_checkpointer: BaseCheckpointSaver,
+) -> None:
+    """Unlike a read, an update can't fall back to the parent graph: its nodes
+    and channels would write the wrong state into the subgraph's namespace."""
+    parent, ns_config = await _run_tool_subgraph(async_checkpointer)
+
+    with pytest.raises(ValueError, match="Subgraph tools not found"):
+        await parent.aupdate_state(ns_config, {"messages": [AIMessage(content="edit")]})
+
+    assert len([s async for s in parent.aget_state_history(ns_config)]) == 3
+
+
+@NEEDS_CONTEXTVARS
+@pytest.mark.parametrize(
+    ("parent_schema", "subagent_schema"),
+    [(DeltaMessagesState, NotesState), (MessagesState, DeltaNotesState)],
+    ids=["delta_parent", "delta_subagent"],
+)
+async def test_undeclared_subgraph_read_across_messages_channel_types_async(
+    async_checkpointer: BaseCheckpointSaver,
+    parent_schema: type,
+    subagent_schema: type,
+) -> None:
+    """deepagents' parent stores `messages` in a `DeltaChannel`, and its
+    subagents in a plain `add_messages` channel. Reading a subagent's namespace
+    with the parent's channel restores its messages, either way round."""
+    parent, ns_config = await _run_tool_subgraph(
+        async_checkpointer,
+        parent_schema=parent_schema,
+        subagent_schema=subagent_schema,
+    )
+
+    state = await parent.aget_state(ns_config)
+    history = [s async for s in parent.aget_state_history(ns_config)]
+
+    assert [m.text for m in state.values["messages"]] == [
+        "research apples",
+        "researcher done",
+    ]
+    assert [[m.text for m in s.values.get("messages", [])] for s in history] == [
+        ["research apples", "researcher done"],
+        ["research apples"],
+        [],
+    ]
+
+
+@NEEDS_CONTEXTVARS
+@pytest.mark.parametrize("review_declared", [False, True])
+async def test_prefix_named_namespaces_resolve_to_their_own_graph_async(
+    async_checkpointer: BaseCheckpointSaver, review_declared: bool
+) -> None:
+    """`research_review` starts with `research`, whose subgraph has a `review`
+    node with its own subgraph. Each namespace still reads with its own graph:
+    a declared subgraph's channels, or the parent's when none is declared."""
+    research = (
+        StateGraph(NotesState)
+        .add_node("review", _notes_graph("research"))
+        .add_edge(START, "review")
+        .compile()
+    )
+
+    class ReviewState(NotesState):
+        verdict: str  # a key `research|review`'s graph doesn't have
+
+    def review(state: ReviewState) -> dict:
+        return {
+            "messages": [AIMessage(content="review done")],
+            "notes": "review",
+            "verdict": "approved",
+        }
+
+    # Subgraph detection doesn't follow a subscript, so `call_review` declares
+    # none.
+    subgraphs = {
+        "review": StateGraph(ReviewState)
+        .add_node("work", review)
+        .add_edge(START, "work")
+        .compile()
+    }
+
+    async def call_review(state: MessagesState) -> dict:
+        await subgraphs["review"].ainvoke(
+            {"messages": [HumanMessage(content="check it")]}
+        )
+        return {}
+
+    parent = (
+        StateGraph(MessagesState)
+        .add_node("research", research)
+        .add_node(
+            "research_review",
+            subgraphs["review"] if review_declared else call_review,
+        )
+        .add_edge(START, "research")
+        .add_edge("research", "research_review")
+        .compile(checkpointer=async_checkpointer)
+    )
+    configs = await _namespace_configs(parent)
+
+    async def values(name: str) -> dict:
+        state = await parent.aget_state(configs[name])
+        return {
+            k: [m.text for m in v] if k == "messages" else v
+            for k, v in state.values.items()
+        }
+
+    assert await values("research") == {
+        "messages": ["go", "research done"],
+        "notes": "research",
+    }
+    if review_declared:
+        assert await values("research_review") == {
+            "messages": ["go", "research done", "review done"],
+            "notes": "review",
+            "verdict": "approved",
+        }
+    else:
+        # The parent's channels have no `notes` or `verdict`.
+        assert await values("research_review") == {
+            "messages": ["check it", "review done"]
+        }