From 501f47ba2fa59ad950c4e122fe9cfddf890ec106 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 11:31:35 -0700 Subject: [PATCH] fix(proxy): run the failure hook when a pass-through stream dies mid-body --- .../pass_through_endpoints.py | 84 +++++++++++++++---- .../test_pass_through_endpoints.py | 72 ++++++++++++++++ 2 files changed, 140 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 79d5d0a016f..37c8cfb09d6 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -868,6 +868,38 @@ async def _log_passthrough_upstream_failure( ) +async def _relay_reporting_failures( + stream: AsyncGenerator[bytes, None], + upstream_status: int, + user_api_key_dict: UserAPIKeyAuth, + request_payload: dict, # mutable-ok: post_call_failure_hook lifts fields onto request_data in place +) -> AsyncGenerator[bytes, None]: + """An upstream that dies mid-stream leaves the client a truncated body and the proxy no record, so run + ``post_call_failure_hook`` (spend row, alerting, failure metric) the way the unified endpoints' generators do. + Error statuses were already reported by ``_log_passthrough_upstream_failure`` and relay untouched.""" + from litellm.proxy.proxy_server import proxy_logging_obj + + try: + async for chunk in stream: + yield chunk + except Exception as e: + if upstream_status >= 400: + raise + try: + await proxy_logging_obj.post_call_failure_hook( + user_api_key_dict=user_api_key_dict, + original_exception=e, + request_data=request_payload, + traceback_str=traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG), + ) + except Exception: # noqa: BLE001 - a failing logging callback must never mask the upstream error + verbose_proxy_logger.warning( + "pass_through_endpoint: post_call_failure_hook raised for a mid-stream upstream error", + exc_info=True, + ) + raise + + from litellm.passthrough.timeout_utils import ( DEFAULT_PASS_THROUGH_REQUEST_TIMEOUT_SECONDS, # noqa: F401 - re-exported for backward compat resolve_llm_passthrough_timeout, # noqa: F401 - re-exported for backward compat @@ -1291,14 +1323,24 @@ async def pass_through_request( return StreamingResponse( wrap_passthrough_sse_bytes_with_keepalive_pings( stream=_own_streamed_managed_ids( - stream=PassThroughStreamingHandler.chunk_processor( - response=response, - request_body=_parsed_body, - litellm_logging_obj=logging_obj, - endpoint_type=endpoint_type, - start_time=start_time, - passthrough_success_handler_obj=pass_through_endpoint_logging, - url_route=str(url), + stream=_relay_reporting_failures( + stream=PassThroughStreamingHandler.chunk_processor( + response=response, + request_body=_parsed_body, + litellm_logging_obj=logging_obj, + endpoint_type=endpoint_type, + start_time=start_time, + passthrough_success_handler_obj=pass_through_endpoint_logging, + url_route=str(url), + ), + upstream_status=response.status_code, + user_api_key_dict=user_api_key_dict, + request_payload=_build_passthrough_failure_request_payload( + parsed_body=_parsed_body, + kwargs=kwargs, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + ), ), managed_id_provider=_managed_id_provider, request=request, @@ -1372,14 +1414,24 @@ async def pass_through_request( return StreamingResponse( wrap_passthrough_sse_bytes_with_keepalive_pings( stream=_own_streamed_managed_ids( - stream=PassThroughStreamingHandler.chunk_processor( - response=response, - request_body=_parsed_body, - litellm_logging_obj=logging_obj, - endpoint_type=endpoint_type, - start_time=start_time, - passthrough_success_handler_obj=pass_through_endpoint_logging, - url_route=str(url), + stream=_relay_reporting_failures( + stream=PassThroughStreamingHandler.chunk_processor( + response=response, + request_body=_parsed_body, + litellm_logging_obj=logging_obj, + endpoint_type=endpoint_type, + start_time=start_time, + passthrough_success_handler_obj=pass_through_endpoint_logging, + url_route=str(url), + ), + upstream_status=response.status_code, + user_api_key_dict=user_api_key_dict, + request_payload=_build_passthrough_failure_request_payload( + parsed_body=_parsed_body, + kwargs=kwargs, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + ), ), managed_id_provider=_managed_id_provider, request=request, diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index d3f17c73499..442adbca08c 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -3988,6 +3988,78 @@ async def test_pass_through_request_streaming_upstream_error_returned_unchanged( assert failure_call_kwargs["original_exception"].status_code == 403 +class _UpstreamDroppingMidStream(httpx.AsyncByteStream): + async def __aiter__(self): + yield b'data: {"id": "chatcmpl-1", "choices": [{"delta": {"content": "hi"}}]}\n\n' + raise httpx.ReadError("upstream dropped the connection mid-stream") + + +async def _relay_everything(body_iterator) -> list: + return [chunk async for chunk in body_iterator] + + +@pytest.mark.asyncio +async def test_pass_through_request_mid_stream_upstream_drop_fires_failure_hook(): + """ + Regression: a 200 stream whose upstream dies mid-body used to end with no + proxy-level failure hook at all, so the request left no spend row, no + failure metric, and no alert; the pre-stream 4xx/5xx path already fires it. + """ + from litellm.llms.custom_httpx.http_handler import get_async_httpx_client + from litellm.types.llms.custom_http import httpxSpecialProvider + + def transport_handler(upstream_request: httpx.Request) -> httpx.Response: + return httpx.Response(200, stream=_UpstreamDroppingMidStream(), headers={"content-type": "text/event-stream"}) + + real_handler = get_async_httpx_client( + llm_provider=httpxSpecialProvider.PassThroughEndpoint, + params={"timeout": resolve_pass_through_request_timeout(None)}, + ) + cache_dict = litellm.in_memory_llm_clients_cache.cache_dict + cache_key = next(key for key, cached in cache_dict.items() if cached is real_handler) + cache_dict[cache_key] = SimpleNamespace(client=httpx.AsyncClient(transport=httpx.MockTransport(transport_handler))) + + mock_proxy_logging = MagicMock() + mock_proxy_logging.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None) + mock_proxy_logging.get_proxy_hook = MagicMock(return_value=None) + + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.scope = {"path": "/relay-chat"} + mock_request.url = MagicMock() + mock_request.url.path = "/relay-chat" + mock_request.body = AsyncMock(return_value=b'{"model": "gpt-5.6", "stream": true}') + mock_request.headers = Headers({"content-type": "application/json"}) + mock_request.query_params = QueryParams({}) + + try: + with patch( # test-quality-ok: proxy_logging_obj is a proxy_server module global read inside pass_through_request; there is no injection seam + "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging + ): + response = await pass_through_request( + request=mock_request, + target="http://target-api.com/v1/chat/completions", + custom_headers={}, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + stream=True, + ) + with pytest.raises(httpx.ReadError): + await _relay_everything(response.body_iterator) + await asyncio.sleep(0) + finally: + cache_dict[cache_key] = real_handler + + mock_proxy_logging.post_call_failure_hook.assert_awaited_once() + failure_call_kwargs = mock_proxy_logging.post_call_failure_hook.call_args.kwargs + assert isinstance(failure_call_kwargs["original_exception"], httpx.ReadError) + request_data = failure_call_kwargs["request_data"] + assert request_data["litellm_call_id"] + assert request_data["model"] == "gpt-5.6" + assert isinstance(request_data["litellm_logging_obj"], LiteLLMLoggingObj) + + @pytest.mark.asyncio async def test_pass_through_request_non_streaming_success_unchanged(): """Success (2xx) passthrough behavior must remain unchanged by the error fix."""