Test fixes for ttls

Author: abrookins

.github/workflows/python-tests.yml

@@ -35,14 +35,14 @@ jobs:
       fail-fast: false
       matrix:
         python-version: [3.12]
-        redis-version: ['6.2.6-v9', '8.0.3', 'latest']
+        redis-version: ['redis-stack:6.2.6-v9', 'redis:8.0.3', 'redis:latest']
 
     steps:
     - uses: actions/checkout@v3
 
     - name: Set Redis image name
       run: |
-          echo "REDIS_IMAGE=redis:${{ matrix.redis-version }}" >> $GITHUB_ENV
+          echo "REDIS_IMAGE=${{ matrix.redis-version }}" >> $GITHUB_ENV
 
     - name: Set up Python
       uses: actions/setup-python@v4

agent_memory_server/working_memory.py

@@ -188,23 +188,24 @@ async def set_working_memory(
     }
 
     try:
-        if working_memory.ttl_seconds:
+        if working_memory.ttl_seconds is not None:
             # Store with TTL
             await redis_client.setex(
                 key,
                 working_memory.ttl_seconds,
-                json.dumps(
-                    data, default=json_datetime_handler
-                ),  # Add custom handler for any remaining datetime objects
+                json.dumps(data, default=json_datetime_handler),
+            )
+            logger.info(
+                f"Set working memory for session {working_memory.session_id} with TTL {working_memory.ttl_seconds}s"
             )
         else:
             await redis_client.set(
                 key,
                 json.dumps(data, default=json_datetime_handler),
             )
-        logger.info(
-            f"Set working memory for session {working_memory.session_id} with TTL {working_memory.ttl_seconds}s"
-        )
+            logger.info(
+                f"Set working memory for session {working_memory.session_id} with no TTL"
+            )
     except Exception as e:
         logger.error(
             f"Error setting working memory for session {working_memory.session_id}: {e}"

tests/test_working_memory.py

@@ -1,9 +1,12 @@
 """Tests for working memory functionality."""
 
+import asyncio
+
 import pytest
 from pydantic import ValidationError
 
 from agent_memory_server.models import MemoryRecord, MemoryTypeEnum, WorkingMemory
+from agent_memory_server.utils.keys import Keys
 from agent_memory_server.working_memory import (
     delete_working_memory,
     get_working_memory,
@@ -60,6 +63,7 @@ async def test_set_and_get_working_memory(self, async_redis_client):
         assert retrieved_mem.memories[0].id == "client-1"
         assert retrieved_mem.memories[1].text == "User is working on a Python project"
         assert retrieved_mem.memories[1].id == "client-2"
+        assert retrieved_mem.ttl_seconds == 1800  # Verify TTL is preserved
 
     @pytest.mark.asyncio
     async def test_get_nonexistent_working_memory(self, async_redis_client):
@@ -155,3 +159,183 @@ async def test_working_memory_validation(self, async_redis_client):
         assert retrieved is not None
         assert len(retrieved.memories) == 1
         assert retrieved.memories[0].id == "test-memory-1"
+
+    @pytest.mark.asyncio
+    async def test_working_memory_ttl_none(self, async_redis_client):
+        """Test working memory without TTL (persistent)"""
+        session_id = "test-session-no-ttl"
+        namespace = "test-namespace"
+
+        memories = [
+            MemoryRecord(
+                text="Persistent memory",
+                id="persistent-1",
+                memory_type=MemoryTypeEnum.SEMANTIC,
+            ),
+        ]
+
+        working_mem = WorkingMemory(
+            memories=memories,
+            session_id=session_id,
+            namespace=namespace,
+            ttl_seconds=None,  # No TTL - should be persistent
+        )
+
+        await set_working_memory(working_mem, redis_client=async_redis_client)
+
+        # Get working memory and verify TTL is None
+        retrieved_mem = await get_working_memory(
+            session_id=session_id,
+            namespace=namespace,
+            redis_client=async_redis_client,
+        )
+
+        assert retrieved_mem is not None
+        assert retrieved_mem.ttl_seconds is None
+
+        # Verify the Redis key has no TTL set (-1 means no TTL)
+        key = Keys.working_memory_key(
+            session_id=session_id,
+            namespace=namespace,
+        )
+        ttl = await async_redis_client.ttl(key)
+        assert ttl == -1  # No TTL set
+
+    @pytest.mark.asyncio
+    async def test_working_memory_ttl_set(self, async_redis_client):
+        """Test working memory with TTL set"""
+        session_id = "test-session-with-ttl"
+        namespace = "test-namespace"
+
+        memories = [
+            MemoryRecord(
+                text="Memory with TTL",
+                id="ttl-memory-1",
+                memory_type=MemoryTypeEnum.SEMANTIC,
+            ),
+        ]
+
+        ttl_seconds = 60  # 1 minute
+        working_mem = WorkingMemory(
+            memories=memories,
+            session_id=session_id,
+            namespace=namespace,
+            ttl_seconds=ttl_seconds,
+        )
+
+        await set_working_memory(working_mem, redis_client=async_redis_client)
+
+        # Get working memory and verify TTL is preserved
+        retrieved_mem = await get_working_memory(
+            session_id=session_id,
+            namespace=namespace,
+            redis_client=async_redis_client,
+        )
+
+        assert retrieved_mem is not None
+        assert retrieved_mem.ttl_seconds == ttl_seconds
+
+        # Verify the Redis key has TTL set (should be <= 60 seconds)
+        key = Keys.working_memory_key(
+            session_id=session_id,
+            namespace=namespace,
+        )
+        ttl = await async_redis_client.ttl(key)
+        assert 0 < ttl <= ttl_seconds
+
+    @pytest.mark.asyncio
+    async def test_working_memory_ttl_expiration(self, async_redis_client):
+        """Test working memory expires after TTL"""
+        session_id = "test-session-expire"
+        namespace = "test-namespace"
+
+        memories = [
+            MemoryRecord(
+                text="Memory that expires",
+                id="expire-memory-1",
+                memory_type=MemoryTypeEnum.SEMANTIC,
+            ),
+        ]
+
+        ttl_seconds = 1  # 1 second
+        working_mem = WorkingMemory(
+            memories=memories,
+            session_id=session_id,
+            namespace=namespace,
+            ttl_seconds=ttl_seconds,
+        )
+
+        await set_working_memory(working_mem, redis_client=async_redis_client)
+
+        # Verify it exists immediately
+        retrieved_mem = await get_working_memory(
+            session_id=session_id,
+            namespace=namespace,
+            redis_client=async_redis_client,
+        )
+        assert retrieved_mem is not None
+
+        # Wait for TTL to expire
+        await asyncio.sleep(1.1)
+
+        # Verify it's gone after TTL
+        retrieved_mem = await get_working_memory(
+            session_id=session_id,
+            namespace=namespace,
+            redis_client=async_redis_client,
+        )
+        assert retrieved_mem is None
+
+    @pytest.mark.asyncio
+    async def test_working_memory_ttl_update_preserves_ttl(self, async_redis_client):
+        """Test that updating working memory preserves TTL"""
+        session_id = "test-session-update-ttl"
+        namespace = "test-namespace"
+
+        memories = [
+            MemoryRecord(
+                text="Original memory",
+                id="original-memory-1",
+                memory_type=MemoryTypeEnum.SEMANTIC,
+            ),
+        ]
+
+        ttl_seconds = 120  # 2 minutes
+        working_mem = WorkingMemory(
+            memories=memories,
+            session_id=session_id,
+            namespace=namespace,
+            ttl_seconds=ttl_seconds,
+        )
+
+        await set_working_memory(working_mem, redis_client=async_redis_client)
+
+        # Update the working memory
+        working_mem.memories.append(
+            MemoryRecord(
+                text="Updated memory",
+                id="updated-memory-1",
+                memory_type=MemoryTypeEnum.SEMANTIC,
+            )
+        )
+
+        await set_working_memory(working_mem, redis_client=async_redis_client)
+
+        # Get updated working memory and verify TTL is preserved
+        retrieved_mem = await get_working_memory(
+            session_id=session_id,
+            namespace=namespace,
+            redis_client=async_redis_client,
+        )
+
+        assert retrieved_mem is not None
+        assert retrieved_mem.ttl_seconds == ttl_seconds
+        assert len(retrieved_mem.memories) == 2
+
+        # Verify the Redis key still has TTL set
+        key = Keys.working_memory_key(
+            session_id=session_id,
+            namespace=namespace,
+        )
+        ttl = await async_redis_client.ttl(key)
+        assert 0 < ttl <= ttl_seconds