From 546984a2b219c3cafd15e4dc0f4217f28d276a34 Mon Sep 17 00:00:00 2001 From: Vaishnavi220506 Date: Fri, 2 Oct 2026 10:24:16 +0530 Subject: [PATCH] fix(routing): bound failed session snapshot retention --- .../router_strategy/adaptive_router/README.md | 11 +- .../adaptive_router/update_queue.py | 13 ++- .../adaptive_router/test_update_queue.py | 104 ++++++++++++++++++ 3 files changed, 124 insertions(+), 4 deletions(-) diff --git a/litellm/router_strategy/adaptive_router/README.md b/litellm/router_strategy/adaptive_router/README.md index 0bf6cae3bf7..9bde6272861 100644 --- a/litellm/router_strategy/adaptive_router/README.md +++ b/litellm/router_strategy/adaptive_router/README.md @@ -68,7 +68,8 @@ Callers may pass header `x-litellm-min-quality-tier: 3` (or metadata key `SIGNAL_GATE_MIN_MESSAGES` messages. - **Persistence.** Bandit cells: aggregated deltas, eventually consistent. Session rows: last-write-wins snapshots. Failed or cancelled flushes retain - unacknowledged rows for the next flush. Retried deltas merge with new feedback; + unacknowledged rows for the next flush within the retention limit below. + Retried deltas merge with new feedback; a newer queued session snapshot replaces a failed older snapshot. Flushes of the same queue are serialized within one process, and their return values count acknowledged writes only @@ -80,8 +81,12 @@ Callers may pass header `x-litellm-min-quality-tier: 3` (or metadata key to distinguish that outcome from a failed write. The queues are in memory, so process termination still loses pending updates. Session write ordering across workers, or server queries continuing after client cancellation, is - not guaranteed. A prolonged database outage retains pending session keys - in memory until a later flush succeeds + not guaranteed. Failed or cancelled session batches retain at most 1,024 + snapshots, prioritizing snapshots queued during the flush, then the most + recently queued failed rows. Older snapshots beyond that bound are dropped. + Successful flushes do not impose this retention limit on newly queued rows. + Session keys containing NUL characters are discarded before enqueueing because + PostgreSQL cannot store them in text columns - **Latency is not in the score.** Quality + cost only. A pathologically slow model can still be picked. diff --git a/litellm/router_strategy/adaptive_router/update_queue.py b/litellm/router_strategy/adaptive_router/update_queue.py index 23aa28c3df6..8893e3d2f21 100644 --- a/litellm/router_strategy/adaptive_router/update_queue.py +++ b/litellm/router_strategy/adaptive_router/update_queue.py @@ -30,6 +30,7 @@ from litellm.repositories.table_repositories import ( StateKey = tuple[str, str, str] # (router_name, request_type, model_name) SessionKey = tuple[str, str, str] # (session_id, router_name, model_name) +_MAX_SESSION_RETRY_ENTRIES: Final[int] = 1024 class AdaptiveRouterUpdateQueue: @@ -87,8 +88,12 @@ class AdaptiveRouterUpdateQueue: SessionState (signals counts + bookkeeping fields). The flusher will upsert this into LiteLLM_AdaptiveRouterSession. """ + if any("\0" in part for part in (session_id, router_name, model_name)): + verbose_router_logger.warning("AdaptiveRouterUpdateQueue: session key cannot be stored in PostgreSQL") + return key: Final[SessionKey] = (session_id, router_name, model_name) async with self._lock: + self._session_agg.pop(key, None) self._session_agg[key] = state_dict self._max_session_size_seen = max(self._max_session_size_seen, len(self._session_agg)) @@ -228,8 +233,14 @@ class AdaptiveRouterUpdateQueue: self._max_state_size_seen = max(self._max_state_size_seen, len(self._state_agg)) async def _restore_session_batch(self, pending: Mapping[SessionKey, Mapping[str, object]]) -> None: + if not pending: + return async with self._lock: - self._session_agg = {**pending, **self._session_agg} + restored: Final = dict(pending) + for key in self._session_agg: + restored.pop(key, None) + merged: Final = {**restored, **self._session_agg} + self._session_agg = dict(tuple(merged.items())[-_MAX_SESSION_RETRY_ENTRIES:]) self._max_session_size_seen = max(self._max_session_size_seen, len(self._session_agg)) # ---- Observability --------------------------------------------------- diff --git a/tests/unit/router_strategy/adaptive_router/test_update_queue.py b/tests/unit/router_strategy/adaptive_router/test_update_queue.py index d7ab46d64d4..4b3dbb52565 100644 --- a/tests/unit/router_strategy/adaptive_router/test_update_queue.py +++ b/tests/unit/router_strategy/adaptive_router/test_update_queue.py @@ -305,3 +305,107 @@ async def test_queue_high_water_mark_includes_restored_rows_and_new_keys(queue, sizes = await queue.queue_size() assert sizes[f"{kind}_pending"] == 3 assert sizes[f"max_{kind}_seen"] == 3 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("invalid_field", ["session_id", "router_name", "model_name"]) +async def test_session_keys_that_postgres_cannot_store_are_not_queued(queue, mock_prisma, invalid_field): + key = {"session_id": "session", "router_name": "router", "model_name": "model"} + key[invalid_field] += "\0" + + await queue.add_session_state(**key, state_dict={"turn_count": 1}) + + assert (await queue.queue_size())["session_pending"] == 0 + assert await queue.flush_session_to_db(mock_prisma) == 0 + mock_prisma.db.litellm_adaptiveroutersession.upsert.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_failed_session_retention_keeps_only_the_newest_rows(queue, mock_prisma, monkeypatch): + monkeypatch.setattr( + "litellm.router_strategy.adaptive_router.update_queue._MAX_SESSION_RETRY_ENTRIES", 3, raising=False + ) + table = mock_prisma.db.litellm_adaptiveroutersession + table.upsert.side_effect = RuntimeError("database unavailable") + for index in range(6): + await queue.add_session_state(f"session-{index}", "router", "model", {"turn_count": index}) + await queue.add_session_state("session-0", "router", "model", {"turn_count": 9}) + + assert await queue.flush_session_to_db(mock_prisma) == 0 + assert (await queue.queue_size())["session_pending"] == 3 + table.reset_mock() + table.upsert.side_effect = None + + assert await queue.flush_session_to_db(mock_prisma) == 3 + assert [ + (call.kwargs["data"]["create"]["session_id"], call.kwargs["data"]["update"]["turn_count"]) + for call in table.upsert.await_args_list + ] == [("session-0", 9), ("session-4", 4), ("session-5", 5)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cancelled", [False, True]) +async def test_session_retention_prioritizes_new_snapshots_during_a_failed_flush( + queue, mock_prisma, monkeypatch, cancelled +): + monkeypatch.setattr( + "litellm.router_strategy.adaptive_router.update_queue._MAX_SESSION_RETRY_ENTRIES", 2, raising=False + ) + started = asyncio.Event() + release = asyncio.Event() + + async def fail_write(**kwargs): + started.set() + await release.wait() + raise RuntimeError("database unavailable") + + table = mock_prisma.db.litellm_adaptiveroutersession + table.upsert.side_effect = fail_write + for session in ("a", "b", "c"): + await queue.add_session_state(session, "router", "model", {"turn_count": 1}) + task = asyncio.create_task(queue.flush_session_to_db(mock_prisma)) + await started.wait() + await queue.add_session_state("a", "router", "model", {"turn_count": 2}) + await queue.add_session_state("d", "router", "model", {"turn_count": 3}) + if cancelled: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + else: + release.set() + assert await task == 0 + assert (await queue.queue_size())["session_pending"] == 2 + table.reset_mock() + table.upsert.side_effect = None + assert await queue.flush_session_to_db(mock_prisma) == 2 + assert [ + (call.kwargs["data"]["create"]["session_id"], call.kwargs["data"]["update"]["turn_count"]) + for call in table.upsert.await_args_list + ] == [("a", 2), ("d", 3)] + + +@pytest.mark.asyncio +async def test_successful_flush_does_not_cap_newly_queued_sessions(queue, mock_prisma, monkeypatch): + monkeypatch.setattr( + "litellm.router_strategy.adaptive_router.update_queue._MAX_SESSION_RETRY_ENTRIES", 1, raising=False + ) + started = asyncio.Event() + release = asyncio.Event() + + async def write(**kwargs): + started.set() + await release.wait() + + table = mock_prisma.db.litellm_adaptiveroutersession + table.upsert.side_effect = write + await queue.add_session_state("old", "router", "model", {"turn_count": 1}) + task = asyncio.create_task(queue.flush_session_to_db(mock_prisma)) + await started.wait() + for session in ("a", "b"): + await queue.add_session_state(session, "router", "model", {"turn_count": 1}) + release.set() + + assert await task == 1 + assert (await queue.queue_size())["session_pending"] == 2 + table.upsert.side_effect = None + assert await queue.flush_session_to_db(mock_prisma) == 2