diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py index e0291975699..de808afbe52 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py @@ -63,6 +63,10 @@ def _metadata_bucket(request_data: Mapping[str, object], key: str) -> Mapping[st return bucket if isinstance(bucket, Mapping) else {} +def _identity_field(request_data: Mapping[str, object], metadata: Mapping[str, object], key: str) -> object: + return metadata[key] if key in metadata else request_data.get(key) + + class CustomCodeGuardrailError(Exception): """Raised when custom code guardrail execution fails.""" @@ -291,9 +295,9 @@ class CustomCodeGuardrail(CustomGuardrail): } return { "model": request_data.get("model"), - "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"), + "user_id": _identity_field(request_data, metadata, "user_api_key_user_id"), + "team_id": _identity_field(request_data, metadata, "user_api_key_team_id"), + "end_user_id": _identity_field(request_data, metadata, "user_api_key_end_user_id"), "metadata": metadata, } diff --git a/tests/test_litellm/proxy/guardrails/test_custom_code_security.py b/tests/test_litellm/proxy/guardrails/test_custom_code_security.py index 59011cea1c3..4e8801041c2 100644 --- a/tests/test_litellm/proxy/guardrails/test_custom_code_security.py +++ b/tests/test_litellm/proxy/guardrails/test_custom_code_security.py @@ -302,6 +302,43 @@ async def test_custom_code_sandbox_merges_caller_metadata_with_litellm_metadata( } +@pytest.mark.asyncio +async def test_custom_code_sandbox_falls_back_to_top_level_identity_for_mcp_calls(): + """MCP pre-call hooks put user_api_key_* at the top level of the synthetic request, with a + metadata bucket that only carries headers; those ids must still reach the sandbox.""" + guardrail = _compile(IDENTITY_ECHO_CODE) + request_data = {"model": "mcp-tool-call", **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"] == ["someone@example.com", "team-1", "end-user-1"] + + +@pytest.mark.asyncio +async def test_custom_code_sandbox_ignores_top_level_identity_when_proxy_bucket_has_it(): + """A caller cannot forge ids through top-level body fields on LLM routes: the proxy bucket + carries every user_api_key_* key (even when None) and it wins over the top level.""" + 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", + "user_api_key_user_id": "forged", + "user_api_key_team_id": "forged-team", + "user_api_key_end_user_id": "forged-end-user", + "metadata": {**CALLER_IDENTITY, "user_api_key_team_id": None}, + } + + 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"] == "['someone@example.com', None, 'end-user-1']" + + @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"