From 038f26472d65fd04a7b521762906cc60e44a00f5 Mon Sep 17 00:00:00 2001 From: Sydney Runkle Date: Tue, 28 Apr 2026 15:52:36 -0400 Subject: [PATCH] fix(checkpoint): prune preserves checkpoint_writes needed for DeltaChannel reconstruction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit InMemorySaver.prune now walks the parent chain from the latest checkpoint and keeps all ancestors whose writes are still needed for DeltaChannel reconstruction — i.e. ancestors with DELTA_SENTINEL blobs where no _DeltaSnapshot has been written yet in the kept ancestry. Once a _DeltaSnapshot blob exists for all sentinel channels (provided by snapshot_frequency), the walk terminates and older checkpoints plus their writes and blobs are safely deleted. With snapshot_frequency=None (pure delta), all intermediate checkpoints are retained since the full write history is required. Setting snapshot_frequency=N bounds the kept ancestry to N steps. Co-Authored-By: Claude Sonnet 4.6 (1M context) --- .../langgraph/checkpoint/memory/__init__.py | 93 +++++++++++ libs/checkpoint/tests/test_memory.py | 145 ++++++++++++++++++ 2 files changed, 238 insertions(+) diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index d9d704336..a45ad697d 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -525,6 +525,99 @@ class InMemorySaver( task_path, ) + def prune( + self, + thread_ids: Sequence[str], + *, + strategy: str = "keep_latest", + ) -> None: + """Prune checkpoints for the given threads. + + For DeltaChannel channels, a checkpoint is only deleted if the walk + from the latest checkpoint would not need to traverse it — i.e., a + `_DeltaSnapshot` blob exists in the kept ancestry that covers all + sentinel channels. Checkpoints that are still in the active walk + chain (because no snapshot has been taken yet, e.g. with + `snapshot_frequency=None`) are retained. + + Args: + thread_ids: Thread IDs to prune. + strategy: ``"keep_latest"`` keeps only the most recent checkpoint + per namespace; ``"delete"`` removes all checkpoints. + """ + for thread_id in thread_ids: + if strategy == "delete": + self.delete_thread(thread_id) + continue + + if strategy != "keep_latest": + raise ValueError( + f"Unknown pruning strategy {strategy!r}. " + "Expected 'keep_latest' or 'delete'." + ) + + for checkpoint_ns, ns_storage in list( + self.storage.get(thread_id, {}).items() + ): + if not ns_storage: + continue + + # Latest checkpoint (uuid6 IDs are lexicographically monotonic) + latest_id = max(ns_storage.keys()) + latest_data, _, _ = ns_storage[latest_id] + latest_cp = self.serde.loads_typed(latest_data) + + # Which channels in the latest checkpoint still have sentinels? + sentinel_channels: set[str] = set() + for ch, ver in latest_cp.get("channel_versions", {}).items(): + blob = self.blobs.get((thread_id, checkpoint_ns, ch, ver)) + if blob is not None and blob[0] == "delta": + sentinel_channels.add(ch) + + # Walk the parent chain to find the oldest ancestor still needed. + # We stop (and mark "safe to prune before here") when all + # sentinel channels are covered by a non-sentinel blob. + required_ids: set[str] = {latest_id} + if sentinel_channels: + _, _, parent_id = ns_storage[latest_id] + remaining = set(sentinel_channels) + while parent_id is not None and remaining: + entry = ns_storage.get(parent_id) + if entry is None: + break + required_ids.add(parent_id) + cp_data, _, grandparent_id = entry + cp = self.serde.loads_typed(cp_data) + resolved: set[str] = set() + for ch in remaining: + ver = cp.get("channel_versions", {}).get(ch) + if ver is None: + continue + blob = self.blobs.get((thread_id, checkpoint_ns, ch, ver)) + if blob is not None and blob[0] not in ("delta", "empty"): + resolved.add(ch) + remaining -= resolved + parent_id = grandparent_id + + # Delete everything outside the required set + for cp_id in list(ns_storage.keys()): + if cp_id in required_ids: + continue + cp_data, _, _ = ns_storage.pop(cp_id) + self.writes.pop((thread_id, checkpoint_ns, cp_id), None) + + # Clean up blobs no longer referenced by any kept checkpoint + live: set[tuple[str, str, str, Any]] = set() + for cp_data, _, _ in ns_storage.values(): + cp = self.serde.loads_typed(cp_data) + for ch, ver in cp.get("channel_versions", {}).items(): + live.add((thread_id, checkpoint_ns, ch, ver)) + for key in [ + k for k in self.blobs if k[:2] == (thread_id, checkpoint_ns) + ]: + if key not in live: + del self.blobs[key] + def delete_thread(self, thread_id: str) -> None: """Delete all checkpoints and writes associated with a thread ID. diff --git a/libs/checkpoint/tests/test_memory.py b/libs/checkpoint/tests/test_memory.py index 617b23f68..7c2e9974c 100644 --- a/libs/checkpoint/tests/test_memory.py +++ b/libs/checkpoint/tests/test_memory.py @@ -663,3 +663,148 @@ class TestPreDeltaBlobTerminator: assert "PRE-DELTA-WRITE" not in values # And the pending write at the target is never folded in. assert "PENDING-AT-TARGET" not in values + + +class TestInMemorySaverPrune: + """Tests for InMemorySaver.prune with DeltaChannel awareness.""" + + def _build_chain( + self, + saver: InMemorySaver, + thread_id: str, + ns: str, + channel: str, + n: int, + *, + snapshot_at: set[int] | None = None, + ) -> list[str]: + """Build a chain of n checkpoints with DELTA_SENTINEL blobs. + + If snapshot_at is provided, writes a _DeltaSnapshot blob at those steps. + Returns list of checkpoint IDs in order (oldest first). + """ + from langgraph.checkpoint.serde.types import _DeltaSnapshot + + serde = saver.serde + cp_ids = [] + parent_id = None + ver_base = "0000000000000000000000000000000{i}.0000000000000000" + + for i in range(n): + cp_id = f"cp{i:04d}" + ver = ver_base.format(i=i) + cp = empty_checkpoint() + cp["id"] = cp_id + cp["channel_versions"][channel] = ver + + if snapshot_at and i in snapshot_at: + blob = serde.dumps_typed(_DeltaSnapshot(value=[f"msg{i}"])) + else: + blob = serde.dumps_typed(DELTA_SENTINEL) + + saver.blobs[(thread_id, ns, channel, ver)] = blob + saver.storage[thread_id][ns][cp_id] = ( + serde.dumps_typed(cp), + serde.dumps_typed({}), + parent_id, + ) + # Add a dummy write for this checkpoint + saver.writes[(thread_id, ns, cp_id)][("task", i)] = ( + "task", + channel, + serde.dumps_typed(f"write{i}"), + "", + ) + cp_ids.append(cp_id) + parent_id = cp_id + + return cp_ids + + def test_prune_pure_delta_keeps_all(self) -> None: + """With no snapshots, all checkpoints are required for reconstruction.""" + saver = InMemorySaver() + thread_id, ns, channel = "t1", "", "messages" + cp_ids = self._build_chain(saver, thread_id, ns, channel, n=5) + + saver.prune([thread_id], strategy="keep_latest") + + # All checkpoints must be retained (walk needs the full chain) + remaining = set(saver.storage[thread_id][ns].keys()) + assert remaining == set(cp_ids) + + def test_prune_with_snapshot_removes_pre_snapshot_checkpoints(self) -> None: + """Checkpoints older than the nearest snapshot can be safely pruned.""" + saver = InMemorySaver() + thread_id, ns, channel = "t1", "", "messages" + # Snapshot at step 2; steps 3 and 4 are sentinels + cp_ids = self._build_chain(saver, thread_id, ns, channel, n=5, snapshot_at={2}) + + saver.prune([thread_id], strategy="keep_latest") + + remaining = set(saver.storage[thread_id][ns].keys()) + # cp0, cp1 (before snapshot) must be gone; cp2, cp3, cp4 must remain + assert cp_ids[0] not in remaining # pre-snapshot + assert cp_ids[1] not in remaining # pre-snapshot + assert cp_ids[2] in remaining # the snapshot itself + assert cp_ids[3] in remaining # sentinel after snapshot + assert cp_ids[4] in remaining # latest + + def test_prune_removes_orphaned_blobs(self) -> None: + """Blob entries for pruned checkpoints are cleaned up.""" + saver = InMemorySaver() + thread_id, ns, channel = "t1", "", "messages" + self._build_chain(saver, thread_id, ns, channel, n=4, snapshot_at={1}) + saver.prune([thread_id], strategy="keep_latest") + + # cp0 blob should be gone (pruned); cp1, cp2, cp3 blobs remain + assert ( + thread_id, + ns, + channel, + f"0000000000000000000000000000000{0}.0000000000000000", + ) not in saver.blobs + for i in range(1, 4): + ver = f"0000000000000000000000000000000{i}.0000000000000000" + assert (thread_id, ns, channel, ver) in saver.blobs + + def test_prune_removes_writes_for_pruned_checkpoints(self) -> None: + """checkpoint_writes for pruned checkpoints are deleted.""" + saver = InMemorySaver() + thread_id, ns, channel = "t1", "", "messages" + cp_ids = self._build_chain(saver, thread_id, ns, channel, n=4, snapshot_at={1}) + + saver.prune([thread_id], strategy="keep_latest") + + # writes for cp0 must be gone + assert (thread_id, ns, cp_ids[0]) not in saver.writes + # writes for cp1+ must remain (they're in the walk chain) + for cp_id in cp_ids[1:]: + assert (thread_id, ns, cp_id) in saver.writes + + def test_prune_delete_strategy_removes_everything(self) -> None: + """strategy='delete' removes all checkpoints for the thread.""" + saver = InMemorySaver() + thread_id, ns, channel = "t1", "", "messages" + self._build_chain(saver, thread_id, ns, channel, n=3) + + saver.prune([thread_id], strategy="delete") + + assert not saver.storage.get(thread_id, {}).get(ns) + assert not any(k[0] == thread_id for k in saver.writes) + assert not any(k[0] == thread_id for k in saver.blobs) + + def test_prune_non_delta_channel_always_pruneable(self) -> None: + """A channel with full snapshot blobs (no sentinels) allows full prune.""" + + saver = InMemorySaver() + thread_id, ns, channel = "t1", "", "messages" + # All snapshots, no sentinels + cp_ids = self._build_chain( + saver, thread_id, ns, channel, n=4, snapshot_at={0, 1, 2, 3} + ) + + saver.prune([thread_id], strategy="keep_latest") + + remaining = set(saver.storage[thread_id][ns].keys()) + # Only the latest checkpoint is needed (all blobs are snapshots) + assert remaining == {cp_ids[-1]}