scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 124

[Fix] Preserve upstream Content-Type when caching HTTP logo

Browse source 
When a logo is fetched from an HTTP URL, the downstream Content-Type
(e.g. image/png for a PNG) was discarded and the cached file was always
served as image/jpeg. This caused the same Content-Type mismatch fixed
for local files in the previous commit.

- Save the upstream Content-Type to cached_logo_type.txt alongside the
  cached image bytes so it survives server restarts
- Read that file when serving from cache instead of hardcoding image/jpeg
- Validate that the persisted type starts with image/ as a safety check
- Add test covering the PNG download + cache roundtrip

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
yuneng-jiang 2026-03-04 21:04:35 -08:00
parent 1f6f16d035
commit 5d6ea96a5d
2 changed files with 94 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

30
litellm/proxy/proxy_server.py
View file

@ -11058,6 +11058,7 @@ async def get_image():
cache_dir = assets_dir if os.access(assets_dir, os.W_OK) else current_dir
cache_path = os.path.join(cache_dir, "cached_logo.jpg")
cache_type_path = os.path.join(cache_dir, "cached_logo_type.txt")
logo_path = os.getenv("UI_LOGO_PATH", default_logo)
verbose_proxy_logger.debug("Reading logo from path: %s", logo_path)
@ -11068,6 +11069,17 @@ async def get_image():
return media_type
return "image/jpeg"
def _read_cached_media_type() -> str:
if os.path.exists(cache_type_path):
try:
with open(cache_type_path) as f:
cached_type = f.read().strip()
if cached_type.startswith("image/"):
return cached_type
except OSError:
pass
return "image/jpeg"
# If UI_LOGO_PATH points to a local file, serve it directly (skip cache)
if logo_path != default_logo and not logo_path.startswith(("http://", "https://")):
if os.path.exists(logo_path):
@ -11080,7 +11092,7 @@ async def get_image():
# [OPTIMIZATION] For HTTP URLs and default logo, check if the cached image exists
if os.path.exists(cache_path):
return FileResponse(cache_path, media_type="image/jpeg")
return FileResponse(cache_path, media_type=_read_cached_media_type())
# Check if the logo path is an HTTP/HTTPS URL
if logo_path.startswith(("http://", "https://")):
@ -11099,8 +11111,22 @@ async def get_image():
with open(cache_path, "wb") as f:
f.write(response.content)
# Persist the upstream Content-Type so it survives restarts
content_type = (
response.headers.get("content-type", "image/jpeg")
.split(";")[0]
.strip()
)
if not content_type.startswith("image/"):
content_type = "image/jpeg"
try:
with open(cache_type_path, "w") as f:
f.write(content_type)
except OSError:
pass
# Return the cached image as a FileResponse
return FileResponse(cache_path, media_type="image/jpeg")
return FileResponse(cache_path, media_type=content_type)
else:
# Handle the case when the image cannot be downloaded
return FileResponse(default_logo, media_type="image/jpeg")

69
tests/proxy_unit_tests/test_get_image.py
View file

@ -45,6 +45,10 @@ async def test_get_image_error_handling():
assert response.headers["content-type"] == "image/jpeg"
def _get_cache_dir(proxy_dir: str) -> str:
return os.path.join(proxy_dir, "proxy")
@pytest.mark.asyncio
async def test_get_image_cache_logic():
"""
@ -60,14 +64,18 @@ async def test_get_image_cache_logic():
else "litellm/proxy/proxy_server.py"
)
)
cache_path = os.path.join(parent_dir, "proxy", "cached_logo.jpg")
if os.path.exists(cache_path):
os.remove(cache_path)
cache_dir = _get_cache_dir(parent_dir)
cache_path = os.path.join(cache_dir, "cached_logo.jpg")
cache_type_path = os.path.join(cache_dir, "cached_logo_type.txt")
for p in (cache_path, cache_type_path):
if os.path.exists(p):
os.remove(p)
# Mock response
mock_response = mock.Mock()
mock_response.status_code = 200
mock_response.content = b"fake image data"
mock_response.headers = {"content-type": "image/jpeg"}
with mock.patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get"
@ -87,3 +95,58 @@ async def test_get_image_cache_logic():
assert response2.status_code == 200
# If cache works, mock_get shouldn't be called again
assert mock_get.call_count == 1
@pytest.mark.asyncio
async def test_get_image_preserves_content_type_from_http_url():
"""
When downloading a logo from an HTTP URL, the Content-Type from the upstream
response should be persisted and used when serving from cache, not hardcoded
to image/jpeg. This matters for PNG/SVG logos served through strict enterprise
proxies that validate Content-Type.
"""
os.environ["UI_LOGO_PATH"] = "http://example.com/logo.png"
parent_dir = os.path.dirname(
os.path.dirname(
app.__file__
if hasattr(app, "__file__")
else "litellm/proxy/proxy_server.py"
)
)
cache_dir = _get_cache_dir(parent_dir)
cache_path = os.path.join(cache_dir, "cached_logo.jpg")
cache_type_path = os.path.join(cache_dir, "cached_logo_type.txt")
for p in (cache_path, cache_type_path):
if os.path.exists(p):
os.remove(p)
mock_response = mock.Mock()
mock_response.status_code = 200
mock_response.content = b"\x89PNG\r\n\x1a\n" # PNG magic bytes
mock_response.headers = {"content-type": "image/png"}
with mock.patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get"
) as mock_get:
mock_get.return_value = mock_response
async with httpx.AsyncClient(
transport=httpx.ASGITransport(app=app), base_url="http://testserver"
) as ac:
# First request: download and cache
response1 = await ac.get("/get_image")
assert response1.status_code == 200
assert response1.headers["content-type"] == "image/png"
# Second request: served from cache with the persisted content-type
response2 = await ac.get("/get_image")
assert response2.status_code == 200
assert response2.headers["content-type"] == "image/png"
# Upstream should only have been hit once
assert mock_get.call_count == 1
# Clean up
for p in (cache_path, cache_type_path):
if os.path.exists(p):
os.remove(p)