fix: resolve HITL interrupt persistence in shallow savers (#133)

The shallow savers were prematurely cleaning up writes when a new
checkpoint was saved, causing Human-in-the-Loop interrupts to be
deleted before they could be consumed during resume.

Changes:
- Remove aggressive write cleanup from aput/put methods in shallow savers
- Fix registry key consistency using to_storage_safe_str() for checkpoint_ns
- Add comprehensive HITL integration tests (7 new tests)
- Update existing tests to reflect new write persistence behavior

Writes are now cleaned up only via delete_thread, TTL expiration,
or explicit overwrite - not during checkpoint transitions.

Author: bsbodden

langgraph/checkpoint/redis/ashallow.py

@@ -229,9 +229,6 @@ async def aput(
             # Create pipeline for all operations
             pipeline = self._redis.pipeline(transaction=False)
 
-            # Get the previous checkpoint ID to potentially clean up its writes
-            pipeline.json().get(checkpoint_key)
-
             # Set the new checkpoint data
             pipeline.json().set(checkpoint_key, "$", checkpoint_data)
 
@@ -240,41 +237,21 @@ async def aput(
                 ttl_seconds = int(self.ttl_config.get("default_ttl") * 60)
                 pipeline.expire(checkpoint_key, ttl_seconds)
 
-            # Execute pipeline to get prev data and set new data
-            results = await pipeline.execute()
-            prev_checkpoint_data = results[0]
-
-            # Check if we need to clean up old writes
-            prev_checkpoint_id = None
-            if prev_checkpoint_data and isinstance(prev_checkpoint_data, dict):
-                prev_checkpoint_id = prev_checkpoint_data.get("checkpoint_id")
-
-            # If checkpoint changed, clean up old writes in a second pipeline
-            if prev_checkpoint_id and prev_checkpoint_id != checkpoint["id"]:
-                thread_zset_key = f"write_keys_zset:{thread_id}:{checkpoint_ns}:shallow"
-
-                # Create cleanup pipeline
-                cleanup_pipeline = self._redis.pipeline(transaction=False)
-
-                # Get all existing write keys
-                cleanup_pipeline.zrange(thread_zset_key, 0, -1)
-
-                # Delete the registry
-                cleanup_pipeline.delete(thread_zset_key)
-
-                # Execute to get keys and delete registry
-                cleanup_results = await cleanup_pipeline.execute()
-                existing_write_keys = cleanup_results[0]
+            # Execute pipeline to set new checkpoint data
+            await pipeline.execute()
 
-                # If there are keys to delete, do it in another pipeline
-                if existing_write_keys:
-                    delete_pipeline = self._redis.pipeline(transaction=False)
-                    for old_key in existing_write_keys:
-                        old_key_str = (
-                            old_key.decode() if isinstance(old_key, bytes) else old_key
-                        )
-                        delete_pipeline.delete(old_key_str)
-                    await delete_pipeline.execute()
+            # NOTE: We intentionally do NOT clean up old writes here.
+            # In the HITL (Human-in-the-Loop) flow, interrupt writes are saved via
+            # put_writes BEFORE the new checkpoint is saved. If we clean up writes
+            # when the checkpoint changes, we would delete the interrupt writes
+            # before they can be consumed when resuming.
+            #
+            # Writes are cleaned up in the following scenarios:
+            # 1. When delete_thread is called
+            # 2. When TTL expires (if configured)
+            # 3. When put_writes is called again for the same task/idx (overwrites)
+            #
+            # See Issue #133 for details on this bug fix.
 
             return next_config
 
@@ -505,7 +482,11 @@ async def aput_writes(
                 writes_objects.append(write_obj)
 
             # Thread-level sorted set for write keys
-            thread_zset_key = f"write_keys_zset:{thread_id}:{checkpoint_ns}:shallow"
+            # Use to_storage_safe_str for consistent key naming
+            safe_checkpoint_ns = to_storage_safe_str(checkpoint_ns)
+            thread_zset_key = (
+                f"write_keys_zset:{thread_id}:{safe_checkpoint_ns}:shallow"
+            )
 
             # Collect all write keys
             write_keys = []
@@ -631,7 +612,9 @@ async def _aload_pending_writes(
             return []
 
         # Use thread-level sorted set
-        thread_zset_key = f"write_keys_zset:{thread_id}:{checkpoint_ns}:shallow"
+        # Use to_storage_safe_str for consistent key naming
+        safe_checkpoint_ns = to_storage_safe_str(checkpoint_ns)
+        thread_zset_key = f"write_keys_zset:{thread_id}:{safe_checkpoint_ns}:shallow"
 
         try:
             # Check if we have any writes in the thread sorted set

langgraph/checkpoint/redis/shallow.py

@@ -200,42 +200,29 @@ def put(
             thread_id, checkpoint_ns
         )
 
-        # Get the previous checkpoint ID to clean up its writes
-        prev_checkpoint_data = self._redis.json().get(checkpoint_key)
-        prev_checkpoint_id = None
-        if prev_checkpoint_data and isinstance(prev_checkpoint_data, dict):
-            prev_checkpoint_id = prev_checkpoint_data.get("checkpoint_id")
-
         with self._redis.pipeline(transaction=False) as pipeline:
             pipeline.json().set(checkpoint_key, "$", checkpoint_data)
 
-            # If checkpoint changed, clean up old writes
-            if prev_checkpoint_id and prev_checkpoint_id != checkpoint["id"]:
-                # Clean up writes from the previous checkpoint
-                thread_write_registry_key = (
-                    f"write_registry:{thread_id}:{checkpoint_ns}:shallow"
-                )
-
-                # Get all existing write keys and delete them
-                existing_write_keys = self._redis.zrange(
-                    thread_write_registry_key, 0, -1
-                )
-                for old_key in existing_write_keys:
-                    old_key_str = (
-                        old_key.decode() if isinstance(old_key, bytes) else old_key
-                    )
-                    pipeline.delete(old_key_str)
-
-                # Clear the registry
-                pipeline.delete(thread_write_registry_key)
-
             # Apply TTL if configured
             if self.ttl_config and "default_ttl" in self.ttl_config:
                 ttl_seconds = int(self.ttl_config.get("default_ttl") * 60)
                 pipeline.expire(checkpoint_key, ttl_seconds)
 
             pipeline.execute()
 
+        # NOTE: We intentionally do NOT clean up old writes here.
+        # In the HITL (Human-in-the-Loop) flow, interrupt writes are saved via
+        # put_writes BEFORE the new checkpoint is saved. If we clean up writes
+        # when the checkpoint changes, we would delete the interrupt writes
+        # before they can be consumed when resuming.
+        #
+        # Writes are cleaned up in the following scenarios:
+        # 1. When delete_thread is called
+        # 2. When TTL expires (if configured)
+        # 3. When put_writes is called again for the same task/idx (overwrites)
+        #
+        # See Issue #133 for details on this bug fix.
+
         return next_config
 
     def list(
@@ -501,8 +488,10 @@ def put_writes(
             writes_objects.append(write_obj)
 
         # THREAD-LEVEL REGISTRY: Only keep writes for the current checkpoint
+        # Use to_storage_safe_str for consistent key naming with delete_thread
+        safe_checkpoint_ns = to_storage_safe_str(checkpoint_ns)
         thread_write_registry_key = (
-            f"write_registry:{thread_id}:{checkpoint_ns}:shallow"
+            f"write_registry:{thread_id}:{safe_checkpoint_ns}:shallow"
         )
 
         # Collect all write keys
@@ -550,8 +539,10 @@ def _load_pending_writes(
 
         # Use thread-level registry that only contains current checkpoint writes
         # All writes belong to the current checkpoint
+        # Use to_storage_safe_str for consistent key naming with delete_thread
+        safe_checkpoint_ns = to_storage_safe_str(checkpoint_ns)
         thread_write_registry_key = (
-            f"write_registry:{thread_id}:{checkpoint_ns}:shallow"
+            f"write_registry:{thread_id}:{safe_checkpoint_ns}:shallow"
         )
 
         # Get all write keys from the thread's registry (already sorted by index)