feat(proxy): enforce trusted safety identifiers

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.

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 [],
     )

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: 
+  # REDIS_PASSWORD:

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()