diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index 602230dd2..8e86e1410 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -1133,35 +1133,40 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): return parent_saved.parent_config[CONF].get(CONFIG_KEY_CHECKPOINT_ID) return parent_checkpoint_id - def _get_checkpoint_before_parent(self) -> CheckpointTuple | None: - """Find the subgraph checkpoint that was current at the parent's - checkpoint time, using the parent checkpoint_id as an upper bound. + def _get_checkpoint_after_parent(self) -> CheckpointTuple | None: + """Find the right subgraph checkpoint to restore when the parent replays. - We return a checkpoint with the historical channel_values but empty - channel_versions and versions_seen (via `empty_checkpoint()`). This is - intentional: the node scheduling logic in `_triggers()` decides whether - to run a node by comparing channel_versions against that node's - versions_seen. By clearing both, every channel looks "new" to every - node, which forces all nodes to re-trigger. Without this, the restored - checkpoint's versions would show all nodes as up-to-date and nothing - would be scheduled to run. + Each time the parent invokes a subgraph, the subgraph creates a series + of checkpoints. Every checkpoint records which parent checkpoint was + active when it was created (in `metadata["parents"]`). The first + checkpoint in each invocation has `source="input"` and contains the + accumulated channel_values from prior invocations but hasn't run any + nodes yet. + + We query for `source="input"` + `parents={parent_ns: parent_id}` to + find the starting checkpoint from the invocation that ran under the + given parent checkpoint — one bounded query, one result. + + We then clear `versions_seen` so all nodes re-trigger from that state. + The existing `is_replaying` logic in `_first()` handles dropping any + cached RESUME writes so that interrupts re-fire. Returns None to start fresh if no such checkpoint exists.""" parent_checkpoint_id = self._get_parent_checkpoint_id() if parent_checkpoint_id and self.checkpointer: - before_config: RunnableConfig = { - CONF: {"checkpoint_id": parent_checkpoint_id} - } + parent_ns = ( + NS_SEP.join(self.checkpoint_ns[:-1]) if self.checkpoint_ns else "" + ) for saved in self.checkpointer.list( - self.checkpoint_config, before=before_config, limit=1 + self.checkpoint_config, + filter={ + "source": "input", + "parents": {parent_ns: parent_checkpoint_id}, + }, + limit=1, ): - checkpoint = empty_checkpoint() - checkpoint["channel_values"] = saved.checkpoint.get( - "channel_values", {} - ) - return CheckpointTuple( - self.checkpoint_config, checkpoint, {"step": -2}, None, [] - ) + saved.checkpoint["versions_seen"] = {} + return saved return None # context manager @@ -1171,16 +1176,13 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): saved = self.checkpointer.get_tuple(self.checkpoint_config) else: saved = None - # When replaying a subgraph that wasn't in the checkpoint map - # (parent checkpoint predates this subgraph), start fresh. - # For stateful subgraphs (checkpointer=True), find the checkpoint - # that was current at the parent's checkpoint time. - if ( - saved is not None - and self.config[CONF].get(CONFIG_KEY_REPLAYING) - and not self.checkpoint_config.get(CONF, {}).get(CONFIG_KEY_CHECKPOINT_ID) - ): - saved = self._get_checkpoint_before_parent() + # When replaying a subgraph, find the checkpoint that was current + # at the parent's checkpoint time. For stateless subgraphs (no + # checkpointer), this returns None and we start fresh as usual. + if self.config[CONF].get( + CONFIG_KEY_REPLAYING + ) and not self.checkpoint_config.get(CONF, {}).get(CONFIG_KEY_CHECKPOINT_ID): + saved = self._get_checkpoint_after_parent() if saved is None: saved = CheckpointTuple( self.checkpoint_config, empty_checkpoint(), {"step": -2}, None, [] @@ -1372,29 +1374,23 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): return parent_saved.parent_config[CONF].get(CONFIG_KEY_CHECKPOINT_ID) return parent_checkpoint_id - async def _aget_checkpoint_before_parent(self) -> CheckpointTuple | None: - """Async version of `_get_checkpoint_before_parent`. See that method's - docstring for details on why we use empty_checkpoint() with only - channel_values restored.""" + async def _aget_checkpoint_after_parent(self) -> CheckpointTuple | None: + """Async version of `_get_checkpoint_after_parent`.""" parent_checkpoint_id = await self._aget_parent_checkpoint_id() if parent_checkpoint_id and self.checkpointer: - before_config: RunnableConfig = { - CONF: {"checkpoint_id": parent_checkpoint_id} - } + parent_ns = ( + NS_SEP.join(self.checkpoint_ns[:-1]) if self.checkpoint_ns else "" + ) async for saved in self.checkpointer.alist( - self.checkpoint_config, before=before_config, limit=1 + self.checkpoint_config, + filter={ + "source": "input", + "parents": {parent_ns: parent_checkpoint_id}, + }, + limit=1, ): - checkpoint = empty_checkpoint() - checkpoint["channel_values"] = saved.checkpoint.get( - "channel_values", {} - ) - return CheckpointTuple( - self.checkpoint_config, - checkpoint, - {"step": -2}, - None, - [], - ) + saved.checkpoint["versions_seen"] = {} + return saved return None # context manager @@ -1404,16 +1400,13 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): saved = await self.checkpointer.aget_tuple(self.checkpoint_config) else: saved = None - # When replaying a subgraph that wasn't in the checkpoint map - # (parent checkpoint predates this subgraph), start fresh. - # For stateful subgraphs (checkpointer=True), find the checkpoint - # that was current at the parent's checkpoint time. - if ( - saved is not None - and self.config[CONF].get(CONFIG_KEY_REPLAYING) - and not self.checkpoint_config.get(CONF, {}).get(CONFIG_KEY_CHECKPOINT_ID) - ): - saved = await self._aget_checkpoint_before_parent() + # When replaying a subgraph, find the checkpoint that was current + # at the parent's checkpoint time. For stateless subgraphs (no + # checkpointer), this returns None and we start fresh as usual. + if self.config[CONF].get( + CONFIG_KEY_REPLAYING + ) and not self.checkpoint_config.get(CONF, {}).get(CONFIG_KEY_CHECKPOINT_ID): + saved = await self._aget_checkpoint_after_parent() if saved is None: saved = CheckpointTuple( self.checkpoint_config, empty_checkpoint(), {"step": -2}, None, [] diff --git a/libs/langgraph/tests/test_time_travel.py b/libs/langgraph/tests/test_time_travel.py index feb80ea8c..c0eab5531 100644 --- a/libs/langgraph/tests/test_time_travel.py +++ b/libs/langgraph/tests/test_time_travel.py @@ -1570,6 +1570,7 @@ def test_stateful_subgraph_retains_state_on_parent_replay( ) -> None: """Stateful subgraph (checkpointer=True) remembers accumulated state from prior invocations when the parent replays.""" + started: list[tuple[str, dict]] = [] observed: list[tuple[str, dict]] = [] class SubState(TypedDict): @@ -1582,13 +1583,15 @@ def test_stateful_subgraph_retains_state_on_parent_replay( return {"results": ["p"]} def step_a(state: SubState) -> SubState: - observed.append(("step_a", dict(state))) + started.append(("step_a", dict(state))) answer = interrupt("question_a") + observed.append(("step_a", dict(state))) return {"value": [f"a:{answer}"]} def step_b(state: SubState) -> SubState: - observed.append(("step_b", dict(state))) + started.append(("step_b", dict(state))) answer = interrupt("question_b") + observed.append(("step_b", dict(state))) return {"value": [f"b:{answer}"]} sub = ( @@ -1612,9 +1615,9 @@ def test_stateful_subgraph_retains_state_on_parent_replay( config = {"configurable": {"thread_id": "1"}} # === 1st invocation: answer "a1" and "b1" === - graph.invoke({"results": []}, config) # hits step_a interrupt - graph.invoke(Command(resume="a1"), config) # hits step_b interrupt - graph.invoke(Command(resume="b1"), config) # completes + graph.invoke({"results": []}, config) # hits step_a interrupt + graph.invoke(Command(resume="a1"), config) # hits step_b interrupt + graph.invoke(Command(resume="b1"), config) # completes # step_a saw empty state (fresh subgraph) assert observed[0] == ("step_a", {"value": []}) @@ -1623,9 +1626,9 @@ def test_stateful_subgraph_retains_state_on_parent_replay( # === 2nd invocation: answer "a2" and "b2" === observed.clear() - graph.invoke({"results": []}, config) # hits step_a interrupt - graph.invoke(Command(resume="a2"), config) # hits step_b interrupt - graph.invoke(Command(resume="b2"), config) # completes + graph.invoke({"results": []}, config) # hits step_a interrupt + graph.invoke(Command(resume="a2"), config) # hits step_b interrupt + graph.invoke(Command(resume="b2"), config) # completes # Stateful subgraph retained state from 1st invocation assert observed[0] == ("step_a", {"value": ["a:a1", "b:b1"]}) @@ -1636,12 +1639,12 @@ def test_stateful_subgraph_retains_state_on_parent_replay( # History is newest-first, so first match = 2nd invocation before_sub_2nd = [s for s in history if s.next == ("sub_node",)][0] - observed.clear() + started.clear() replay = graph.invoke(None, before_sub_2nd.config) assert "__interrupt__" in replay # Replay sees 1st invocation's final state, NOT 2nd invocation's - assert observed[0] == ("step_a", {"value": ["a:a1", "b:b1"]}) + assert started[0] == ("step_a", {"value": ["a:a1", "b:b1"]}) def test_stateful_subgraph_retains_state_on_parent_fork( @@ -1691,14 +1694,14 @@ def test_stateful_subgraph_retains_state_on_parent_fork( config = {"configurable": {"thread_id": "1"}} # === 1st invocation: answer "a1" and "b1" === - graph.invoke({"results": []}, config) # hits step_a interrupt - graph.invoke(Command(resume="a1"), config) # hits step_b interrupt - graph.invoke(Command(resume="b1"), config) # completes + graph.invoke({"results": []}, config) # hits step_a interrupt + graph.invoke(Command(resume="a1"), config) # hits step_b interrupt + graph.invoke(Command(resume="b1"), config) # completes # === 2nd invocation: answer "a2" and "b2" === - graph.invoke({"results": []}, config) # hits step_a interrupt - graph.invoke(Command(resume="a2"), config) # hits step_b interrupt - graph.invoke(Command(resume="b2"), config) # completes + graph.invoke({"results": []}, config) # hits step_a interrupt + graph.invoke(Command(resume="a2"), config) # hits step_b interrupt + graph.invoke(Command(resume="b2"), config) # completes # === Fork from checkpoint before sub_node in 2nd invocation === history = list(graph.get_state_history(config)) @@ -1718,6 +1721,7 @@ def test_stateless_subgraph_starts_fresh_on_parent_replay( ) -> None: """Stateless subgraph (no checkpointer) always starts with empty state, even after prior invocations have completed.""" + started: list[tuple[str, dict]] = [] observed: list[tuple[str, dict]] = [] class SubState(TypedDict): @@ -1730,13 +1734,15 @@ def test_stateless_subgraph_starts_fresh_on_parent_replay( return {"results": ["p"]} def step_a(state: SubState) -> SubState: - observed.append(("step_a", dict(state))) + started.append(("step_a", dict(state))) answer = interrupt("question_a") + observed.append(("step_a", dict(state))) return {"value": [f"a:{answer}"]} def step_b(state: SubState) -> SubState: - observed.append(("step_b", dict(state))) + started.append(("step_b", dict(state))) answer = interrupt("question_b") + observed.append(("step_b", dict(state))) return {"value": [f"b:{answer}"]} sub = ( @@ -1760,9 +1766,9 @@ def test_stateless_subgraph_starts_fresh_on_parent_replay( config = {"configurable": {"thread_id": "1"}} # === 1st invocation: answer "a1" and "b1" === - graph.invoke({"results": []}, config) # hits step_a interrupt - graph.invoke(Command(resume="a1"), config) # hits step_b interrupt - graph.invoke(Command(resume="b1"), config) # completes + graph.invoke({"results": []}, config) # hits step_a interrupt + graph.invoke(Command(resume="a1"), config) # hits step_b interrupt + graph.invoke(Command(resume="b1"), config) # completes # step_a saw empty state, step_b saw only step_a's answer assert observed[0] == ("step_a", {"value": []}) @@ -1770,9 +1776,9 @@ def test_stateless_subgraph_starts_fresh_on_parent_replay( # === 2nd invocation: answer "a2" and "b2" === observed.clear() - graph.invoke({"results": []}, config) # hits step_a interrupt - graph.invoke(Command(resume="a2"), config) # hits step_b interrupt - graph.invoke(Command(resume="b2"), config) # completes + graph.invoke({"results": []}, config) # hits step_a interrupt + graph.invoke(Command(resume="a2"), config) # hits step_b interrupt + graph.invoke(Command(resume="b2"), config) # completes # Stateless subgraph starts fresh — no memory of 1st invocation assert observed[0] == ("step_a", {"value": []}) @@ -1782,9 +1788,9 @@ def test_stateless_subgraph_starts_fresh_on_parent_replay( history = list(graph.get_state_history(config)) before_sub_2nd = [s for s in history if s.next == ("sub_node",)][0] - observed.clear() + started.clear() replay = graph.invoke(None, before_sub_2nd.config) assert "__interrupt__" in replay # Stateless subgraph starts completely fresh on replay - assert observed[0] == ("step_a", {"value": []}) + assert started[0] == ("step_a", {"value": []}) diff --git a/libs/langgraph/tests/test_time_travel_async.py b/libs/langgraph/tests/test_time_travel_async.py index a36898f78..71645d29d 100644 --- a/libs/langgraph/tests/test_time_travel_async.py +++ b/libs/langgraph/tests/test_time_travel_async.py @@ -1603,6 +1603,7 @@ async def test_stateful_subgraph_retains_state_on_parent_replay( ) -> None: """Stateful subgraph (checkpointer=True) remembers accumulated state from prior invocations when the parent replays.""" + started: list[tuple[str, dict]] = [] observed: list[tuple[str, dict]] = [] class SubState(TypedDict): @@ -1615,13 +1616,15 @@ async def test_stateful_subgraph_retains_state_on_parent_replay( return {"results": ["p"]} def step_a(state: SubState) -> SubState: - observed.append(("step_a", dict(state))) + started.append(("step_a", dict(state))) answer = interrupt("question_a") + observed.append(("step_a", dict(state))) return {"value": [f"a:{answer}"]} def step_b(state: SubState) -> SubState: - observed.append(("step_b", dict(state))) + started.append(("step_b", dict(state))) answer = interrupt("question_b") + observed.append(("step_b", dict(state))) return {"value": [f"b:{answer}"]} sub = ( @@ -1645,9 +1648,9 @@ async def test_stateful_subgraph_retains_state_on_parent_replay( config = {"configurable": {"thread_id": "1"}} # === 1st invocation: answer "a1" and "b1" === - await graph.ainvoke({"results": []}, config) # hits step_a interrupt - await graph.ainvoke(Command(resume="a1"), config) # hits step_b interrupt - await graph.ainvoke(Command(resume="b1"), config) # completes + await graph.ainvoke({"results": []}, config) # hits step_a interrupt + await graph.ainvoke(Command(resume="a1"), config) # hits step_b interrupt + await graph.ainvoke(Command(resume="b1"), config) # completes # step_a saw empty state (fresh subgraph) assert observed[0] == ("step_a", {"value": []}) @@ -1656,9 +1659,9 @@ async def test_stateful_subgraph_retains_state_on_parent_replay( # === 2nd invocation: answer "a2" and "b2" === observed.clear() - await graph.ainvoke({"results": []}, config) # hits step_a interrupt - await graph.ainvoke(Command(resume="a2"), config) # hits step_b interrupt - await graph.ainvoke(Command(resume="b2"), config) # completes + await graph.ainvoke({"results": []}, config) # hits step_a interrupt + await graph.ainvoke(Command(resume="a2"), config) # hits step_b interrupt + await graph.ainvoke(Command(resume="b2"), config) # completes # Stateful subgraph retained state from 1st invocation assert observed[0] == ("step_a", {"value": ["a:a1", "b:b1"]}) @@ -1669,12 +1672,12 @@ async def test_stateful_subgraph_retains_state_on_parent_replay( # History is newest-first, so first match = 2nd invocation before_sub_2nd = [s for s in history if s.next == ("sub_node",)][0] - observed.clear() + started.clear() replay = await graph.ainvoke(None, before_sub_2nd.config) assert "__interrupt__" in replay # Replay sees 1st invocation's final state, NOT 2nd invocation's - assert observed[0] == ("step_a", {"value": ["a:a1", "b:b1"]}) + assert started[0] == ("step_a", {"value": ["a:a1", "b:b1"]}) @pytest.mark.skipif( @@ -1728,14 +1731,14 @@ async def test_stateful_subgraph_retains_state_on_parent_fork( config = {"configurable": {"thread_id": "1"}} # === 1st invocation: answer "a1" and "b1" === - await graph.ainvoke({"results": []}, config) # hits step_a interrupt - await graph.ainvoke(Command(resume="a1"), config) # hits step_b interrupt - await graph.ainvoke(Command(resume="b1"), config) # completes + await graph.ainvoke({"results": []}, config) # hits step_a interrupt + await graph.ainvoke(Command(resume="a1"), config) # hits step_b interrupt + await graph.ainvoke(Command(resume="b1"), config) # completes # === 2nd invocation: answer "a2" and "b2" === - await graph.ainvoke({"results": []}, config) # hits step_a interrupt - await graph.ainvoke(Command(resume="a2"), config) # hits step_b interrupt - await graph.ainvoke(Command(resume="b2"), config) # completes + await graph.ainvoke({"results": []}, config) # hits step_a interrupt + await graph.ainvoke(Command(resume="a2"), config) # hits step_b interrupt + await graph.ainvoke(Command(resume="b2"), config) # completes # === Fork from checkpoint before sub_node in 2nd invocation === history = [s async for s in graph.aget_state_history(config)] @@ -1762,6 +1765,7 @@ async def test_stateless_subgraph_starts_fresh_on_parent_replay( ) -> None: """Stateless subgraph (no checkpointer) always starts with empty state, even after prior invocations have completed.""" + started: list[tuple[str, dict]] = [] observed: list[tuple[str, dict]] = [] class SubState(TypedDict): @@ -1774,13 +1778,15 @@ async def test_stateless_subgraph_starts_fresh_on_parent_replay( return {"results": ["p"]} def step_a(state: SubState) -> SubState: - observed.append(("step_a", dict(state))) + started.append(("step_a", dict(state))) answer = interrupt("question_a") + observed.append(("step_a", dict(state))) return {"value": [f"a:{answer}"]} def step_b(state: SubState) -> SubState: - observed.append(("step_b", dict(state))) + started.append(("step_b", dict(state))) answer = interrupt("question_b") + observed.append(("step_b", dict(state))) return {"value": [f"b:{answer}"]} sub = ( @@ -1804,9 +1810,9 @@ async def test_stateless_subgraph_starts_fresh_on_parent_replay( config = {"configurable": {"thread_id": "1"}} # === 1st invocation: answer "a1" and "b1" === - await graph.ainvoke({"results": []}, config) # hits step_a interrupt - await graph.ainvoke(Command(resume="a1"), config) # hits step_b interrupt - await graph.ainvoke(Command(resume="b1"), config) # completes + await graph.ainvoke({"results": []}, config) # hits step_a interrupt + await graph.ainvoke(Command(resume="a1"), config) # hits step_b interrupt + await graph.ainvoke(Command(resume="b1"), config) # completes # step_a saw empty state, step_b saw only step_a's answer assert observed[0] == ("step_a", {"value": []}) @@ -1814,9 +1820,9 @@ async def test_stateless_subgraph_starts_fresh_on_parent_replay( # === 2nd invocation: answer "a2" and "b2" === observed.clear() - await graph.ainvoke({"results": []}, config) # hits step_a interrupt - await graph.ainvoke(Command(resume="a2"), config) # hits step_b interrupt - await graph.ainvoke(Command(resume="b2"), config) # completes + await graph.ainvoke({"results": []}, config) # hits step_a interrupt + await graph.ainvoke(Command(resume="a2"), config) # hits step_b interrupt + await graph.ainvoke(Command(resume="b2"), config) # completes # Stateless subgraph starts fresh — no memory of 1st invocation assert observed[0] == ("step_a", {"value": []}) @@ -1826,9 +1832,9 @@ async def test_stateless_subgraph_starts_fresh_on_parent_replay( history = [s async for s in graph.aget_state_history(config)] before_sub_2nd = [s for s in history if s.next == ("sub_node",)][0] - observed.clear() + started.clear() replay = await graph.ainvoke(None, before_sub_2nd.config) assert "__interrupt__" in replay # Stateless subgraph starts completely fresh on replay - assert observed[0] == ("step_a", {"value": []}) + assert started[0] == ("step_a", {"value": []})