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

litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000000_daily_tag_spend_add_team_id/migration.sql

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

litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000001_daily_tag_spend_team_id_unique_index/migration.sql

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

litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000003_daily_tag_spend_drop_tag_only_unique_index/migration.sql

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

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 '';

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])

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,

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

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

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])

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])

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

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 (

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, ...]]] = []

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