diff --git a/litellm/proxy/db/db_transaction_queue/adaptive_router_update_queue.py b/litellm/proxy/db/db_transaction_queue/adaptive_router_update_queue.py index 7f5d9f78541..c76ca16aa35 100644 --- a/litellm/proxy/db/db_transaction_queue/adaptive_router_update_queue.py +++ b/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: diff --git a/litellm/router_strategy/adaptive_router/adaptive_router.py b/litellm/router_strategy/adaptive_router/adaptive_router.py index ae5e39d2ee0..1e8d02185d7 100644 --- a/litellm/router_strategy/adaptive_router/adaptive_router.py +++ b/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 ) diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_adaptive_router.py b/tests/test_litellm/router_strategy/adaptive_router/test_adaptive_router.py index aed217cdc21..93c4db90dad 100644 --- a/tests/test_litellm/router_strategy/adaptive_router/test_adaptive_router.py +++ b/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():