From 6d9f9f53264255072e05128165d09eb7f13e71d7 Mon Sep 17 00:00:00 2001 From: Hunter Lovell Date: Tue, 23 Jun 2026 13:43:04 -0700 Subject: [PATCH] fix(pregel): persist pending writes on interrupt --- libs/langgraph/langgraph/pregel/_runner.py | 27 +++++ libs/langgraph/tests/test_pregel.py | 121 ++++++++++++++++++++- 2 files changed, 145 insertions(+), 3 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/_runner.py b/libs/langgraph/langgraph/pregel/_runner.py index 4d53b9f9d..159510658 100644 --- a/libs/langgraph/langgraph/pregel/_runner.py +++ b/libs/langgraph/langgraph/pregel/_runner.py @@ -34,6 +34,8 @@ from langgraph._internal._constants import ( ERROR_SOURCE_NODE, INTERRUPT, NO_WRITES, + NULL_TASK_ID, + PUSH, RESUME, RETURN, ) @@ -585,6 +587,12 @@ class PregelRunner: if isinstance(exception, GraphInterrupt): # save interrupt to checkpointer if exception.args[0]: + if pending_writes := _writes_to_persist_on_interrupt(task.writes): + # GraphInterrupt is a controlled suspension point, not a + # task failure. Persist writes emitted before the + # suspension without associating them with task + # completion, so the interrupted task remains resumable. + self.put_writes()(NULL_TASK_ID, pending_writes) # type: ignore[misc] writes = [(INTERRUPT, exception.args[0])] if resumes := [w for w in task.writes if w[0] == RESUME]: writes.extend(resumes) @@ -939,3 +947,22 @@ async def _acall_impl( destination.set_exception(RuntimeError("Task not scheduled")) except Exception as exc: destination.set_exception(exc) + + +def _writes_to_persist_on_interrupt( + writes: Iterable[tuple[str, Any]], +) -> list[tuple[str, Any]]: + return [ + write + for write in writes + if write[0] + not in ( + ERROR, + ERROR_SOURCE_NODE, + INTERRUPT, + NO_WRITES, + PUSH, + RESUME, + RETURN, + ) + ] diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 0aae1318e..d043c578f 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -40,15 +40,25 @@ from pytest_mock import MockerFixture from syrupy import SnapshotAssertion from typing_extensions import NotRequired, TypedDict -from langgraph._internal._constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL +from langgraph._internal._constants import ( + CONFIG_KEY_NODE_FINISHED, + CONFIG_KEY_SEND, + ERROR, + PULL, +) from langgraph.channels.binop import BinaryOperatorAggregate from langgraph.channels.delta import DeltaChannel from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic from langgraph.channels.untracked_value import UntrackedValue -from langgraph.config import get_stream_writer -from langgraph.errors import GraphRecursionError, InvalidUpdateError, ParentCommand +from langgraph.config import get_config, get_stream_writer +from langgraph.errors import ( + GraphInterrupt, + GraphRecursionError, + InvalidUpdateError, + ParentCommand, +) from langgraph.func import entrypoint, task from langgraph.graph import END, START, StateGraph from langgraph.graph.message import MessagesState, _messages_delta_reducer, add_messages @@ -4820,6 +4830,111 @@ def test_parent_command( ) +def test_pending_send_write_persists_on_interrupt( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + class State(TypedDict): + quickjs_checkpoint: dict + result: str + + def node(state: State) -> dict[str, Any]: + send = get_config()["configurable"][CONFIG_KEY_SEND] + send([("quickjs_checkpoint", {"snapshot": "abc"})]) + value = interrupt("trip") + return {"result": value} + + builder = StateGraph(State) + builder.add_node("node", node) + builder.add_edge(START, "node") + builder.add_edge("node", END) + graph = builder.compile(checkpointer=sync_checkpointer) + config = {"configurable": {"thread_id": "1"}} + + first = graph.invoke({"quickjs_checkpoint": {}, "result": ""}, config) + assert "__interrupt__" in first + + snapshot = graph.get_state(config) + assert snapshot.values["quickjs_checkpoint"] == {"snapshot": "abc"} + assert snapshot.next == ("node",) + assert snapshot.tasks[0].interrupts + + second = graph.invoke(Command(resume="ok"), config) + assert second["quickjs_checkpoint"] == {"snapshot": "abc"} + assert second["result"] == "ok" + assert graph.get_state(config).values["quickjs_checkpoint"] == {"snapshot": "abc"} + + +def test_pending_send_write_persists_on_tool_subgraph_interrupt( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + from langchain_core.tools import tool + + class SubgraphState(TypedDict): + value: str + + def subgraph_node(state: SubgraphState) -> dict[str, Any]: + value = interrupt("subgraph trip") + return {"value": value} + + subgraph_builder = StateGraph(SubgraphState) + subgraph_builder.add_node("subgraph_node", subgraph_node) + subgraph_builder.add_edge(START, "subgraph_node") + subgraph = subgraph_builder.compile(checkpointer=True) + + class ParentState(TypedDict): + messages: Annotated[list[AnyMessage], add_messages] + quickjs_checkpoint: dict + + @tool + def task_tool() -> str: + """Run the interrupting subgraph task.""" + send = get_config()["configurable"][CONFIG_KEY_SEND] + try: + response = subgraph.invoke({"value": ""}) + except GraphInterrupt: + send([("quickjs_checkpoint", {"snapshot": "tool"})]) + raise + return response["value"] + + def call_tool(state: ParentState) -> dict[str, Any]: + return { + "messages": AIMessage( + content="", + tool_calls=[ + { + "name": "task_tool", + "args": {}, + "id": "tool_call_1", + } + ], + ) + } + + builder = StateGraph(ParentState) + builder.add_node("agent", call_tool) + builder.add_node("tools", ToolNode([task_tool])) + builder.add_edge(START, "agent") + builder.add_edge("agent", "tools") + builder.add_edge("tools", END) + graph = builder.compile(checkpointer=sync_checkpointer) + config = {"configurable": {"thread_id": "1"}} + + first = graph.invoke( + {"messages": [HumanMessage(content="start")], "quickjs_checkpoint": {}}, config + ) + assert "__interrupt__" in first + + snapshot = graph.get_state(config) + assert snapshot.values["quickjs_checkpoint"] == {"snapshot": "tool"} + assert snapshot.next == ("tools",) + assert snapshot.tasks[0].interrupts + + second = graph.invoke(Command(resume="ok"), config) + assert second["quickjs_checkpoint"] == {"snapshot": "tool"} + assert second["messages"][-1].content == "ok" + assert graph.get_state(config).values["quickjs_checkpoint"] == {"snapshot": "tool"} + + def test_interrupt_subgraph(sync_checkpointer: BaseCheckpointSaver): class State(TypedDict): baz: str