Merge fcdff6ed11 into 559247fa84

.github/workflows/test-unit-proxy-db.yml

@@ -195,6 +195,7 @@ jobs:
               tests/proxy_unit_tests/test_check_responses_cost.py
               tests/proxy_unit_tests/test_response_polling_handler.py
               tests/proxy_unit_tests/test_response_polling_pre_call_checks.py
+              tests/proxy_unit_tests/test_safety_identifier.py
               tests/proxy_unit_tests/test_realtime_cache.py
               tests/proxy_unit_tests/test_proxy_exception_mapping.py
               tests/proxy_unit_tests/test_custom_tokenizer_bug.py

litellm/litellm_core_utils/safety_identifier.py

@@ -0,0 +1,32 @@
+import hashlib
+from typing import Final, Protocol
+
+
+class _SafetyIdentifierPayload(Protocol):
+    def get(self, key: str, /) -> object | None: ...
+
+    def __setitem__(self, key: str, value: object, /) -> None: ...
+
+    def __contains__(self, key: object, /) -> bool: ...
+
+    def pop(self, key: str, default: object | None = None, /) -> object | None: ...
+
+
+def enforce_safety_identifier(
+    *,
+    data: _SafetyIdentifierPayload,
+    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

@@ -3,7 +3,8 @@ import contextlib
 import json
 import logging
 import math
-from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Mapping, Sequence
+import os
+from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Mapping, MutableMapping, Sequence
 from datetime import datetime
 from functools import lru_cache
 from types import MappingProxyType
@@ -44,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,
 )
@@ -73,6 +75,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
 
@@ -1514,6 +1517,23 @@ class ProxyBaseLLMRequestProcessing:
     def __init__(self, data: dict):
         self.data = data
 
+    @staticmethod
+    def _enforce_safety_identifier(
+        *,
+        data: MutableMapping[str, object],
+        route_type: ProxyRouteType,
+        user_api_key_dict: UserAPIKeyAuth,
+    ) -> None:
+        if route_type not in ("acompletion", "aresponses"):
+            return
+        if str_to_bool(os.getenv("LITELLM_ENFORCE_SAFETY_IDENTIFIER")) is not True:
+            return
+        enforce_safety_identifier(
+            data=data,
+            user_id=user_api_key_dict.user_id,
+            enabled=True,
+        )
+
     @staticmethod
     def _merge_passthrough_streaming_headers(
         response_headers: httpx.Headers | dict | None,
@@ -1986,6 +2006,12 @@ class ProxyBaseLLMRequestProcessing:
             trust_client_model_info=False,
         )
 
+        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
@@ -2004,6 +2030,12 @@ class ProxyBaseLLMRequestProcessing:
                 llm_router=llm_router,
             )
 
+        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/responses/streaming_iterator.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import asyncio
 import json
+import os
 import time
 import traceback
 import uuid
@@ -30,9 +31,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,
@@ -93,6 +96,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:
+    user_id: Final[str | None] = user_api_key_dict.user_id if user_api_key_dict is not None else None
+    enabled: Final[bool] = str_to_bool(os.getenv("LITELLM_ENFORCE_SAFETY_IDENTIFIER")) is True
+    modified = enforce_safety_identifier(data=msg_obj, user_id=user_id, enabled=enabled)
+    nested_candidate: Final = msg_obj.get("response")
+    if _is_json_object(nested_candidate):
+        modified = enforce_safety_identifier(data=nested_candidate, user_id=user_id, enabled=enabled) or modified
+    return modified
+
+
 class _MutableJsonObject(Protocol):
     @overload
     def get(self, key: str, /) -> object | None: ...
@@ -1766,16 +1782,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)
@@ -2543,6 +2561,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

@@ -216,6 +216,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

@@ -4319,6 +4319,9 @@ def get_optional_params(
     allowed_openai_params = allowed_openai_params or []
     supported_params.extend(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,157 @@
+import hashlib
+from typing import Literal
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+import litellm
+from litellm.litellm_core_utils.safety_identifier import enforce_safety_identifier
+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: pytest.MonkeyPatch):
+    monkeypatch.setenv("LITELLM_ENFORCE_SAFETY_IDENTIFIER", "true")
+
+    data = {"safety_identifier": "caller-value"}
+    ProxyBaseLLMRequestProcessing._enforce_safety_identifier(
+        data=data,
+        route_type="acompletion",
+        user_api_key_dict=UserAPIKeyAuth(user_id="user-123"),
+    )
+
+    assert data["safety_identifier"] == hashlib.sha256(b"user-123").hexdigest()
+
+
+def test_enforce_safety_identifier_is_idempotent():
+    safety_identifier = hashlib.sha256(b"user-123").hexdigest()
+    data = {"safety_identifier": safety_identifier}
+
+    modified = enforce_safety_identifier(data=data, user_id="user-123", enabled=True)
+
+    assert modified is False
+    assert data == {"safety_identifier": safety_identifier}
+
+
+@pytest.mark.parametrize("setting", [None, "false"])
+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"}
+
+    ProxyBaseLLMRequestProcessing._enforce_safety_identifier(
+        data=data,
+        route_type="acompletion",
+        user_api_key_dict=UserAPIKeyAuth(user_id="user-123"),
+    )
+
+    assert data == {"safety_identifier": "caller-value"}
+
+
+def test_enforce_safety_identifier_removes_untrusted_identifier(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setenv("LITELLM_ENFORCE_SAFETY_IDENTIFIER", "true")
+    data = {"safety_identifier": "caller-value"}
+
+    ProxyBaseLLMRequestProcessing._enforce_safety_identifier(
+        data=data,
+        route_type="aresponses",
+        user_api_key_dict=UserAPIKeyAuth(user_id=None),
+    )
+
+    assert data == {}
+
+
+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"}
+
+    ProxyBaseLLMRequestProcessing._enforce_safety_identifier(
+        data=data,
+        route_type="aembedding",
+        user_api_key_dict=UserAPIKeyAuth(user_id="user-123"),
+    )
+
+    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: str, model: str):
+    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"
+
+
+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: pytest.MonkeyPatch, route_type: Literal["acompletion", "aresponses"]
+):
+    monkeypatch.setenv("LITELLM_ENFORCE_SAFETY_IDENTIFIER", "true")
+    request = MagicMock()
+    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()

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
 
@@ -1205,6 +1206,130 @@ 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_nested_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", "response": {"safety_identifier": "caller-value"}})
+        )
+
+        assert json.loads(masked)["response"]["safety_identifier"] == hashlib.sha256(b"user-123").hexdigest()
+
+    @pytest.mark.asyncio
+    async def test_nested_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", "response": {"safety_identifier": "caller-value"}})
+        )
+
+        assert "safety_identifier" not in json.loads(masked)["response"]
+
+    @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_managed_nested_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", "response": {"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