From b2a946562f5bf9f7b10c1ff7b41d76d01bbaa931 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 14 Sep 2026 21:03:17 +0000 Subject: [PATCH] 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"