diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index 0b8593c3e..4cf1f80fc 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -509,106 +509,6 @@ 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 self.serde.loads_typed(blob) is DELTA_SENTINEL - ): - 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] != "empty" - and self.serde.loads_typed(blob) is not DELTA_SENTINEL - ): - 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 7c2e9974c..16de77aca 100644 --- a/libs/checkpoint/tests/test_memory.py +++ b/libs/checkpoint/tests/test_memory.py @@ -664,147 +664,3 @@ class TestPreDeltaBlobTerminator: # 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]}