fix(proxy): close safety identifier enforcement gaps

litellm/litellm_core_utils/safety_identifier.py

@@ -0,0 +1,23 @@
+import hashlib
+from collections.abc import MutableMapping
+from typing import Final
+
+
+def enforce_safety_identifier(
+    *,
+    data: MutableMapping[str, object],  # mutable-ok: trusted enforcement rewrites the request payload in place
+    user_id: str | None,
+    enabled: bool,
+) -> bool:
+    if not enabled:
+        return False
+    if user_id:
+        safety_identifier: Final = hashlib.sha256(user_id.encode("utf-8")).hexdigest()
+        if data.get("safety_identifier") == safety_identifier:
+            return False
+        data["safety_identifier"] = safety_identifier  # rebind-ok: enforce the trusted request identity in place
+        return True
+    if "safety_identifier" not in data:
+        return False
+    data.pop("safety_identifier", None)  # rebind-ok: remove the untrusted client value when no identity exists
+    return True

litellm/proxy/common_request_processing.py

@@ -1,11 +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 collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Mapping, MutableMapping, Sequence
 from datetime import datetime
 from functools import lru_cache
 from types import MappingProxyType
@@ -46,6 +45,7 @@ from litellm.litellm_core_utils.llm_response_utils.get_headers import (
     get_response_headers,
 )
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
+from litellm.litellm_core_utils.safety_identifier import enforce_safety_identifier
 from litellm.litellm_core_utils.streaming_handler import (
     backfill_missing_cache_usage_fields,
 )
@@ -1538,22 +1538,19 @@ class ProxyBaseLLMRequestProcessing:
     @staticmethod
     def _enforce_safety_identifier(
         *,
-        data: dict[str, Any],
+        data: MutableMapping[str, object],
         route_type: ProxyRouteType,
         user_api_key_dict: UserAPIKeyAuth,
-    ) -> dict[str, Any]:
+    ) -> None:
         if route_type not in ("acompletion", "aresponses"):
-            return data
+            return
         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 {  # mutable-ok: downstream request processing mutates payloads
-            **data,
-            "safety_identifier": safety_identifier,
-        }
+            return
+        enforce_safety_identifier(
+            data=data,
+            user_id=user_api_key_dict.user_id,
+            enabled=True,
+        )
 
     @staticmethod
     def _merge_passthrough_streaming_headers(
@@ -2028,7 +2025,7 @@ class ProxyBaseLLMRequestProcessing:
             trust_client_model_info=False,
         )
 
-        self.data = self._enforce_safety_identifier(
+        self._enforce_safety_identifier(
             data=self.data,
             route_type=route_type,
             user_api_key_dict=user_api_key_dict,
@@ -2046,7 +2043,7 @@ class ProxyBaseLLMRequestProcessing:
             call_type=route_type,
         )
 
-        self.data = self._enforce_safety_identifier(
+        self._enforce_safety_identifier(
             data=self.data,
             route_type=route_type,
             user_api_key_dict=user_api_key_dict,

litellm/responses/streaming_iterator.py

@@ -2,14 +2,15 @@ from __future__ import annotations
 
 import asyncio
 import json
+import os
 import time
 import traceback
 import uuid
-from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
+from collections.abc import Awaitable, Callable, Iterable, Mapping, MutableMapping, Sequence
 from datetime import datetime
 from functools import lru_cache
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, overload, runtime_checkable
+from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, overload, runtime_checkable
 
 import httpx
 from openai._streaming import SSEDecoder
@@ -29,9 +30,11 @@ from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_b
 from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
     update_response_metadata,
 )
+from litellm.litellm_core_utils.safety_identifier import enforce_safety_identifier
 from litellm.litellm_core_utils.thread_pool_executor import executor
 from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
 from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils
+from litellm.secret_managers.main import str_to_bool
 from litellm.types.llms.openai import (
     PART_UNION_TYPES,
     ResponseAPIUsage,
@@ -92,6 +95,19 @@ def _is_str_mapping(value: object) -> TypeIs[dict[str, str]]:  # guard-ok: verif
     return _is_json_object(value) and all(isinstance(item, str) for item in value.values())
 
 
+def _enforce_responses_ws_safety_identifier(
+    msg_obj: _MutableJsonObject,
+    user_api_key_dict: UserAPIKeyAuth | None,
+) -> bool:
+    return enforce_safety_identifier(
+        data=cast(  # cast-ok: JSON protocol is backed by a mutable response.create dictionary
+            MutableMapping[str, object], msg_obj
+        ),
+        user_id=user_api_key_dict.user_id if user_api_key_dict is not None else None,
+        enabled=str_to_bool(os.getenv("LITELLM_ENFORCE_SAFETY_IDENTIFIER")) is True,
+    )
+
+
 class _MutableJsonObject(Protocol):
     @overload
     def get(self, key: str, /) -> object | None: ...
@@ -1744,16 +1760,18 @@ class ResponsesWebSocketStreaming:
         if msg_obj.get("type") != "response.create":
             return message
 
+        safety_identifier_modified: Final = _enforce_responses_ws_safety_identifier(msg_obj, self.user_api_key_dict)
+
         # Always enforce the authorized model, even when PII masking is off.
         model_modified: Final = self._enforce_authorized_model(msg_obj)
 
         if not self.guardrail_callbacks:
-            return json.dumps(msg_obj) if model_modified else message
+            return json.dumps(msg_obj) if model_modified or safety_identifier_modified else message
 
         if "metadata" not in self.request_data:
             self.request_data["metadata"] = {}
 
-        modified = model_modified
+        modified = model_modified or safety_identifier_modified
         guardrail_cbs: Final[tuple[PresidioGuardrailCallback, ...]] = tuple(self.guardrail_callbacks)
         for cb in guardrail_cbs:
             presidio_config = cb.get_presidio_settings_from_request_data(self.request_data)
@@ -2521,6 +2539,8 @@ class ManagedResponsesWebSocketHandler:
         if msg_obj is None:
             return
 
+        _enforce_responses_ws_safety_identifier(msg_obj, self.user_api_key_dict)
+
         # generate=false is a prompt-cache warmup hint (sent by codex prewarm).
         # Native provider sockets handle it server-side, but there is no HTTP
         # equivalent and the frame carries empty input. Managed providers must

litellm/responses/utils.py

@@ -210,6 +210,12 @@ class ResponsesAPIRequestUtils:
         should_drop_params: Final = litellm.drop_params or drop_params is True
 
         non_default_params: Final = cast(dict, response_api_optional_params)
+        if (
+            "safety_identifier" in non_default_params
+            and "safety_identifier" not in supported_params
+            and (allowed_openai_params is None or "safety_identifier" not in allowed_openai_params)
+        ):
+            non_default_params.pop("safety_identifier")
         # Check for unsupported parameters
         ResponsesAPIRequestUtils._check_valid_arg(
             supported_params=supported_params + (allowed_openai_params or []),

litellm/utils.py

@@ -4254,11 +4254,6 @@ 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")
 

tests/proxy_unit_tests/test_safety_identifier.py

@@ -1,74 +1,78 @@
 import hashlib
+from typing import Literal
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 from fastapi import Request
 
 import litellm
+from litellm.llms.perplexity.responses.transformation import PerplexityResponsesConfig
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
+from litellm.responses.utils import ResponsesAPIRequestUtils
 
 
-def test_enforce_safety_identifier_hashes_authenticated_user(monkeypatch):
+def test_enforce_safety_identifier_hashes_authenticated_user(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setenv("LITELLM_ENFORCE_SAFETY_IDENTIFIER", "true")
 
-    result = ProxyBaseLLMRequestProcessing._enforce_safety_identifier(
-        data={"safety_identifier": "caller-value"},
+    data = {"safety_identifier": "caller-value"}
+    ProxyBaseLLMRequestProcessing._enforce_safety_identifier(
+        data=data,
         route_type="acompletion",
         user_api_key_dict=UserAPIKeyAuth(user_id="user-123"),
     )
 
-    assert result["safety_identifier"] == hashlib.sha256(b"user-123").hexdigest()
+    assert data["safety_identifier"] == hashlib.sha256(b"user-123").hexdigest()
 
 
 @pytest.mark.parametrize("setting", [None, "false"])
-def test_enforce_safety_identifier_is_opt_in(monkeypatch, setting):
+def test_enforce_safety_identifier_is_opt_in(monkeypatch: pytest.MonkeyPatch, setting: str | None):
     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(
+    ProxyBaseLLMRequestProcessing._enforce_safety_identifier(
         data=data,
         route_type="acompletion",
         user_api_key_dict=UserAPIKeyAuth(user_id="user-123"),
     )
 
-    assert result == data
+    assert data == {"safety_identifier": "caller-value"}
 
 
-def test_enforce_safety_identifier_skips_missing_user_id(monkeypatch):
+def test_enforce_safety_identifier_removes_untrusted_identifier(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setenv("LITELLM_ENFORCE_SAFETY_IDENTIFIER", "true")
     data = {"safety_identifier": "caller-value"}
 
-    result = ProxyBaseLLMRequestProcessing._enforce_safety_identifier(
+    ProxyBaseLLMRequestProcessing._enforce_safety_identifier(
         data=data,
         route_type="aresponses",
         user_api_key_dict=UserAPIKeyAuth(user_id=None),
     )
 
-    assert result == data
+    assert data == {}
 
 
-def test_enforce_safety_identifier_only_applies_to_openai_generation_routes(monkeypatch):
+def test_enforce_safety_identifier_only_applies_to_openai_generation_routes(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setenv("LITELLM_ENFORCE_SAFETY_IDENTIFIER", "true")
     data = {"safety_identifier": "caller-value"}
 
-    result = ProxyBaseLLMRequestProcessing._enforce_safety_identifier(
+    ProxyBaseLLMRequestProcessing._enforce_safety_identifier(
         data=data,
         route_type="aembedding",
         user_api_key_dict=UserAPIKeyAuth(user_id="user-123"),
     )
 
-    assert result == data
+    assert data == {"safety_identifier": "caller-value"}
 
 
 @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):
+def test_unsupported_safety_identifier_is_dropped_by_provider_translation(provider: str, model: str):
     result = litellm.get_optional_params(
         model=model,
         custom_llm_provider=provider,
@@ -88,9 +92,21 @@ def test_supported_safety_identifier_is_preserved_by_provider_translation():
     assert result["safety_identifier"] == "trusted-value"
 
 
+def test_unsupported_safety_identifier_is_dropped_by_responses_translation():
+    result = ResponsesAPIRequestUtils.get_optional_params_responses_api(
+        model="sonar",
+        responses_api_provider_config=PerplexityResponsesConfig(),
+        response_api_optional_params={"safety_identifier": "trusted-value"},
+    )
+
+    assert "safety_identifier" not in result
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("route_type", ["acompletion", "aresponses"])
-async def test_pre_call_hook_cannot_override_enforced_safety_identifier(monkeypatch, route_type):
+async def test_pre_call_hook_cannot_override_enforced_safety_identifier(
+    monkeypatch: pytest.MonkeyPatch, route_type: Literal["acompletion", "aresponses"]
+):
     monkeypatch.setenv("LITELLM_ENFORCE_SAFETY_IDENTIFIER", "true")
     request = MagicMock(spec=Request)
     request.headers.get.return_value = "call-id"

tests/test_litellm/responses/test_responses_websocket_all_providers.py

@@ -7,8 +7,9 @@ Tests that:
 3. Providers without native websocket support use ManagedResponsesWebSocketHandler
 """
 
+import hashlib
 import json
-from unittest.mock import MagicMock
+from unittest.mock import AsyncMock, MagicMock
 
 import pytest
 
@@ -1196,6 +1197,66 @@ class TestWebSocketProjectQuotaEnforcement:
 
 
 class TestNativeWebSocketGuardrails:
+    @pytest.mark.asyncio
+    async def test_response_create_overwrites_safety_identifier(self, monkeypatch: pytest.MonkeyPatch):
+        from litellm.proxy._types import UserAPIKeyAuth
+        from litellm.responses.streaming_iterator import ResponsesWebSocketStreaming
+
+        monkeypatch.setenv("LITELLM_ENFORCE_SAFETY_IDENTIFIER", "true")
+        handler = ResponsesWebSocketStreaming(
+            websocket=MagicMock(),
+            backend_ws=MagicMock(),
+            logging_obj=MagicMock(),
+            user_api_key_dict=UserAPIKeyAuth(user_id="user-123"),
+        )
+
+        masked = await handler._mask_response_create(
+            json.dumps({"type": "response.create", "safety_identifier": "caller-value"})
+        )
+
+        assert json.loads(masked)["safety_identifier"] == hashlib.sha256(b"user-123").hexdigest()
+
+    @pytest.mark.asyncio
+    async def test_response_create_removes_safety_identifier_without_user_id(self, monkeypatch: pytest.MonkeyPatch):
+        from litellm.proxy._types import UserAPIKeyAuth
+        from litellm.responses.streaming_iterator import ResponsesWebSocketStreaming
+
+        monkeypatch.setenv("LITELLM_ENFORCE_SAFETY_IDENTIFIER", "true")
+        handler = ResponsesWebSocketStreaming(
+            websocket=MagicMock(),
+            backend_ws=MagicMock(),
+            logging_obj=MagicMock(),
+            user_api_key_dict=UserAPIKeyAuth(user_id=None),
+        )
+
+        masked = await handler._mask_response_create(
+            json.dumps({"type": "response.create", "safety_identifier": "caller-value"})
+        )
+
+        assert "safety_identifier" not in json.loads(masked)
+
+    @pytest.mark.asyncio
+    async def test_managed_response_create_forwards_trusted_safety_identifier(self, monkeypatch: pytest.MonkeyPatch):
+        from litellm.proxy._types import UserAPIKeyAuth
+        from litellm.responses.streaming_iterator import ManagedResponsesWebSocketHandler
+
+        monkeypatch.setenv("LITELLM_ENFORCE_SAFETY_IDENTIFIER", "true")
+        handler = ManagedResponsesWebSocketHandler(
+            websocket=MagicMock(),
+            model="gpt-4o",
+            logging_obj=MagicMock(),
+            user_api_key_dict=UserAPIKeyAuth(user_id="user-123"),
+        )
+        stream_and_forward = AsyncMock(return_value=None)
+        monkeypatch.setattr(handler, "_stream_and_forward", stream_and_forward)
+
+        await handler._process_response_create(
+            json.dumps({"type": "response.create", "input": "hi", "safety_identifier": "caller-value"})
+        )
+
+        call_kwargs = stream_and_forward.call_args.args[1]
+        assert call_kwargs["safety_identifier"] == hashlib.sha256(b"user-123").hexdigest()
+
     @pytest.mark.asyncio
     async def test_response_create_injects_authorized_model(self):
         import json