From 1508cbf33c436066175d2a1e056192e2bdd79d70 Mon Sep 17 00:00:00 2001 From: Caspar Broekhuizen Date: Mon, 29 Sep 2025 19:13:32 -0700 Subject: [PATCH] fix(langgraph): fix nested subgraph checkpoint replay --- libs/langgraph/langgraph/pregel/_loop.py | 60 ++++++++++++++- libs/langgraph/tests/test_pregel.py | 91 +++++++++++++++++++++++ libs/langgraph/tests/test_pregel_async.py | 91 +++++++++++++++++++++++ 3 files changed, 240 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index 023ac5545..5e5d2382c 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -1048,7 +1048,35 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): def __enter__(self) -> Self: if self.checkpointer: - saved = self.checkpointer.get_tuple(self.checkpoint_config) + if not self.is_nested: + saved = self.checkpointer.get_tuple(self.checkpoint_config) + else: + # when rewinding into a nested graph, ignore checkpoints newer than the + # boundary id so the subgraph reruns instead of replaying future state. + boundary_checkpoint_id = self.checkpoint_config["metadata"].get( + "checkpoint_id" + ) + checkpoint_ns: str = self.checkpoint_config["configurable"].get( + "checkpoint_ns", "" + ) + + config = self.checkpoint_config.copy() + config["configurable"]["checkpoint_ns"] = checkpoint_ns + + checkpoints = self.checkpointer.list(config) + + # find the first checkpoint where checkpoint id <= boundary id. + # if boundary_checkpoint_id=None, e.g. when not time-traveling, + # stop at the first (max) checkpoint in the namespace. + saved = None + for checkpoint in checkpoints: + if ( + boundary_checkpoint_id + and checkpoint.checkpoint["id"] > boundary_checkpoint_id + ): + continue + saved = checkpoint + break else: saved = None if saved is None: @@ -1227,7 +1255,35 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): async def __aenter__(self) -> Self: if self.checkpointer: - saved = await self.checkpointer.aget_tuple(self.checkpoint_config) + if not self.is_nested: + saved = await self.checkpointer.aget_tuple(self.checkpoint_config) + else: + # when rewinding into a nested graph, ignore checkpoints newer than the + # boundary id so the subgraph reruns instead of replaying future state. + boundary_checkpoint_id = self.checkpoint_config["metadata"].get( + "checkpoint_id" + ) + checkpoint_ns: str = self.checkpoint_config["configurable"].get( + "checkpoint_ns", "" + ) + + config = self.checkpoint_config.copy() + config["configurable"]["checkpoint_ns"] = checkpoint_ns + + checkpoints = self.checkpointer.alist(config) + + # find the first checkpoint where checkpoint id <= boundary id. + # if boundary_checkpoint_id=None, e.g. when not time-traveling, + # stop at the first (max) checkpoint in the namespace. + saved = None + async for checkpoint in checkpoints: + if ( + boundary_checkpoint_id + and checkpoint.checkpoint["id"] > boundary_checkpoint_id + ): + continue + saved = checkpoint + break else: saved = None if saved is None: diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 09274782c..292f8efd3 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -8516,3 +8516,94 @@ def test_interrupt_stream_mode_values(): result = [*app.stream(State(), stream_mode="values")] assert "__interrupt__" in result[-1] + + +def test_subgraph_resume_reexecutes_from_valid_checkpoint( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + class InnerState(TypedDict, total=False): + input: str + output: str + + class OuterState(TypedDict, total=False): + input: str + output: str + + inner_model_calls = 0 + inner_tool_calls = 0 + outer_tool_calls = 0 + + def inner_model(state: InnerState) -> None: + nonlocal inner_model_calls + inner_model_calls += 1 + + def inner_tool(state: InnerState) -> dict[str, str]: + nonlocal inner_tool_calls + inner_tool_calls += 1 + return {"output": f"inner {inner_tool_calls}"} + + def should_continue_inner(state: InnerState) -> str: + return END if "output" in state else "inner_tool" + + subgraph_builder = StateGraph(state_schema=InnerState) + subgraph_builder.add_node("model", inner_model) + subgraph_builder.add_node("inner_tool", inner_tool) + subgraph_builder.add_edge("inner_tool", "model") + subgraph_builder.add_conditional_edges( + "model", + should_continue_inner, + path_map=["inner_tool", END], + ) + subgraph_builder.set_entry_point("model") + subgraph = subgraph_builder.compile(checkpointer=True) + + outer_model_calls = 0 + + def outer_model(state: OuterState) -> Command | None: + nonlocal outer_model_calls + outer_model_calls += 1 + if "output" in state: + return Command(goto=END) + return None + + def outer_tool(state: OuterState) -> dict[str, str]: + nonlocal outer_tool_calls + outer_tool_calls += 1 + result = subgraph.invoke(state) + return {"output": f"subgraph result: {result['output']}"} + + def should_continue(state: OuterState) -> str: + return END if "output" in state else "tools" + + graph_builder = StateGraph(state_schema=OuterState) + graph_builder.add_node("model", outer_model) + graph_builder.add_node("tools", outer_tool) + graph_builder.add_edge("tools", "model") + graph_builder.add_conditional_edges( + "model", + should_continue, + path_map=["tools", END], + ) + graph_builder.set_entry_point("model") + graph = graph_builder.compile(checkpointer=sync_checkpointer) + + config: RunnableConfig = {"configurable": {"thread_id": str(uuid.uuid4())}} + + graph.invoke({"input": "hello"}, config=config, interrupt_after=["tools"]) + + assert inner_tool_calls == 1 + assert outer_tool_calls == 1 + + history = list(graph.get_state_history(config)) + resume_state = next( + state + for state in history + if state.next == ("model",) and "output" not in state.values + ) + + result = graph.invoke(Command(resume="resume"), resume_state.config) + + assert result["output"] == "subgraph result: inner 2" + assert inner_tool_calls == 2 + assert outer_tool_calls == 2 + assert outer_model_calls == 3 diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index a6296a0c4..109cf67e3 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -9211,3 +9211,94 @@ async def test_astream_waiter_cleanup_on_cancel( assert recorded_tasks, "expected stream.wait() task to be created" assert set(finished_tasks) == set(recorded_tasks) assert all(t.done() for t in recorded_tasks) + + +async def test_subgraph_resume_reexecutes_from_valid_checkpoint( + async_checkpointer: BaseCheckpointSaver, +) -> None: + class InnerState(TypedDict, total=False): + input: str + output: str + + class OuterState(TypedDict, total=False): + input: str + output: str + + inner_model_calls = 0 + inner_tool_calls = 0 + outer_tool_calls = 0 + + async def inner_model(state: InnerState) -> None: + nonlocal inner_model_calls + inner_model_calls += 1 + + async def inner_tool(state: InnerState) -> dict[str, str]: + nonlocal inner_tool_calls + inner_tool_calls += 1 + return {"output": f"inner {inner_tool_calls}"} + + def should_continue_inner(state: InnerState) -> str: + return END if "output" in state else "inner_tool" + + subgraph_builder = StateGraph(state_schema=InnerState) + subgraph_builder.add_node("model", inner_model) + subgraph_builder.add_node("inner_tool", inner_tool) + subgraph_builder.add_edge("inner_tool", "model") + subgraph_builder.add_conditional_edges( + "model", + should_continue_inner, + path_map=["inner_tool", END], + ) + subgraph_builder.set_entry_point("model") + subgraph = subgraph_builder.compile(checkpointer=True) + + outer_model_calls = 0 + + async def outer_model(state: OuterState) -> Command | None: + nonlocal outer_model_calls + outer_model_calls += 1 + if "output" in state: + return Command(goto=END) + return None + + async def outer_tool(state: OuterState) -> dict[str, str]: + nonlocal outer_tool_calls + outer_tool_calls += 1 + result = await subgraph.ainvoke(state) + return {"output": f"subgraph result: {result['output']}"} + + def should_continue(state: OuterState) -> str: + return END if "output" in state else "tools" + + graph_builder = StateGraph(state_schema=OuterState) + graph_builder.add_node("model", outer_model) + graph_builder.add_node("tools", outer_tool) + graph_builder.add_edge("tools", "model") + graph_builder.add_conditional_edges( + "model", + should_continue, + path_map=["tools", END], + ) + graph_builder.set_entry_point("model") + graph = graph_builder.compile(checkpointer=async_checkpointer) + + config: RunnableConfig = {"configurable": {"thread_id": str(uuid.uuid4())}} + + await graph.ainvoke({"input": "hello"}, config=config, interrupt_after=["tools"]) + + assert inner_tool_calls == 1 + assert outer_tool_calls == 1 + + history = [state async for state in graph.aget_state_history(config)] + resume_state = next( + state + for state in history + if state.next == ("model",) and "output" not in state.values + ) + + result = await graph.ainvoke(Command(resume="resume"), resume_state.config) + + assert result["output"] == "subgraph result: inner 2" + assert inner_tool_calls == 2 + assert outer_tool_calls == 2 + assert outer_model_calls >= 3