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