diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index aeb502599dc..91c07036bde 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1396,6 +1396,35 @@ app.add_middleware( app.add_middleware(PrometheusAuthMiddleware) +@app.middleware("http") +async def rewrite_redirect_location(request: Request, call_next: Any) -> Response: + """ + Rewrite Location headers in redirect responses to use the public + hostname from X-Forwarded-Host / X-Forwarded-Proto when available. + + This handles cases where the reverse proxy does not preserve the + original Host header (e.g. Kong with preserve_host: false), which + causes Starlette's Router.redirect_slashes to build redirect URLs + using the internal pod IP instead of the public hostname. + + Only activates when X-Forwarded-Host is present, so it is a no-op + for local development and environments without a reverse proxy. + """ + response = await call_next(request) + if response.status_code in (301, 302, 307, 308): + location = response.headers.get("location", "") + fwd_host = request.headers.get("x-forwarded-host", "") + fwd_proto = request.headers.get("x-forwarded-proto", "https") + if fwd_host and location: + from urllib.parse import urlparse, urlunparse + + parsed = urlparse(location) + if parsed.netloc: # only rewrite absolute URLs + new = parsed._replace(scheme=fwd_proto, netloc=fwd_host) + response.headers["location"] = urlunparse(new) + return response + + def mount_swagger_ui(): swagger_directory = os.path.join(current_dir, "swagger") swagger_path = "/" if server_root_path is None else server_root_path diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 7deda21c215..7e9b0f97b5b 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -135,6 +135,45 @@ def test_join_paths_nested_path(): assert result == "http://0.0.0.0:4000/v1/chat/completions" +@pytest.mark.asyncio +async def test_rewrite_redirect_location_with_forwarded_host(): + """Test that redirect Location headers are rewritten using X-Forwarded-Host""" + from starlette.testclient import TestClient + from starlette.responses import RedirectResponse + from litellm.proxy.proxy_server import app + + # Create a test client that sends X-Forwarded-Host + client = TestClient(app) + # Hit /ui which will trigger a trailing-slash redirect to /ui/ + response = client.get( + "/ui", + headers={ + "x-forwarded-host": "external.company.com", + "x-forwarded-proto": "https", + }, + follow_redirects=False, + ) + if response.status_code in (301, 302, 307, 308): + location = response.headers.get("location", "") + # The Location should use the forwarded host, not an internal IP + assert "external.company.com" in location + assert location.startswith("https://") + + +@pytest.mark.asyncio +async def test_rewrite_redirect_location_no_forwarded_host(): + """Test that redirect Location headers are NOT rewritten without X-Forwarded-Host""" + from starlette.testclient import TestClient + from litellm.proxy.proxy_server import app + + client = TestClient(app) + response = client.get("/ui", follow_redirects=False) + if response.status_code in (301, 302, 307, 308): + location = response.headers.get("location", "") + # Without X-Forwarded-Host, the location should use the original host + assert "external.company.com" not in location + + def _patch_today(monkeypatch, year, month, day): class PatchedDate(real_datetime.date): @classmethod