Merge pull request #41126 from BerriAI/litellm_custom_code_guardrail_identity

fix(guardrails): resolve caller identity from metadata buckets in custom code guardrail

litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py

@@ -58,6 +58,11 @@ if TYPE_CHECKING:
     from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 
 
+def _metadata_bucket(request_data: Mapping[str, object], key: str) -> Mapping[str, object]:
+    bucket: Final = request_data.get(key)
+    return bucket if isinstance(bucket, Mapping) else {}
+
+
 class CustomCodeGuardrailError(Exception):
     """Raised when custom code guardrail execution fails."""
 
@@ -280,12 +285,16 @@ class CustomCodeGuardrail(CustomGuardrail):
         Returns:
             Safe subset of request data
         """
+        metadata: Final = {
+            **_metadata_bucket(request_data, "metadata"),
+            **_metadata_bucket(request_data, "litellm_metadata"),
+        }
         return {
             "model": request_data.get("model"),
-            "user_id": request_data.get("user_api_key_user_id"),
-            "team_id": request_data.get("user_api_key_team_id"),
-            "end_user_id": request_data.get("user_api_key_end_user_id"),
-            "metadata": request_data.get("metadata", {}),
+            "user_id": metadata.get("user_api_key_user_id"),
+            "team_id": metadata.get("user_api_key_team_id"),
+            "end_user_id": metadata.get("user_api_key_end_user_id"),
+            "metadata": metadata,
         }
 
     def _process_result(

litellm/proxy/utils.py

@@ -1269,7 +1269,12 @@ class ProxyLogging:
             # (e.g. MCPJWTSigner) to independently verify the caller's identity
             # before re-signing an outbound token (FR-5 verify+re-sign).
             "incoming_bearer_token": kwargs.get("incoming_bearer_token"),
-            "metadata": {"headers": kwargs.get("headers") or {}},
+            "metadata": {
+                "headers": kwargs.get("headers") or {},
+                "user_api_key_user_id": kwargs.get("user_api_key_user_id"),
+                "user_api_key_team_id": kwargs.get("user_api_key_team_id"),
+                "user_api_key_end_user_id": kwargs.get("user_api_key_end_user_id"),
+            },
         }
         user_api_key_auth: Final = kwargs.get("user_api_key_auth")
         if isinstance(user_api_key_auth, UserAPIKeyAuth):

tests/test_litellm/proxy/guardrails/test_custom_code_security.py

@@ -250,6 +250,76 @@ async def test_custom_code_flag_default_reason_and_empty_metadata():
     }
 
 
+IDENTITY_ECHO_CODE = (
+    "def apply_guardrail(inputs, request_data, input_type):\n"
+    "    return flag('identity', metadata={\n"
+    "        'ids': [request_data['user_id'], request_data['team_id'], request_data['end_user_id']],\n"
+    "        'metadata_keys': sorted(request_data['metadata'].keys()),\n"
+    "    })\n"
+)
+CALLER_IDENTITY = {
+    "user_api_key_user_id": "someone@example.com",
+    "user_api_key_team_id": "team-1",
+    "user_api_key_end_user_id": "end-user-1",
+    "user_api_key_alias": "guardrail-repro-key",
+}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"])
+async def test_custom_code_sandbox_sees_caller_identity_from_proxy_metadata_bucket(metadata_key):
+    """LIT-6609: the proxy writes user_api_key_* into `metadata` (chat) or `litellm_metadata`
+    (/v1/messages, responses, batches, files); the sandbox must resolve ids from either."""
+    guardrail = _compile(IDENTITY_ECHO_CODE)
+    request_data = {"model": "m", metadata_key: dict(CALLER_IDENTITY)}
+
+    await guardrail.apply_guardrail(inputs={"texts": ["x"]}, request_data=request_data, input_type="request")
+
+    entry = request_data[metadata_key]["standard_logging_guardrail_information"][0]
+    assert entry["guardrail_response"]["metadata"] == {
+        "ids": ["someone@example.com", "team-1", "end-user-1"],
+        "metadata_keys": sorted(CALLER_IDENTITY),
+    }
+
+
+@pytest.mark.asyncio
+async def test_custom_code_sandbox_merges_caller_metadata_with_litellm_metadata():
+    """On litellm_metadata routes the caller's own `metadata` field must stay visible next to
+    the proxy identity block, and the proxy block wins on key collisions."""
+    guardrail = _compile(IDENTITY_ECHO_CODE)
+    request_data = {
+        "model": "m",
+        "metadata": {"trace_id": "abc", "user_api_key_user_id": "forged"},
+        "litellm_metadata": dict(CALLER_IDENTITY),
+    }
+
+    await guardrail.apply_guardrail(inputs={"texts": ["x"]}, request_data=request_data, input_type="request")
+
+    entry = request_data["litellm_metadata"]["standard_logging_guardrail_information"][0]
+    assert entry["guardrail_response"]["metadata"] == {
+        "ids": ["someone@example.com", "team-1", "end-user-1"],
+        "metadata_keys": sorted([*CALLER_IDENTITY, "trace_id"]),
+    }
+
+
+@pytest.mark.asyncio
+async def test_custom_code_sandbox_ignores_top_level_identity_fields():
+    """Only the proxy-owned metadata buckets carry identity; user_api_key_* keys at the top level
+    of the request body are caller-controlled on ordinary routes and must never become ids."""
+    code = (
+        "def apply_guardrail(inputs, request_data, input_type):\n"
+        "    ids = [request_data['user_id'], request_data['team_id'], request_data['end_user_id']]\n"
+        "    return flag('identity', metadata={'ids': str(ids)})\n"
+    )
+    guardrail = _compile(code)
+    request_data = {"model": "m", **CALLER_IDENTITY, "metadata": {"headers": {}}}
+
+    await guardrail.apply_guardrail(inputs={"texts": ["x"]}, request_data=request_data, input_type="request")
+
+    entry = request_data["metadata"]["standard_logging_guardrail_information"][0]
+    assert entry["guardrail_response"]["metadata"]["ids"] == "[None, None, None]"
+
+
 @pytest.mark.asyncio
 async def test_custom_code_allow_still_records_success_not_flagged():
     code = "def apply_guardrail(inputs, request_data, input_type):\n    return allow()\n"

tests/test_litellm/proxy/utils/proxy_logging/test_mcp_bridging.py

@@ -87,6 +87,27 @@ def test_convert_mcp_to_llm_format_exposes_headers_on_metadata(proxy_logging, ma
     assert out["metadata"]["headers"] == {"x-nuid": "nuid-1"}
 
 
+def test_convert_mcp_to_llm_format_exposes_caller_identity_on_metadata(proxy_logging, make_mcp_request_obj):
+    """Custom code guardrails resolve user_id/team_id/end_user_id from the proxy-owned metadata
+    bucket on every route, so the MCP bridge has to write the authenticated ids there too."""
+    req = make_mcp_request_obj()
+    out = proxy_logging._convert_mcp_to_llm_format(
+        request_obj=req,
+        kwargs={
+            "user_api_key_user_id": "u-1",
+            "user_api_key_team_id": "t-1",
+            "user_api_key_end_user_id": "eu-1",
+            "headers": {"x-nuid": "nuid-1"},
+        },
+    )
+    assert out["metadata"] == {
+        "headers": {"x-nuid": "nuid-1"},
+        "user_api_key_user_id": "u-1",
+        "user_api_key_team_id": "t-1",
+        "user_api_key_end_user_id": "eu-1",
+    }
+
+
 def test_convert_mcp_to_llm_format_defaults_headers_to_empty(proxy_logging, make_mcp_request_obj):
     req = make_mcp_request_obj()
     out = proxy_logging._convert_mcp_to_llm_format(request_obj=req, kwargs={})