From 58f8810b622b44f1cb7ff5ad72fd5d6cd678207c Mon Sep 17 00:00:00 2001 From: Mrinal Chanshetty Date: Thu, 29 Oct 2026 11:29:14 -0700 Subject: [PATCH] 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 --- .../managed_object_attribution.py | 3 +- litellm/constants.py | 1 + litellm/proxy/db/db_spend_update_writer.py | 16 ++++---- .../proxy/hooks/proxy_track_cost_callback.py | 9 ++++- litellm/proxy/litellm_pre_call_utils.py | 11 +++++ scripts/check_schema_compatibility.py | 36 +++++++++++++++-- .../common_utils/test_check_batch_cost.py | 15 +++++++ .../proxy/db/test_db_spend_update_writer.py | 38 ++++++++---------- .../hooks/test_proxy_track_cost_callback.py | 30 ++++++++++++++ .../proxy/spend_tracking/test_spend_event.py | 17 ++++++++ tests/unit/proxy/test_db_schema_changes.py | 34 ++++++++++++++-- .../unit/proxy/test_litellm_pre_call_utils.py | 40 +++++++++++++++++++ 12 files changed, 212 insertions(+), 38 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py b/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py index a4a40a7a3d7..b0300e1653c 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py +++ b/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), diff --git a/litellm/constants.py b/litellm/constants.py index e21751c7fd1..8e101950cf3 100644 --- a/litellm/constants.py +++ b/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" diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 59278bd93ec..a95675c79f0 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/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: diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 980f0a5d410..30a9e72ccef 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/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, diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 4444d6d25f9..4a2e6ac35a5 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/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) diff --git a/scripts/check_schema_compatibility.py b/scripts/check_schema_compatibility.py index 5e98db3c7e2..541201fede8 100644 --- a/scripts/check_schema_compatibility.py +++ b/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") diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py index 5b37cee88c5..faec25d68cd 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/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 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 7a4dc7d30bc..4dfbca7db2c 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/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 diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index a4e78a48a1a..dbd191f78bb 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/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(): """ diff --git a/tests/unit/proxy/spend_tracking/test_spend_event.py b/tests/unit/proxy/spend_tracking/test_spend_event.py index ff449235582..e65d834e41d 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_event.py +++ b/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())) diff --git a/tests/unit/proxy/test_db_schema_changes.py b/tests/unit/proxy/test_db_schema_changes.py index 49dfaec1f16..e6f9eb20dcc 100644 --- a/tests/unit/proxy/test_db_schema_changes.py +++ b/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: diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 91c87c63067..dbb06c53c7c 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/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