scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

fix(spend): key the session rollup on the caller's credential as well

Browse source 
session_id is caller-controlled, and this row sums spend, baseline spend and the
savings between them. Keyed on the session alone, one caller reusing another's
id folds their turns into that tenant's dollars and merges two conversations
into one benchmark row.

The codebase has two keying conventions and this table was following the wrong
one. Session-scoped state is keyed on the session and the router that owns it:
LiteLLM_AdaptiveRouterSession is @@id([session_id, router_name, model_name]) and
the deployment-affinity cache key is model_group plus session_id, neither
carrying a tenant. Spend accounting is keyed on the credential: every daily
table is @@unique([<entity>, date, api_key, model, ...]). The discriminator is
not whether a row is per-session, it is whether the row adds up money, and this
one does.

So the key gains api_key, matching the accounting tables it belongs with. A
credential rotated mid-conversation now splits that session's rollup in two,
which undercounts one row rather than merging two tenants into one wrong number.
This commit is contained in:
Tin Chi Lo 2026-08-04 14:22:49 -07:00
parent 71f2e85fbe
commit 3e50a52ecb
9 changed files with 109 additions and 43 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm-proxy-extras/litellm_proxy_extras/migrations/20260801000000_add_auto_router_session_rollup/migration.sql
View file

@ -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

3
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -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")
}

5
litellm/proxy/db/db_spend_update_writer.py
View file

@ -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

2
litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py
View file

@ -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",
)

3
litellm/proxy/schema.prisma
View file

@ -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")
}

26
litellm/proxy/spend_tracking/auto_router_session_queue.py
View file

@ -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(

3
schema.prisma
View file

@ -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")
}

3
tests/test_litellm/proxy/db/test_db_spend_update_writer.py
View file

@ -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):

104
tests/test_litellm/proxy/spend_tracking/test_auto_router_sessions.py
View file

@ -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