From 313e1495fa2156e03ab50be784a0ed806d3e7e42 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 16 Sep 2026 01:31:27 +0000 Subject: [PATCH] fix(openai): sanitize image generation failure telemetry Exception post_call paths were logging the raw request dict, which could leak extra_headers and nested bearer tokens. Use the sanitized logged body on every failure path and keep extra_body/extra_query keys that are not headers. Co-authored-by: Zsanz3 --- litellm/llms/openai/openai.py | 27 ++++--- ...t_openai_image_generation_extra_headers.py | 78 ++++++++++++++++++- 2 files changed, 92 insertions(+), 13 deletions(-) diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 0c6cff5bca5..69ff2e9b8f9 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -362,13 +362,22 @@ def _embedding_request_without_sdk_defaults( return body, options +def _image_generation_logged_body(data: Mapping[str, object]) -> dict[str, object]: + extra_body: Final = { + key: item for key, item in _as_mapping(data.get("extra_body")).items() if key != "extra_headers" + } + extra_query: Final = dict(_as_mapping(data.get("extra_query"))) + return { # mutable-ok: loggers isinstance-check this payload as a dict + **{key: item for key, item in data.items() if key not in _SDK_OPTION_KEYS}, + **({} if not extra_query else {"extra_query": extra_query}), + **({} if not extra_body else {"extra_body": extra_body}), + } + + def _image_generation_request_data( data: Mapping[str, object], headers: Mapping[str, object] | None ) -> tuple[dict[str, object], Mapping[str, object]]: extra_body_source: Final = _as_mapping(data.get("extra_body")) - extra_body: Final = types.MappingProxyType( - {key: item for key, item in extra_body_source.items() if key != "extra_headers"} - ) extra_headers: Final = types.MappingProxyType( { **_as_mapping(extra_body_source.get("extra_headers")), @@ -376,17 +385,12 @@ def _image_generation_request_data( **_as_mapping(headers), } ) - extra_query: Final = _as_mapping(data.get("extra_query")) - logged_body: Final = { # mutable-ok: loggers isinstance-check this payload as a dict - key: item for key, item in data.items() if key not in _SDK_OPTION_KEYS - } + logged_body: Final = _image_generation_logged_body(data) request_data: Final = types.MappingProxyType( dict( ( *logged_body.items(), *((("extra_headers", extra_headers),) if extra_headers else ()), - *((("extra_query", extra_query),) if extra_query else ()), - *((("extra_body", extra_body),) if extra_body else ()), ) ) ) @@ -1483,6 +1487,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): logging_obj.post_call( input=prompt, api_key=api_key, + additional_args={"complete_input_dict": _image_generation_logged_body(data)}, original_response=str(e), ) raise e @@ -1569,7 +1574,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): logging_obj.post_call( input=prompt, api_key=api_key, - additional_args={"complete_input_dict": data}, + additional_args={"complete_input_dict": _image_generation_logged_body(data)}, original_response=str(e), ) raise e @@ -1578,7 +1583,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): logging_obj.post_call( input=prompt, api_key=api_key, - additional_args={"complete_input_dict": data}, + additional_args={"complete_input_dict": _image_generation_logged_body(data)}, original_response=str(e), ) if hasattr(e, "status_code"): diff --git a/tests/test_litellm/llms/openai/image_generation/test_openai_image_generation_extra_headers.py b/tests/test_litellm/llms/openai/image_generation/test_openai_image_generation_extra_headers.py index b422e62b9c2..c7378a447a0 100644 --- a/tests/test_litellm/llms/openai/image_generation/test_openai_image_generation_extra_headers.py +++ b/tests/test_litellm/llms/openai/image_generation/test_openai_image_generation_extra_headers.py @@ -10,8 +10,24 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from litellm.llms.openai.common_utils import OpenAIError from litellm.llms.openai.openai import OpenAIChatCompletion +_SECRET_HEADERS = {"Authorization": "Bearer leaked-token", "cf-aig-authorization": "Bearer custom-token"} + + +def _assert_logged_body_omits_credentials(logged_body: object) -> None: + assert isinstance(logged_body, dict) + assert "extra_headers" not in logged_body + assert "Authorization" not in logged_body + assert "cf-aig-authorization" not in logged_body + extra_body = logged_body.get("extra_body") or {} + assert "extra_headers" not in extra_body + assert "Bearer leaked-token" not in str(logged_body) + assert extra_body.get("foo") == "bar" + assert logged_body.get("extra_query") == {"user": "alice"} + assert logged_body.get("size") == "1024x1024" + @pytest.fixture def openai_chat_completions(): @@ -235,6 +251,7 @@ class TestImageGenerationExtraHeaders: "quality": "low", "extra_headers": test_headers, "extra_body": {"extra_headers": test_headers, "foo": "bar"}, + "extra_query": {"user": "alice"}, } if is_async: @@ -267,10 +284,10 @@ class TestImageGenerationExtraHeaders: extra_body = kwargs.get("extra_body") or {} assert "extra_headers" not in extra_body assert extra_body.get("foo") == "bar" + assert kwargs.get("extra_query") == {"user": "alice"} logged_body = mock_logging_obj.pre_call.call_args[1]["additional_args"]["complete_input_dict"] - assert "extra_headers" not in logged_body - assert "extra_headers" not in (logged_body.get("extra_body") or {}) + _assert_logged_body_omits_credentials(logged_body) def test_optional_params_headers_merge_prefers_caller_headers(self, openai_chat_completions, mock_logging_obj): """When both optional_params and the headers arg carry extra_headers, @@ -302,6 +319,63 @@ class TestImageGenerationExtraHeaders: _, kwargs = mock_openai_client.images.generate.call_args assert kwargs.get("extra_headers") == {"cf-aig-authorization": "Bearer from-headers"} + @pytest.mark.parametrize( + "is_async,raised", + [ + (False, OpenAIError(status_code=400, message="bad image params")), + (False, RuntimeError("forced failure")), + (True, OpenAIError(status_code=400, message="bad image params")), + (True, RuntimeError("forced failure")), + ], + ids=["sync-openai-error", "sync-runtime-error", "async-openai-error", "async-runtime-error"], + ) + @pytest.mark.asyncio + async def test_failure_telemetry_omits_extra_headers( + self, openai_chat_completions, mock_logging_obj, is_async, raised + ): + """post_call on a generate() failure must use the sanitized body, not raw data.""" + mock_openai_client = MagicMock() + mock_openai_client.api_key = "test-key" + mock_openai_client._base_url._uri_reference = "https://api.openai.com" + + optional_params = { + "size": "1024x1024", + "extra_headers": _SECRET_HEADERS, + "extra_body": {"extra_headers": _SECRET_HEADERS, "foo": "bar"}, + "extra_query": {"user": "alice"}, + } + + if is_async: + mock_openai_client.images.generate = AsyncMock(side_effect=raised) + with pytest.raises(type(raised)): + await openai_chat_completions.aimage_generation( + prompt="A white cat", + data={"model": "gpt-image-2", "prompt": "A white cat", **optional_params}, + model_response=MagicMock(), + timeout=60.0, + logging_obj=mock_logging_obj, + api_key="test-key", + headers=_SECRET_HEADERS, + client=mock_openai_client, + ) + else: + mock_openai_client.images.generate.side_effect = raised + with pytest.raises(OpenAIError): + openai_chat_completions.image_generation( + model="gpt-image-2", + prompt="A white cat", + timeout=60.0, + optional_params=optional_params, + logging_obj=mock_logging_obj, + api_key="test-key", + headers=_SECRET_HEADERS, + client=mock_openai_client, + ) + + mock_openai_client.images.generate.assert_called_once() + logged_body = mock_logging_obj.post_call.call_args[1]["additional_args"]["complete_input_dict"] + _assert_logged_body_omits_credentials(logged_body) + class TestImageGenerationEntryPointHeaders: """Test that litellm.image_generation() passes headers through to the OpenAI provider."""