optimize session_index

Author: zhilingluo

src/agentscope_runtime/engine/services/session_history/redis_session_history_service.py

@@ -91,8 +91,9 @@ async def health(self) -> bool:
     def _session_key(self, user_id: str, session_id: str):
         return f"session:{user_id}:{session_id}"
 
-    def _index_key(self, user_id: str):
-        return f"session_index:{user_id}"
+    def _session_pattern(self, user_id: str):
+        """Generate the pattern for scanning session keys for a user."""
+        return f"session:{user_id}:*"
 
     def _session_to_json(self, session: Session) -> str:
         return session.model_dump_json()
@@ -114,7 +115,6 @@ async def create_session(
         key = self._session_key(user_id, sid)
 
         await self._redis.set(key, self._session_to_json(session))
-        await self._redis.sadd(self._index_key(user_id), sid)
 
         # Set TTL for the session key if configured
         if self._ttl_seconds is not None:
@@ -130,15 +130,7 @@ async def get_session(
         key = self._session_key(user_id, session_id)
         session_json = await self._redis.get(key)
         if session_json is None:
-            session = Session(id=session_id, user_id=user_id)
-            await self._redis.set(key, self._session_to_json(session))
-            await self._redis.sadd(self._index_key(user_id), session_id)
-
-            # Set TTL for the session key if configured
-            if self._ttl_seconds is not None:
-                await self._redis.expire(key, self._ttl_seconds)
-
-            return session
+            return None
 
         session = self._session_from_json(session_json)
 
@@ -151,19 +143,33 @@ async def get_session(
     async def delete_session(self, user_id: str, session_id: str):
         key = self._session_key(user_id, session_id)
         await self._redis.delete(key)
-        await self._redis.srem(self._index_key(user_id), session_id)
 
     async def list_sessions(self, user_id: str) -> list[Session]:
-        idx_key = self._index_key(user_id)
-        session_ids = await self._redis.smembers(idx_key)
+        """List all sessions for a user by scanning session keys.
+
+        Uses SCAN to find all session:{user_id}:* keys. Expired sessions
+        naturally disappear as their keys expire, avoiding stale entries.
+        """
+        pattern = self._session_pattern(user_id)
         sessions = []
-        for sid in session_ids:
-            key = self._session_key(user_id, sid)
-            session_json = await self._redis.get(key)
-            if session_json:
-                session = self._session_from_json(session_json)
-                session.messages = []
-                sessions.append(session)
+        cursor = 0
+
+        while True:
+            cursor, keys = await self._redis.scan(
+                cursor,
+                match=pattern,
+                count=100,
+            )
+            for key in keys:
+                session_json = await self._redis.get(key)
+                if session_json:
+                    session = self._session_from_json(session_json)
+                    session.messages = []
+                    sessions.append(session)
+
+            if cursor == 0:
+                break
+
         return sessions
 
     async def append_message(
@@ -192,49 +198,54 @@ async def append_message(
         key = self._session_key(user_id, session_id)
 
         session_json = await self._redis.get(key)
-        if session_json:
-            stored_session = self._session_from_json(session_json)
-            stored_session.messages.extend(norm_message)
-
-            # Limit the number of messages per session to prevent memory issues
-            if self._max_messages_per_session is not None:
-                if (
-                    len(stored_session.messages)
-                    > self._max_messages_per_session
-                ):
-                    # Keep only the most recent messages
-                    stored_session.messages = stored_session.messages[
-                        -self._max_messages_per_session :
-                    ]
-
-            await self._redis.set(key, self._session_to_json(stored_session))
-            await self._redis.sadd(self._index_key(user_id), session_id)
-
-            # Set TTL for the session key if configured
-            if self._ttl_seconds is not None:
-                await self._redis.expire(key, self._ttl_seconds)
-        else:
-            print(
-                f"Warning: Session {session.id} not found in storage for "
-                f"append_message.",
+        if session_json is None:
+            raise RuntimeError(
+                f"Session {session_id} not found or has expired for user "
+                f"{user_id}. Previous memory/state has been lost. "
+                f"Please create a new session.",
             )
 
+        stored_session = self._session_from_json(session_json)
+        stored_session.messages.extend(norm_message)
+
+        # Limit the number of messages per session to prevent memory issues
+        if self._max_messages_per_session is not None:
+            if len(stored_session.messages) > self._max_messages_per_session:
+                # Keep only the most recent messages
+                stored_session.messages = stored_session.messages[
+                    -self._max_messages_per_session :
+                ]
+
+        await self._redis.set(key, self._session_to_json(stored_session))
+
+        # Set TTL for the session key if configured
+        if self._ttl_seconds is not None:
+            await self._redis.expire(key, self._ttl_seconds)
+
     async def delete_user_sessions(self, user_id: str) -> None:
         """
         Deletes all session history data for a specific user.
 
+        Uses SCAN to find all session keys for the user and deletes them.
+
         Args:
             user_id (str): The ID of the user whose session history data should
              be deleted
         """
         if not self._redis:
             raise RuntimeError("Redis connection is not available")
 
-        index_key = self._index_key(user_id)
-        session_ids = await self._redis.smembers(index_key)
+        pattern = self._session_pattern(user_id)
+        cursor = 0
 
-        for session_id in session_ids:
-            key = self._session_key(user_id, session_id)
-            await self._redis.delete(key)
+        while True:
+            cursor, keys = await self._redis.scan(
+                cursor,
+                match=pattern,
+                count=100,
+            )
+            if keys:
+                await self._redis.delete(*keys)
 
-        await self._redis.delete(index_key)
+            if cursor == 0:
+                break

tests/unit/test_redis_session_history_service.py

@@ -118,25 +118,19 @@ async def test_get_session(
     assert refetched_session is not None
     assert refetched_session.messages == []
 
-    # Test getting a non-existent session (should create a new one)
+    # Test getting a non-existent session (should return None)
     non_existent_session = await session_history_service.get_session(
         user_id,
         "non_existent_id",
     )
-    assert non_existent_session is not None
-    assert non_existent_session.id == "non_existent_id"
-    assert non_existent_session.user_id == user_id
-    assert non_existent_session.messages == []
+    assert non_existent_session is None
 
-    # Test getting a session for a different user (should create a new one)
+    # Test getting a session for a different user (should return None)
     other_user_session = await session_history_service.get_session(
         "other_user",
         created_session.id,
     )
-    assert other_user_session is not None
-    assert other_user_session.id == created_session.id
-    assert other_user_session.user_id == "other_user"
-    assert other_user_session.messages == []
+    assert other_user_session is None
 
 
 @pytest.mark.asyncio
@@ -156,17 +150,12 @@ async def test_delete_session(
 
     await session_history_service.delete_session(user_id, session.id)
 
-    # Ensure session is deleted - get_session will create a new empty session
+    # Ensure session is deleted - get_session should return None
     retrieved_session = await session_history_service.get_session(
         user_id,
         session.id,
     )
-    assert retrieved_session is not None
-    assert retrieved_session.id == session.id
-    assert retrieved_session.user_id == user_id
-    assert (
-        retrieved_session.messages == []
-    )  # Should be empty as it's a new session
+    assert retrieved_session is None
 
     # Test deleting a non-existent session (should not raise error)
     await session_history_service.delete_session(user_id, "non_existent_id")
@@ -275,26 +264,19 @@ async def test_append_message(
     for i, msg in enumerate(stored_session.messages[2:]):
         assert msg.content == messages3[i].get("content")
 
-    # Test appending to a non-existent session
+    # Test appending to a non-existent session (should raise RuntimeError)
     non_existent_session = Session(
         id="non_existent",
         user_id=user_id,
         messages=[],
     )
-    # This should not raise an error, but print a warning.
-    await session_history_service.append_message(
-        non_existent_session,
-        message1,
-    )
-    # get_session will create a new session, not the one we tried to append to
-    retrieved_session = await session_history_service.get_session(
-        user_id,
-        "non_existent",
-    )
-    assert retrieved_session is not None
-    assert (
-        retrieved_session.messages == []
-    )  # Empty as it's a newly created session
+    # This should raise a RuntimeError indicating
+    # the session is missing/expired
+    with pytest.raises(RuntimeError, match="not found or has expired"):
+        await session_history_service.append_message(
+            non_existent_session,
+            message1,
+        )
 
 
 @pytest.mark.asyncio
@@ -347,12 +329,11 @@ async def test_ttl_expiration():
         key_exists = await fake_redis.exists(key)
         assert key_exists == 0, "Key should be expired and deleted"
 
-        # Verify get_session creates a new session after expiry
+        # Verify get_session returns None after expiry
         retrieved_after_expiry = await service.get_session(user_id, session.id)
-        assert retrieved_after_expiry is not None
         assert (
-            retrieved_after_expiry.messages == []
-        ), "Session should be empty after expiry"
+            retrieved_after_expiry is None
+        ), "Session should return None after expiry"
 
     finally:
         await service.stop()