From 4b468053b792814cddaff58b0a2d74eba343cda1 Mon Sep 17 00:00:00 2001 From: Mrinal Chanshetty Date: Mon, 19 Oct 2026 12:55:30 -0700 Subject: [PATCH 1/7] feat(proxy): record the sending team on daily tag spend rows LiteLLM_DailyTagSpend gains team_id TEXT NOT NULL DEFAULT '' ("" = the request had no team) and the unique key becomes (tag, team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint), so one tag used by several teams is kept as one row per team. A (team_id, date, tag) index serves team-scoped tag reads. Migrations add the column, build the new unique index and the team index CONCURRENTLY, then drop the old unique index CONCURRENTLY. The spend writer stamps the request's team_id on each daily tag transaction and keys the in-memory/Redis merge on it; the bulk upsert adds team_id to the tag table's insert columns and ON CONFLICT target. Autorouter baseline adjustments carry team_id too. Entries queued by an older pod, and baseline records stored before this change, carry no team_id and land under "". Upgrade: an old-version writer targets the dropped unique index and fails every flush, so drain and stop old instances before running these migrations. LIT-8516 --- .../migration.sql | 8 ++++++++ .../migration.sql | 14 ++++++++++++++ .../migration.sql | 12 ++++++++++++ .../migration.sql | 15 +++++++++++++++ .../litellm_proxy_extras/schema.prisma | 4 +++- litellm/proxy/_types.py | 4 ++++ litellm/proxy/db/baseline_accounting.py | 4 ++++ litellm/proxy/db/daily_spend_bulk_upsert.py | 19 ++++++++++++++----- litellm/proxy/db/db_spend_update_writer.py | 8 ++++++-- litellm/proxy/schema.prisma | 4 +++- schema.prisma | 4 +++- .../proxy/db/test_daily_spend_bulk_upsert.py | 10 +++++----- 12 files changed, 91 insertions(+), 15 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000000_daily_tag_spend_add_team_id/migration.sql create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000001_daily_tag_spend_team_id_unique_index/migration.sql create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000002_daily_tag_spend_team_id_date_tag_index/migration.sql create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000003_daily_tag_spend_drop_tag_only_unique_index/migration.sql 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 new file mode 100644 index 00000000000..1a8322fa254 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000000_daily_tag_spend_add_team_id/migration.sql @@ -0,0 +1,8 @@ +-- 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 new file mode 100644 index 00000000000..1816ff63ce6 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000001_daily_tag_spend_team_id_unique_index/migration.sql @@ -0,0 +1,14 @@ +-- 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/20261017000002_daily_tag_spend_team_id_date_tag_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000002_daily_tag_spend_team_id_date_tag_index/migration.sql new file mode 100644 index 00000000000..474d5311305 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000002_daily_tag_spend_team_id_date_tag_index/migration.sql @@ -0,0 +1,12 @@ +-- CreateIndex (CONCURRENTLY) +-- +-- Serves team-scoped tag reads: WHERE team_id = ANY(...) AND date BETWEEN ... grouped by tag. +-- +-- 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 INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_DailyTagSpend_team_id_date_tag_idx" ON "LiteLLM_DailyTagSpend"("team_id", "date", "tag"); 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 new file mode 100644 index 00000000000..7db51879c01 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000003_daily_tag_spend_drop_tag_only_unique_index/migration.sql @@ -0,0 +1,15 @@ +-- 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. +-- +-- 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/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 83a0e6f9e05..c84b9f810fe 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1086,6 +1086,7 @@ model LiteLLM_DailyTagSpend { id String @id @default(uuid()) request_id String? tag String? + team_id String @default("") // "" = the request had no team date String api_key String model String? @@ -1111,9 +1112,10 @@ model LiteLLM_DailyTagSpend { created_at DateTime @default(now()) updated_at DateTime @updatedAt - @@unique([tag, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) + @@unique([tag, team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint], map: "LiteLLM_DailyTagSpend_tag_team_id_dimensions_key") @@index([date]) @@index([tag, date]) + @@index([team_id, date, tag]) @@index([api_key]) @@index([model]) @@index([mcp_namespaced_tool_name]) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index dc7383569dc..5d09280dc85 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -5159,6 +5159,10 @@ class DailyEndUserSpendTransaction(BaseDailySpendTransaction): class DailyTagSpendTransaction(BaseDailySpendTransaction): request_id: str | None tag: str + # "" when the request had no team. Not required: entries queued by a pod running the + # previous release, or replayed from the Redis buffer across an upgrade, carry no such key, + # and the upsert reads a missing value as "". + team_id: NotRequired[str] class DailyAgentSpendTransaction(BaseDailySpendTransaction): diff --git a/litellm/proxy/db/baseline_accounting.py b/litellm/proxy/db/baseline_accounting.py index f7a2be654c8..5aca1ebaa50 100644 --- a/litellm/proxy/db/baseline_accounting.py +++ b/litellm/proxy/db/baseline_accounting.py @@ -58,6 +58,9 @@ class DailyBaselineAttribution(LiteLLMBaseModel): model_group: str | None = None endpoint: str | None = None mcp_namespaced_tool_name: str | None = None + # The request's team, for the team-scoped tag rows. Defaulted so records stored before + # the column existed still validate; they land under "" like any pre-upgrade tag spend. + team_id: str = "" targets: tuple[DailyBaselineTarget, ...] = () def adjustment(self, target: DailyBaselineTarget, savings_delta: float, request_id: str) -> SpendRow: @@ -71,6 +74,7 @@ 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 eb130a5196f..b66d4cbc014 100644 --- a/litellm/proxy/db/daily_spend_bulk_upsert.py +++ b/litellm/proxy/db/daily_spend_bulk_upsert.py @@ -32,6 +32,9 @@ class DailySpendTable: name: str entity_id_column: str carries_request_id: bool = False + # Rows also keyed by the team that sent the request, so one entity id shared by several + # teams is kept as one row per team. + team_scoped: bool = False DAILY_SPEND_TABLES: Final[Mapping[DailySpendEntity, DailySpendTable]] = MappingProxyType( @@ -41,7 +44,9 @@ DAILY_SPEND_TABLES: Final[Mapping[DailySpendEntity, DailySpendTable]] = MappingP "org": DailySpendTable(name="LiteLLM_DailyOrganizationSpend", entity_id_column="organization_id"), "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), + "tag": DailySpendTable( + name="LiteLLM_DailyTagSpend", entity_id_column="tag", carries_request_id=True, team_scoped=True + ), } ) @@ -124,9 +129,14 @@ 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 (table.entity_id_column, *_KEY_COLUMNS)) + return tuple(_as_text(transaction.get(column)) for column in _conflict_columns(table)) def _merge(group: Sequence[SpendRow]) -> SpendRow: @@ -174,8 +184,7 @@ def _row_params( def _insert_columns(table: DailySpendTable) -> tuple[str, ...]: return ( "id", - table.entity_id_column, - *_KEY_COLUMNS, + *_conflict_columns(table), "model_group", *_COUNTER_COLUMNS, *_SPEND_COLUMNS, @@ -213,7 +222,7 @@ def build_bulk_upsert( sql: Final = ( f'INSERT INTO {quoted_table} ({_quoted(columns)}, "updated_at")\n' f"VALUES {rows}\n" - f"ON CONFLICT ({_quoted((table.entity_id_column, *_KEY_COLUMNS))}) DO UPDATE SET\n" + f"ON CONFLICT ({_quoted(_conflict_columns(table))}) DO UPDATE SET\n" f" {increments}{request_id_update},\n" f" \"updated_at\" = (NOW() AT TIME ZONE 'UTC')" ) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 166a42adfc0..70106ab6ba7 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -918,6 +918,7 @@ 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 "", targets=tuple( DailyBaselineTarget(entity=entity, entity_id=identity) for entity, values in bases @@ -3034,13 +3035,16 @@ class DBSpendUpdateWriter: return request_tags: Final = daily_spend_entity_ids(payload, "tag") + team_id: Final = payload.get("team_id") or "" for tag in request_tags: if tag is None: continue endpoint_str = base_daily_transaction.get("endpoint") or "" - daily_transaction_key = f"{tag}_{base_daily_transaction['date']}_{payload['api_key']}_{payload['model']}_{payload['custom_llm_provider']}_{endpoint_str}" + # team_id is part of the key so the queue never folds two teams' spend on one tag + # into a single entry before the upsert can keep them apart. + daily_transaction_key = f"{tag}_{team_id}_{base_daily_transaction['date']}_{payload['api_key']}_{payload['model']}_{payload['custom_llm_provider']}_{endpoint_str}" daily_transaction = DailyTagSpendTransaction( - tag=tag, **base_daily_transaction, request_id=payload["request_id"] + tag=tag, team_id=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 83a0e6f9e05..c84b9f810fe 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1086,6 +1086,7 @@ model LiteLLM_DailyTagSpend { id String @id @default(uuid()) request_id String? tag String? + team_id String @default("") // "" = the request had no team date String api_key String model String? @@ -1111,9 +1112,10 @@ model LiteLLM_DailyTagSpend { created_at DateTime @default(now()) updated_at DateTime @updatedAt - @@unique([tag, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) + @@unique([tag, team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint], map: "LiteLLM_DailyTagSpend_tag_team_id_dimensions_key") @@index([date]) @@index([tag, date]) + @@index([team_id, date, tag]) @@index([api_key]) @@index([model]) @@index([mcp_namespaced_tool_name]) diff --git a/schema.prisma b/schema.prisma index 83a0e6f9e05..c84b9f810fe 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1086,6 +1086,7 @@ model LiteLLM_DailyTagSpend { id String @id @default(uuid()) request_id String? tag String? + team_id String @default("") // "" = the request had no team date String api_key String model String? @@ -1111,9 +1112,10 @@ model LiteLLM_DailyTagSpend { created_at DateTime @default(now()) updated_at DateTime @updatedAt - @@unique([tag, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) + @@unique([tag, team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint], map: "LiteLLM_DailyTagSpend_tag_team_id_dimensions_key") @@index([date]) @@index([tag, date]) + @@index([team_id, date, tag]) @@index([api_key]) @@index([model]) @@index([mcp_namespaced_tool_name]) 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 5bce9137666..0912c0c406d 100644 --- a/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py +++ b/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py @@ -88,10 +88,10 @@ 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 - # 25 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 * 25 - assert "$2500::text" in sql + # 26 bound columns per row (team_id included) 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", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"' + '"tag", "team_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"' ) From 5cbe0eb02e6dea20346803596b7fefe03f4637ee Mon Sep 17 00:00:00 2001 From: Mrinal Chanshetty Date: Mon, 19 Oct 2026 13:41:40 -0700 Subject: [PATCH 2/7] test(proxy): cover team-attributed daily tag spend Unit: one tag from two teams stays two upsert rows; an entry with no team_id (queued by an older pod) folds into the "" row; only the tag table is keyed by team_id. The writer queues one tag from two teams as separate team-attributed entries, read back through the real queue. Integration: two teams send the same tag through the proxy and LiteLLM_DailyTagSpend holds one row per team with the hand-computed spend. Mutations (each red, then restored): - team_id dropped from the upsert conflict key: the two-team and legacy-entry upsert tests fail - team_id added to every table's conflict key: all five per-table tests fail - team_id dropped from the writer's queue key: the writer test fails, with all three requests merged into one entry credited to team-a - the writer stops recording the team: the integration test fails, with both teams' spend under team_id '' LIT-8516 --- .../spend/test_team_tag_spend_attribution.py | 44 +++++++++++++++++++ .../proxy/db/test_daily_spend_bulk_upsert.py | 39 ++++++++++++++++ .../proxy/db/test_db_spend_update_writer.py | 39 ++++++++++++++++ 3 files changed, 122 insertions(+) create mode 100644 tests/integration/spend/test_team_tag_spend_attribution.py diff --git a/tests/integration/spend/test_team_tag_spend_attribution.py b/tests/integration/spend/test_team_tag_spend_attribution.py new file mode 100644 index 00000000000..b0bd50ca44e --- /dev/null +++ b/tests/integration/spend/test_team_tag_spend_attribution.py @@ -0,0 +1,44 @@ +import uuid +from typing import Final + +import pytest +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows + +_SPEND_PER_CALL: Final = 20 * 0.001 + 20 * 0.002 + + +def test_one_tag_sent_by_two_teams_is_recorded_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} + for team, calls in calls_per_team.items(): + key = scenario.key(team_id=team, models=[model]) + 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 + + rows: Final = eventually( + lambda: 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, + ) + + assert sorted((row["team_id"], int(row["api_requests"])) for row in rows) == sorted(calls_per_team.items()), ( + rows + ) + for row in rows: + assert float(row["spend"]) == pytest.approx(calls_per_team[row["team_id"]] * _SPEND_PER_CALL), rows 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 0912c0c406d..9886ec91cda 100644 --- a/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py +++ b/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py @@ -144,6 +144,45 @@ 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(): + merged = 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(): + """Entries queued by an older pod, or replayed from Redis across the upgrade, carry no + team_id key. They arbitrate to the same row as an explicit "" rather than a second one.""" + legacy = tag_txn(spend=0.25) + unattributed = tag_txn(team_id="", spend=0.75) + + merged = 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): + """Only LiteLLM_DailyTagSpend has the extra team_id column; adding it to any other + table's conflict target names a column that table does not have.""" + table = DAILY_SPEND_TABLES[entity] + row = {**tag_txn(), table.entity_id_column: "id-1"} + + sql, _ = build_bulk_upsert(table, merge_by_conflict_key(table, (row,))) + + conflict_target = 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 1ec65a4fdcd..af7a00bae25 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -1488,6 +1488,45 @@ async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_i ) +@pytest.mark.asyncio +async def test_one_tag_from_two_teams_is_queued_as_two_team_attributed_entries(): + """The queue merges entries by key before the upsert ever sees them, so a key without the + team would fold both teams' spend on a shared tag into one entry.""" + writer = DBSpendUpdateWriter() + prisma = MagicMock() + prisma.get_request_status = MagicMock(return_value="success") + + def tagged(team_id: str | None, spend: float) -> dict: + return { + "request_id": f"req-{team_id}", + "request_tags": '["cost-center:cc-42"]', + "team_id": team_id, + "user": "user-1", + "startTime": "2026-10-17T00:00:00", + "api_key": "key-1", + "model": "gpt-4o", + "custom_llm_provider": "openai", + "model_group": "gpt-4o", + "prompt_tokens": 10, + "completion_tokens": 5, + "spend": spend, + "metadata": '{"usage_object": {}}', + } + + for team_id, spend in (("team-a", 1.0), ("team-b", 2.0), (None, 0.5)): + await writer.add_spend_log_transaction_to_daily_tag_transaction( + payload=tagged(team_id, spend), prisma_client=prisma + ) + + queued = await writer.daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + + assert sorted((txn["tag"], txn["team_id"], txn["spend"]) for txn in queued.values()) == [ + ("cost-center:cc-42", "", 0.5), + ("cost-center:cc-42", "team-a", 1.0), + ("cost-center:cc-42", "team-b", 2.0), + ], queued + + @pytest.mark.asyncio async def test_add_spend_log_transaction_to_daily_org_transaction_injects_org_id_and_queues_update(): """ From 9e2e2eb6dc6e32ef2e6725469f99182b53a5a564 Mon Sep 17 00:00:00 2001 From: Mrinal Chanshetty Date: Mon, 19 Oct 2026 15:38:23 -0700 Subject: [PATCH 3/7] fix(proxy): satisfy type gates for team-attributed tag spend DailyTagSpendTransaction.team_id is ReadOnly (LIT012), since nothing writes it after the transaction is built. The two payload.get("team_id") reads keep .get rather than indexing and carry the file's existing reportUnknownMemberType suppression for TypedDict.get. The caller swallows exceptions at debug level, so indexing would silently drop a request's tag spend whenever a payload lacked team_id, instead of recording it as unattributed. LIT-8516 --- litellm/proxy/_types.py | 2 +- litellm/proxy/db/db_spend_update_writer.py | 6 ++++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 5d09280dc85..11ef5d67a4b 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -5162,7 +5162,7 @@ class DailyTagSpendTransaction(BaseDailySpendTransaction): # "" when the request had no team. Not required: entries queued by a pod running the # previous release, or replayed from the Redis buffer across an upgrade, carry no such key, # and the upsert reads a missing value as "". - team_id: NotRequired[str] + team_id: ReadOnly[NotRequired[str]] class DailyAgentSpendTransaction(BaseDailySpendTransaction): diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 70106ab6ba7..614870ad133 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -918,7 +918,7 @@ 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 "", + 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 @@ -3035,7 +3035,9 @@ class DBSpendUpdateWriter: return request_tags: Final = daily_spend_entity_ids(payload, "tag") - team_id: Final = payload.get("team_id") or "" + # .get, not indexing: the caller swallows exceptions at debug level, so a payload without + # team_id must still record its tag spend (as unattributed) rather than drop it silently. + 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 From 79fd5378bc55b8ea527f785ee9a067067e6e8d88 Mon Sep 17 00:00:00 2001 From: Mrinal Chanshetty Date: Mon, 19 Oct 2026 20:05:20 -0700 Subject: [PATCH 4/7] fix(proxy): set lock_timeout on team_id migration and allow CONCURRENTLY in DROP INDEX check The team_id column migration runs a plain ALTER TABLE, so it now starts with SET LOCAL lock_timeout like the other recent migrations. test_drop_index_uses_if_exists rejected DROP INDEX CONCURRENTLY IF EXISTS; it now accepts an optional CONCURRENTLY, matching the CREATE INDEX check. --- .../20261017000000_daily_tag_spend_add_team_id/migration.sql | 2 ++ .../litellm_proxy_extras/test_litellm_proxy_extras_utils.py | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) 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 index 1a8322fa254..0b73df2d05b 100644 --- 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 @@ -1,3 +1,5 @@ +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). 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 02bdf3429ea..ad889af1873 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+IF\s+EXISTS", line, re.IGNORECASE + r"DROP\s+INDEX\s+(?:CONCURRENTLY\s+)?IF\s+EXISTS", line, re.IGNORECASE ): violations.append(f" {migration_name}:{line_num}: {line.strip()}") assert ( From 0e5f178cfe5f308c0bade19f85ec01248864b167 Mon Sep 17 00:00:00 2001 From: Mrinal Chanshetty Date: Tue, 20 Oct 2026 14:33:35 -0700 Subject: [PATCH 5/7] test(proxy): prove team tag rows with one key shared by two teams The integration test gave each team its own key, so api_key alone kept the rows apart and dropping team_id from the tag queue key still passed. One key now moves between the two teams through /key/update, the shape JWT callers (all recorded as api_key "") and moved keys produce. The upgrade note in the drop-index migration warns operators not to run the schema-sync SQL an old instance prints, which drops team_id and both new indexes. Trims comments that restated the code. --- .../migration.sql | 3 +++ litellm/proxy/_types.py | 3 --- litellm/proxy/db/baseline_accounting.py | 2 -- litellm/proxy/db/daily_spend_bulk_upsert.py | 2 -- litellm/proxy/db/db_spend_update_writer.py | 5 +---- tests/integration/spend/test_team_tag_spend_attribution.py | 5 +++-- tests/unit/proxy/db/test_daily_spend_bulk_upsert.py | 6 ------ tests/unit/proxy/db/test_db_spend_update_writer.py | 2 -- 8 files changed, 7 insertions(+), 21 deletions(-) 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 index 7db51879c01..ebc4f7ac457 100644 --- 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 @@ -6,6 +6,9 @@ -- -- 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 diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 11ef5d67a4b..e34b1d7b9a4 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -5159,9 +5159,6 @@ class DailyEndUserSpendTransaction(BaseDailySpendTransaction): class DailyTagSpendTransaction(BaseDailySpendTransaction): request_id: str | None tag: str - # "" when the request had no team. Not required: entries queued by a pod running the - # previous release, or replayed from the Redis buffer across an upgrade, carry no such key, - # and the upsert reads a missing value as "". team_id: ReadOnly[NotRequired[str]] diff --git a/litellm/proxy/db/baseline_accounting.py b/litellm/proxy/db/baseline_accounting.py index 5aca1ebaa50..5e817410d2b 100644 --- a/litellm/proxy/db/baseline_accounting.py +++ b/litellm/proxy/db/baseline_accounting.py @@ -58,8 +58,6 @@ class DailyBaselineAttribution(LiteLLMBaseModel): model_group: str | None = None endpoint: str | None = None mcp_namespaced_tool_name: str | None = None - # The request's team, for the team-scoped tag rows. Defaulted so records stored before - # the column existed still validate; they land under "" like any pre-upgrade tag spend. team_id: str = "" targets: tuple[DailyBaselineTarget, ...] = () diff --git a/litellm/proxy/db/daily_spend_bulk_upsert.py b/litellm/proxy/db/daily_spend_bulk_upsert.py index b66d4cbc014..dcc2b738a51 100644 --- a/litellm/proxy/db/daily_spend_bulk_upsert.py +++ b/litellm/proxy/db/daily_spend_bulk_upsert.py @@ -32,8 +32,6 @@ class DailySpendTable: name: str entity_id_column: str carries_request_id: bool = False - # Rows also keyed by the team that sent the request, so one entity id shared by several - # teams is kept as one row per team. team_scoped: bool = False diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 614870ad133..c7b31e4cc0f 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -3035,15 +3035,12 @@ class DBSpendUpdateWriter: return request_tags: Final = daily_spend_entity_ids(payload, "tag") - # .get, not indexing: the caller swallows exceptions at debug level, so a payload without - # team_id must still record its tag spend (as unattributed) rather than drop it silently. + # .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 "" - # team_id is part of the key so the queue never folds two teams' spend on one tag - # into a single entry before the upsert can keep them apart. daily_transaction_key = f"{tag}_{team_id}_{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"] diff --git a/tests/integration/spend/test_team_tag_spend_attribution.py b/tests/integration/spend/test_team_tag_spend_attribution.py index b0bd50ca44e..d82394e143b 100644 --- a/tests/integration/spend/test_team_tag_spend_attribution.py +++ b/tests/integration/spend/test_team_tag_spend_attribution.py @@ -8,14 +8,15 @@ from integration._support.database import read_rows _SPEND_PER_CALL: Final = 20 * 0.001 + 20 * 0.002 -def test_one_tag_sent_by_two_teams_is_recorded_once_per_team(gateway: Gateway) -> None: +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(): - key = scenario.key(team_id=team, models=[model]) + gateway.post("/key/update", {"key": key, "team_id": team}) for _ in range(calls): reply = gateway.post( "/v1/chat/completions", 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 9886ec91cda..dd4267906ed 100644 --- a/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py +++ b/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py @@ -88,8 +88,6 @@ 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 (team_id included) 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 @@ -153,8 +151,6 @@ def test_one_tag_from_two_teams_stays_two_rows(): def test_an_entry_without_team_id_folds_into_the_unattributed_row(): - """Entries queued by an older pod, or replayed from Redis across the upgrade, carry no - team_id key. They arbitrate to the same row as an explicit "" rather than a second one.""" legacy = tag_txn(spend=0.25) unattributed = tag_txn(team_id="", spend=0.75) @@ -168,8 +164,6 @@ def test_an_entry_without_team_id_folds_into_the_unattributed_row(): @pytest.mark.parametrize("entity", ["user", "team", "org", "end_user", "agent"]) def test_only_the_tag_table_is_keyed_by_team(entity): - """Only LiteLLM_DailyTagSpend has the extra team_id column; adding it to any other - table's conflict target names a column that table does not have.""" table = DAILY_SPEND_TABLES[entity] row = {**tag_txn(), table.entity_id_column: "id-1"} 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 af7a00bae25..d9acb8035d1 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -1490,8 +1490,6 @@ async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_i @pytest.mark.asyncio async def test_one_tag_from_two_teams_is_queued_as_two_team_attributed_entries(): - """The queue merges entries by key before the upsert ever sees them, so a key without the - team would fold both teams' spend on a shared tag into one entry.""" writer = DBSpendUpdateWriter() prisma = MagicMock() prisma.get_request_status = MagicMock(return_value="success") From f9e70654a2a4bf30ee68d7404da31918c8548d4d Mon Sep 17 00:00:00 2001 From: Mrinal Chanshetty Date: Tue, 20 Oct 2026 16:08:21 -0700 Subject: [PATCH 6/7] fix(ci): stop the schema compatibility check reading block attributes as fields parse_model_fields split every model line on whitespace and keyed it by the first token, so @@unique([tag, date, ...]) became a field named "@@unique([tag," and its column list its type. Changing a composite unique key failed the check as a type change, and @@ lines sharing that first token overwrote each other, so most of them were never compared. Comment lines were read as fields too. Block attributes and comments are now skipped, so the check compares columns, which is what its messages and tests describe. --- scripts/check_schema_compatibility.py | 6 +++++- tests/unit/proxy/test_db_schema_changes.py | 22 ++++++++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/scripts/check_schema_compatibility.py b/scripts/check_schema_compatibility.py index 898c65dff2c..5e98db3c7e2 100644 --- a/scripts/check_schema_compatibility.py +++ b/scripts/check_schema_compatibility.py @@ -15,10 +15,14 @@ Schema = Mapping[str, Mapping[str, str]] MODELS: Final = re.compile(r"^[ \t]*model\s+(\w+)\s*\{(.*?)^[ \t]*\}", re.MULTILINE | re.DOTALL) +def _is_field(part: list[str]) -> bool: + return len(part) >= 2 and not part[0].startswith(("@@", "//")) + + def parse_model_fields(schema: str) -> Schema: def fields(body: str) -> Mapping[str, str]: parts: Final = tuple(line.split() for line in body.splitlines()) - return MappingProxyType({part[0]: " ".join(part[1:]) for part in parts if len(part) >= 2}) + return MappingProxyType({part[0]: " ".join(part[1:]) for part in parts if _is_field(part)}) return MappingProxyType({match[1]: fields(match[2]) for match in MODELS.finditer(schema)}) diff --git a/tests/unit/proxy/test_db_schema_changes.py b/tests/unit/proxy/test_db_schema_changes.py index 3a5d0fb1a0a..49dfaec1f16 100644 --- a/tests/unit/proxy/test_db_schema_changes.py +++ b/tests/unit/proxy/test_db_schema_changes.py @@ -16,6 +16,28 @@ def test_parse_model_fields_retains_types_and_attributes() -> None: assert parse_model_fields(ORIGINAL) == {"Item": {"id": "String @id", "label": "String?"}} +def test_block_attributes_and_comments_are_not_fields() -> None: + schema: Final = ( + "model Item {\n id String @id\n label String?\n // free text\n" + ' @@unique([label, id], map: "item_label_id_key")\n @@index([label, id])\n @@index([id])\n}\n' + ) + assert parse_model_fields(schema) == {"Item": {"id": "String @id", "label": "String?"}} + + +def test_changing_a_composite_unique_key_is_not_a_field_change() -> None: + before: Final = "model Item {\n id String @id\n label String?\n @@unique([label, id])\n}\n" + after: Final = "model Item {\n id String @id\n label String?\n @@unique([label, id, tenant])\n}\n" + assert check_breaking_changes(parse_model_fields(before), parse_model_fields(after)) == () + + +def test_a_narrowed_column_is_still_caught_next_to_block_attributes() -> None: + before: Final = "model Item {\n id String @id\n label String?\n @@index([label, id])\n}\n" + after: Final = "model Item {\n id String @id\n label String\n @@index([label, id])\n}\n" + assert check_breaking_changes(parse_model_fields(before), parse_model_fields(after)) == ( + "Breaking: Field Item.label changed from optional to required", + ) + + def test_parse_model_fields_keeps_json_defaults_and_empty_models() -> None: schema: Final = 'model Item {\n data Json @default("{}")\n}\nmodel Empty {\n}\n' assert parse_model_fields(schema) == {"Item": {"data": 'Json @default("{}")'}, "Empty": {}} From 3982bcc9d7b00f0ddab470571d96a8c6490001d5 Mon Sep 17 00:00:00 2001 From: Mrinal Chanshetty Date: Sat, 24 Oct 2026 20:18:20 -0500 Subject: [PATCH 7/7] fix(proxy): keep tag and team apart in the daily tag spend queue key The tag queue key joined its fields with "_", and the new team_id sat next to the free-text tag, so tag "cost_center" from team "west" and tag "cost" from team "center_west" got the same key. The queue then summed both into one entry under the first team, before the upsert could keep them apart. The key is now a JSON list of the same fields. Also types the new tests (Final, return types, the entity literal). --- litellm/proxy/db/db_spend_update_writer.py | 12 +++- .../proxy/db/test_daily_spend_bulk_upsert.py | 22 ++++--- .../proxy/db/test_db_spend_update_writer.py | 64 +++++++++++++------ 3 files changed, 67 insertions(+), 31 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 2e11fdd9c9d..22540866a77 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -3080,7 +3080,17 @@ class DBSpendUpdateWriter: if tag is None: continue endpoint_str = base_daily_transaction.get("endpoint") or "" - daily_transaction_key = f"{tag}_{team_id}_{base_daily_transaction['date']}_{payload['api_key']}_{payload['model']}_{payload['custom_llm_provider']}_{endpoint_str}" + 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 = DailyTagSpendTransaction( tag=tag, team_id=team_id, **base_daily_transaction, request_id=payload["request_id"] ) 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 dd4267906ed..61a7d7dac7f 100644 --- a/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py +++ b/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py @@ -3,11 +3,13 @@ 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, @@ -142,19 +144,19 @@ 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(): - merged = merge_by_conflict_key( +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(): - legacy = tag_txn(spend=0.25) - unattributed = tag_txn(team_id="", spend=0.75) +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 = merge_by_conflict_key(TAG_TABLE, (legacy, unattributed)) + merged: Final = merge_by_conflict_key(TAG_TABLE, (legacy, unattributed)) assert len(merged) == 1, merged key, folded = merged[0] @@ -163,13 +165,13 @@ def test_an_entry_without_team_id_folds_into_the_unattributed_row(): @pytest.mark.parametrize("entity", ["user", "team", "org", "end_user", "agent"]) -def test_only_the_tag_table_is_keyed_by_team(entity): - table = DAILY_SPEND_TABLES[entity] - row = {**tag_txn(), table.entity_id_column: "id-1"} +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 = re.search(r"ON CONFLICT \(([^)]*)\)", sql) + 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", ' 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 ec3620a32e4..ac235902925 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -20,7 +20,7 @@ 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, SpendUpdateQueueItem +from litellm.proxy._types import DailyTagSpendTransaction, Litellm_EntityType, SpendLogsPayload, SpendUpdateQueueItem from litellm.proxy.db.db_spend_update_writer import ( _TEAM_ADVISORY_LOCK_SQL, _TEAM_MEMBER_SPEND_SQL, @@ -1490,16 +1490,12 @@ async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_i ) -@pytest.mark.asyncio -async def test_one_tag_from_two_teams_is_queued_as_two_team_attributed_entries(): - writer = DBSpendUpdateWriter() - prisma = MagicMock() - prisma.get_request_status = MagicMock(return_value="success") - - def tagged(team_id: str | None, spend: float) -> dict: - return { - "request_id": f"req-{team_id}", - "request_tags": '["cost-center:cc-42"]', +def _tagged_spend_log(tag: str, team_id: str | None, spend: float) -> SpendLogsPayload: + return cast( + SpendLogsPayload, + { + "request_id": f"req-{tag}-{team_id}", + "request_tags": json.dumps([tag]), "team_id": team_id, "user": "user-1", "startTime": "2026-10-17T00:00:00", @@ -1511,20 +1507,48 @@ async def test_one_tag_from_two_teams_is_queued_as_two_team_attributed_entries() "completion_tokens": 5, "spend": spend, "metadata": '{"usage_object": {}}', - } + }, + ) - for team_id, spend in (("team-a", 1.0), ("team-b", 2.0), (None, 0.5)): + +@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(team_id, spend), prisma_client=prisma + payload=_tagged_spend_log(tag, team_id, spend), prisma_client=prisma ) - queued = await writer.daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + queued: Final = await writer.daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() - assert sorted((txn["tag"], txn["team_id"], txn["spend"]) for txn in queued.values()) == [ - ("cost-center:cc-42", "", 0.5), - ("cost-center:cc-42", "team-a", 1.0), - ("cost-center:cc-42", "team-b", 2.0), - ], queued + assert sorted((txn["tag"], txn["team_id"], txn["spend"]) for txn in queued.values()) == expected, queued @pytest.mark.asyncio