Merge 0142382f46 into 5c60824024

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

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:

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

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