fix(proxy): gate redaction opt out controls

litellm/proxy/litellm_pre_call_utils.py

@@ -149,9 +149,16 @@ _UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS = frozenset(
 )
 _CLIENT_MOCK_CONTROL_FIELDS = frozenset({"mock_response", "mock_tool_calls"})
 _ALLOW_CLIENT_MOCK_RESPONSE_METADATA_KEY = "allow_client_mock_response"
+_ALLOW_CLIENT_MESSAGE_REDACTION_OPT_OUT_METADATA_KEY = (
+    "allow_client_message_redaction_opt_out"
+)
 
 
-def _strip_untrusted_request_header_controls(headers: Any) -> None:
+def _strip_untrusted_request_header_controls(
+    headers: Any,
+    *,
+    allow_client_message_redaction_opt_out: bool = False,
+) -> None:
     if not isinstance(headers, dict):
         return
 
@@ -160,6 +167,8 @@ def _strip_untrusted_request_header_controls(headers: Any) -> None:
             isinstance(header_name, str)
             and header_name.lower() in _UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS
         ):
+            if allow_client_message_redaction_opt_out:
+                continue
             headers.pop(header_name, None)
 
 
@@ -171,18 +180,37 @@ def _is_false_like(value: Any) -> bool:
     return False
 
 
-def _key_or_team_allows_client_mock_response(
+def _key_or_team_metadata_flag_is_true(
     user_api_key_dict: UserAPIKeyAuth,
+    metadata_key: str,
 ) -> bool:
     for admin_metadata in (user_api_key_dict.metadata, user_api_key_dict.team_metadata):
         if (
             isinstance(admin_metadata, dict)
-            and admin_metadata.get(_ALLOW_CLIENT_MOCK_RESPONSE_METADATA_KEY) is True
+            and admin_metadata.get(metadata_key) is True
         ):
             return True
     return False
 
 
+def _key_or_team_allows_client_mock_response(
+    user_api_key_dict: UserAPIKeyAuth,
+) -> bool:
+    return _key_or_team_metadata_flag_is_true(
+        user_api_key_dict=user_api_key_dict,
+        metadata_key=_ALLOW_CLIENT_MOCK_RESPONSE_METADATA_KEY,
+    )
+
+
+def _key_or_team_allows_client_message_redaction_opt_out(
+    user_api_key_dict: UserAPIKeyAuth,
+) -> bool:
+    return _key_or_team_metadata_flag_is_true(
+        user_api_key_dict=user_api_key_dict,
+        metadata_key=_ALLOW_CLIENT_MESSAGE_REDACTION_OPT_OUT_METADATA_KEY,
+    )
+
+
 def _get_metadata_variable_name(request: Request) -> str:
     """
     Helper to return what the "metadata" field should be called in the request data
@@ -1043,6 +1071,9 @@ async def add_litellm_data_to_request(  # noqa: PLR0915
     _allow_client_mock_response = _key_or_team_allows_client_mock_response(
         user_api_key_dict
     )
+    _allow_client_message_redaction_opt_out = (
+        _key_or_team_allows_client_message_redaction_opt_out(user_api_key_dict)
+    )
     for _internal_key in _UNTRUSTED_ROOT_CONTROL_FIELDS:
         if _allow_client_mock_response and _internal_key in _CLIENT_MOCK_CONTROL_FIELDS:
             continue
@@ -1086,9 +1117,13 @@ async def add_litellm_data_to_request(  # noqa: PLR0915
         forward_llm_provider_auth_headers=forward_llm_auth,
         authenticated_with_header=authenticated_with_header,
     )
-    _strip_untrusted_request_header_controls(_headers)
+    _strip_untrusted_request_header_controls(
+        _headers,
+        allow_client_message_redaction_opt_out=_allow_client_message_redaction_opt_out,
+    )
     if (
-        litellm.turn_off_message_logging is True
+        not _allow_client_message_redaction_opt_out
+        and litellm.turn_off_message_logging is True
         and "turn_off_message_logging" in data
         and _is_false_like(data["turn_off_message_logging"])
     ):
@@ -1230,7 +1265,12 @@ async def add_litellm_data_to_request(  # noqa: PLR0915
     for _meta_key in ("metadata", "litellm_metadata"):
         _user_meta = data.get(_meta_key)
         if isinstance(_user_meta, dict):
-            _strip_untrusted_request_header_controls(_user_meta.get("headers"))
+            _strip_untrusted_request_header_controls(
+                _user_meta.get("headers"),
+                allow_client_message_redaction_opt_out=(
+                    _allow_client_message_redaction_opt_out
+                ),
+            )
             for _k in [
                 k
                 for k in _user_meta

tests/test_litellm/proxy/test_litellm_pre_call_utils.py

@@ -686,6 +686,73 @@ async def test_add_litellm_data_to_request_strips_client_redaction_bypass_contro
     }
 
 
+@pytest.mark.parametrize(
+    "auth_kwargs",
+    [
+        {"metadata": {"allow_client_message_redaction_opt_out": True}},
+        {"team_metadata": {"allow_client_message_redaction_opt_out": True}},
+    ],
+)
+@pytest.mark.asyncio
+async def test_add_litellm_data_to_request_allows_redaction_opt_out_with_admin_opt_in(
+    auth_kwargs,
+):
+    request_mock = MagicMock(spec=Request)
+    request_mock.url.path = "/v1/chat/completions"
+    request_mock.url = MagicMock()
+    request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions"
+    request_mock.method = "POST"
+    request_mock.query_params = {}
+    request_mock.headers = {
+        "Content-Type": "application/json",
+        "litellm-disable-message-redaction": "true",
+    }
+    request_mock.client = MagicMock()
+    request_mock.client.host = "127.0.0.1"
+
+    original_turn_off_message_logging = litellm.turn_off_message_logging
+    litellm.turn_off_message_logging = True
+    try:
+        updated = await add_litellm_data_to_request(
+            data={
+                "model": "gpt-3.5-turbo",
+                "messages": [{"role": "user", "content": "hello"}],
+                "turn_off_message_logging": False,
+                "metadata": {"headers": {"litellm-disable-message-redaction": "true"}},
+                "litellm_metadata": json.dumps(
+                    {"headers": {"LiteLLM-Disable-Message-Redaction": "true"}}
+                ),
+            },
+            request=request_mock,
+            user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key", **auth_kwargs),
+            proxy_config=MagicMock(),
+            general_settings={},
+            version="test-version",
+        )
+    finally:
+        litellm.turn_off_message_logging = original_turn_off_message_logging
+
+    assert updated["turn_off_message_logging"] is False
+    assert "litellm-disable-message-redaction" in {
+        header.lower() for header in updated["metadata"]["headers"]
+    }
+    assert "litellm-disable-message-redaction" in {
+        header.lower()
+        for header in updated["metadata"]["requester_metadata"].get("headers", {})
+    }
+    assert "litellm-disable-message-redaction" in {
+        header.lower() for header in updated["proxy_server_request"]["headers"]
+    }
+    assert "litellm-disable-message-redaction" in {
+        header.lower()
+        for header in updated["proxy_server_request"]["body"]["metadata"]["headers"]
+    }
+    assert "litellm-disable-message-redaction" in {
+        header.lower()
+        for header in (updated.get("litellm_metadata") or {}).get("headers", {})
+    }
+
+
 @pytest.mark.asyncio
 async def test_add_litellm_data_to_request_ignores_x_litellm_tags_header_without_permission():
     """Regression: the `x-litellm-tags` header bypassed the body-metadata