fix(tracker): bound per-session bookkeeping dicts with LRU eviction (#18)

beliefstate/config.py

@@ -295,6 +295,10 @@ def validate_resolution_strategy(cls, v: str) -> str:
         default=300,
         description="Maximum tokens reserved for belief injection in prompts. If belief summary exceeds this, use relevance-based filtering.",
     )
+    max_tracked_sessions: int = Field(
+        default=10000,
+        description="Maximum number of sessions kept in the in-memory turn/state/provider bookkeeping. Least-recently-used sessions are evicted beyond this to bound memory in long-running servers. Set to 0 to disable eviction (unbounded).",
+    )
 
     # Context injection filtering
     exclude_sources: List[str] = Field(

beliefstate/tracker.py

@@ -8,7 +8,7 @@
 from datetime import datetime, timezone
 from functools import wraps
 from contextvars import ContextVar
-from collections import deque
+from collections import deque, OrderedDict
 
 from beliefstate.config import TrackerConfig
 from beliefstate.call import LLMCall, LLMResponse
@@ -429,9 +429,12 @@ def __init__(
             strategy=self.config.resolution_strategy,
             respect_strategy_for_updates=self.config.respect_strategy_for_updates,
         )
-        self._session_turn_counters: Dict[str, int] = {}
-        self._session_turn_states: Dict[str, int] = {}
-        self._session_providers: Dict[str, str] = {}
+        # LRU-ordered so long-running servers with many unique sessions don't
+        # grow these unbounded; least-recently-used sessions are evicted once
+        # config.max_tracked_sessions is exceeded (see _touch_session).
+        self._session_turn_counters: "OrderedDict[str, int]" = OrderedDict()
+        self._session_turn_states: "OrderedDict[str, int]" = OrderedDict()
+        self._session_providers: "OrderedDict[str, str]" = OrderedDict()
         self._stats = TrackerStats()
         self._pending_tasks: Set[asyncio.Task[None]] = set()
         self._pending_conflict_notes: Dict[str, List[str]] = {}
@@ -515,6 +518,27 @@ def _dispatch(self, coro: Any) -> None:
         task.add_done_callback(self._pending_tasks.discard)
         self._sweep_stale_locks()
 
+    def _touch_session(self, session_id: str) -> None:
+        """Mark ``session_id`` as most-recently-used and evict LRU sessions.
+
+        Keeps the per-session bookkeeping dicts bounded so a server handling
+        many unique sessions over a long uptime does not leak memory. The turn
+        counter dict is the LRU authority (it is written every tracked turn);
+        evicted sessions are dropped from the state/provider dicts and the lock
+        registry too.
+        """
+        counters = self._session_turn_counters
+        if session_id in counters:
+            counters.move_to_end(session_id)
+        limit = self.config.max_tracked_sessions
+        if limit <= 0:
+            return
+        while len(counters) > limit:
+            oldest_sid, _ = counters.popitem(last=False)
+            self._session_turn_states.pop(oldest_sid, None)
+            self._session_providers.pop(oldest_sid, None)
+            _session_locks.pop(oldest_sid, None)
+
     def _sweep_stale_locks(self) -> None:
         """Remove session locks for sessions no longer tracked."""
         active_sessions = set(self._session_turn_counters.keys())
@@ -1237,6 +1261,7 @@ async def wrapper(*args: Any, **kwargs: Any) -> Any:
                         self._session_turn_counters.get(session_id, 0) + 1
                     )
                     current_turn = self._session_turn_counters[session_id]
+                    self._touch_session(session_id)
 
                 if auto_inject:
                     last_user_msg = ""

tests/test_config.py

@@ -65,6 +65,10 @@ def test_default_token_aware_injection(self):
         assert c.enable_token_aware_injection is True
         assert c.belief_budget_tokens == 300
 
+    def test_default_max_tracked_sessions(self):
+        c = TrackerConfig()
+        assert c.max_tracked_sessions == 10000
+
 
 class TestTrackerConfigOverrides:
     """Verify custom overrides are accepted."""

tests/test_tracker_advanced.py

@@ -129,6 +129,44 @@ def test_set_session(self):
         session_context.set("default")
 
 
+class TestSessionEviction:
+    """LRU eviction of the in-memory session bookkeeping dicts (issue #18)."""
+
+    def _make_tracker(self, max_tracked_sessions):
+        config = make_config(max_tracked_sessions=max_tracked_sessions)
+        return BeliefTracker(config=config, adapter=MagicMock())
+
+    def _add(self, tracker, sid):
+        tracker._session_turn_counters[sid] = 1
+        tracker._session_turn_states[sid] = 1
+        tracker._session_providers[sid] = "openai"
+        tracker._touch_session(sid)
+
+    def test_evicts_lru_beyond_limit(self):
+        tracker = self._make_tracker(3)
+        for i in range(5):
+            self._add(tracker, f"s{i}")
+        # Only the 3 most-recently-used sessions survive, in lockstep.
+        assert set(tracker._session_turn_counters) == {"s2", "s3", "s4"}
+        assert set(tracker._session_turn_states) == {"s2", "s3", "s4"}
+        assert set(tracker._session_providers) == {"s2", "s3", "s4"}
+
+    def test_touch_refreshes_recency(self):
+        tracker = self._make_tracker(3)
+        for sid in ("a", "b", "c"):
+            self._add(tracker, sid)
+        # Re-touch "a" so "b" becomes the LRU; adding "d" must evict "b".
+        tracker._touch_session("a")
+        self._add(tracker, "d")
+        assert set(tracker._session_turn_counters) == {"a", "c", "d"}
+
+    def test_zero_disables_eviction(self):
+        tracker = self._make_tracker(0)
+        for i in range(50):
+            self._add(tracker, f"s{i}")
+        assert len(tracker._session_turn_counters) == 50
+
+
 # ── get_beliefs / get_stats / get_summary ────────────────────────────────