fix(proxy): attribute tag spend to the virtual key's team on every spend path

The tag writer read the in-memory auth object, which the spend collector, batch and background
cost polling, and the failure hook all drop, so their tag rows never got a team. The proxy now
stamps user_api_key_is_virtual_key in request metadata from the unforgeable via_virtual_key
marker (client-sent user_api_key_* fields are already stripped), the failure hook and batch
attribution set it too, and the writer reads it with user_api_key_team_id

The schema compatibility check flags a removed or changed @@unique again, since an older
proxy's ON CONFLICT names it, and ignores trailing // comments on field lines

enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py

@@ -15,7 +15,7 @@ from typing import TYPE_CHECKING, Final, Protocol, TypeGuard, cast
 from pydantic import TypeAdapter
 
 from litellm._logging import verbose_proxy_logger
-from litellm.constants import CLI_SESSION_KEY_PREFIX
+from litellm.constants import CLI_SESSION_KEY_PREFIX, USER_API_KEY_IS_VIRTUAL_KEY_METADATA_KEY
 from litellm.repositories.table_repositories import (
     DeletedVerificationTokenRepository,
     DeprecatedVerificationTokenRepository,
@@ -310,6 +310,7 @@ def _shared_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str,
         "user_api_key": api_key,
         "user_api_key_hash": api_key,
         "user_api_key_team_id": team_id,
+        USER_API_KEY_IS_VIRTUAL_KEY_METADATA_KEY: rows.key is not None,
         **(
             {
                 "user_api_key_user_email": getattr(rows.user, "user_email", None),

litellm/constants.py

@@ -1693,6 +1693,7 @@ ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: Final = "_litellm_router_usage_counted
 INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
 RESPONSES_WEBSOCKET_CALL_ID_METADATA_KEY: Final = "responses_websocket_call_id"
 BILLED_RETRY_ATTEMPT_METADATA_KEY: Final = "billed_retry_attempt"
+USER_API_KEY_IS_VIRTUAL_KEY_METADATA_KEY: Final = "user_api_key_is_virtual_key"
 MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: Final = "_managed_background_response"
 SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated"
 SESSION_ID_OMITTED_METADATA_KEY: Final = "litellm_session_id_omitted"

litellm/proxy/db/db_spend_update_writer.py

@@ -33,6 +33,7 @@ from litellm.constants import (
     DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME,
     DB_SPEND_UPDATE_JOB_NAME,
     INTERNAL_CALL_ORIGIN_METADATA_KEY,
+    USER_API_KEY_IS_VIRTUAL_KEY_METADATA_KEY,
 )
 from litellm.litellm_core_utils.core_helpers import finite_cost, single_log_line
 from litellm.litellm_core_utils.litellm_logging import coerce_model_access_groups
@@ -52,7 +53,6 @@ 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 (
@@ -356,14 +356,14 @@ 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")
+    buckets: Final = tuple(
+        _SPEND_METADATA_ADAPTER.validate_python(params.get(bucket) or {}) 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 ""
+    marked: Final = next(
+        (bucket for bucket in buckets if bucket.get(USER_API_KEY_IS_VIRTUAL_KEY_METADATA_KEY) is True), None
+    )
+    team_id: Final = marked.get("user_api_key_team_id") if marked is not None else None
+    return team_id if isinstance(team_id, str) else ""
 
 
 def _is_billed_retry_attempt(metadata: SpendLogsMetadata) -> bool:

litellm/proxy/hooks/proxy_track_cost_callback.py

@@ -13,7 +13,11 @@ import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.batches.batch_utils import batch_cost_is_final
 from litellm.caching.redis_batch import flush_post_call_redis_batches
-from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED, BILLED_RETRY_ATTEMPT_METADATA_KEY
+from litellm.constants import (
+    BACKGROUND_INTERACTION_COST_POLLING_ENABLED,
+    BILLED_RETRY_ATTEMPT_METADATA_KEY,
+    USER_API_KEY_IS_VIRTUAL_KEY_METADATA_KEY,
+)
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.litellm_core_utils.core_helpers import (
     budget_reservation_from_metadata,
@@ -480,6 +484,9 @@ class _ProxyDBLogger(CustomLogger):
             LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_dict)
         )
         _metadata["user_api_key"] = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict)
+        _metadata[USER_API_KEY_IS_VIRTUAL_KEY_METADATA_KEY] = LiteLLMProxyRequestSetup.authenticated_with_virtual_key(
+            user_api_key_dict
+        )
         _metadata["status"] = "failure"
         _error_information = StandardLoggingPayloadSetup.get_error_information(
             original_exception=original_exception,

litellm/proxy/litellm_pre_call_utils.py

@@ -45,6 +45,7 @@ from litellm.constants import (
     SESSION_ID_GENERATED_METADATA_KEY,
     SESSION_ID_OMITTED_METADATA_KEY,
     TYPESAFE_DANGER_ZONE_HEADER_PREFIX,
+    USER_API_KEY_IS_VIRTUAL_KEY_METADATA_KEY,
     VECTOR_STORE_CALLER_METADATA_KWARG,
     X_LITELLM_DISABLE_CALLBACKS,
 )
@@ -2093,6 +2094,13 @@ class LiteLLMProxyRequestSetup:
         project_tags: Final = (user_api_key_dict.project_metadata or {}).get("tags") or ()
         return tuple(dict.fromkeys((*key_tags, *team_tags, *project_tags)))
 
+    @staticmethod
+    def authenticated_with_virtual_key(user_api_key_dict: UserAPIKeyAuth) -> bool:
+        """A serializable copy of the unforgeable via_virtual_key marker, carried in request metadata so
+        spend paths that drop the auth object (collector, batch and background polling, failure logging)
+        still know the request came through a virtual key."""
+        return user_api_key_dict.via_virtual_key and not user_api_key_dict.is_session_token
+
     @staticmethod
     def add_user_api_key_auth_to_request_metadata(
         data: dict,
@@ -2107,6 +2115,9 @@ class LiteLLMProxyRequestSetup:
         )
         data[_metadata_variable_name].update(user_api_key_logged_metadata)
         data[_metadata_variable_name]["user_api_key"] = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict)
+        data[_metadata_variable_name][USER_API_KEY_IS_VIRTUAL_KEY_METADATA_KEY] = (
+            LiteLLMProxyRequestSetup.authenticated_with_virtual_key(user_api_key_dict)
+        )
 
         # Key-owned agent_id for spend attribution; keep existing (e.g. from header) if key has none
         _key_agent_id: Final = getattr(user_api_key_dict, "agent_id", None)

scripts/check_schema_compatibility.py

@@ -12,21 +12,47 @@ from typing import Final
 from pydantic import TypeAdapter
 
 Schema = Mapping[str, Mapping[str, str]]
+UniqueKeys = Mapping[str, frozenset[str]]
 MODELS: Final = re.compile(r"^[ \t]*model\s+(\w+)\s*\{(.*?)^[ \t]*\}", re.MULTILINE | re.DOTALL)
 
 
+def _without_comment(tokens: list[str]) -> list[str]:
+    return tokens[: next((i for i, token in enumerate(tokens) if token.startswith("//")), len(tokens))]
+
+
 def _is_field(part: list[str]) -> bool:
-    return len(part) >= 2 and not part[0].startswith(("@@", "//"))
+    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())
+        parts: Final = tuple(_without_comment(line.split()) for line in body.splitlines())
         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)})
 
 
+def parse_unique_keys(schema: str) -> UniqueKeys:
+    def unique_keys(body: str) -> frozenset[str]:
+        parts: Final = tuple(_without_comment(line.split()) for line in body.splitlines())
+        return frozenset(" ".join(part) for part in parts if part and part[0].startswith("@@unique"))
+
+    return MappingProxyType({match[1]: unique_keys(match[2]) for match in MODELS.finditer(schema)})
+
+
+def check_unique_key_changes(old_keys: UniqueKeys, new_keys: UniqueKeys) -> tuple[str, ...]:
+    """An older proxy's upsert names its unique key in ON CONFLICT, so dropping or changing one breaks it."""
+
+    def changes() -> Iterator[str]:
+        for model, keys in old_keys.items():
+            if model not in new_keys:
+                continue
+            for key in sorted(keys - new_keys[model]):
+                yield f"Breaking: Unique key {model}.{key} was removed or changed"
+
+    return tuple(changes())
+
+
 def _field_changes(model: str, field: str, old_type: str, new_fields: Mapping[str, str]) -> Iterator[str]:
     if field not in new_fields:
         yield f"Breaking: Field {model}.{field} was removed"
@@ -67,8 +93,10 @@ def main() -> int:
     previous: Final = subprocess.run(
         ("git", "show", f"{base}:schema.prisma"), capture_output=True, text=True, check=True
     ).stdout
-    changes: Final = check_breaking_changes(
-        parse_model_fields(previous), parse_model_fields(Path("schema.prisma").read_text())
+    current: Final = Path("schema.prisma").read_text()
+    changes: Final = (
+        *check_breaking_changes(parse_model_fields(previous), parse_model_fields(current)),
+        *check_unique_key_changes(parse_unique_keys(previous), parse_unique_keys(current)),
     )
     if changes:
         sys.stderr.write("\n".join(changes) + "\n")

tests/unit/proxy/common_utils/test_check_batch_cost.py

@@ -2486,6 +2486,21 @@ class TestBatchCostAttribution:
         assert metadata["user_api_key_team_alias"] == "Team Alpha"
         assert metadata["tags"] == ["env:prod"]
 
+    @pytest.mark.parametrize(("key_row_exists", "expected"), [(True, "team-alpha"), (False, "")])
+    @pytest.mark.asyncio
+    async def test_tag_spend_team_needs_the_batch_to_come_from_a_virtual_key(self, key_row_exists, expected):
+        """Tag spend is attributed to a team only when the creating key is a real virtual key row,
+        not a JWT, a CLI session or a deleted key."""
+        from types import SimpleNamespace
+
+        from litellm.proxy.db.db_spend_update_writer import _virtual_key_team_id
+
+        instance = self._instance(key_row=SimpleNamespace(key_alias="prod-key") if key_row_exists else None)
+
+        metadata = await instance.batch_metadata(self._job(), "batch-1")
+
+        assert _virtual_key_team_id({"litellm_params": {"metadata": metadata}}) == expected
+
     @pytest.mark.asyncio
     async def test_metadata_tolerates_legacy_row_without_columns(self):
         """Rows created before the columns existed carry only created_by/team_id and must

tests/unit/proxy/db/test_db_spend_update_writer.py

@@ -20,13 +20,12 @@ from redis.exceptions import DataError
 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.constants import SPEND_ROLLUP_LOCK_TIMEOUT_MS, USER_API_KEY_IS_VIRTUAL_KEY_METADATA_KEY
 from litellm.proxy._types import (
     DailyTagSpendTransaction,
     Litellm_EntityType,
     SpendLogsPayload,
     SpendUpdateQueueItem,
-    UserAPIKeyAuth,
 )
 from litellm.proxy.db.db_spend_update_writer import (
     _TEAM_ADVISORY_LOCK_SQL,
@@ -1501,37 +1500,34 @@ async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_i
             )
 
 
-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"),
+    ("bucket", "marker", "team_id", "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, ""),
+        ("metadata", True, "team-a", "team-a"),
+        ("litellm_metadata", True, "team-a", "team-a"),
+        ("metadata", True, None, ""),
+        ("metadata", False, "team-a", ""),
+        ("metadata", "true", "team-a", ""),
+        ("metadata", None, "team-a", ""),
     ],
     ids=[
         "virtual-key",
         "virtual-key-litellm-metadata",
         "virtual-key-without-team",
         "jwt-or-custom-auth",
-        "ui-session",
-        "forged-metadata",
-        "no-auth",
+        "non-boolean-marker",
+        "no-marker",
     ],
 )
-def test_tag_team_comes_only_from_a_proxy_validated_team_virtual_key(
-    bucket: str, auth: UserAPIKeyAuth | dict[str, object] | None, expected: str
+def test_tag_team_comes_only_from_a_request_marked_as_a_virtual_key(
+    bucket: str, marker: object, team_id: str | None, expected: str
 ) -> None:
-    kwargs: Final = {"litellm_params": {bucket: {"user_api_key_auth": auth, "user_api_key_team_id": "team-a"}}}
+    metadata: Final = {
+        "user_api_key_team_id": team_id,
+        **({USER_API_KEY_IS_VIRTUAL_KEY_METADATA_KEY: marker} if marker is not None else {}),
+    }
 
-    assert _virtual_key_team_id(kwargs) == expected
+    assert _virtual_key_team_id({"litellm_params": {bucket: metadata}}) == expected
 
 
 @pytest.mark.asyncio

tests/unit/proxy/hooks/test_proxy_track_cost_callback.py

@@ -287,6 +287,36 @@ async def test_attempt_rows_add_spend_but_not_requests_to_daily_spend():
     assert [usage.spend for usage in model_usage if usage is not None] == [7.2e-05, 7.2e-05, 4.8e-05]
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(("via_virtual_key", "expected"), [(True, "team-a"), (False, "")])
+async def test_failed_request_keeps_the_virtual_key_team_for_tag_spend(via_virtual_key: bool, expected: str):
+    """On /v1/messages and /v1/responses the auth object lives in litellm_metadata, which the failure
+    row does not copy, so the tag writer reads the virtual-key marker from the failure metadata."""
+    from litellm.proxy.db.db_spend_update_writer import _virtual_key_team_id
+
+    auth: Final = UserAPIKeyAuth(api_key="a" * 64, team_id="team-a").model_copy(
+        update={"via_virtual_key": via_virtual_key}
+    )
+    request_data: Final = {
+        "model": "claude-sonnet",
+        "messages": [{"role": "user", "content": "Hello"}],
+        "litellm_metadata": {"user_api_key_auth": auth, "tags": ["cost-center:cc-42"]},
+        "proxy_server_request": {"request_id": "test_request_id"},
+    }
+
+    with patch(
+        "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database",
+        new_callable=AsyncMock,
+    ) as mock_update_database:
+        await ProxyDBLogger().async_post_call_failure_hook(
+            request_data=request_data,
+            original_exception=Exception("upstream 500"),
+            user_api_key_dict=auth,
+        )
+
+        assert _virtual_key_team_id(mock_update_database.call_args[1]["kwargs"]) == expected
+
+
 @pytest.mark.asyncio
 async def test_async_post_call_failure_hook_carries_guardrail_info_from_litellm_metadata():
     """

tests/unit/proxy/spend_tracking/test_spend_event.py

@@ -6,7 +6,9 @@ import pytest
 
 import litellm
 from litellm.caching.caching import Cache
+from litellm.constants import USER_API_KEY_IS_VIRTUAL_KEY_METADATA_KEY
 from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy.db.db_spend_update_writer import _virtual_key_team_id
 from litellm.proxy.spend_tracking.spend_event import (
     CACHE_OFF_KEY,
     SpendEventBuildError,
@@ -151,6 +153,21 @@ def test_round_trip_preserves_identity_usage_reservation_and_tools():
     assert "complete_streaming_response" not in args.kwargs
 
 
+def test_round_trip_keeps_the_virtual_key_team_for_tag_spend():
+    base: Final = _success_kwargs()
+    kwargs: Final = {
+        **base,
+        "litellm_params": {
+            **base["litellm_params"],
+            "metadata": {**base["litellm_params"]["metadata"], USER_API_KEY_IS_VIRTUAL_KEY_METADATA_KEY: True},
+        },
+    }
+    event: Final = decode_spend_event(_build(kwargs, _response()))
+    assert not isinstance(event, SpendEventDecodeError)
+
+    assert _virtual_key_team_id(spend_event_callback_args(event).kwargs) == "team-1"
+
+
 def test_streaming_event_reconstructs_complete_streaming_response():
     kwargs: Final = {**_success_kwargs(), "stream": True, "complete_streaming_response": _response()}
     event: Final = decode_spend_event(_build(kwargs, _response()))

tests/unit/proxy/test_db_schema_changes.py

@@ -2,7 +2,12 @@ from typing import Final
 
 import pytest
 
-from scripts.check_schema_compatibility import check_breaking_changes, parse_model_fields
+from scripts.check_schema_compatibility import (
+    check_breaking_changes,
+    check_unique_key_changes,
+    parse_model_fields,
+    parse_unique_keys,
+)
 
 ORIGINAL: Final = "model Item {\n id String @id\n label String?\n}\n"
 
@@ -24,10 +29,33 @@ def test_block_attributes_and_comments_are_not_fields() -> None:
     assert parse_model_fields(schema) == {"Item": {"id": "String @id", "label": "String?"}}
 
 
-def test_changing_a_composite_unique_key_is_not_a_field_change() -> None:
+@pytest.mark.parametrize(
+    "after",
+    [
+        "model Item {\n id String @id\n label String?\n @@unique([label, id, tenant])\n}\n",
+        "model Item {\n id String @id\n label String?\n}\n",
+    ],
+    ids=["changed", "removed"],
+)
+def test_changing_or_removing_a_unique_key_is_breaking(after: str) -> 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)) == ()
+    assert check_unique_key_changes(parse_unique_keys(before), parse_unique_keys(after)) == (
+        "Breaking: Unique key Item.@@unique([label, id]) was removed or changed",
+    )
+
+
+def test_adding_an_index_or_rewording_a_comment_is_not_breaking() -> None:
+    before: Final = (
+        'model Item {\n id String @id\n label String @default("") // old note\n'
+        " @@unique([label, id])\n @@index([id])\n}\n"
+    )
+    after: Final = (
+        'model Item {\n id String @id\n label String @default("") // reworded note\n'
+        " @@unique([label,  id]) // same key\n @@index([id])\n @@index([label, id])\n}\n"
+    )
+    assert check_breaking_changes(parse_model_fields(before), parse_model_fields(after)) == ()
+    assert check_unique_key_changes(parse_unique_keys(before), parse_unique_keys(after)) == ()
 
 
 def test_a_narrowed_column_is_still_caught_next_to_block_attributes() -> None:

tests/unit/proxy/test_litellm_pre_call_utils.py

@@ -597,6 +597,46 @@ async def test_add_litellm_data_to_request_strips_all_user_api_key_prefix_keys()
     assert attacker_leaks == [], f"Unexpected leaked keys: {attacker_leaks}"
 
 
+@pytest.mark.parametrize(
+    ("auth_fields", "expected"),
+    [
+        ({"via_virtual_key": True}, True),
+        ({}, False),
+        ({"via_virtual_key": True, "is_session_token": True}, False),
+    ],
+    ids=["virtual-key", "jwt-or-custom-auth", "ui-session"],
+)
+@pytest.mark.asyncio
+async def test_only_the_proxy_marks_a_request_as_made_with_a_virtual_key(
+    auth_fields: dict[str, bool], expected: bool
+) -> None:
+    from litellm.constants import USER_API_KEY_IS_VIRTUAL_KEY_METADATA_KEY
+    from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
+
+    request_mock = MagicMock(spec=Request)
+    request_mock.url = MagicMock()
+    request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions"
+    request_mock.method = "POST"
+    request_mock.query_params = {}
+    request_mock.headers = {"Content-Type": "application/json"}
+    request_mock.client = MagicMock()
+    request_mock.client.host = "127.0.0.1"
+    forged: Final = {USER_API_KEY_IS_VIRTUAL_KEY_METADATA_KEY: True, "user_api_key_team_id": "victim-team"}
+
+    updated: Final = await add_litellm_data_to_request(
+        data={"model": "gpt-3.5-turbo", "metadata": {**forged}, "litellm_metadata": {**forged}},
+        request=request_mock,
+        user_api_key_dict=UserAPIKeyAuth(api_key="a" * 64, team_id="team-a").model_copy(update=auth_fields),
+        proxy_config=MagicMock(),
+        general_settings={},
+        version="test-version",
+    )
+
+    assert updated["metadata"][USER_API_KEY_IS_VIRTUAL_KEY_METADATA_KEY] is expected
+    assert updated["metadata"]["user_api_key_team_id"] == "team-a"
+    assert USER_API_KEY_IS_VIRTUAL_KEY_METADATA_KEY not in (updated.get("litellm_metadata") or {})
+
+
 @pytest.mark.asyncio
 async def test_add_litellm_data_to_request_string_metadata_does_not_crash():
     """Regression: pre-strip code that pre-populated data['metadata'][k]=v