diff --git a/libs/langgraph/langgraph/pregel/_checkpoint.py b/libs/langgraph/langgraph/pregel/_checkpoint.py index ba0c6178a..127700728 100644 --- a/libs/langgraph/langgraph/pregel/_checkpoint.py +++ b/libs/langgraph/langgraph/pregel/_checkpoint.py @@ -80,6 +80,8 @@ def get_updated_channels_from_tasks( def get_delta_channels_from_all_channels( channels: Mapping[str, BaseChannel], + *, + include_unavailable: bool = False, ) -> set[str]: """Every available DeltaChannel. @@ -91,7 +93,7 @@ def get_delta_channels_from_all_channels( return { k for k, ch in channels.items() - if isinstance(ch, DeltaChannel) and ch.is_available() + if isinstance(ch, DeltaChannel) and (include_unavailable or ch.is_available()) } @@ -146,7 +148,9 @@ def create_checkpoint_plan_for_update_state_api( "parents": parents, } if is_fresh_thread or is_fork: - return get_delta_channels_from_all_channels(channels), metadata + return get_delta_channels_from_all_channels( + channels, include_unavailable=is_fork + ), metadata new_counters = create_metadata_for_update_state_api( channels, @@ -193,7 +197,9 @@ def create_fork_checkpoint( channels, step, get_next_version=get_next_version, - channels_to_snapshot=get_delta_channels_from_all_channels(channels), + channels_to_snapshot=get_delta_channels_from_all_channels( + channels, include_unavailable=True + ), ) @@ -225,9 +231,20 @@ def create_checkpoint( values = {} channel_versions = dict(checkpoint["channel_versions"]) for k in channels: - if k not in channel_versions: - continue ch = channels[k] + if k not in channel_versions: + # Nothing was ever written to this channel on this branch, so + # it has no version and `put` would drop any blob stored for + # it. A *forced* snapshot still has to land: it is the only + # thing that stops the ancestor walk running past this + # checkpoint into a fork base that holds another branch's + # writes. Mint a first version so the blob survives. + if k in channels_to_snapshot and get_next_version is not None: + channel_versions[k] = get_next_version(None, None) + values[k] = _DeltaSnapshot( + ch.get() if ch.is_available() else ch.typ() + ) + continue if k in channels_to_snapshot: # Callers force a full snapshot blob here: exit mode when a # delta channel reaches its snapshot cadence, update_state on diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index 8bf03f680..21989fdf3 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -1174,7 +1174,7 @@ class PregelLoop: self._delta_channels_forced_snapshot.update( k for k in self._delta_channels_awaiting_fork_snapshot - if (ch := self.channels.get(k)) is not None and ch.is_available() + if k in self.channels ) # create new checkpoint channels_to_snapshot = ( diff --git a/libs/langgraph/tests/test_delta_channel_fork.py b/libs/langgraph/tests/test_delta_channel_fork.py index bb991e3f8..d53a86689 100644 --- a/libs/langgraph/tests/test_delta_channel_fork.py +++ b/libs/langgraph/tests/test_delta_channel_fork.py @@ -12,7 +12,8 @@ replay, so the plain channel is the oracle: after a fork the two must agree. Coverage: fork by ``invoke`` with new input (sync/async, all durabilities), fork off the checkpoint that predates the thread's first input (sync/async), -fork by ``update_state`` / ``aupdate_state``, and guards that neither an +fork by ``update_state`` / ``aupdate_state``, fork before the delta channel +ever had a value, and guards that neither an unaddressed run nor an unaddressed multi-superstep ``bulk_update_state`` departs from the normal ``snapshot_frequency`` cadence. """ @@ -46,6 +47,9 @@ def _append(current: list | None, writes: Sequence[Any]) -> list: class _State(TypedDict): log: Annotated[list, DeltaChannel(_append, snapshot_frequency=1000)] plain: Annotated[list, add] + # Written only by the tests that fork before `log` ever has a value, so the + # fork advances without touching the delta channel. + other: Annotated[list, add] def _build(checkpointer: BaseCheckpointSaver, tag: str) -> Any: @@ -65,6 +69,19 @@ def _build(checkpointer: BaseCheckpointSaver, tag: str) -> Any: return builder.compile(checkpointer=checkpointer) +def _build_without_delta_writes(checkpointer: BaseCheckpointSaver, tag: str) -> Any: + """Compile a graph whose node writes only ``other``, never the delta channel.""" + + def node(state: _State) -> dict: + return {"other": [f"{tag}-other"]} + + builder = StateGraph(_State) + builder.add_node("n", node) + builder.set_entry_point("n") + builder.set_finish_point("n") + return builder.compile(checkpointer=checkpointer) + + def _thread(thread_id: str) -> RunnableConfig: return {"configurable": {"thread_id": thread_id}} @@ -257,6 +274,81 @@ def test_unaddressed_run_keeps_snapshot_cadence( assert not _snapshotted_checkpoints(sync_checkpointer, config) +def test_fork_before_first_value_when_fork_never_writes_the_channel( + sync_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + """Seal the fork even when the channel has no value to snapshot. + + Forking before ``log`` was ever written leaves nothing to copy into the + fork's first checkpoint, so without a minted version and an empty blob the + boundary goes unrecorded and the walk runs into the base. The fork here + never writes ``log`` at all, so no later superstep can seal it either. + """ + config = _thread("t") + graph = _build(sync_checkpointer, "first") + graph.invoke(_input("in-1"), config, durability=durability) + + root = list(graph.get_state_history(config))[-1] + assert root.values["log"] == [] + + _build_without_delta_writes(sync_checkpointer, "third").invoke( + {"other": ["in-9"]}, _at(config, root), durability=durability + ) + + state = graph.get_state(config) + _assert_fork_is_clean(state, "in-1") + assert state.values["log"] == [] + + +async def test_afork_before_first_value_when_fork_never_writes_the_channel( + async_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + """Async twin of ``test_fork_before_first_value_when_fork_never_writes_the_channel``.""" + config = _thread("t") + graph = _build(async_checkpointer, "first") + await graph.ainvoke(_input("in-1"), config, durability=durability) + + root = [snapshot async for snapshot in graph.aget_state_history(config)][-1] + assert root.values["log"] == [] + + await _build_without_delta_writes(async_checkpointer, "third").ainvoke( + {"other": ["in-9"]}, _at(config, root), durability=durability + ) + + state = await graph.aget_state(config) + _assert_fork_is_clean(state, "in-1") + assert state.values["log"] == [] + + +def test_fork_before_first_value_by_bulk_update( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + """The fork's first superstep touches another key, the delta key comes later. + + The first checkpoint has to seal the boundary on its own. Deferring until + the superstep that finally writes ``log`` is too late, because that + superstep reconstructs through the unsealed checkpoint first. + """ + config = _thread("t") + graph = _build(sync_checkpointer, "first") + graph.invoke(_input("in-1"), config) + + root = list(graph.get_state_history(config))[-1] + assert root.values["log"] == [] + + forked = graph.bulk_update_state( + _at(config, root), + [ + [StateUpdate({"other": ["s1"]}, "n")], + [StateUpdate(_input("s2"), "n")], + ], + ) + + state = graph.get_state(forked) + _assert_fork_is_clean(state, "in-1") + assert state.values["log"] == ["s2"] + + @pytest.mark.parametrize("first_as_node", [INPUT, END, "__copy__"]) def test_fork_by_bulk_update_whose_first_superstep_skips_the_plan( sync_checkpointer: BaseCheckpointSaver, first_as_node: str