fix(langgraph): add async language

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:

libs/langgraph/tests/test_pregel.py

@@ -8516,3 +8516,112 @@ 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:
+    """Verify that a resumed subgraph doesn't replay future state during a time-jump.
+
+    After rewinding to a pre-tool checkpoint, the entire subgraph should execute again rather than replaying state from the previous execution.
+    """
+
+    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) -> Optional[Command]:
+        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))
+    # i: history[i].next, history[i].values:
+    # 0: ('model',) {'input': 'hello', 'output': 'subgraph result: inner 1'}
+    # 1: ('tools',) {'input': 'hello'}
+    # 2: ('model',) {'input': 'hello'}
+    # 3: ('__start__',) {}
+
+    # resume from tools (node that executes subgraph)
+    resume_state = history[1]
+
+    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 == 2
+
+    # resume from model
+    resume_state = history[2]
+
+    result = graph.invoke(Command(resume="resume"), resume_state.config)
+
+    assert result["output"] == "subgraph result: inner 3"
+    assert inner_tool_calls == 3
+    assert outer_tool_calls == 3
+    assert outer_model_calls == 4

libs/langgraph/tests/test_pregel_async.py

@@ -9211,3 +9211,112 @@ 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:
+    """Verify that a resumed subgraph doesn't replay future state during a time-jump.
+
+    After rewinding to a pre-tool checkpoint, the entire subgraph should execute again rather than replaying state from the previous execution.
+    """
+
+    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) -> Optional[Command]:
+        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)]
+    # i: history[i].next, history[i].values:
+    # 0: ('model',) {'input': 'hello', 'output': 'subgraph result: inner 1'}
+    # 1: ('tools',) {'input': 'hello'}
+    # 2: ('model',) {'input': 'hello'}
+    # 3: ('__start__',) {}
+
+    # resume from tools (node that executes subgraph)
+    resume_state = history[1]
+
+    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 == 2
+
+    # resume from model
+    resume_state = history[2]
+
+    result = await graph.ainvoke(Command(resume="resume"), resume_state.config)
+
+    assert result["output"] == "subgraph result: inner 3"
+    assert inner_tool_calls == 3
+    assert outer_tool_calls == 3
+    assert outer_model_calls == 4