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 ()
@@ -3277,6 +3294,7 @@ class Pregel(
 
         config = ensure_config(self.config, config)
         run_manager = None
+        checkpointer = None
         try:
             # assign defaults
             (
@@ -3609,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"]