diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260801000000_add_auto_router_session_rollup/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260801000000_add_auto_router_session_rollup/migration.sql index e30748c83f0..ec1472b57cf 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260801000000_add_auto_router_session_rollup/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260801000000_add_auto_router_session_rollup/migration.sql @@ -1,5 +1,6 @@ -- CreateTable CREATE TABLE IF NOT EXISTS "LiteLLM_AutoRouterSession" ( + "api_key" TEXT NOT NULL, "session_id" TEXT NOT NULL, "model_group" TEXT NOT NULL, "router_kind" TEXT NOT NULL, @@ -27,7 +28,7 @@ CREATE TABLE IF NOT EXISTS "LiteLLM_AutoRouterSession" ( "model_state" JSONB NOT NULL DEFAULT '{}', "updated_at" TIMESTAMP(3) NOT NULL, - CONSTRAINT "LiteLLM_AutoRouterSession_pkey" PRIMARY KEY ("session_id","model_group") + CONSTRAINT "LiteLLM_AutoRouterSession_pkey" PRIMARY KEY ("api_key","session_id","model_group") ); -- CreateIndex diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 502541f36cd..28990f01b4b 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1393,6 +1393,7 @@ model LiteLLM_AdaptiveRouterSession { // Per-(session, auto-router) rollup behind the auto-router benchmarks dashboard. model LiteLLM_AutoRouterSession { + api_key String session_id String model_group String router_kind String @@ -1426,7 +1427,7 @@ model LiteLLM_AutoRouterSession { model_state Json @default("{}") updated_at DateTime @default(now()) @updatedAt - @@id([session_id, model_group]) + @@id([api_key, session_id, model_group]) @@index([model_group, last_turn_at], map: "idx_auto_router_session_group_activity") @@index([last_turn_at], map: "idx_auto_router_session_last_turn") } diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 23e8a132f94..e3ac340916c 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -303,7 +303,8 @@ class DBSpendUpdateWriter: session_id = payload.get("session_id") model_group = payload.get("model_group") model = payload.get("model") - if prisma_client is None or llm_router is None or not session_id or not model_group or not model: + api_key = payload.get("api_key") + if prisma_client is None or llm_router is None or not session_id or not model_group or not model or not api_key: return if not serves_an_auto_router(llm_router, model_group): return @@ -334,7 +335,7 @@ class DBSpendUpdateWriter: cost_breakdown=_metadata.get("cost_breakdown"), ) await self.auto_router_session_queue.record_turn( - key=(session_id, model_group), + key=(api_key, session_id, model_group), router_kind=router_kind, # The same setting savings.py priced this turn against, stored on the row so # the dashboard names the baseline the numbers were actually computed with diff --git a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py index c2baa7f6da1..6596703d69d 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -194,7 +194,7 @@ class SpendLogCleanup: prisma_client, cutoff_date, table_name="LiteLLM_AutoRouterSession", - key_columns=("session_id", "model_group"), + key_columns=("api_key", "session_id", "model_group"), time_column="last_turn_at", ) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 502541f36cd..28990f01b4b 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1393,6 +1393,7 @@ model LiteLLM_AdaptiveRouterSession { // Per-(session, auto-router) rollup behind the auto-router benchmarks dashboard. model LiteLLM_AutoRouterSession { + api_key String session_id String model_group String router_kind String @@ -1426,7 +1427,7 @@ model LiteLLM_AutoRouterSession { model_state Json @default("{}") updated_at DateTime @default(now()) @updatedAt - @@id([session_id, model_group]) + @@id([api_key, session_id, model_group]) @@index([model_group, last_turn_at], map: "idx_auto_router_session_group_activity") @@index([last_turn_at], map: "idx_auto_router_session_last_turn") } diff --git a/litellm/proxy/spend_tracking/auto_router_session_queue.py b/litellm/proxy/spend_tracking/auto_router_session_queue.py index 528f78ba94e..011ce93ec40 100644 --- a/litellm/proxy/spend_tracking/auto_router_session_queue.py +++ b/litellm/proxy/spend_tracking/auto_router_session_queue.py @@ -41,7 +41,15 @@ from litellm.repositories.table_repositories import AutoRouterSessionRepository if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient -SessionKey = tuple[str, str] +SessionKey = tuple[str, str, str] +"""``(api_key, session_id, model_group)``. + +``session_id`` is caller-controlled, so on its own it lets one caller write into +another's rollup by reusing their id. The row sums spend, baseline spend and the +savings between them, which puts it in the same family as the daily spend tables, +and every one of those carries ``api_key`` in its key for exactly this reason. A +key rotated mid-conversation splits that session's rollup in two, which +undercounts one row rather than merging two tenants into one wrong number.""" DEFAULT_MAX_STAGED_TURNS = 50_000 MAX_CACHED_SESSIONS = 10_000 @@ -94,14 +102,18 @@ def _chunked(batch: _Chunk, size: int) -> tuple[_Chunk, ...]: def _key_fields(key: SessionKey) -> Mapping[str, str]: - """One composite primary key as the two columns that make it up.""" - session_id, model_group = key - return {"session_id": session_id, "model_group": model_group} # mutable-ok: prisma's query API takes dict payloads + """One composite primary key as the three columns that make it up.""" + api_key, session_id, model_group = key + return { # mutable-ok: prisma's query API takes dict payloads + "api_key": api_key, + "session_id": session_id, + "model_group": model_group, + } def _unique_where(key: SessionKey) -> Mapping[str, Mapping[str, str]]: """The composite primary key, under the name prisma gives the ``@@id`` selector.""" - return {"session_id_model_group": _key_fields(key)} # mutable-ok: prisma's query API takes dict payloads + return {"api_key_session_id_model_group": _key_fields(key)} # mutable-ok: prisma's query API takes dict payloads def _increment(value: float) -> Mapping[str, float]: @@ -279,7 +291,9 @@ class AutoRouterSessionQueue: verbose_proxy_logger.warning("auto_router_sessions: could not load %d session states (%s)", len(missing), e) return StateUnavailable() stored = { # mutable-ok: built once from the rows this read returned - (row.session_id, row.model_group): state_from_row(row.last_model, row.last_turn_at, row.model_state) + (row.api_key, row.session_id, row.model_group): state_from_row( + row.last_model, row.last_turn_at, row.model_state + ) for row in rows } return MappingProxyType( diff --git a/schema.prisma b/schema.prisma index 502541f36cd..28990f01b4b 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1393,6 +1393,7 @@ model LiteLLM_AdaptiveRouterSession { // Per-(session, auto-router) rollup behind the auto-router benchmarks dashboard. model LiteLLM_AutoRouterSession { + api_key String session_id String model_group String router_kind String @@ -1426,7 +1427,7 @@ model LiteLLM_AutoRouterSession { model_state Json @default("{}") updated_at DateTime @default(now()) @updatedAt - @@id([session_id, model_group]) + @@id([api_key, session_id, model_group]) @@index([model_group, last_turn_at], map: "idx_auto_router_session_group_activity") @@index([last_turn_at], map: "idx_auto_router_session_last_turn") } diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 73c09593bc0..e6a4aa4cf47 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -2285,6 +2285,7 @@ def _turn_payload( "routed_model": "anthropic/claude-haiku-4-5", } return { + "api_key": "key-1", "session_id": session_id, "model_group": model_group, "model": "anthropic/claude-haiku-4-5", @@ -2316,7 +2317,7 @@ class TestRecordingAnAutoRouterTurn: queue = await _record(router, _turn_payload("smart-router")) - assert queue.staged == [(("session-1", "smart-router"), "complexity", "anthropic/claude-haiku-4-5")] + assert queue.staged == [(("key-1", "session-1", "smart-router"), "complexity", "anthropic/claude-haiku-4-5")] @pytest.mark.parametrize("kind", ["complexity", "quality", "adaptive", "semantic"]) async def test_a_turn_is_folded_under_the_kind_the_router_recorded(self, kind): diff --git a/tests/test_litellm/proxy/spend_tracking/test_auto_router_sessions.py b/tests/test_litellm/proxy/spend_tracking/test_auto_router_sessions.py index c99f3f0e21f..e20ba655346 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_auto_router_sessions.py +++ b/tests/test_litellm/proxy/spend_tracking/test_auto_router_sessions.py @@ -370,7 +370,16 @@ def test_state_from_row_without_a_timestamp_starts_at_the_epoch(): class _StoredRow: """A session rollup as prisma hands it back.""" - def __init__(self, session_id: str, model_group: str, last_model: str, last_turn_at: float, model_state: dict): + def __init__( + self, + session_id: str, + model_group: str, + last_model: str, + last_turn_at: float, + model_state: dict, + api_key: str = "k1", + ): + self.api_key = api_key self.session_id = session_id self.model_group = model_group self.last_model = last_model @@ -397,8 +406,8 @@ class _RecordingTable: self.reads += 1 if self.fail_read: raise RuntimeError("transient database fault") - wanted = {(pair["session_id"], pair["model_group"]) for pair in where["OR"]} - return [row for row in self.rows if (row.session_id, row.model_group) in wanted] + wanted = {(pair["api_key"], pair["session_id"], pair["model_group"]) for pair in where["OR"]} + return [row for row in self.rows if (row.api_key, row.session_id, row.model_group) in wanted] def commit(self, statements): self.transactions += 1 @@ -474,9 +483,10 @@ def _turn_at(at: float, model: str = MODEL_A) -> TurnFacts: return turn(model, at=at, created=5000) -def _stored_session(session_id: str = "s1", model_group: str = "g") -> _StoredRow: +def _stored_session(session_id: str = "s1", model_group: str = "g", api_key: str = "k1") -> _StoredRow: """A session last served on MODEL_B that has already used MODEL_A.""" return _StoredRow( + api_key=api_key, session_id=session_id, model_group=model_group, last_model=MODEL_B, @@ -497,7 +507,7 @@ class TestTheLoggingPathNeverTouchesTheDatabase: queue = AutoRouterSessionQueue() for at in (0, 60, 120): - await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(at)) + await queue.record_turn(("k1", "s1", "g"), "complexity", None, _turn_at(at)) table = _RecordingTable() assert await queue.flush(_RecordingPrisma(table)) == 1 @@ -510,12 +520,12 @@ class TestTheLoggingPathNeverTouchesTheDatabase: prisma = _RecordingPrisma(table) queue = AutoRouterSessionQueue() for at in (0, 60, 120): - await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(at)) + await queue.record_turn(("k1", "s1", "g"), "complexity", None, _turn_at(at)) await queue.flush(prisma) assert table.reads == 1 - await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(180)) + await queue.record_turn(("k1", "s1", "g"), "complexity", None, _turn_at(180)) await queue.flush(prisma) assert table.reads == 1 @@ -530,7 +540,7 @@ class TestFlushDurability: table = _RecordingTable(fail_write=True) prisma = _RecordingPrisma(table) queue = AutoRouterSessionQueue() - await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(0)) + await queue.record_turn(("k1", "s1", "g"), "complexity", None, _turn_at(0)) assert await queue.flush(prisma) == 0 @@ -544,11 +554,11 @@ class TestFlushDurability: table = _RecordingTable(fail_write=True) prisma = _RecordingPrisma(table) queue = AutoRouterSessionQueue() - await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(0)) + await queue.record_turn(("k1", "s1", "g"), "complexity", None, _turn_at(0)) await queue.flush(prisma) table.fail_write = False - await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(60)) + await queue.record_turn(("k1", "s1", "g"), "complexity", None, _turn_at(60)) assert await queue.flush(prisma) == 1 assert table.upserts[0][1]["create"]["turns"] == 2 @@ -559,7 +569,7 @@ class TestFlushDurability: table = _RecordingTable(fail_read=True, rows=(_stored_session(),)) prisma = _RecordingPrisma(table) queue = AutoRouterSessionQueue() - await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(120)) + await queue.record_turn(("k1", "s1", "g"), "complexity", None, _turn_at(120)) assert await queue.flush(prisma) == 0 assert table.upserts == [] @@ -571,7 +581,7 @@ class TestFlushDurability: table = _RecordingTable(fail_read=True, rows=(_stored_session(),)) prisma = _RecordingPrisma(table) queue = AutoRouterSessionQueue() - await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(120)) + await queue.record_turn(("k1", "s1", "g"), "complexity", None, _turn_at(120)) await queue.flush(prisma) table.fail_read = False @@ -586,11 +596,47 @@ class TestFlushDurability: prisma = _RecordingPrisma(_RecordingTable()) queue = AutoRouterSessionQueue() - await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(0)) + await queue.record_turn(("k1", "s1", "g"), "complexity", None, _turn_at(0)) assert await queue.flush(prisma) == 1 assert await queue.flush(prisma) == 0 +@pytest.mark.asyncio +class TestOneCallerCannotWriteIntoAnothersRollup: + """`session_id` is caller-controlled, and the row sums spend and savings. + + Keyed on the session alone, a caller reusing someone else's id would fold + their turns into that tenant's dollars. The key carries the trusted + credential for the same reason every daily spend table does. + """ + + async def test_the_same_session_id_under_two_keys_is_two_rollups(self): + from litellm.proxy.spend_tracking.auto_router_session_queue import AutoRouterSessionQueue + + table = _RecordingTable() + queue = AutoRouterSessionQueue() + await queue.record_turn(("victim-key", "shared-id", "g"), "complexity", None, _turn_at(0)) + await queue.record_turn(("attacker-key", "shared-id", "g"), "complexity", None, _turn_at(60)) + + assert await queue.flush(_RecordingPrisma(table)) == 2 + keys = {where["api_key_session_id_model_group"]["api_key"] for where, _ in table.upserts} + assert keys == {"victim-key", "attacker-key"} + assert all(data["create"]["turns"] == 1 for _, data in table.upserts) + + async def test_history_is_looked_up_under_the_callers_own_key(self): + """Reusing the id must not even read the other tenant's state, let alone fold onto it.""" + from litellm.proxy.spend_tracking.auto_router_session_queue import AutoRouterSessionQueue + + victim = _stored_session(api_key="victim-key", session_id="shared-id") + table = _RecordingTable(rows=(victim,)) + queue = AutoRouterSessionQueue() + await queue.record_turn(("attacker-key", "shared-id", "g"), "complexity", None, _turn_at(300)) + + assert await queue.flush(_RecordingPrisma(table)) == 1 + assert table.upserts[0][1]["create"]["first_visit_turns"] == 1 + assert table.upserts[0][1]["create"]["return_turns"] == 0 + + @pytest.mark.asyncio class TestTheRowRecordsTheFoldsAnswer: """Session progress on the row comes from the folded state, not the staged turns. @@ -604,7 +650,7 @@ class TestTheRowRecordsTheFoldsAnswer: table = _RecordingTable(rows=(_stored_session(),)) queue = AutoRouterSessionQueue() - await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(30)) + await queue.record_turn(("k1", "s1", "g"), "complexity", None, _turn_at(30)) assert await queue.flush(_RecordingPrisma(table)) == 1 update = table.upserts[0][1]["update"] @@ -616,7 +662,7 @@ class TestTheRowRecordsTheFoldsAnswer: table = _RecordingTable(rows=(_stored_session(),)) queue = AutoRouterSessionQueue() - await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(300)) + await queue.record_turn(("k1", "s1", "g"), "complexity", None, _turn_at(300)) assert await queue.flush(_RecordingPrisma(table)) == 1 update = table.upserts[0][1]["update"] @@ -636,7 +682,7 @@ class TestStagingCostsTheSamePerTurn: from litellm.proxy.spend_tracking.auto_router_session_queue import AutoRouterSessionQueue queue = AutoRouterSessionQueue() - key = ("s1", "g") + key = ("k1", "s1", "g") await queue.record_turn(key, "complexity", None, _turn_at(0)) buffer = queue._pending[key].turns @@ -652,7 +698,7 @@ class TestStagingCostsTheSamePerTurn: table = _RecordingTable() queue = AutoRouterSessionQueue() for at in range(500): - await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(at * 60)) + await queue.record_turn(("k1", "s1", "g"), "complexity", None, _turn_at(at * 60)) assert await queue.flush(_RecordingPrisma(table)) == 1 assert table.upserts[0][1]["create"]["turns"] == 500 @@ -668,7 +714,7 @@ class TestStagingIsBounded: table = _RecordingTable() queue = AutoRouterSessionQueue(max_staged_turns=2) for i in range(5): - await queue.record_turn((f"s{i}", "g"), "complexity", None, _turn_at(0)) + await queue.record_turn(("k1", f"s{i}", "g"), "complexity", None, _turn_at(0)) assert await queue.flush(_RecordingPrisma(table)) == 2 @@ -679,7 +725,7 @@ class TestStagingIsBounded: table = _RecordingTable() queue = AutoRouterSessionQueue(max_staged_turns=2) for at in (0, 60, 120, 180): - await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(at)) + await queue.record_turn(("k1", "s1", "g"), "complexity", None, _turn_at(at)) await queue.flush(_RecordingPrisma(table)) assert table.upserts[0][1]["create"]["turns"] == 2 @@ -690,10 +736,10 @@ class TestStagingIsBounded: table = _RecordingTable() prisma = _RecordingPrisma(table) queue = AutoRouterSessionQueue(max_staged_turns=1) - await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(0)) + await queue.record_turn(("k1", "s1", "g"), "complexity", None, _turn_at(0)) await queue.flush(prisma) - await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(60)) + await queue.record_turn(("k1", "s1", "g"), "complexity", None, _turn_at(60)) assert await queue.flush(prisma) == 1 assert table.upserts[1][1]["update"]["turns"] == {"increment": 1} @@ -711,11 +757,11 @@ class TestStagingIsBounded: async def arrive_mid_flush(): for at in (200, 260): - await queue.record_turn(("s2", "g"), "complexity", None, _turn_at(at)) + await queue.record_turn(("k1", "s2", "g"), "complexity", None, _turn_at(at)) table = _RefillingTable(arrive_mid_flush, fail_write=True) for at in (0, 60): - await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(at)) + await queue.record_turn(("k1", "s1", "g"), "complexity", None, _turn_at(at)) await queue.flush(_RecordingPrisma(table)) @@ -744,7 +790,7 @@ class TestFlushCostDoesNotGrowWithSessionCount: table = _RecordingTable() queue = AutoRouterSessionQueue() for i in range(250): - await queue.record_turn((f"s{i}", "g"), "complexity", None, _turn_at(0)) + await queue.record_turn(("k1", f"s{i}", "g"), "complexity", None, _turn_at(0)) assert await queue.flush(_RecordingPrisma(table)) == 250 assert (table.reads, table.transactions) == (1, 1) @@ -757,11 +803,11 @@ class TestFlushCostDoesNotGrowWithSessionCount: prisma = _RecordingPrisma(table) queue = AutoRouterSessionQueue() for i in range(50): - await queue.record_turn((f"s{i}", "g"), "complexity", None, _turn_at(0)) + await queue.record_turn(("k1", f"s{i}", "g"), "complexity", None, _turn_at(0)) await queue.flush(prisma) for i in range(50): - await queue.record_turn((f"s{i}", "g"), "complexity", None, _turn_at(60)) + await queue.record_turn(("k1", f"s{i}", "g"), "complexity", None, _turn_at(60)) assert await queue.flush(prisma) == 50 assert (table.reads, table.transactions) == (1, 2) @@ -773,14 +819,14 @@ class TestFlushCostDoesNotGrowWithSessionCount: prisma = _RecordingPrisma(table) queue = AutoRouterSessionQueue() for i in range(5): - await queue.record_turn((f"s{i}", "g"), "complexity", None, _turn_at(0)) + await queue.record_turn(("k1", f"s{i}", "g"), "complexity", None, _turn_at(0)) assert await queue.flush(prisma) == 0 assert table.upserts == [] table.fail_write = False assert await queue.flush(prisma) == 5 - written = {where["session_id_model_group"]["session_id"] for where, _ in table.upserts} + written = {where["api_key_session_id_model_group"]["session_id"] for where, _ in table.upserts} assert written == {f"s{i}" for i in range(5)} async def test_one_read_serves_a_chunk_of_sessions_that_all_have_history(self): @@ -790,7 +836,7 @@ class TestFlushCostDoesNotGrowWithSessionCount: table = _RecordingTable(rows=tuple(_stored_session(session_id=f"s{i}") for i in range(5))) queue = AutoRouterSessionQueue() for i in range(5): - await queue.record_turn((f"s{i}", "g"), "complexity", None, _turn_at(120)) + await queue.record_turn(("k1", f"s{i}", "g"), "complexity", None, _turn_at(120)) assert await queue.flush(_RecordingPrisma(table)) == 5 assert table.reads == 1