From 232007e7f92a037bfb48f16edde32e02b403f6a2 Mon Sep 17 00:00:00 2001 From: Dominic White Date: Mon, 7 Sep 2026 09:03:45 +0200 Subject: [PATCH] feat(proxy): enforce trusted safety identifiers --- litellm/proxy/common_request_processing.py | 32 +++++ litellm/utils.py | 8 ++ proxy_server_config.yaml | 3 +- .../test_safety_identifier.py | 131 ++++++++++++++++++ 4 files changed, 173 insertions(+), 1 deletion(-) create mode 100644 tests/proxy_unit_tests/test_safety_identifier.py diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 5e6c9b34332..074782c3ec7 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1,8 +1,10 @@ import asyncio import contextlib +import hashlib import json import logging import math +import os from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Mapping, Sequence from datetime import datetime from functools import lru_cache @@ -67,6 +69,7 @@ from litellm.proxy.utils import ProxyLogging, _check_and_merge_model_level_guard from litellm.router import Router from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict from litellm.router_utils.common_utils import resolve_model_group_alias +from litellm.secret_managers.main import str_to_bool from litellm.types.guardrails import GuardrailEventHooks from litellm.types.router import RouterRateLimitError @@ -1532,6 +1535,23 @@ class ProxyBaseLLMRequestProcessing: def __init__(self, data: dict): self.data = data + @staticmethod + def _enforce_safety_identifier( + *, + data: dict[str, object], + route_type: ProxyRouteType, + user_api_key_dict: UserAPIKeyAuth, + ) -> dict[str, object]: + if route_type not in {"acompletion", "aresponses"}: + return data + if str_to_bool(os.getenv("LITELLM_ENFORCE_SAFETY_IDENTIFIER")) is not True: + return data + user_id: Final = user_api_key_dict.user_id + if not user_id: + return data + safety_identifier: Final = hashlib.sha256(user_id.encode("utf-8")).hexdigest() + return {**data, "safety_identifier": safety_identifier} + @staticmethod def _merge_passthrough_streaming_headers( response_headers: httpx.Headers | dict | None, @@ -2005,6 +2025,12 @@ class ProxyBaseLLMRequestProcessing: trust_client_model_info=False, ) + self.data = self._enforce_safety_identifier( + data=self.data, + route_type=route_type, + user_api_key_dict=user_api_key_dict, + ) + # An auto router with its own compression policy is authoritative for this # request: suppress every other compression guardrail and arm whichever one # the policy names for the model call, before those guardrails get a chance @@ -2017,6 +2043,12 @@ class ProxyBaseLLMRequestProcessing: call_type=route_type, ) + self.data = self._enforce_safety_identifier( + data=self.data, + route_type=route_type, + user_api_key_dict=user_api_key_dict, + ) + # Refresh AFTER pre_call_hook: guardrails (e.g. Presidio PII masking) may # have mutated `self.data` in place, and the audit-trail snapshot taken in # add_litellm_data_to_request predates that mutation. diff --git a/litellm/utils.py b/litellm/utils.py index bc2f4a86f12..941854e3075 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -4254,6 +4254,14 @@ def get_optional_params( allowed_openai_params = allowed_openai_params or [] supported_params.extend(allowed_openai_params) + # safety_identifier is injected by the proxy for trusted attribution. It is + # optional and provider-specific, so do not make providers that do not + # advertise it reject the entire request. Providers that support it still + # receive it through their normal parameter mapping, and callers can opt + # into an unlisted provider parameter via allowed_openai_params. + if "safety_identifier" in non_default_params and "safety_identifier" not in supported_params: + non_default_params.pop("safety_identifier") + _check_valid_arg( supported_params=supported_params or [], ) diff --git a/proxy_server_config.yaml b/proxy_server_config.yaml index 73990153227..28346279b24 100644 --- a/proxy_server_config.yaml +++ b/proxy_server_config.yaml @@ -246,7 +246,8 @@ general_settings: forward_headers: True # environment_variables: + # LITELLM_ENFORCE_SAFETY_IDENTIFIER: "true" # Hash the authenticated user_id and overwrite client safety_identifier values on chat/responses requests # settings for using redis caching # REDIS_HOST: redis-16337.c322.us-east-1-2.ec2.cloud.redislabs.com # REDIS_PORT: "16337" - # REDIS_PASSWORD: \ No newline at end of file + # REDIS_PASSWORD: diff --git a/tests/proxy_unit_tests/test_safety_identifier.py b/tests/proxy_unit_tests/test_safety_identifier.py new file mode 100644 index 00000000000..83ede9be79e --- /dev/null +++ b/tests/proxy_unit_tests/test_safety_identifier.py @@ -0,0 +1,131 @@ +import hashlib +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import Request + +import litellm +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + +def test_enforce_safety_identifier_hashes_authenticated_user(monkeypatch): + monkeypatch.setenv("LITELLM_ENFORCE_SAFETY_IDENTIFIER", "true") + + result = ProxyBaseLLMRequestProcessing._enforce_safety_identifier( + data={"safety_identifier": "caller-value"}, + route_type="acompletion", + user_api_key_dict=UserAPIKeyAuth(user_id="user-123"), + ) + + assert result["safety_identifier"] == hashlib.sha256(b"user-123").hexdigest() + + +@pytest.mark.parametrize("setting", [None, "false"]) +def test_enforce_safety_identifier_is_opt_in(monkeypatch, setting): + if setting is None: + monkeypatch.delenv("LITELLM_ENFORCE_SAFETY_IDENTIFIER", raising=False) + else: + monkeypatch.setenv("LITELLM_ENFORCE_SAFETY_IDENTIFIER", setting) + data = {"safety_identifier": "caller-value"} + + result = ProxyBaseLLMRequestProcessing._enforce_safety_identifier( + data=data, + route_type="acompletion", + user_api_key_dict=UserAPIKeyAuth(user_id="user-123"), + ) + + assert result == data + + +def test_enforce_safety_identifier_skips_missing_user_id(monkeypatch): + monkeypatch.setenv("LITELLM_ENFORCE_SAFETY_IDENTIFIER", "true") + data = {"safety_identifier": "caller-value"} + + result = ProxyBaseLLMRequestProcessing._enforce_safety_identifier( + data=data, + route_type="aresponses", + user_api_key_dict=UserAPIKeyAuth(user_id=None), + ) + + assert result == data + + +def test_enforce_safety_identifier_only_applies_to_openai_generation_routes(monkeypatch): + monkeypatch.setenv("LITELLM_ENFORCE_SAFETY_IDENTIFIER", "true") + data = {"safety_identifier": "caller-value"} + + result = ProxyBaseLLMRequestProcessing._enforce_safety_identifier( + data=data, + route_type="aembedding", + user_api_key_dict=UserAPIKeyAuth(user_id="user-123"), + ) + + assert result == data + + +@pytest.mark.parametrize( + ("provider", "model"), + [("anthropic", "claude-3-5-sonnet-20241022"), ("gemini", "gemini-2.0-flash")], +) +def test_unsupported_safety_identifier_is_dropped_by_provider_translation(provider, model): + result = litellm.get_optional_params( + model=model, + custom_llm_provider=provider, + safety_identifier="trusted-value", + ) + + assert "safety_identifier" not in result + + +def test_supported_safety_identifier_is_preserved_by_provider_translation(): + result = litellm.get_optional_params( + model="gpt-4o", + custom_llm_provider="openai", + safety_identifier="trusted-value", + ) + + assert result["safety_identifier"] == "trusted-value" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route_type", ["acompletion", "aresponses"]) +async def test_pre_call_hook_cannot_override_enforced_safety_identifier(monkeypatch, route_type): + monkeypatch.setenv("LITELLM_ENFORCE_SAFETY_IDENTIFIER", "true") + request = MagicMock(spec=Request) + request.headers.get.return_value = "call-id" + logging_obj = MagicMock() + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(return_value={"model": "gpt-5", "safety_identifier": "hook-value"}) + user_api_key_dict = UserAPIKeyAuth(user_id="user-123") + processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-5", "safety_identifier": "caller-value"}) + + with ( + patch( # test-quality-ok: isolate shared pre-call ordering without making an upstream request + "litellm.proxy.common_request_processing.add_litellm_data_to_request", + new=AsyncMock(side_effect=lambda **kwargs: kwargs["data"]), + ), + patch( # test-quality-ok: isolate shared pre-call ordering without initializing logging callbacks + "litellm.proxy.common_request_processing.litellm.utils.function_setup", + return_value=(logging_obj, processor.data), + ), + patch( # test-quality-ok: isolate shared pre-call ordering from router configuration + "litellm.proxy.common_request_processing._check_and_merge_model_level_guardrails", + side_effect=lambda **kwargs: kwargs["data"], + ), + patch( # test-quality-ok: isolate shared pre-call ordering from optional compression hooks + "litellm.proxy.common_request_processing._arm_auto_router_compression", + new=AsyncMock(), + ), + ): + result, _ = await processor.common_processing_pre_call_logic( + request=request, + general_settings={}, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + proxy_config=MagicMock(), + route_type=route_type, + version="test", + ) + + assert result["safety_identifier"] == hashlib.sha256(b"user-123").hexdigest()