fix(adaptive_router): bound owner cache, drop PK from upsert update, redact PII

- _owner_cache now opportunistically sweeps expired entries past
  _OWNER_CACHE_SWEEP_THRESHOLD live entries. Previously sessions that never
  came back piled up forever.
- flush_session_to_db strips session_id/router_name/model_name from the update
  payload. Prisma rejects writes to @@id fields.
- record_turn no longer persists last_user_content / last_assistant_content /
  tool_call_history / pending_tool_calls. Those are needed only in-memory for
  the next turn's signal detection; writing user prompts and tool payloads to
  the DB would store PII for every conversation.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

litellm/proxy/db/db_transaction_queue/adaptive_router_update_queue.py

@@ -166,6 +166,14 @@ class AdaptiveRouterUpdateQueue:
                 # NOTE: Prisma client lower-cases model names, so
                 # `LiteLLM_AdaptiveRouterSession` -> `litellm_adaptiveroutersession`
                 # (single 's', not 'litellm_adaptiverouterssession').
+                # Strip PK fields from the update payload — Prisma rejects
+                # writes to fields that are part of the @@id. asdict(state)
+                # always carries them, so build a separate update dict.
+                update_payload = {
+                    k: v
+                    for k, v in payload.items()
+                    if k not in ("session_id", "router_name", "model_name")
+                }
                 await prisma_client.db.litellm_adaptiveroutersession.upsert(
                     where={
                         "session_id_router_name_model_name": {
@@ -179,9 +187,9 @@ class AdaptiveRouterUpdateQueue:
                             "session_id": session_id,
                             "router_name": router,
                             "model_name": model,
-                            **payload,
+                            **update_payload,
                         },
-                        "update": payload,
+                        "update": update_payload,
                     },
                 )
             except Exception as e:

litellm/router_strategy/adaptive_router/adaptive_router.py

@@ -47,6 +47,8 @@ from litellm.router_strategy.adaptive_router.config import (
 # Sweep session-state cache when it exceeds this many live entries. Expired
 # entries are dropped in bulk; amortizes to O(1) per insert.
 _SESSION_STATE_SWEEP_THRESHOLD: int = 1024
+# Same pattern for the owner cache.
+_OWNER_CACHE_SWEEP_THRESHOLD: int = 1024
 from litellm.router_strategy.adaptive_router.signals import (
     SessionState,
     SignalDelta,
@@ -228,6 +230,12 @@ class AdaptiveRouter:
             self._skipped_updates_total += 1
             return False
 
+        # Opportunistic bulk sweep — sessions that never come back would
+        # otherwise pile up here forever. Same threshold pattern as the
+        # session-state cache.
+        if len(self._owner_cache) >= _OWNER_CACHE_SWEEP_THRESHOLD:
+            self._evict_expired_owner_cache(now)
+
         # No live owner -> claim for current_model.
         self._owner_cache[session_key] = (
             current_model,
@@ -235,6 +243,11 @@ class AdaptiveRouter:
         )
         return True
 
+    def _evict_expired_owner_cache(self, now: float) -> None:
+        expired = [k for k, (_, exp) in self._owner_cache.items() if exp <= now]
+        for k in expired:
+            self._owner_cache.pop(k, None)
+
     async def get_state_snapshot(self) -> Dict[str, Any]:
         """In-memory snapshot for the introspection endpoint. Cheap; no DB hit."""
         cells = []
@@ -361,7 +374,20 @@ class AdaptiveRouter:
             "AdaptiveRouter[%s]: record_turn delta=%s", self.router_name, delta
         )
 
+        # Strip the raw conversation content before persisting. The
+        # last_user/assistant_content and tool_call_history fields are only
+        # needed in-memory for the next turn's incremental signal detection;
+        # writing user prompts and tool payloads to the DB would store PII
+        # for every adaptive-router conversation. Counts + bookkeeping is
+        # all the persisted row needs.
         snapshot = asdict(state)
+        for sensitive in (
+            "last_user_content",
+            "last_assistant_content",
+            "tool_call_history",
+            "pending_tool_calls",
+        ):
+            snapshot.pop(sensitive, None)
         await self.queue.add_session_state(
             session_id, self.router_name, model_name, snapshot
         )

tests/test_litellm/router_strategy/adaptive_router/test_adaptive_router.py

@@ -118,6 +118,25 @@ def test_claim_or_check_owner_expired_owner_reclaims_for_new_model(monkeypatch):
     assert r._skipped_updates_total == 0
 
 
+def test_owner_cache_evicts_expired_entries_when_threshold_crossed(monkeypatch):
+    """Past _OWNER_CACHE_SWEEP_THRESHOLD live entries, new claims sweep stale."""
+    r = _make_router()
+    monkeypatch.setattr(ar_module, "_OWNER_CACHE_SWEEP_THRESHOLD", 5)
+    monkeypatch.setattr(ar_module.time, "time", lambda: 1_000.0)
+    for i in range(5):
+        r.claim_or_check_owner(f"old-{i}", "fast")
+    assert len(r._owner_cache) == 5
+
+    # Jump past TTL so all "old-*" entries are now expired.
+    monkeypatch.setattr(
+        ar_module.time, "time", lambda: 1_000.0 + OWNER_CACHE_TTL_SECONDS + 1
+    )
+    r.claim_or_check_owner("new-1", "fast")
+    # Sweep ran -> only the new entry remains.
+    assert "new-1" in r._owner_cache
+    assert all(k.startswith("new-") for k in r._owner_cache)
+
+
 # ---- record_turn --------------------------------------------------------
 
 
@@ -149,6 +168,16 @@ async def test_record_turn_pushes_to_queue():
     # satisfaction fired -> alpha delta -> add_state_delta called
     r.queue.add_state_delta.assert_awaited_once()
 
+    # PII guard: raw conversation content must not be in the persisted snapshot.
+    snapshot = r.queue.add_session_state.call_args.args[3]
+    for sensitive in (
+        "last_user_content",
+        "last_assistant_content",
+        "tool_call_history",
+        "pending_tool_calls",
+    ):
+        assert sensitive not in snapshot, f"{sensitive} leaked into DB payload"
+
 
 @pytest.mark.asyncio
 async def test_record_turn_satisfaction_increments_alpha():