fix: prevent blob and write accumulation in ShallowRedisSaver classes (#13)

  Add cleanup logic to AsyncShallowRedisSaver and ShallowRedisSaver to delete
  old blobs and writes when storing new checkpoints. This prevents memory bloat
  when using shallow savers, which should only keep the latest checkpoint state.
  Add comprehensive test to verify the fix works correctly.

Author: bsbodden

langgraph/checkpoint/redis/ashallow.py

@@ -172,7 +172,7 @@ async def aput(
         metadata: CheckpointMetadata,
         new_versions: ChannelVersions,
     ) -> RunnableConfig:
-        """Store only the latest checkpoint asynchronously."""
+        """Store only the latest checkpoint asynchronously and clean up old blobs."""
         configurable = config["configurable"].copy()
         thread_id = configurable.pop("thread_id")
         checkpoint_ns = configurable.pop("checkpoint_ns")
@@ -200,6 +200,10 @@ async def aput(
             checkpoint_data["source"] = metadata["source"]
             checkpoint_data["step"] = metadata["step"]
 
+        # Note: Need to keep track of the current versions to keep
+        current_channel_versions = new_versions.copy()
+        
+        # Store the new checkpoint, which replaces any existing one due to the shallow key
         await self.checkpoints_index.load(
             [checkpoint_data],
             keys=[
@@ -209,7 +213,39 @@ async def aput(
             ],
         )
 
-        # Store blob values
+        # Before storing the new blobs, clean up old ones that won't be needed
+        # - Get a list of all blob keys for this thread_id and checkpoint_ns
+        # - Then delete the ones that aren't in new_versions
+        cleanup_pipeline = self._redis.pipeline()
+        
+        # Get all blob keys for this thread/namespace
+        blob_key_pattern = AsyncShallowRedisSaver._make_shallow_redis_checkpoint_blob_key_pattern(
+            thread_id, checkpoint_ns
+        )
+        existing_blob_keys = await self._redis.keys(blob_key_pattern)
+        
+        # Process each existing blob key to determine if it should be kept or deleted
+        if existing_blob_keys:
+            for blob_key in existing_blob_keys:
+                key_parts = blob_key.decode().split(REDIS_KEY_SEPARATOR)
+                # The key format is checkpoint_blob:thread_id:checkpoint_ns:channel:version
+                if len(key_parts) >= 5:
+                    channel = key_parts[3]
+                    version = key_parts[4]
+                    
+                    # Only keep the blob if it's referenced by the current versions
+                    if (channel in current_channel_versions and 
+                        current_channel_versions[channel] == version):
+                        # This is a current version, keep it
+                        continue
+                    else:
+                        # This is an old version, delete it
+                        cleanup_pipeline.delete(blob_key)
+            
+            # Execute the cleanup
+            await cleanup_pipeline.execute()
+        
+        # Store the new blob values
         blobs = self._dump_blobs(
             thread_id,
             checkpoint_ns,
@@ -385,7 +421,7 @@ async def aput_writes(
         task_id: str,
         task_path: str = "",
     ) -> None:
-        """Store intermediate writes for the latest checkpoint.
+        """Store intermediate writes for the latest checkpoint and clean up old writes.
 
         Args:
             config (RunnableConfig): Configuration of the related checkpoint.
@@ -413,23 +449,48 @@ async def aput_writes(
                 "blob": blob,
             }
             writes_objects.append(write_obj)
-
-            upsert_case = all(w[0] in WRITES_IDX_MAP for w in writes)
-            for write_obj in writes_objects:
-                key = self._make_redis_checkpoint_writes_key(
-                    thread_id,
-                    checkpoint_ns,
-                    checkpoint_id,
-                    task_id,
-                    write_obj["idx"],
-                )
-                if upsert_case:
-                    tx = partial(_write_obj_tx, key=key, write_obj=write_obj)
-                    await self._redis.transaction(tx, key)
-                else:
-                    # Unlike AsyncRedisSaver, the shallow implementation always overwrites
-                    # the full object, so we don't check for key existence here.
-                    await self._redis.json().set(key, "$", write_obj)
+        
+        # First clean up old writes for this thread and namespace if they're for a different checkpoint_id
+        cleanup_pipeline = self._redis.pipeline()
+        
+        # Get all writes keys for this thread/namespace
+        writes_key_pattern = AsyncShallowRedisSaver._make_shallow_redis_checkpoint_writes_key_pattern(
+            thread_id, checkpoint_ns
+        )
+        existing_writes_keys = await self._redis.keys(writes_key_pattern)
+        
+        # Process each existing writes key to determine if it should be kept or deleted
+        if existing_writes_keys:
+            for write_key in existing_writes_keys:
+                key_parts = write_key.decode().split(REDIS_KEY_SEPARATOR)
+                # The key format is checkpoint_write:thread_id:checkpoint_ns:checkpoint_id:task_id:idx
+                if len(key_parts) >= 5:
+                    key_checkpoint_id = key_parts[3]
+                    
+                    # If the write is for a different checkpoint_id, delete it
+                    if key_checkpoint_id != checkpoint_id:
+                        cleanup_pipeline.delete(write_key)
+            
+            # Execute the cleanup
+            await cleanup_pipeline.execute()
+
+        # Now store the new writes
+        upsert_case = all(w[0] in WRITES_IDX_MAP for w in writes)
+        for write_obj in writes_objects:
+            key = self._make_redis_checkpoint_writes_key(
+                thread_id,
+                checkpoint_ns,
+                checkpoint_id,
+                task_id,
+                write_obj["idx"],
+            )
+            if upsert_case:
+                tx = partial(_write_obj_tx, key=key, write_obj=write_obj)
+                await self._redis.transaction(tx, key)
+            else:
+                # Unlike AsyncRedisSaver, the shallow implementation always overwrites
+                # the full object, so we don't check for key existence here.
+                await self._redis.json().set(key, "$", write_obj)
 
     async def aget_channel_values(
         self, thread_id: str, checkpoint_ns: str, checkpoint_id: str
@@ -622,4 +683,15 @@ def put_writes(
 
     @staticmethod
     def _make_shallow_redis_checkpoint_key(thread_id: str, checkpoint_ns: str) -> str:
+        """Create a key for shallow checkpoints using only thread_id and checkpoint_ns."""
         return REDIS_KEY_SEPARATOR.join([CHECKPOINT_PREFIX, thread_id, checkpoint_ns])
+        
+    @staticmethod
+    def _make_shallow_redis_checkpoint_blob_key_pattern(thread_id: str, checkpoint_ns: str) -> str:
+        """Create a pattern to match all blob keys for a thread and namespace."""
+        return REDIS_KEY_SEPARATOR.join([CHECKPOINT_BLOB_PREFIX, thread_id, checkpoint_ns]) + ":*"
+        
+    @staticmethod
+    def _make_shallow_redis_checkpoint_writes_key_pattern(thread_id: str, checkpoint_ns: str) -> str:
+        """Create a pattern to match all writes keys for a thread and namespace."""
+        return REDIS_KEY_SEPARATOR.join([CHECKPOINT_WRITE_PREFIX, thread_id, checkpoint_ns]) + ":*"

langgraph/checkpoint/redis/shallow.py

@@ -119,7 +119,7 @@ def put(
         metadata: CheckpointMetadata,
         new_versions: ChannelVersions,
     ) -> RunnableConfig:
-        """Store only the latest checkpoint."""
+        """Store only the latest checkpoint and clean up old blobs."""
         configurable = config["configurable"].copy()
         thread_id = configurable.pop("thread_id")
         checkpoint_ns = configurable.pop("checkpoint_ns")
@@ -146,6 +146,9 @@ def put(
         if all(key in metadata for key in ["source", "step"]):
             checkpoint_data["source"] = metadata["source"]
             checkpoint_data["step"] = metadata["step"]
+            
+        # Note: Need to keep track of the current versions to keep
+        current_channel_versions = new_versions.copy()
 
         self.checkpoints_index.load(
             [checkpoint_data],
@@ -155,6 +158,38 @@ def put(
                 )
             ],
         )
+        
+        # Before storing the new blobs, clean up old ones that won't be needed
+        # - Get a list of all blob keys for this thread_id and checkpoint_ns
+        # - Then delete the ones that aren't in new_versions
+        cleanup_pipeline = self._redis.json().pipeline(transaction=False)
+        
+        # Get all blob keys for this thread/namespace
+        blob_key_pattern = ShallowRedisSaver._make_shallow_redis_checkpoint_blob_key_pattern(
+            thread_id, checkpoint_ns
+        )
+        existing_blob_keys = self._redis.keys(blob_key_pattern)
+        
+        # Process each existing blob key to determine if it should be kept or deleted
+        if existing_blob_keys:
+            for blob_key in existing_blob_keys:
+                key_parts = blob_key.decode().split(REDIS_KEY_SEPARATOR)
+                # The key format is checkpoint_blob:thread_id:checkpoint_ns:channel:version
+                if len(key_parts) >= 5:
+                    channel = key_parts[3]
+                    version = key_parts[4]
+                    
+                    # Only keep the blob if it's referenced by the current versions
+                    if (channel in current_channel_versions and 
+                        current_channel_versions[channel] == version):
+                        # This is a current version, keep it
+                        continue
+                    else:
+                        # This is an old version, delete it
+                        cleanup_pipeline.delete(blob_key)
+            
+            # Execute the cleanup
+            cleanup_pipeline.execute()
 
         # Store blob values
         blobs = self._dump_blobs(
@@ -408,7 +443,7 @@ def put_writes(
         task_id: str,
         task_path: str = "",
     ) -> None:
-        """Store intermediate writes linked to a checkpoint.
+        """Store intermediate writes linked to a checkpoint and clean up old writes.
 
         Args:
             config: Configuration of the related checkpoint.
@@ -436,6 +471,30 @@ def put_writes(
                 "blob": blob,
             }
             writes_objects.append(write_obj)
+            
+        # First clean up old writes for this thread and namespace if they're for a different checkpoint_id
+        cleanup_pipeline = self._redis.json().pipeline(transaction=False)
+        
+        # Get all writes keys for this thread/namespace
+        writes_key_pattern = ShallowRedisSaver._make_shallow_redis_checkpoint_writes_key_pattern(
+            thread_id, checkpoint_ns
+        )
+        existing_writes_keys = self._redis.keys(writes_key_pattern)
+        
+        # Process each existing writes key to determine if it should be kept or deleted
+        if existing_writes_keys:
+            for write_key in existing_writes_keys:
+                key_parts = write_key.decode().split(REDIS_KEY_SEPARATOR)
+                # The key format is checkpoint_write:thread_id:checkpoint_ns:checkpoint_id:task_id:idx
+                if len(key_parts) >= 5:
+                    key_checkpoint_id = key_parts[3]
+                    
+                    # If the write is for a different checkpoint_id, delete it
+                    if key_checkpoint_id != checkpoint_id:
+                        cleanup_pipeline.delete(write_key)
+            
+            # Execute the cleanup
+            cleanup_pipeline.execute()
 
         # For each write, check existence and then perform appropriate operation
         with self._redis.json().pipeline(transaction=False) as pipeline:
@@ -470,18 +529,25 @@ def _dump_blobs(
         values: dict[str, Any],
         versions: ChannelVersions,
     ) -> List[Tuple[str, dict[str, Any]]]:
+        """Convert blob data for Redis storage.
+        
+        In the shallow implementation, we use the version in the key to allow
+        storing multiple versions without conflicts and to facilitate cleanup.
+        """
         if not versions:
             return []
 
         return [
             (
-                ShallowRedisSaver._make_shallow_redis_checkpoint_blob_key(
-                    thread_id, checkpoint_ns, k
+                # Use the base Redis checkpoint blob key to include version, enabling version tracking
+                BaseRedisSaver._make_redis_checkpoint_blob_key(
+                    thread_id, checkpoint_ns, k, ver
                 ),
                 {
                     "thread_id": thread_id,
                     "checkpoint_ns": checkpoint_ns,
                     "channel": k,
+                    "version": ver, # Include version in the data as well
                     "type": (
                         self._get_type_and_blob(values[k])[0]
                         if k in values
@@ -581,12 +647,24 @@ def _load_pending_sends(
 
     @staticmethod
     def _make_shallow_redis_checkpoint_key(thread_id: str, checkpoint_ns: str) -> str:
+        """Create a key for shallow checkpoints using only thread_id and checkpoint_ns."""
         return REDIS_KEY_SEPARATOR.join([CHECKPOINT_PREFIX, thread_id, checkpoint_ns])
 
     @staticmethod
     def _make_shallow_redis_checkpoint_blob_key(
         thread_id: str, checkpoint_ns: str, channel: str
     ) -> str:
+        """Create a key for a blob in a shallow checkpoint."""
         return REDIS_KEY_SEPARATOR.join(
             [CHECKPOINT_BLOB_PREFIX, thread_id, checkpoint_ns, channel]
         )
+        
+    @staticmethod
+    def _make_shallow_redis_checkpoint_blob_key_pattern(thread_id: str, checkpoint_ns: str) -> str:
+        """Create a pattern to match all blob keys for a thread and namespace."""
+        return REDIS_KEY_SEPARATOR.join([CHECKPOINT_BLOB_PREFIX, thread_id, checkpoint_ns]) + ":*"
+        
+    @staticmethod
+    def _make_shallow_redis_checkpoint_writes_key_pattern(thread_id: str, checkpoint_ns: str) -> str:
+        """Create a pattern to match all writes keys for a thread and namespace."""
+        return REDIS_KEY_SEPARATOR.join([CHECKPOINT_WRITE_PREFIX, thread_id, checkpoint_ns]) + ":*"