diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f837e75039e..4973d277c2b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2545,7 +2545,9 @@ async def otel_unhandled_exception_handler(request: Request, exc: Exception): if isinstance(exc, (ProxyException, HTTPException, RequestValidationError)): raise exc response: Final = await _unhandled_exception_response(request, exc) - response.headers.update(trace_response_headers(getattr(request.state, "parent_otel_span", None))) + response.headers.update( + {**trace_response_headers(getattr(request.state, "parent_otel_span", None)), "connection": "close"} + ) return response diff --git a/tests/unit/proxy/proxy_server/test_exception_handlers.py b/tests/unit/proxy/proxy_server/test_exception_handlers.py index 341875c7cca..40016b8bb85 100644 --- a/tests/unit/proxy/proxy_server/test_exception_handlers.py +++ b/tests/unit/proxy/proxy_server/test_exception_handlers.py @@ -643,6 +643,31 @@ async def test_otel_unhandled_exception_handler_answers_a_db_outage_with_503_no_ assert json.loads(response.body) == _DB_OUTAGE_503_BODY +@pytest.mark.parametrize( + ("path", "exc", "status_code"), + [ + ("/v1/messages/count_tokens", RuntimeError("kaboom"), 500), + ("/v1/messages/count_tokens", httpx.ConnectError("All connection attempts failed"), 503), + ("/v1/traces", RuntimeError("kaboom"), 500), + ], + ids=["unhandled_500", "db_outage_503", "otlp_500"], +) +def test_an_unhandled_exception_response_tells_the_client_to_drop_the_connection( + path: str, exc: Exception, status_code: int +) -> None: + async def fail() -> None: + raise exc + + app: Final = FastAPI() + app.exception_handler(Exception)(otel_unhandled_exception_handler) + app.add_api_route(path, fail, methods=["POST"]) + + response: Final = TestClient(app, raise_server_exceptions=False).post(path) + + assert response.status_code == status_code + assert response.headers["connection"] == "close" + + @pytest.mark.asyncio @pytest.mark.parametrize("path", ("/v1/messages", "/claude_code_gateway/v1/messages", "/anthropic/v1/messages")) async def test_otel_unhandled_exception_handler_answers_the_anthropic_api_error_on_the_messages_routes(