diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7566e02baea..6d7dfb2cf93 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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") diff --git a/tests/proxy_unit_tests/test_get_image.py b/tests/proxy_unit_tests/test_get_image.py index ad8c2672754..1c160f38eae 100644 --- a/tests/proxy_unit_tests/test_get_image.py +++ b/tests/proxy_unit_tests/test_get_image.py @@ -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)