From ce2dd040ee576fbaa09fd1c157fc53dc12adfa4c Mon Sep 17 00:00:00 2001 From: Mrinal Chanshetty Date: Thu, 29 Oct 2026 08:57:23 -0700 Subject: [PATCH] fix(proxy): keep the daily tag spend unique key and record the virtual key's team Swapping the unique key on LiteLLM_DailyTagSpend broke older proxies during a rolling upgrade, because their ON CONFLICT target no longer matched an index. The key now stays as it is and team_id is a plain column, stamped only when a team's virtual key authenticated the request. That key already pins each row to one team through api_key. JWT, custom auth, OAuth2, UI sessions, the master key and team-less keys get '' On conflict a non-empty team overwrites the stored one, so rows an older proxy wrote are filled in by the next write and a key moved mid-day lands under its latest team. The two migrations that built the new unique index and dropped the old one are removed, and the remaining two are renamed to sort after main's latest migration --- .../migration.sql | 10 -- .../migration.sql | 14 -- .../migration.sql | 18 --- .../migration.sql | 8 ++ .../migration.sql | 0 .../litellm_proxy_extras/schema.prisma | 4 +- litellm/proxy/db/baseline_accounting.py | 2 - litellm/proxy/db/daily_spend_bulk_upsert.py | 27 ++-- litellm/proxy/db/db_spend_update_writer.py | 39 ++--- litellm/proxy/schema.prisma | 4 +- schema.prisma | 4 +- .../spend/test_team_tag_spend_attribution.py | 133 +++++++++++++----- .../test_litellm_proxy_extras_utils.py | 2 +- .../proxy/db/test_daily_spend_bulk_upsert.py | 41 +----- .../proxy/db/test_db_spend_update_writer.py | 107 ++++++++------ 15 files changed, 218 insertions(+), 195 deletions(-) delete mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000000_daily_tag_spend_add_team_id/migration.sql delete mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000001_daily_tag_spend_team_id_unique_index/migration.sql delete mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000003_daily_tag_spend_drop_tag_only_unique_index/migration.sql create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20261029000000_daily_tag_spend_add_team_id/migration.sql rename litellm-proxy-extras/litellm_proxy_extras/migrations/{20261017000002_daily_tag_spend_team_id_date_tag_index => 20261029000001_daily_tag_spend_team_id_date_tag_index}/migration.sql (100%) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000000_daily_tag_spend_add_team_id/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000000_daily_tag_spend_add_team_id/migration.sql deleted file mode 100644 index 0b73df2d05b..00000000000 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000000_daily_tag_spend_add_team_id/migration.sql +++ /dev/null @@ -1,10 +0,0 @@ -SET LOCAL lock_timeout = '10s'; - --- Records the team that sent each request on its daily tag spend rows, so one tag used by --- several teams lands in one row per team. "" means the request had no team, and every row --- written before this migration keeps "" (attribution is not retroactive). --- --- NOT NULL DEFAULT '' rather than nullable: the daily spend upsert normalizes every conflict --- column to '' because a NULL never matches itself in a unique index. On PostgreSQL 11+ adding --- a column with a constant default is a metadata-only change, so this does not rewrite the table. -ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN IF NOT EXISTS "team_id" TEXT NOT NULL DEFAULT ''; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000001_daily_tag_spend_team_id_unique_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000001_daily_tag_spend_team_id_unique_index/migration.sql deleted file mode 100644 index 1816ff63ce6..00000000000 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000001_daily_tag_spend_team_id_unique_index/migration.sql +++ /dev/null @@ -1,14 +0,0 @@ --- CreateIndex (CONCURRENTLY) --- --- The new conflict target for the daily tag spend upsert: the old unique key plus team_id. --- Built before the old unique index is dropped (next migrations), so the table is never --- without a unique index the writer can arbitrate on. --- --- Disclaimer: --- - CREATE INDEX CONCURRENTLY cannot run inside a transaction. This migration must stay a --- single statement so Prisma Migrate on PostgreSQL can apply it outside a transaction. --- - If the build is interrupted, Postgres may leave an INVALID index that must be dropped and --- recreated. --- - Do not edit this file after it has been applied to any database: Prisma checksums --- migrations; add a new migration instead. -CREATE UNIQUE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_DailyTagSpend_tag_team_id_dimensions_key" ON "LiteLLM_DailyTagSpend"("tag", "team_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000003_daily_tag_spend_drop_tag_only_unique_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000003_daily_tag_spend_drop_tag_only_unique_index/migration.sql deleted file mode 100644 index ebc4f7ac457..00000000000 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000003_daily_tag_spend_drop_tag_only_unique_index/migration.sql +++ /dev/null @@ -1,18 +0,0 @@ --- DropIndex (CONCURRENTLY) --- --- Drops the unique key without team_id. With it in place, two teams sending the same tag on --- the same day for the same key and model would collide on it, so it has to go once the --- team-aware unique index exists. --- --- Upgrade note: an old-version writer still targets this index in its ON CONFLICT clause and --- fails every flush once it is gone. Drain and stop old instances before running migrations. --- An old instance booting against the migrated database also logs "prisma schema out of sync" --- with SQL that drops team_id and both new indexes. Do not run it: it deletes team attribution --- and its CREATE UNIQUE INDEX fails once two teams share a key. --- --- Disclaimer: --- - DROP INDEX CONCURRENTLY cannot run inside a transaction. This migration must stay a --- single statement so Prisma Migrate on PostgreSQL can apply it outside a transaction. --- - Do not edit this file after it has been applied to any database: Prisma checksums --- migrations; add a new migration instead. -DROP INDEX CONCURRENTLY IF EXISTS "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key"; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261029000000_daily_tag_spend_add_team_id/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261029000000_daily_tag_spend_add_team_id/migration.sql new file mode 100644 index 00000000000..0e15d9745b5 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261029000000_daily_tag_spend_add_team_id/migration.sql @@ -0,0 +1,8 @@ +SET LOCAL lock_timeout = '10s'; + +-- The team that owns the virtual key a tag row was spent through. "" when the request did not +-- come through a virtual key with a team (JWT, custom auth, OAuth2, team-less keys), and on every +-- row written before this migration. Not part of the unique key: a team's virtual key already +-- maps the row to one team, so older proxies keep upserting on the unchanged key during a rolling +-- upgrade. On PostgreSQL 11+ adding a column with a constant default does not rewrite the table. +ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN IF NOT EXISTS "team_id" TEXT NOT NULL DEFAULT ''; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000002_daily_tag_spend_team_id_date_tag_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261029000001_daily_tag_spend_team_id_date_tag_index/migration.sql similarity index 100% rename from litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000002_daily_tag_spend_team_id_date_tag_index/migration.sql rename to litellm-proxy-extras/litellm_proxy_extras/migrations/20261029000001_daily_tag_spend_team_id_date_tag_index/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index b018cf4788c..fd772d13023 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1086,7 +1086,7 @@ model LiteLLM_DailyTagSpend { id String @id @default(uuid()) request_id String? tag String? - team_id String @default("") // "" = the request had no team + team_id String @default("") // "" unless spent through a team's virtual key date String api_key String model String? @@ -1112,7 +1112,7 @@ model LiteLLM_DailyTagSpend { created_at DateTime @default(now()) updated_at DateTime @updatedAt - @@unique([tag, team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint], map: "LiteLLM_DailyTagSpend_tag_team_id_dimensions_key") + @@unique([tag, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) @@index([date]) @@index([tag, date]) @@index([team_id, date, tag]) diff --git a/litellm/proxy/db/baseline_accounting.py b/litellm/proxy/db/baseline_accounting.py index 5e817410d2b..f7a2be654c8 100644 --- a/litellm/proxy/db/baseline_accounting.py +++ b/litellm/proxy/db/baseline_accounting.py @@ -58,7 +58,6 @@ class DailyBaselineAttribution(LiteLLMBaseModel): model_group: str | None = None endpoint: str | None = None mcp_namespaced_tool_name: str | None = None - team_id: str = "" targets: tuple[DailyBaselineTarget, ...] = () def adjustment(self, target: DailyBaselineTarget, savings_delta: float, request_id: str) -> SpendRow: @@ -72,7 +71,6 @@ class DailyBaselineAttribution(LiteLLMBaseModel): "model_group": self.model_group, "endpoint": self.endpoint, "mcp_namespaced_tool_name": self.mcp_namespaced_tool_name, - **({"team_id": self.team_id} if table.team_scoped else {}), table.entity_id_column: target.entity_id, "request_id": request_id, "autorouter_savings_spend": savings_delta, diff --git a/litellm/proxy/db/daily_spend_bulk_upsert.py b/litellm/proxy/db/daily_spend_bulk_upsert.py index dcc2b738a51..74d4ee29029 100644 --- a/litellm/proxy/db/daily_spend_bulk_upsert.py +++ b/litellm/proxy/db/daily_spend_bulk_upsert.py @@ -32,7 +32,7 @@ class DailySpendTable: name: str entity_id_column: str carries_request_id: bool = False - team_scoped: bool = False + carries_team_id: bool = False DAILY_SPEND_TABLES: Final[Mapping[DailySpendEntity, DailySpendTable]] = MappingProxyType( @@ -43,7 +43,7 @@ DAILY_SPEND_TABLES: Final[Mapping[DailySpendEntity, DailySpendTable]] = MappingP "end_user": DailySpendTable(name="LiteLLM_DailyEndUserSpend", entity_id_column="end_user_id"), "agent": DailySpendTable(name="LiteLLM_DailyAgentSpend", entity_id_column="agent_id"), "tag": DailySpendTable( - name="LiteLLM_DailyTagSpend", entity_id_column="tag", carries_request_id=True, team_scoped=True + name="LiteLLM_DailyTagSpend", entity_id_column="tag", carries_request_id=True, carries_team_id=True ), } ) @@ -127,14 +127,9 @@ def _as_float(value: object) -> float: return float(value) if isinstance(value, (int, float)) else 0.0 -def _conflict_columns(table: DailySpendTable) -> tuple[str, ...]: - """The table's unique constraint, in constraint order.""" - return (table.entity_id_column, *(("team_id",) if table.team_scoped else ()), *_KEY_COLUMNS) - - def conflict_key(table: DailySpendTable, transaction: SpendRow) -> tuple[str, ...]: """The tuple the database arbitrates the upsert on, normalized free of NULLs.""" - return tuple(_as_text(transaction.get(column)) for column in _conflict_columns(table)) + return tuple(_as_text(transaction.get(column)) for column in (table.entity_id_column, *_KEY_COLUMNS)) def _merge(group: Sequence[SpendRow]) -> SpendRow: @@ -176,17 +171,20 @@ def _row_params( *(_as_int(transaction.get(column)) for column in _COUNTER_COLUMNS), *(_as_float(transaction.get(column)) for column in _SPEND_COLUMNS), *((None if request_id is None else _as_text(request_id),) if table.carries_request_id else ()), + *((_as_text(transaction.get("team_id")),) if table.carries_team_id else ()), ) def _insert_columns(table: DailySpendTable) -> tuple[str, ...]: return ( "id", - *_conflict_columns(table), + table.entity_id_column, + *_KEY_COLUMNS, "model_group", *_COUNTER_COLUMNS, *_SPEND_COLUMNS, *(("request_id",) if table.carries_request_id else ()), + *(("team_id",) if table.carries_team_id else ()), ) @@ -217,11 +215,18 @@ def build_bulk_upsert( if table.carries_request_id else "" ) + # Rows written by an older proxy, or by a request without a team's virtual key, carry "", so + # only a known team overwrites the stored one. + team_id_update: Final = ( + f', "team_id" = COALESCE(NULLIF(EXCLUDED."team_id", \'\'), {quoted_table}."team_id")' + if table.carries_team_id + else "" + ) sql: Final = ( f'INSERT INTO {quoted_table} ({_quoted(columns)}, "updated_at")\n' f"VALUES {rows}\n" - f"ON CONFLICT ({_quoted(_conflict_columns(table))}) DO UPDATE SET\n" - f" {increments}{request_id_update},\n" + f"ON CONFLICT ({_quoted((table.entity_id_column, *_KEY_COLUMNS))}) DO UPDATE SET\n" + f" {increments}{request_id_update}{team_id_update},\n" f" \"updated_at\" = (NOW() AT TIME ZONE 'UTC')" ) return sql, tuple(value for key, transaction in batch for value in _row_params(table, key, transaction)) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 9e3b0fd3163..59278bd93ec 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -52,6 +52,7 @@ from litellm.proxy._types import ( SpendLogsPayload, SpendUpdateQueueItem, ToolDiscoveryQueueItem, + UserAPIKeyAuth, ) from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET, project_cache_key from litellm.proxy.db.daily_spend_bulk_upsert import ( @@ -351,6 +352,20 @@ def _parse_spend_log_metadata(payload: Mapping[str, object]) -> SpendLogsMetadat return parsed +def _virtual_key_team_id(kwargs: Mapping[str, object]) -> str: + """The team of the virtual key that authenticated the request, or "" for JWT, custom auth, OAuth2, + UI sessions, the master key and team-less keys, whose api_key does not pin a tag row to one team.""" + params: Final = _SPEND_METADATA_ADAPTER.validate_python(kwargs.get("litellm_params") or {}) + candidates: Final = tuple( + _SPEND_METADATA_ADAPTER.validate_python(params.get(bucket) or {}).get("user_api_key_auth") + for bucket in ("litellm_metadata", "metadata") + ) + auth: Final = next((found for found in candidates if isinstance(found, UserAPIKeyAuth)), None) + if auth is None or not auth.via_virtual_key or auth.is_session_token: + return "" + return auth.team_id or "" + + def _is_billed_retry_attempt(metadata: SpendLogsMetadata) -> bool: """An earlier billed retry attempt's own spend row, not a request the caller sent.""" return BILLED_RETRY_ATTEMPT_METADATA_KEY in metadata and metadata[BILLED_RETRY_ATTEMPT_METADATA_KEY] is True @@ -721,6 +736,7 @@ class DBSpendUpdateWriter: litellm_proxy_budget_name=litellm_proxy_budget_name, payload=payload, request_model_access_groups=get_request_model_access_groups(kwargs), + virtual_key_team_id=_virtual_key_team_id(_SPEND_METADATA_ADAPTER.validate_python(kwargs or {})), ), task_name="spend rollups", ) @@ -1040,7 +1056,6 @@ class DBSpendUpdateWriter: model_group=base.get("model_group"), endpoint=base.get("endpoint"), mcp_namespaced_tool_name=base.get("mcp_namespaced_tool_name"), - team_id=payload.get("team_id") or "", # pyright: ignore[reportUnknownMemberType] # TypedDict.get overloads carry Any defaults targets=tuple( DailyBaselineTarget(entity=entity, entity_id=identity) for entity, values in bases @@ -1146,6 +1161,7 @@ class DBSpendUpdateWriter: payload: SpendLogsPayload, request_model_access_groups: Sequence[str] = (), project_id: str | None = None, + virtual_key_team_id: str = "", ): """ Runs every spend-update helper for one request. Each helper goes through @@ -1245,7 +1261,9 @@ class DBSpendUpdateWriter: ) await self._run_spend_update_step( "add_spend_log_transaction_to_daily_tag_transaction", - self.add_spend_log_transaction_to_daily_tag_transaction(payload=payload_copy, prisma_client=prisma_client), + self.add_spend_log_transaction_to_daily_tag_transaction( + payload=payload_copy, prisma_client=prisma_client, virtual_key_team_id=virtual_key_team_id + ), payload=payload_copy, ) await self._enqueue_model_usage_transaction(payload=payload_copy, prisma_client=prisma_client) @@ -3080,6 +3098,7 @@ class DBSpendUpdateWriter: self, payload: SpendLogsPayload, prisma_client: PrismaClient | None = None, + virtual_key_team_id: str = "", ) -> None: if prisma_client is None: verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.") @@ -3095,25 +3114,13 @@ class DBSpendUpdateWriter: return request_tags: Final = daily_spend_entity_ids(payload, "tag") - # .get, not indexing: the caller swallows exceptions, so a KeyError would silently drop tag spend - team_id: Final = payload.get("team_id") or "" # pyright: ignore[reportUnknownMemberType] # TypedDict.get overloads carry Any defaults for tag in request_tags: if tag is None: continue endpoint_str = base_daily_transaction.get("endpoint") or "" - daily_transaction_key = json.dumps( - [ - tag, - team_id, - base_daily_transaction["date"], - payload["api_key"], - payload["model"], - payload["custom_llm_provider"], - endpoint_str, - ] - ) + daily_transaction_key = f"{tag}_{base_daily_transaction['date']}_{payload['api_key']}_{payload['model']}_{payload['custom_llm_provider']}_{endpoint_str}" daily_transaction = DailyTagSpendTransaction( - tag=tag, team_id=team_id, **base_daily_transaction, request_id=payload["request_id"] + tag=tag, team_id=virtual_key_team_id, **base_daily_transaction, request_id=payload["request_id"] ) await self.daily_tag_spend_update_queue.add_update(update={daily_transaction_key: daily_transaction}) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index b018cf4788c..fd772d13023 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1086,7 +1086,7 @@ model LiteLLM_DailyTagSpend { id String @id @default(uuid()) request_id String? tag String? - team_id String @default("") // "" = the request had no team + team_id String @default("") // "" unless spent through a team's virtual key date String api_key String model String? @@ -1112,7 +1112,7 @@ model LiteLLM_DailyTagSpend { created_at DateTime @default(now()) updated_at DateTime @updatedAt - @@unique([tag, team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint], map: "LiteLLM_DailyTagSpend_tag_team_id_dimensions_key") + @@unique([tag, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) @@index([date]) @@index([tag, date]) @@index([team_id, date, tag]) diff --git a/schema.prisma b/schema.prisma index b018cf4788c..fd772d13023 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1086,7 +1086,7 @@ model LiteLLM_DailyTagSpend { id String @id @default(uuid()) request_id String? tag String? - team_id String @default("") // "" = the request had no team + team_id String @default("") // "" unless spent through a team's virtual key date String api_key String model String? @@ -1112,7 +1112,7 @@ model LiteLLM_DailyTagSpend { created_at DateTime @default(now()) updated_at DateTime @updatedAt - @@unique([tag, team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint], map: "LiteLLM_DailyTagSpend_tag_team_id_dimensions_key") + @@unique([tag, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) @@index([date]) @@index([tag, date]) @@index([team_id, date, tag]) diff --git a/tests/integration/spend/test_team_tag_spend_attribution.py b/tests/integration/spend/test_team_tag_spend_attribution.py index d82394e143b..558735f282e 100644 --- a/tests/integration/spend/test_team_tag_spend_attribution.py +++ b/tests/integration/spend/test_team_tag_spend_attribution.py @@ -1,45 +1,112 @@ +import math import uuid from typing import Final -import pytest -from integration._support.client import Gateway, eventually -from integration._support.database import read_rows +from integration._support.client import Gateway, Scenario, eventually +from integration._support.database import read_rows, write_rows +from pydantic import BaseModel, ConfigDict, TypeAdapter _SPEND_PER_CALL: Final = 20 * 0.001 + 20 * 0.002 -def test_one_key_moved_between_teams_records_its_tag_once_per_team(gateway: Gateway) -> None: - with gateway.scenario() as scenario: - model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) - tag: Final = f"cost-center:{uuid.uuid4().hex}" - teams: Final = tuple(scenario.team(models=[model]) for _ in range(2)) - calls_per_team: Final = {teams[0]: 1, teams[1]: 2} - key: Final = scenario.key(team_id=teams[0], models=[model]) - for team, calls in calls_per_team.items(): - gateway.post("/key/update", {"key": key, "team_id": team}) - for _ in range(calls): - reply = gateway.post( - "/v1/chat/completions", - { - "model": model, - "messages": [{"role": "user", "content": f"tagged {uuid.uuid4().hex}"}], - "metadata": {"tags": [tag]}, - }, - key=key, - ) - assert reply["usage"]["total_tokens"] == 40, reply +class _Usage(BaseModel): + model_config = ConfigDict(frozen=True) - rows: Final = eventually( - lambda: read_rows( + total_tokens: int + + +class _Reply(BaseModel): + model_config = ConfigDict(frozen=True) + + usage: _Usage + + +class _TagRow(BaseModel): + model_config = ConfigDict(frozen=True) + + team_id: str + spend: float + api_requests: int + + +_TAG_ROWS: Final = TypeAdapter(tuple[_TagRow, ...]) + + +def _send_tagged(gateway: Gateway, *, key: str, model: str, tag: str, calls: int) -> None: + for _ in range(calls): + reply = gateway.post( + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"tagged {uuid.uuid4().hex}"}], + "metadata": {"tags": [tag]}, + }, + key=key, + ) + assert _Reply.model_validate(reply).usage.total_tokens == 40, reply + + +def _tag_rows(tag: str, *, requests: int) -> tuple[_TagRow, ...]: + return eventually( + lambda: _TAG_ROWS.validate_python( + read_rows( 'SELECT team_id, spend, api_requests FROM "LiteLLM_DailyTagSpend" WHERE tag = %s ORDER BY team_id', (tag,), - ), - lambda values: sum(int(row["api_requests"]) for row in values) == 3, - seconds=70, - ) + ) + ), + lambda rows: sum(row.api_requests for row in rows) == requests, + seconds=70, + ) - assert sorted((row["team_id"], int(row["api_requests"])) for row in rows) == sorted(calls_per_team.items()), ( - rows - ) + +def _model(scenario: Scenario) -> str: + return scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + + +def test_a_shared_tag_is_split_by_the_team_of_each_virtual_key(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = _model(scenario) + tag: Final = f"cost-center:{uuid.uuid4().hex}" + teams: Final = tuple(scenario.team(models=[model]) for _ in range(2)) + calls: Final = {teams[0]: 1, teams[1]: 2, "": 1} + for team, count in calls.items(): + key = scenario.key(team_id=team, models=[model]) if team else scenario.key(models=[model]) + _send_tagged(gateway, key=key, model=model, tag=tag, calls=count) + + rows: Final = _tag_rows(tag, requests=sum(calls.values())) + + assert sorted((row.team_id, row.api_requests) for row in rows) == sorted(calls.items()), rows for row in rows: - assert float(row["spend"]) == pytest.approx(calls_per_team[row["team_id"]] * _SPEND_PER_CALL), rows + assert math.isclose(row.spend, calls[row.team_id] * _SPEND_PER_CALL), rows + + +def test_a_key_moved_between_teams_keeps_one_row_under_its_latest_team(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = _model(scenario) + tag: Final = f"moved-key:{uuid.uuid4().hex}" + teams: Final = tuple(scenario.team(models=[model]) for _ in range(2)) + key: Final = scenario.key(team_id=teams[0], models=[model]) + _send_tagged(gateway, key=key, model=model, tag=tag, calls=1) + _tag_rows(tag, requests=1) + gateway.post("/key/update", {"key": key, "team_id": teams[1]}) + _send_tagged(gateway, key=key, model=model, tag=tag, calls=2) + + rows: Final = _tag_rows(tag, requests=3) + + assert [(row.team_id, row.api_requests) for row in rows] == [(teams[1], 3)], rows + + +def test_a_row_an_older_proxy_wrote_without_a_team_gets_the_team_on_the_next_write(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = _model(scenario) + tag: Final = f"rolling-upgrade:{uuid.uuid4().hex}" + team: Final = scenario.team(models=[model]) + key: Final = scenario.key(team_id=team, models=[model]) + _send_tagged(gateway, key=key, model=model, tag=tag, calls=1) + _tag_rows(tag, requests=1) + write_rows("UPDATE \"LiteLLM_DailyTagSpend\" SET team_id = '' WHERE tag = %s", (tag,)) + _send_tagged(gateway, key=key, model=model, tag=tag, calls=1) + + rows: Final = _tag_rows(tag, requests=2) + + assert [(row.team_id, row.api_requests) for row in rows] == [(team, 2)], rows diff --git a/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py b/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py index ad889af1873..02bdf3429ea 100644 --- a/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py +++ b/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py @@ -435,7 +435,7 @@ class TestMigrationSQLIdempotency: for migration_name, sql in _guarded_migrations(all_migrations): for line_num, line in _statements(sql): if re.search(r"DROP\s+INDEX\s+", line, re.IGNORECASE) and not re.search( - r"DROP\s+INDEX\s+(?:CONCURRENTLY\s+)?IF\s+EXISTS", line, re.IGNORECASE + r"DROP\s+INDEX\s+IF\s+EXISTS", line, re.IGNORECASE ): violations.append(f" {migration_name}:{line_num}: {line.strip()}") assert ( diff --git a/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py b/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py index 61a7d7dac7f..1f80e184516 100644 --- a/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py +++ b/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py @@ -3,13 +3,11 @@ import re from collections.abc import AsyncIterator from contextlib import AbstractAsyncContextManager, asynccontextmanager -from typing import Final import pytest from litellm.proxy.db.daily_spend_bulk_upsert import ( DAILY_SPEND_TABLES, - DailySpendEntity, build_bulk_upsert, conflict_key, merge_by_conflict_key, @@ -90,6 +88,8 @@ def test_one_statement_carries_every_row_in_the_batch(): assert sql.count("INSERT INTO") == 1 assert len(re.findall(r"ON CONFLICT", sql)) == 1 + # 26 bound columns per row plus the inlined updated_at, so the row count is what + # separates one multi-row statement from a hundred single-row ones. assert len(params) == 100 * 26 assert "$2600::text" in sql assert sql.count("(NOW() AT TIME ZONE 'UTC')") == 100 + 1 @@ -101,7 +101,7 @@ def test_conflict_target_is_the_full_unique_constraint(): conflict_target = re.search(r"ON CONFLICT \(([^)]*)\)", sql) assert conflict_target is not None assert conflict_target.group(1) == ( - '"tag", "team_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"' + '"tag", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"' ) @@ -144,41 +144,6 @@ def test_non_tag_tables_carry_no_request_id_column(): assert '"user_id"' in sql -def test_one_tag_from_two_teams_stays_two_rows() -> None: - merged: Final = merge_by_conflict_key( - TAG_TABLE, (tag_txn(team_id="team-b", spend=2.0), tag_txn(team_id="team-a", spend=1.0)) - ) - - assert [(txn["team_id"], txn["spend"]) for _, txn in merged] == [("team-a", 1.0), ("team-b", 2.0)], merged - - -def test_an_entry_without_team_id_folds_into_the_unattributed_row() -> None: - legacy: Final = tag_txn(spend=0.25) - unattributed: Final = tag_txn(team_id="", spend=0.75) - - merged: Final = merge_by_conflict_key(TAG_TABLE, (legacy, unattributed)) - - assert len(merged) == 1, merged - key, folded = merged[0] - assert key[:2] == ("team-a", ""), key - assert folded["spend"] == pytest.approx(1.0) - - -@pytest.mark.parametrize("entity", ["user", "team", "org", "end_user", "agent"]) -def test_only_the_tag_table_is_keyed_by_team(entity: DailySpendEntity) -> None: - table: Final = DAILY_SPEND_TABLES[entity] - row: Final = {**tag_txn(), table.entity_id_column: "id-1"} - - sql, _ = build_bulk_upsert(table, merge_by_conflict_key(table, (row,))) - - conflict_target: Final = re.search(r"ON CONFLICT \(([^)]*)\)", sql) - assert conflict_target is not None - assert conflict_target.group(1) == ( - f'"{table.entity_id_column}", "date", "api_key", "model", "custom_llm_provider", ' - '"mcp_namespaced_tool_name", "endpoint"' - ) - - class _RecordingDb: def __init__(self) -> None: self.statements: list[tuple[str, tuple[object, ...]]] = [] diff --git a/tests/unit/proxy/db/test_db_spend_update_writer.py b/tests/unit/proxy/db/test_db_spend_update_writer.py index 474d10ff1bf..7a4dc7d30bc 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -21,13 +21,20 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm._service_logger import ServiceTypes from litellm.constants import SPEND_ROLLUP_LOCK_TIMEOUT_MS -from litellm.proxy._types import DailyTagSpendTransaction, Litellm_EntityType, SpendLogsPayload, SpendUpdateQueueItem +from litellm.proxy._types import ( + DailyTagSpendTransaction, + Litellm_EntityType, + SpendLogsPayload, + SpendUpdateQueueItem, + UserAPIKeyAuth, +) from litellm.proxy.db.db_spend_update_writer import ( _TEAM_ADVISORY_LOCK_SQL, _TEAM_MEMBER_SPEND_SQL, DBSpendUpdateWriter, _spend_tables_left_to_send, _SpendTableName, + _virtual_key_team_id, ) from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import DailySpendUpdateQueue from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer @@ -1494,65 +1501,73 @@ async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_i ) -def _tagged_spend_log(tag: str, team_id: str | None, spend: float) -> SpendLogsPayload: - return cast( +def _validated_key(**fields: str | bool | None) -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="a" * 64, team_id="team-a").model_copy(update={"via_virtual_key": True, **fields}) + + +@pytest.mark.parametrize( + ("bucket", "auth", "expected"), + [ + ("metadata", _validated_key(), "team-a"), + ("litellm_metadata", _validated_key(), "team-a"), + ("metadata", _validated_key(team_id=None), ""), + ("metadata", _validated_key(via_virtual_key=False), ""), + ("metadata", _validated_key(is_session_token=True), ""), + ("metadata", {"via_virtual_key": True, "team_id": "team-a"}, ""), + ("metadata", None, ""), + ], + ids=[ + "virtual-key", + "virtual-key-litellm-metadata", + "virtual-key-without-team", + "jwt-or-custom-auth", + "ui-session", + "forged-metadata", + "no-auth", + ], +) +def test_tag_team_comes_only_from_a_proxy_validated_team_virtual_key( + bucket: str, auth: UserAPIKeyAuth | dict[str, object] | None, expected: str +) -> None: + kwargs: Final = {"litellm_params": {bucket: {"user_api_key_auth": auth, "user_api_key_team_id": "team-a"}}} + + assert _virtual_key_team_id(kwargs) == expected + + +@pytest.mark.asyncio +async def test_every_tag_row_of_a_request_carries_its_virtual_key_team() -> None: + writer: Final = DBSpendUpdateWriter() + prisma: Final = MagicMock() + prisma.get_request_status = MagicMock(return_value="success") + payload: Final = cast( SpendLogsPayload, { - "request_id": f"req-{tag}-{team_id}", - "request_tags": json.dumps([tag]), - "team_id": team_id, + "request_id": "req-1", + "request_tags": '["cost-center:cc-42", "prod"]', + "team_id": "team-from-jwt", "user": "user-1", - "startTime": "2026-10-17T00:00:00", - "api_key": "key-1", + "startTime": "2026-10-29T00:00:00", + "api_key": "a" * 64, "model": "gpt-4o", "custom_llm_provider": "openai", "model_group": "gpt-4o", "prompt_tokens": 10, "completion_tokens": 5, - "spend": spend, + "spend": 0.01, "metadata": '{"usage_object": {}}', }, ) - -@pytest.mark.parametrize( - ("requests", "expected"), - [ - ( - ( - ("cost-center:cc-42", "team-a", 1.0), - ("cost-center:cc-42", "team-b", 2.0), - ("cost-center:cc-42", None, 0.5), - ), - [ - ("cost-center:cc-42", "", 0.5), - ("cost-center:cc-42", "team-a", 1.0), - ("cost-center:cc-42", "team-b", 2.0), - ], - ), - ( - (("cost_center", "west", 1.0), ("cost", "center_west", 2.0)), - [("cost", "center_west", 2.0), ("cost_center", "west", 1.0)], - ), - ], - ids=["one-tag-two-teams", "underscore-joined-names-collide"], -) -async def test_each_tag_and_team_pair_is_queued_as_its_own_entry( - requests: tuple[tuple[str, str | None, float], ...], - expected: list[tuple[str, str, float]], -) -> None: - writer: Final = DBSpendUpdateWriter() - prisma: Final = MagicMock() - prisma.get_request_status = MagicMock(return_value="success") - - for tag, team_id, spend in requests: - await writer.add_spend_log_transaction_to_daily_tag_transaction( - payload=_tagged_spend_log(tag, team_id, spend), prisma_client=prisma - ) - + await writer.add_spend_log_transaction_to_daily_tag_transaction( + payload=payload, prisma_client=prisma, virtual_key_team_id="team-a" + ) queued: Final = await writer.daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + tag_rows: Final = tuple(cast(DailyTagSpendTransaction, txn) for txn in queued.values()) - assert sorted((txn["tag"], txn["team_id"], txn["spend"]) for txn in queued.values()) == expected, queued + assert sorted((txn["tag"], txn.get("team_id")) for txn in tag_rows) == [ + ("cost-center:cc-42", "team-a"), + ("prod", "team-a"), + ], queued @pytest.mark.asyncio