keep TTL during migration

both cli and lazy migration
test

Author: fcenedes

agent_memory_server/cli.py

@@ -188,33 +188,39 @@ async def run_migration():
         for batch_start in range(0, len(string_keys), batch_size):
             batch_keys = string_keys[batch_start : batch_start + batch_size]
 
-            # First, read all string data in a pipeline
+            # Read all string data and TTLs in a pipeline
             read_pipe = redis.pipeline()
             for key in batch_keys:
                 read_pipe.get(key)
-            string_data_list = await read_pipe.execute()
+                read_pipe.ttl(key)
+            results = await read_pipe.execute()
+
+            # Parse results (alternating: data, ttl, data, ttl, ...)
+            migrations = []  # List of (key, data, ttl) tuples
+            for i, key in enumerate(batch_keys):
+                string_data = results[i * 2]
+                ttl = results[i * 2 + 1]
 
-            # Parse and prepare migration
-            migrations = []  # List of (key, data) tuples
-            for key, string_data in zip(batch_keys, string_data_list, strict=False):
                 if string_data is None:
                     continue
 
                 try:
                     if isinstance(string_data, bytes):
                         string_data = string_data.decode("utf-8")
                     data = json_module.loads(string_data)
-                    migrations.append((key, data))
+                    migrations.append((key, data, ttl))
                 except Exception as e:
                     errors += 1
                     logger.error(f"Failed to parse key {key}: {e}")
 
-            # Execute migrations in a pipeline (delete + json.set)
+            # Execute migrations in a pipeline (delete + json.set + expire if needed)
             if migrations:
                 write_pipe = redis.pipeline()
-                for key, data in migrations:
+                for key, data, ttl in migrations:
                     write_pipe.delete(key)
                     write_pipe.json().set(key, "$", data)
+                    if ttl > 0:
+                        write_pipe.expire(key, ttl)
 
                 try:
                     await write_pipe.execute()
@@ -224,10 +230,12 @@ async def run_migration():
                     logger.warning(
                         f"Batch migration failed, retrying individually: {e}"
                     )
-                    for key, data in migrations:
+                    for key, data, ttl in migrations:
                         try:
                             await redis.delete(key)
                             await redis.json().set(key, "$", data)
+                            if ttl > 0:
+                                await redis.expire(key, ttl)
                             migrated += 1
                         except Exception as e2:
                             errors += 1

agent_memory_server/working_memory.py

@@ -181,13 +181,17 @@ async def _migrate_string_to_json(
         logger.info(f"Migrating working memory key {key} from string to JSON format")
 
         # Atomically migrate the key from string to JSON using a Lua script
-        # The script: if key is string, get value, delete, set as JSON; else do nothing
+        # The script: get TTL, get value, delete, set as JSON, restore TTL if > 0
         lua_script = """
         local key = KEYS[1]
         if redis.call('TYPE', key).ok == 'string' then
+            local ttl = redis.call('TTL', key)
             local val = redis.call('GET', key)
             redis.call('DEL', key)
             redis.call('JSON.SET', key, '$', ARGV[1])
+            if ttl > 0 then
+                redis.call('EXPIRE', key, ttl)
+            end
             return val
         else
             return nil
@@ -196,10 +200,6 @@ async def _migrate_string_to_json(
         # Pass the JSON string as ARGV[1]
         await redis_client.eval(lua_script, 1, key, json.dumps(data))
 
-        # Preserve TTL if it was set
-        # Note: TTL is lost during migration since we deleted the key
-        # The next set_working_memory call will restore it if configured
-
         logger.info(f"Successfully migrated working memory key {key} to JSON format")
 
         # Decrement the counter (O(1) instead of re-scanning all keys)

tests/test_working_memory.py

@@ -435,6 +435,67 @@ async def test_backward_compatibility_string_to_json_migration(
         assert retrieved_again is not None
         assert retrieved_again.session_id == session_id
 
+    @pytest.mark.asyncio
+    async def test_migration_preserves_ttl(self, async_redis_client):
+        """Test that TTL is preserved when migrating from string to JSON format."""
+        import json
+
+        from agent_memory_server.working_memory import reset_migration_status
+
+        # Reset migration status to ensure lazy migration is active
+        reset_migration_status()
+
+        session_id = "test-ttl-migration-session"
+        namespace = "test-namespace"
+        ttl_seconds = 3600  # 1 hour
+
+        # Create old-format data with TTL
+        old_format_data = {
+            "messages": [],
+            "memories": [],
+            "session_id": session_id,
+            "namespace": namespace,
+            "context": None,
+            "user_id": None,
+            "tokens": 0,
+            "ttl_seconds": ttl_seconds,
+            "data": {},
+            "long_term_memory_strategy": {"strategy": "discrete"},
+            "last_accessed": 1704067200,
+            "created_at": 1704067200,
+            "updated_at": 1704067200,
+        }
+
+        # Store as old string format with TTL
+        key = Keys.working_memory_key(session_id=session_id, namespace=namespace)
+        await async_redis_client.set(key, json.dumps(old_format_data), ex=ttl_seconds)
+
+        # Verify TTL is set
+        ttl_before = await async_redis_client.ttl(key)
+        assert ttl_before > 0
+        assert ttl_before <= ttl_seconds
+
+        # Trigger migration by reading
+        retrieved_mem = await get_working_memory(
+            session_id=session_id,
+            namespace=namespace,
+            redis_client=async_redis_client,
+        )
+        assert retrieved_mem is not None
+
+        # Verify key was migrated to JSON
+        key_type = await async_redis_client.type(key)
+        if isinstance(key_type, bytes):
+            key_type = key_type.decode("utf-8")
+        assert key_type == "ReJSON-RL"
+
+        # Verify TTL was preserved
+        ttl_after = await async_redis_client.ttl(key)
+        assert ttl_after > 0
+        # TTL should be close to original (within a few seconds of test execution)
+        assert ttl_after <= ttl_seconds
+        assert ttl_after >= ttl_before - 5  # Allow 5 seconds for test execution
+
     @pytest.mark.asyncio
     async def test_check_and_set_migration_status_with_no_keys(
         self, async_redis_client