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 de808afbe52..e0291975699 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,10 +63,6 @@ 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.""" @@ -295,9 +291,9 @@ class CustomCodeGuardrail(CustomGuardrail): } return { "model": request_data.get("model"), - "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"), + "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, } diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index c095586b6c9..f523c2a7151 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1253,7 +1253,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): 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 4e8801041c2..068cd0d8ed7 100644 --- a/tests/test_litellm/proxy/guardrails/test_custom_code_security.py +++ b/tests/test_litellm/proxy/guardrails/test_custom_code_security.py @@ -303,40 +303,21 @@ 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.""" +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", - "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}, - } + 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"] == "['someone@example.com', None, 'end-user-1']" + assert entry["guardrail_response"]["metadata"]["ids"] == "[None, None, None]" @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_mcp_bridging.py b/tests/test_litellm/proxy/utils/proxy_logging/test_mcp_bridging.py index 56057dce7e0..438b2351034 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_mcp_bridging.py +++ b/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={})