From b2a946562f5bf9f7b10c1ff7b41d76d01bbaa931 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 14 Sep 2026 21:03:17 +0000 Subject: [PATCH 1/3] fix(guardrails): resolve caller identity from metadata buckets in custom code guardrail The sandbox read user_api_key_* off the top level of request_data, where the proxy never writes them, and only forwarded request_data["metadata"], which is empty on /v1/messages, /v1/responses, batches and files because those routes keep proxy state in litellm_metadata. Merge both buckets (litellm_metadata wins) and resolve ids from the merged dict Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../custom_code/custom_code_guardrail.py | 17 ++++-- .../guardrails/test_custom_code_security.py | 52 +++++++++++++++++++ 2 files changed, 65 insertions(+), 4 deletions(-) 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 d5ef1e949b8..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 @@ -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( 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 7971cf62c9a..59011cea1c3 100644 --- a/tests/test_litellm/proxy/guardrails/test_custom_code_security.py +++ b/tests/test_litellm/proxy/guardrails/test_custom_code_security.py @@ -250,6 +250,58 @@ 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_allow_still_records_success_not_flagged(): code = "def apply_guardrail(inputs, request_data, input_type):\n return allow()\n" From 90ba447974f76018776e52e28883f55d85607c94 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 14 Sep 2026 21:20:46 +0000 Subject: [PATCH 2/3] fix(guardrails): keep top-level caller identity for MCP pre-call custom code guardrails Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../custom_code/custom_code_guardrail.py | 10 +++-- .../guardrails/test_custom_code_security.py | 37 +++++++++++++++++++ 2 files changed, 44 insertions(+), 3 deletions(-) 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" From 87bf5730da0aad31f27f37132131f2766f196d0c Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 14 Sep 2026 21:37:44 +0000 Subject: [PATCH 3/3] fix(guardrails): carry MCP caller identity in synthetic metadata instead of trusting top-level fields Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../custom_code/custom_code_guardrail.py | 10 ++----- litellm/proxy/utils.py | 7 ++++- .../guardrails/test_custom_code_security.py | 29 ++++--------------- .../utils/proxy_logging/test_mcp_bridging.py | 21 ++++++++++++++ 4 files changed, 35 insertions(+), 32 deletions(-) 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={})