From ac9c0f511de11376b2b45f0fb1a4e4c0891d2bf1 Mon Sep 17 00:00:00 2001 From: Hunter Lovell Date: Mon, 30 Mar 2026 18:30:34 -0700 Subject: [PATCH] feat(langgraph): add on_interrupt hook to StateGraph.compile() --- libs/langgraph/langgraph/graph/state.py | 7 + libs/langgraph/langgraph/pregel/_loop.py | 63 ++++++++- libs/langgraph/langgraph/pregel/main.py | 11 ++ libs/langgraph/langgraph/types.py | 15 ++- libs/langgraph/tests/test_pregel.py | 148 ++++++++++++++++++++++ libs/langgraph/tests/test_pregel_async.py | 114 +++++++++++++++++ 6 files changed, 353 insertions(+), 5 deletions(-) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index f3299de92..d622fc789 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -76,6 +76,7 @@ from langgraph.types import ( CachePolicy, Checkpointer, Command, + OnInterruptHook, RetryPolicy, Send, ensure_valid_checkpointer, @@ -831,6 +832,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): interrupt_after: All | list[str] | None = None, debug: bool = False, name: str | None = None, + on_interrupt: OnInterruptHook | None = None, ) -> CompiledStateGraph[StateT, ContextT, InputT, OutputT]: """Compiles the `StateGraph` into a `CompiledStateGraph` object. @@ -850,6 +852,10 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): interrupt_after: An optional list of node names to interrupt after. debug: A flag indicating whether to enable debug mode. name: The name to use for the compiled graph. + on_interrupt: An optional callback that is invoked whenever the graph + execution is interrupted. Called with the list of `Interrupt` objects. + + May be a sync function or an async coroutine function. Returns: CompiledStateGraph: The compiled `StateGraph`. @@ -910,6 +916,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): store=store, cache=cache, name=name or "LangGraph", + on_interrupt=on_interrupt, ) compiled.attach_node(START, None) diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index 79201d0d6..64c6e78b9 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio import binascii import concurrent.futures +import warnings from collections import defaultdict, deque from collections.abc import Callable, Iterator, Mapping, Sequence from contextlib import ( @@ -12,7 +13,7 @@ from contextlib import ( ExitStack, ) from datetime import datetime, timezone -from inspect import signature +from inspect import iscoroutinefunction, signature from types import TracebackType from typing import ( Any, @@ -115,6 +116,8 @@ from langgraph.types import ( CachePolicy, Command, Durability, + Interrupt, + OnInterruptHook, PregelExecutableTask, RetryPolicy, Send, @@ -157,6 +160,7 @@ class PregelLoop: manager: None | AsyncParentRunManager | ParentRunManager interrupt_after: All | Sequence[str] interrupt_before: All | Sequence[str] + on_interrupt: OnInterruptHook | None durability: Durability retry_policy: Sequence[RetryPolicy] cache_policy: CachePolicy | None @@ -226,6 +230,7 @@ class PregelLoop: migrate_checkpoint: Callable[[Checkpoint], None] | None = None, retry_policy: Sequence[RetryPolicy] = (), cache_policy: CachePolicy | None = None, + on_interrupt: OnInterruptHook | None = None, ) -> None: self.stream = stream self.config = config @@ -242,6 +247,7 @@ class PregelLoop: self.stream_keys = stream_keys self.interrupt_after = interrupt_after self.interrupt_before = interrupt_before + self.on_interrupt = on_interrupt self.manager = manager self.is_nested = CONFIG_KEY_TASK_ID in self.config.get(CONF, {}) self.skip_done_tasks = CONFIG_KEY_CHECKPOINT_ID not in config[CONF] @@ -865,14 +871,36 @@ class PregelLoop: [{INTERRUPT: cast(GraphInterrupt, exc_value).args[0]}] ), ) - # save final output + # save final output first, so graph state is consistent even + # if the on_interrupt hook raises self.output = read_channels(self.channels, self.output_keys) + # call on_interrupt hook + if self.on_interrupt is not None: + interrupts: list[Interrupt] = ( + list(cast(GraphInterrupt, exc_value).args[0]) + if exc_value is not None and exc_value.args and exc_value.args[0] + else [] + ) + self._call_on_interrupt(interrupts) # suppress interrupt return True elif exc_type is None: # save final output self.output = read_channels(self.channels, self.output_keys) + def _call_on_interrupt(self, interrupts: list[Interrupt]) -> None: + """Call the on_interrupt hook synchronously.""" + if self.on_interrupt is None: + return + if iscoroutinefunction(self.on_interrupt): + warnings.warn( + "Async on_interrupt hook cannot be called from sync graph execution. " + "Use a sync function or run the graph with astream/ainvoke.", + stacklevel=2, + ) + return + self.on_interrupt(interrupts) + def _emit( self, mode: StreamMode, @@ -985,6 +1013,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): migrate_checkpoint: Callable[[Checkpoint], None] | None = None, retry_policy: Sequence[RetryPolicy] = (), cache_policy: CachePolicy | None = None, + on_interrupt: OnInterruptHook | None = None, ) -> None: super().__init__( input, @@ -1006,6 +1035,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): retry_policy=retry_policy, cache_policy=cache_policy, durability=durability, + on_interrupt=on_interrupt, ) self.stack = ExitStack() if checkpointer: @@ -1161,6 +1191,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): migrate_checkpoint: Callable[[Checkpoint], None] | None = None, retry_policy: Sequence[RetryPolicy] = (), cache_policy: CachePolicy | None = None, + on_interrupt: OnInterruptHook | None = None, ) -> None: super().__init__( input, @@ -1182,6 +1213,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): retry_policy=retry_policy, cache_policy=cache_policy, durability=durability, + on_interrupt=on_interrupt, ) self.stack = AsyncExitStack() if checkpointer: @@ -1257,6 +1289,18 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): }, ) + _deferred_on_interrupt_args: list[Interrupt] | None = None + + def _call_on_interrupt(self, interrupts: list[Interrupt]) -> None: + """Override for async loop: defer async hooks to __aexit__.""" + if self.on_interrupt is None: + return + if iscoroutinefunction(self.on_interrupt): + # Defer async hooks — they will be awaited in __aexit__ + self._deferred_on_interrupt_args = interrupts + else: + self.on_interrupt(interrupts) + # context manager async def __aenter__(self) -> Self: @@ -1315,14 +1359,25 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): exc_value: BaseException | None, traceback: TracebackType | None, ) -> bool | None: - # unwind stack + # unwind stack (calls _suppress_interrupt synchronously) exit_task = asyncio.create_task( self.stack.__aexit__(exc_type, exc_value, traceback) ) try: - return await exit_task + result = await exit_task except asyncio.CancelledError as e: # Bubble up the exit task upon cancellation to permit the API # consumer to await it before e.g., reusing the DB connection. e.args = (*e.args, exit_task) raise + # Await deferred async on_interrupt hook (set by _call_on_interrupt) + if ( + self._deferred_on_interrupt_args is not None + and self.on_interrupt is not None + ): + interrupts = self._deferred_on_interrupt_args + self._deferred_on_interrupt_args = None + coro = self.on_interrupt(interrupts) + if coro is not None: + await coro + return result diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index 37e8125f9..fa5b272d9 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -138,6 +138,7 @@ from langgraph.types import ( Command, Durability, Interrupt, + OnInterruptHook, Send, StateSnapshot, StateUpdate, @@ -622,6 +623,12 @@ class Pregel( context_schema: type[ContextT] | None = None """Specifies the schema for the context object that will be passed to the workflow.""" + on_interrupt: OnInterruptHook | None = None + """Optional callback invoked when the graph execution is interrupted. + + Called with the list of `Interrupt` objects whenever the graph pauses. + May be a sync or async callable.""" + config: RunnableConfig | None = None name: str = "LangGraph" @@ -652,6 +659,7 @@ class Pregel( config: RunnableConfig | None = None, trigger_to_nodes: Mapping[str, Sequence[str]] | None = None, name: str = "LangGraph", + on_interrupt: OnInterruptHook | None = None, **deprecated_kwargs: Unpack[DeprecatedKwargs], ) -> None: if ( @@ -695,6 +703,7 @@ class Pregel( ) self.cache_policy = cache_policy self.context_schema = context_schema + self.on_interrupt = on_interrupt self.config = config self.trigger_to_nodes = trigger_to_nodes or {} self.name = name @@ -2599,6 +2608,7 @@ class Pregel( migrate_checkpoint=self._migrate_checkpoint, retry_policy=self.retry_policy, cache_policy=self.cache_policy, + on_interrupt=self.on_interrupt, ) as loop: # create runner runner = PregelRunner( @@ -2908,6 +2918,7 @@ class Pregel( migrate_checkpoint=self._migrate_checkpoint, retry_policy=self.retry_policy, cache_policy=self.cache_policy, + on_interrupt=self.on_interrupt, ) as loop: # create runner runner = PregelRunner( diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index c8a798a38..740a1cc49 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -2,7 +2,7 @@ from __future__ import annotations import sys from collections import deque -from collections.abc import Callable, Hashable, Sequence +from collections.abc import Awaitable, Callable, Hashable, Sequence from dataclasses import asdict, dataclass from typing import ( TYPE_CHECKING, @@ -56,6 +56,7 @@ __all__ = ( "Durability", "interrupt", "Overwrite", + "OnInterruptHook", "ensure_valid_checkpointer", ) @@ -109,6 +110,18 @@ StreamWriter = Callable[[Any], None] Always injected into nodes if requested as a keyword argument, but it's a no-op when not using `stream_mode="custom"`.""" +OnInterruptHook = ( + Callable[[list["Interrupt"]], None] | Callable[[list["Interrupt"]], Awaitable[None]] +) +"""Callback invoked when a graph execution is interrupted. + +Called with the list of `Interrupt` objects whenever the graph pauses due to +an `interrupt()` call or `interrupt_before`/`interrupt_after` configuration. + +May be a regular function or an async coroutine function. Async hooks are +awaited in async graph execution; in sync execution only sync hooks are called. +""" + _DC_KWARGS = {"kw_only": True, "slots": True, "frozen": True} diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 97928ec6f..461a36141 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -8893,3 +8893,151 @@ def test_fork_does_not_apply_pending_writes( # Should be: 1 (input) + 20 (forked node_a) + 100 (node_b) = 121 assert result == {"value": 121} + + +def test_on_interrupt_hook_with_interrupt_call( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + """Test that on_interrupt hook fires when interrupt() is called in a node.""" + hook_calls: list[list[Interrupt]] = [] + + def my_on_interrupt(interrupts: list[Interrupt]) -> None: + hook_calls.append(interrupts) + + class State(TypedDict): + value: str + + def ask_human(state: State) -> dict: + answer = interrupt("what should I do?") + return {"value": answer} + + builder = StateGraph(State) + builder.add_node("ask", ask_human) + builder.add_edge(START, "ask") + + graph = builder.compile( + checkpointer=sync_checkpointer, + on_interrupt=my_on_interrupt, + ) + + config = {"configurable": {"thread_id": "1"}} + + # First invocation: should trigger interrupt and call the hook + result = list(graph.stream({"value": ""}, config)) + assert len(result) == 1 + assert "__interrupt__" in result[0] + + # Hook should have been called once with the interrupt data + assert len(hook_calls) == 1 + assert len(hook_calls[0]) == 1 + assert hook_calls[0][0].value == "what should I do?" + + # Resume — no new interrupt, hook should not fire again + hook_calls.clear() + result = list(graph.stream(Command(resume="do this"), config)) + assert any("ask" in chunk for chunk in result) + assert len(hook_calls) == 0 + + +def test_on_interrupt_hook_with_interrupt_before( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + """Test that on_interrupt hook fires for interrupt_before config.""" + hook_calls: list[list[Interrupt]] = [] + + def my_on_interrupt(interrupts: list[Interrupt]) -> None: + hook_calls.append(interrupts) + + class State(TypedDict): + value: int + + def add_one(state: State) -> dict: + return {"value": state["value"] + 1} + + builder = StateGraph(State) + builder.add_node("add_one", add_one) + builder.add_edge(START, "add_one") + + graph = builder.compile( + checkpointer=sync_checkpointer, + interrupt_before=["add_one"], + on_interrupt=my_on_interrupt, + ) + + config = {"configurable": {"thread_id": "1"}} + + # Should interrupt before add_one runs + result = list(graph.stream({"value": 0}, config)) + assert any("__interrupt__" in chunk for chunk in result) + + # Hook should have been called (empty interrupt list for config-level interrupts) + assert len(hook_calls) == 1 + assert hook_calls[0] == [] + + +def test_on_interrupt_hook_not_called_without_interrupt( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + """Test that on_interrupt hook is NOT called when graph completes normally.""" + hook_calls: list[list[Interrupt]] = [] + + def my_on_interrupt(interrupts: list[Interrupt]) -> None: + hook_calls.append(interrupts) + + class State(TypedDict): + value: int + + def add_one(state: State) -> dict: + return {"value": state["value"] + 1} + + builder = StateGraph(State) + builder.add_node("add_one", add_one) + builder.add_edge(START, "add_one") + + graph = builder.compile( + checkpointer=sync_checkpointer, + on_interrupt=my_on_interrupt, + ) + + config = {"configurable": {"thread_id": "1"}} + + result = graph.invoke({"value": 0}, config) + assert result == {"value": 1} + + # Hook should NOT have been called + assert len(hook_calls) == 0 + + +def test_on_interrupt_hook_exception_propagates( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + """Test that exceptions in the on_interrupt hook propagate to the caller.""" + + def bad_hook(interrupts: list[Interrupt]) -> None: + raise RuntimeError("hook exploded") + + class State(TypedDict): + value: str + + def ask(state: State) -> dict: + answer = interrupt("question") + return {"value": answer} + + builder = StateGraph(State) + builder.add_node("ask", ask) + builder.add_edge(START, "ask") + + graph = builder.compile( + checkpointer=sync_checkpointer, + on_interrupt=bad_hook, + ) + + config = {"configurable": {"thread_id": "1"}} + + # Hook error should propagate + with pytest.raises(RuntimeError, match="hook exploded"): + list(graph.stream({"value": ""}, config)) + + # Graph state should still be checkpointed and resumable despite the hook error + result = list(graph.stream(Command(resume="answer"), config)) + assert any("ask" in chunk for chunk in result) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index e57c09b3d..c2005b394 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -9345,3 +9345,117 @@ async def test_fork_does_not_apply_pending_writes( # 1 (input) + 20 (forked node_a) + 100 (node_b) = 121 assert result == {"value": 121} + + +async def test_on_interrupt_hook_async_with_interrupt_call( + async_checkpointer: BaseCheckpointSaver, +) -> None: + """Test that an async on_interrupt hook fires when interrupt() is called.""" + hook_calls: list[list[Interrupt]] = [] + + async def my_on_interrupt(interrupts: list[Interrupt]) -> None: + hook_calls.append(interrupts) + + class State(TypedDict): + value: str + + def ask_human(state: State) -> dict: + answer = interrupt("what should I do?") + return {"value": answer} + + builder = StateGraph(State) + builder.add_node("ask", ask_human) + builder.add_edge(START, "ask") + + graph = builder.compile( + checkpointer=async_checkpointer, + on_interrupt=my_on_interrupt, + ) + + config = {"configurable": {"thread_id": "1"}} + + # First invocation: should trigger interrupt and call the async hook + result = [chunk async for chunk in graph.astream({"value": ""}, config)] + assert len(result) == 1 + assert "__interrupt__" in result[0] + + # Hook should have been called once with the interrupt data + assert len(hook_calls) == 1 + assert len(hook_calls[0]) == 1 + assert hook_calls[0][0].value == "what should I do?" + + # Resume — no new interrupt, hook should not fire again + hook_calls.clear() + result = [chunk async for chunk in graph.astream(Command(resume="do this"), config)] + assert any("ask" in chunk for chunk in result) + assert len(hook_calls) == 0 + + +async def test_on_interrupt_hook_sync_in_async_graph( + async_checkpointer: BaseCheckpointSaver, +) -> None: + """Test that a sync on_interrupt hook works in async graph execution.""" + hook_calls: list[list[Interrupt]] = [] + + def my_sync_hook(interrupts: list[Interrupt]) -> None: + hook_calls.append(interrupts) + + class State(TypedDict): + value: str + + def ask_human(state: State) -> dict: + answer = interrupt("question?") + return {"value": answer} + + builder = StateGraph(State) + builder.add_node("ask", ask_human) + builder.add_edge(START, "ask") + + graph = builder.compile( + checkpointer=async_checkpointer, + on_interrupt=my_sync_hook, + ) + + config = {"configurable": {"thread_id": "1"}} + + result = [chunk async for chunk in graph.astream({"value": ""}, config)] + assert "__interrupt__" in result[0] + + # Sync hook should work fine in async execution + assert len(hook_calls) == 1 + assert hook_calls[0][0].value == "question?" + + +async def test_on_interrupt_hook_async_exception_propagates( + async_checkpointer: BaseCheckpointSaver, +) -> None: + """Test that exceptions in the async on_interrupt hook propagate.""" + + async def bad_hook(interrupts: list[Interrupt]) -> None: + raise RuntimeError("async hook exploded") + + class State(TypedDict): + value: str + + def ask(state: State) -> dict: + answer = interrupt("question") + return {"value": answer} + + builder = StateGraph(State) + builder.add_node("ask", ask) + builder.add_edge(START, "ask") + + graph = builder.compile( + checkpointer=async_checkpointer, + on_interrupt=bad_hook, + ) + + config = {"configurable": {"thread_id": "1"}} + + # Hook error should propagate + with pytest.raises(RuntimeError, match="async hook exploded"): + [chunk async for chunk in graph.astream({"value": ""}, config)] + + # Graph state should still be checkpointed and resumable despite the hook error + result = [chunk async for chunk in graph.astream(Command(resume="answer"), config)] + assert any("ask" in chunk for chunk in result)