scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-17 23:51:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 67

fix(guardrails): carry MCP caller identity in synthetic metadata instead of trusting top-level fields

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng 2026-09-14 21:37:44 +00:00
parent 90ba447974
commit 87bf5730da
4 changed files with 35 additions and 32 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py
View file

@ -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,
}

7
litellm/proxy/utils.py
View file

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

29
tests/test_litellm/proxy/guardrails/test_custom_code_security.py
View file

@ -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

21
tests/test_litellm/proxy/utils/proxy_logging/test_mcp_bridging.py
View file

@ -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={})