feat: remain 2000 session in memory using LRU (to remain old active sessions) and remain 20 records per session with FIFO

dori · dori · commit 074bd31f22b5 · 2025-09-07T16:45:44.000+03:00
diff --git a/src/mcp_as_a_judge/constants.py b/src/mcp_as_a_judge/constants.py
@@ -14,4 +14,5 @@
 # Database Configuration
 DATABASE_URL = "sqlite://:memory:"
 MAX_SESSION_RECORDS = 20  # Maximum records to keep per session (FIFO)
-RECORD_RETENTION_DAYS = 1
+MAX_TOTAL_SESSIONS = 2000  # Maximum total sessions to keep (LRU cleanup)
+RECORD_RETENTION_DAYS = 1  # Optional time-based cleanup (fallback)
diff --git a/src/mcp_as_a_judge/db/cleanup_service.py b/src/mcp_as_a_judge/db/cleanup_service.py
@@ -7,10 +7,10 @@
 
 from datetime import datetime, timedelta
 
-from sqlalchemy import Engine
+from sqlalchemy import Engine, func
 from sqlmodel import Session, select
 
-from mcp_as_a_judge.constants import RECORD_RETENTION_DAYS
+from mcp_as_a_judge.constants import MAX_TOTAL_SESSIONS, RECORD_RETENTION_DAYS
 from mcp_as_a_judge.db.interface import ConversationRecord
 from mcp_as_a_judge.logging_config import get_logger
 
@@ -22,8 +22,16 @@ class ConversationCleanupService:
     """
     Service for cleaning up old conversation history records.
 
-    Handles time-based cleanup: Removes records older than retention period.
-    Note: LRU cleanup is handled by the SQLite provider during save operations.
+    Implements session-based LRU cleanup strategy:
+    - Maintains max 2000 sessions by removing least recently used sessions
+    - Runs once per day to avoid performance overhead
+
+    LRU vs FIFO for Better User Experience:
+    - LRU (Least Recently Used): Keeps sessions that users are actively using, even if they're old
+    - FIFO (First In, First Out): Would remove oldest sessions regardless of recent activity
+    - LRU provides better UX because active conversations are preserved longer
+
+    Note: Per-session FIFO cleanup (max 20 records) is handled by the SQLite provider.
     """
 
     def __init__(self, engine: Engine) -> None:
@@ -34,8 +42,10 @@ def __init__(self, engine: Engine) -> None:
             engine: SQLAlchemy engine for database operations
         """
         self.engine = engine
+        self.max_total_sessions = MAX_TOTAL_SESSIONS
         self.retention_days = RECORD_RETENTION_DAYS
         self.last_cleanup_time = datetime.utcnow()
+        self.last_session_cleanup_time = datetime.utcnow()
 
     def cleanup_old_records(self) -> int:
         """
@@ -79,3 +89,136 @@ def cleanup_old_records(self) -> int:
                 f"🧹 Daily cleanup: Deleted {old_count} records older than {self.retention_days} days"
             )
             return old_count
+
+    def get_session_count(self) -> int:
+        """
+        Get the total number of unique sessions in the database.
+
+        Returns:
+            Number of unique sessions
+        """
+        with Session(self.engine) as session:
+            # Count distinct session_ids
+            count_stmt = select(func.count(func.distinct(ConversationRecord.session_id)))
+            result = session.exec(count_stmt).first()
+            return result or 0
+
+    def get_least_recently_used_sessions(self, limit: int) -> list[str]:
+        """
+        Get session IDs of the least recently used sessions.
+
+        Uses LRU strategy: finds sessions with the oldest "last activity" timestamp.
+        Last activity = MAX(timestamp) for each session (most recent record in session).
+
+        Args:
+            limit: Number of session IDs to return
+
+        Returns:
+            List of session IDs ordered by last activity (oldest first)
+        """
+        with Session(self.engine) as session:
+            # Find sessions with oldest last activity (LRU)
+            # GROUP BY session_id, ORDER BY MAX(timestamp) ASC to get least recently used
+            lru_stmt = (
+                select(
+                    ConversationRecord.session_id,
+                    func.max(ConversationRecord.timestamp).label("last_activity")
+                )
+                .group_by(ConversationRecord.session_id)
+                .order_by(func.max(ConversationRecord.timestamp).asc())
+                .limit(limit)
+            )
+
+            results = session.exec(lru_stmt).all()
+            return [result.session_id for result in results]
+
+    def delete_sessions(self, session_ids: list[str]) -> int:
+        """
+        Bulk delete all records for the given session IDs.
+
+        Args:
+            session_ids: List of session IDs to delete
+
+        Returns:
+            Number of records deleted
+        """
+        if not session_ids:
+            return 0
+
+        with Session(self.engine) as session:
+            # Count records before deletion for logging
+            count_stmt = select(ConversationRecord).where(
+                ConversationRecord.session_id.in_(session_ids)
+            )
+            records_to_delete = session.exec(count_stmt).all()
+            delete_count = len(records_to_delete)
+
+            # Bulk delete all records for these sessions
+            for record in records_to_delete:
+                session.delete(record)
+
+            session.commit()
+
+            logger.info(
+                f"🗑️ Deleted {delete_count} records from {len(session_ids)} sessions: "
+                f"{', '.join(session_ids[:3])}{'...' if len(session_ids) > 3 else ''}"
+            )
+
+            return delete_count
+
+    def cleanup_excess_sessions(self) -> int:
+        """
+        Remove least recently used sessions when total sessions exceed MAX_TOTAL_SESSIONS.
+
+        This implements LRU (Least Recently Used) cleanup strategy:
+        - Keeps sessions that users are actively using (better UX than FIFO)
+        - Only runs once per day to avoid excessive cleanup operations
+        - During the day, session count can exceed limit (e.g., 5000 sessions is not a memory issue)
+        - Daily cleanup brings it back to the target limit (2000 sessions)
+        - Removes entire sessions (all records for those session_ids)
+
+        Returns:
+            Number of records deleted
+        """
+        # Only run session cleanup once per day
+        if (datetime.utcnow() - self.last_session_cleanup_time).days < 1:
+            return 0
+
+        current_session_count = self.get_session_count()
+
+        if current_session_count <= self.max_total_sessions:
+            logger.info(
+                f"🧹 Daily session LRU cleanup: {current_session_count} sessions "
+                f"(max: {self.max_total_sessions}) - no cleanup needed"
+            )
+            self.last_session_cleanup_time = datetime.utcnow()
+            return 0
+
+        # Calculate how many sessions to remove
+        sessions_to_remove = current_session_count - self.max_total_sessions
+
+        logger.info(
+            f"🧹 Daily session LRU cleanup: {current_session_count} sessions exceeds limit "
+            f"({self.max_total_sessions}), removing {sessions_to_remove} least recently used sessions"
+        )
+
+        # Get least recently used sessions
+        lru_session_ids = self.get_least_recently_used_sessions(sessions_to_remove)
+
+        if not lru_session_ids:
+            logger.warning("🧹 No sessions found for LRU cleanup")
+            self.last_session_cleanup_time = datetime.utcnow()
+            return 0
+
+        # Delete all records for these sessions
+        deleted_count = self.delete_sessions(lru_session_ids)
+
+        # Reset cleanup tracking
+        self.last_session_cleanup_time = datetime.utcnow()
+
+        logger.info(
+            f"✅ Daily session LRU cleanup completed: removed {sessions_to_remove} sessions, "
+            f"deleted {deleted_count} records"
+        )
+
+        return deleted_count
diff --git a/src/mcp_as_a_judge/db/db_config.py b/src/mcp_as_a_judge/db/db_config.py
@@ -8,6 +8,7 @@
 from mcp_as_a_judge.constants import (
     DATABASE_URL,
     MAX_SESSION_RECORDS,
+    MAX_TOTAL_SESSIONS,
     RECORD_RETENTION_DAYS,
 )
 
@@ -61,6 +62,7 @@ class DatabaseConfig:
     def __init__(self) -> None:
         self.url = DATABASE_URL
         self.max_session_records = MAX_SESSION_RECORDS
+        self.max_total_sessions = MAX_TOTAL_SESSIONS
         self.record_retention_days = RECORD_RETENTION_DAYS
 
 
diff --git a/src/mcp_as_a_judge/db/providers/sqlite_provider.py b/src/mcp_as_a_judge/db/providers/sqlite_provider.py
@@ -29,8 +29,9 @@ class SQLiteProvider(ConversationHistoryDB):
     Features:
     - SQLModel with SQLAlchemy for type safety
     - In-memory or file-based SQLite storage
-    - LRU cleanup per session
-    - Time-based cleanup (configurable retention)
+    - Two-level cleanup strategy:
+      1. Daily session-based LRU cleanup (max 2000 sessions, removes least recently used)
+      2. Per-session FIFO cleanup (max 20 records per session, runs on every save)
     - Session-based conversation retrieval
     """
 
@@ -58,7 +59,9 @@ def __init__(self, max_session_records: int = 20, url: str = "") -> None:
 
         logger.info(
             f"🗄️ SQLModel SQLite provider initialized: {connection_string}, "
-            f"max_records={max_session_records}, retention_days={self._cleanup_service.retention_days}"
+            f"max_records_per_session={max_session_records}, "
+            f"max_total_sessions={self._cleanup_service.max_total_sessions}, "
+            f"retention_days={self._cleanup_service.retention_days}"
         )
 
     def _parse_sqlite_url(self, url: str) -> str:
@@ -86,6 +89,14 @@ def _cleanup_old_records(self) -> int:
         """
         return self._cleanup_service.cleanup_old_records()
 
+    def _cleanup_excess_sessions(self) -> int:
+        """
+        Remove least recently used sessions when total sessions exceed limit.
+        This implements daily LRU cleanup to maintain max 2000 sessions for better memory management.
+        Runs once per day - during the day session count can exceed limit without issues.
+        """
+        return self._cleanup_service.cleanup_excess_sessions()
+
     def _cleanup_old_messages(self, session_id: str) -> int:
         """
         Remove old messages from a session using FIFO strategy.
@@ -164,10 +175,10 @@ async def save_conversation(
 
         logger.info("✅ Successfully inserted record into conversation_history table")
 
-        # Daily cleanup: run once per day to remove old records
-        self._cleanup_old_records()
+        # Daily session LRU cleanup: maintain max 2000 sessions (runs once per day)
+        self._cleanup_excess_sessions()
 
-        # Always perform LRU cleanup for this session (lightweight)
+        # Per-session FIFO cleanup: maintain max 20 records per session (runs on every save)
         self._cleanup_old_messages(session_id)
 
         return record_id
diff --git a/tests/test_conversation_history_lifecycle.py b/tests/test_conversation_history_lifecycle.py
@@ -198,6 +198,126 @@ async def test_time_based_cleanup_integration(self):
             f"✅ After time-based cleanup: {len(records_after)} records (within retention)"
         )
 
+    @pytest.mark.asyncio
+    async def test_lru_session_cleanup_lifecycle(self):
+        """Test LRU session cleanup: keeps most recently used sessions."""
+        # Create provider with small session limit for testing
+        db = SQLiteProvider(max_session_records=20)
+
+        # Override session limit for testing (normally 2000)
+        db._cleanup_service.max_total_sessions = 3
+
+        print("\n🔄 TESTING LRU SESSION CLEANUP LIFECYCLE")
+        print("=" * 60)
+
+        # PHASE 1: Create sessions with different activity patterns
+        print("📝 PHASE 1: Creating sessions with different activity patterns...")
+
+        # Session A: Created first, but will be most recently used
+        await db.save_conversation("session_A", "tool1", "input1", "output1")
+        print("   Session A: Created (oldest creation time)")
+
+        # Session B: Created second
+        await db.save_conversation("session_B", "tool1", "input1", "output1")
+        print("   Session B: Created")
+
+        # Session C: Created third
+        await db.save_conversation("session_C", "tool1", "input1", "output1")
+        print("   Session C: Created")
+
+        # Session D: Created fourth
+        await db.save_conversation("session_D", "tool1", "input1", "output1")
+        print("   Session D: Created")
+
+        # Session E: Created fifth
+        await db.save_conversation("session_E", "tool1", "input1", "output1")
+        print("   Session E: Created")
+
+        # Add recent activity to Session A (making it most recently used)
+        await db.save_conversation("session_A", "tool2", "recent_input", "recent_output")
+        print("   Session A: Updated with recent activity (most recently used)")
+
+        # Verify all sessions exist
+        current_count = db._cleanup_service.get_session_count()
+        assert current_count == 5, f"Expected 5 sessions, got {current_count}"
+        print(f"✅ Phase 1: {current_count} sessions created successfully")
+
+        # PHASE 2: Test LRU detection
+        print("\n🔍 PHASE 2: Testing LRU detection...")
+
+        # Get least recently used sessions (should be B, C - oldest last activity)
+        lru_sessions = db._cleanup_service.get_least_recently_used_sessions(2)
+        print(f"   LRU sessions to remove: {lru_sessions}")
+
+        # Should be sessions B and C (oldest last activity)
+        assert "session_B" in lru_sessions, "Session B should be LRU"
+        assert "session_C" in lru_sessions, "Session C should be LRU"
+        assert "session_A" not in lru_sessions, "Session A should NOT be LRU (recent activity)"
+        print("✅ Phase 2: LRU detection working correctly")
+
+        # PHASE 3: Trigger LRU cleanup
+        print("\n🧹 PHASE 3: Triggering LRU session cleanup...")
+
+        # Force cleanup by mocking old cleanup time
+        old_time = datetime.utcnow() - timedelta(days=2)
+        db._cleanup_service.last_session_cleanup_time = old_time
+
+        # Trigger cleanup
+        deleted_count = db._cleanup_service.cleanup_excess_sessions()
+        print(f"   Deleted {deleted_count} records")
+
+        # Verify session count is now at limit
+        final_count = db._cleanup_service.get_session_count()
+        assert final_count == 3, f"Expected 3 sessions after cleanup, got {final_count}"
+        print(f"✅ Phase 3: Session count reduced to {final_count}")
+
+        # PHASE 4: Verify which sessions remain
+        print("\n🔍 PHASE 4: Verifying remaining sessions...")
+
+        sessions_to_check = ["session_A", "session_B", "session_C", "session_D", "session_E"]
+        remaining_sessions = []
+        deleted_sessions = []
+
+        for session_id in sessions_to_check:
+            records = await db.get_session_conversations(session_id)
+            if records:
+                remaining_sessions.append(session_id)
+                last_activity = max(r.timestamp for r in records)
+                print(f"   ✅ {session_id}: {len(records)} records, last activity: {last_activity}")
+            else:
+                deleted_sessions.append(session_id)
+                print(f"   ❌ {session_id}: DELETED (was least recently used)")
+
+        # Verify correct sessions were kept/deleted
+        assert "session_A" in remaining_sessions, "Session A should remain (most recent activity)"
+        assert "session_D" in remaining_sessions, "Session D should remain (recent creation)"
+        assert "session_E" in remaining_sessions, "Session E should remain (most recent creation)"
+        assert "session_B" in deleted_sessions, "Session B should be deleted (LRU)"
+        assert "session_C" in deleted_sessions, "Session C should be deleted (LRU)"
+
+        print("✅ Phase 4: LRU cleanup preserved active sessions correctly")
+
+        # PHASE 5: Verify Session A has both records (original + recent activity)
+        print("\n📊 PHASE 5: Verifying session record preservation...")
+
+        session_a_records = await db.get_session_conversations("session_A")
+        assert len(session_a_records) == 2, f"Session A should have 2 records, got {len(session_a_records)}"
+
+        # Check that both records exist
+        sources = [r.source for r in session_a_records]
+        assert "tool1" in sources, "Session A should have original tool1 record"
+        assert "tool2" in sources, "Session A should have recent tool2 record"
+
+        print(f"   Session A has {len(session_a_records)} records: {sources}")
+        print("✅ Phase 5: Session record preservation working correctly")
+
+        print("\n🎯 LRU CLEANUP SUMMARY:")
+        print("   - Sessions B, C deleted (least recently used)")
+        print("   - Sessions A, D, E preserved (most recently used)")
+        print("   - Session A preserved despite being oldest (recent activity)")
+        print("   - LRU provides better UX than FIFO!")
+        print("✅ LRU session cleanup lifecycle test PASSED!")
+
         print("✅ Time-based cleanup integration working correctly")
 
     @pytest.mark.asyncio
@@ -266,6 +386,7 @@ async def run_tests():
         await test_instance.test_save_retrieve_fifo_cleanup_lifecycle()
         await test_instance.test_multiple_sessions_isolation()
         await test_instance.test_time_based_cleanup_integration()
+        await test_instance.test_lru_session_cleanup_lifecycle()
         await test_instance.test_edge_cases_and_error_handling()
         print("\n🎉 ALL CONVERSATION HISTORY LIFECYCLE TESTS PASSED!")
 
