diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index 5ce2dfb3b87..703942b3131 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -223,15 +223,25 @@ def _a2a_jsonrpc_error_chunk(exc: HTTPException, request_id: str | None) -> Mapp def _ensure_litellm_metadata(data: dict, user_api_key_dict: UserAPIKeyAuth) -> None: - """Populate data['litellm_metadata'] from user_api_key_dict if absent.""" - if "litellm_metadata" not in data: - from litellm.llms.base_llm.guardrail_translation.base_translation import ( - BaseTranslation, - ) + """Merge the ``user_api_key_*`` metadata into the request's resolved bucket. - user_metadata: Final = BaseTranslation.transform_user_api_key_dict_to_metadata(user_api_key_dict) - if user_metadata: - data["litellm_metadata"] = user_metadata + Not ``data["litellm_metadata"]``: creating that key on a route whose bucket + is ``metadata`` flips every reader of + ``get_metadata_variable_name_from_kwargs`` onto the new dict and silently + strands the tags and router fields (#45543). + ``merge_user_api_key_metadata_into_request`` resolves the bucket and + ``setdefault``s, so an upstream value is never overwritten. + """ + from litellm.llms.base_llm.guardrail_translation.base_translation import ( + BaseTranslation, + ) + + BaseTranslation.merge_user_api_key_metadata_into_request( + cast( # cast-ok: _RequestData is this module's own alias for the bare-dict request payload + _RequestData, data + ), + user_api_key_dict, + ) class UnifiedLLMGuardrails(CustomLogger): diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_grayswan.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_grayswan.py index a5ba9dab0bd..da3fb33bedd 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_grayswan.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_grayswan.py @@ -570,7 +570,18 @@ def test_prepare_payload_includes_litellm_metadata( def test_ensure_litellm_metadata_populates_from_user_api_key_dict() -> None: - """Verify _ensure_litellm_metadata populates litellm_metadata.""" + """Verify _ensure_litellm_metadata populates the request's metadata bucket. + + Updated for #45543. This asserted the literal `litellm_metadata` key, which + is what the old implementation created unconditionally -- and creating it on + a route whose bucket is `metadata` is the bug: every reader resolves the + bucket by presence, so a new key moves them all onto a near-empty dict. The + property being protected, that the key metadata reaches the guardrail, is + unchanged; it is now asserted on the resolved bucket. + """ + from litellm.litellm_core_utils.core_helpers import ( + get_metadata_variable_name_from_kwargs, + ) from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( _ensure_litellm_metadata, ) @@ -580,23 +591,37 @@ def test_ensure_litellm_metadata_populates_from_user_api_key_dict() -> None: _ensure_litellm_metadata(data, user_auth) - assert "litellm_metadata" in data - assert data["litellm_metadata"]["user_api_key_user_id"] == "u1" - assert data["litellm_metadata"]["user_api_key_team_id"] == "t1" + bucket_name = get_metadata_variable_name_from_kwargs(data) + assert bucket_name == "metadata" + assert data[bucket_name]["user_api_key_user_id"] == "u1" + assert data[bucket_name]["user_api_key_team_id"] == "t1" -def test_ensure_litellm_metadata_noop_when_already_present() -> None: - """Verify _ensure_litellm_metadata does not overwrite existing litellm_metadata.""" +def test_ensure_litellm_metadata_does_not_overwrite_existing_values() -> None: + """Verify _ensure_litellm_metadata leaves values already in the bucket alone. + + Updated for #45543. This asserted the bucket stayed *exactly* + `{"existing": "value"}`, because the old implementation returned early + whenever `litellm_metadata` was present and so added nothing at all. The + merge now runs on the resolved bucket via `setdefault`, so the property + worth protecting -- an existing value is never overwritten -- still holds, + and is asserted directly rather than through whole-dict equality. The key + metadata the guardrail needs is added alongside it. + """ from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( _ensure_litellm_metadata, ) user_auth = UserAPIKeyAuth(user_id="should-not-appear") - data: dict = {"litellm_metadata": {"existing": "value"}} + data: dict = {"litellm_metadata": {"existing": "value", "user_api_key_user_id": "already-resolved"}} _ensure_litellm_metadata(data, user_auth) - assert data["litellm_metadata"] == {"existing": "value"} + # Still the `litellm_metadata` bucket: the fix resolves the bucket, it does + # not flip routes that genuinely use this one. + assert "metadata" not in data + assert data["litellm_metadata"]["existing"] == "value" + assert data["litellm_metadata"]["user_api_key_user_id"] == "already-resolved" class _CapturingClient: diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_wingback.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_wingback.py index c9e331de4a2..58a377cb779 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_wingback.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_wingback.py @@ -252,13 +252,29 @@ async def test_proxy_post_call_hook_blocks_on_wingback_verdict(registered: _Wing @pytest.mark.asyncio @pytest.mark.parametrize("registered", [("pre_call", ALLOWED)], indirect=True) async def test_proxy_pre_call_hook_records_one_guardrail_entry_per_check(registered: _Wingback) -> None: + """One guardrail entry per check, recorded in the request's metadata bucket. + + Updated for #45543. This read the literal `litellm_metadata` key, which + existed only because the old `_ensure_litellm_metadata` created it + unconditionally -- and on a chat-completions request, whose bucket is + `metadata`, creating that key is the bug: every reader resolves the bucket + by presence, so a new key moves them all onto a near-empty dict. The + property being protected, one entry per check with its status, is + unchanged; it is now asserted on the resolved bucket. + """ + from litellm.litellm_core_utils.core_helpers import ( + get_metadata_variable_name_from_kwargs, + ) + data: Final = _chat_request("hello") await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook( user_api_key_dict=UserAPIKeyAuth(), data=data, call_type="completion" ) - entries: Final = data["litellm_metadata"]["standard_logging_guardrail_information"] + bucket_name = get_metadata_variable_name_from_kwargs(data) + assert bucket_name == "metadata" + entries: Final = data[bucket_name]["standard_logging_guardrail_information"] assert [(entry["guardrail_name"], entry["guardrail_status"]) for entry in entries] == [ ("wingback-runtime-security", "success") ] diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 84f4d7a4173..89dd7a36766 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -21,6 +21,9 @@ from litellm.integrations.custom_guardrail import ( log_guardrail_information, ) from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route +from litellm.litellm_core_utils.core_helpers import ( + get_metadata_variable_name_from_kwargs, +) from litellm.llms import discover_guardrail_translation_mappings, load_guardrail_translation_mappings from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation from litellm.llms.base_llm.guardrail_translation.utils import ( @@ -3290,4 +3293,75 @@ async def test_pass_through_guardrail_text_is_the_client_body_without_the_inject texts: Final = [call["inputs"]["texts"] for call in guardrail.apply_calls] assert [[json.loads(text) for text in call_texts] for call_texts in texts] == [[body]], texts - assert data["litellm_metadata"]["user_api_key_user_email"] == "someone@example.com", data["litellm_metadata"] + # the key metadata lands in whichever bucket the request resolves to (#45543) + bucket: Final = data[get_metadata_variable_name_from_kwargs(data)] + assert bucket["user_api_key_user_email"] == "someone@example.com", bucket + + +class TestEnsureLitellmMetadataBucket: + """``_ensure_litellm_metadata`` must merge into the bucket the request + already resolves to, not plant ``litellm_metadata``. Creating that key on a + chat-completions request flips every reader onto a dict holding only the + ``user_api_key_*`` keys, so tags and router fields are left unread: tag + routing picks a deployment of another model and SpendLogs lose + ``model_group``/``model_id``, without raising (#45543). Same bucket-flip + shape as TestStreamingGuardrailInformationBucket above.""" + + @pytest.fixture + def user_api_key_dict(self) -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="sk-test-45543", user_id="user-1", team_id="team-1") + + def test_chat_completions_request_keeps_its_metadata_bucket(self, user_api_key_dict: UserAPIKeyAuth) -> None: + data: Final = { + "model": "all-models", + "metadata": {"tags": ["model:glm-5"], "model_group": "all-models"}, + } + + unified_module._ensure_litellm_metadata(data, user_api_key_dict) + + # Creating the key at all is the bug; its presence is what flips the readers. + assert "litellm_metadata" not in data + assert get_metadata_variable_name_from_kwargs(data) == "metadata" + assert data["metadata"]["tags"] == ["model:glm-5"] + assert data["metadata"]["model_group"] == "all-models" + assert data["metadata"]["user_api_key_team_id"] == "team-1" + assert data["metadata"]["user_api_key_user_id"] == "user-1" + + def test_existing_metadata_values_are_not_overwritten(self, user_api_key_dict: UserAPIKeyAuth) -> None: + data: Final = { + "model": "all-models", + "metadata": {"user_api_key_user_id": "resolved-upstream"}, + } + + unified_module._ensure_litellm_metadata(data, user_api_key_dict) + + assert data["metadata"]["user_api_key_user_id"] == "resolved-upstream" + assert data["metadata"]["user_api_key_team_id"] == "team-1" + + def test_litellm_metadata_route_still_uses_litellm_metadata(self, user_api_key_dict: UserAPIKeyAuth) -> None: + """Batch and file routes do keep proxy state in ``litellm_metadata``, so + the bucket is resolved rather than hardcoded in either direction.""" + data: Final = {"model": "m", "litellm_metadata": {"tags": ["keep-me"]}} + + unified_module._ensure_litellm_metadata(data, user_api_key_dict) + + assert "metadata" not in data + assert get_metadata_variable_name_from_kwargs(data) == "litellm_metadata" + assert data["litellm_metadata"]["tags"] == ["keep-me"] + assert data["litellm_metadata"]["user_api_key_team_id"] == "team-1" + + def test_request_without_any_metadata_gets_the_default_bucket(self, user_api_key_dict: UserAPIKeyAuth) -> None: + data: Final = {"model": "m"} + + unified_module._ensure_litellm_metadata(data, user_api_key_dict) + + assert "litellm_metadata" not in data + assert data["metadata"]["user_api_key_team_id"] == "team-1" + + def test_no_user_api_key_metadata_creates_no_bucket(self) -> None: + data: Final = {"model": "m"} + + unified_module._ensure_litellm_metadata(data, None) # pyright: ignore[reportArgumentType] # the hook is reached with no key on unauthenticated routes + + assert "litellm_metadata" not in data + assert "metadata" not in data