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): keep top-level caller identity for MCP pre-call custom code guardrails

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:20:46 +00:00
parent b2a946562f
commit 90ba447974
2 changed files with 44 additions and 3 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,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,
}

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

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