scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-10 08:28:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 108

fix(openai): sanitize image generation failure telemetry

Browse source 
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 <Zsanz3@users.noreply.github.com>
This commit is contained in:
Cursor Agent 2026-09-16 01:31:27 +00:00
parent 11a3499d4d
commit 313e1495fa
No known key found for this signature in database
2 changed files with 92 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

27
litellm/llms/openai/openai.py
View file

@ -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"):

78
tests/test_litellm/llms/openai/image_generation/test_openai_image_generation_extra_headers.py
View file

@ -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."""