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/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 1a67ec2108a..8152020ef4a 100644 --- a/tests/integration/spend/test_team_tag_spend_attribution.py +++ b/tests/integration/spend/test_team_tag_spend_attribution.py @@ -47,14 +47,15 @@ def _tag_breakdown(gateway: Gateway, team: str, *, key: str | None = None) -> di return totals -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/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 ( 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")