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

fix(proxy): run the failure hook when a pass-through stream dies mid-body

Browse source
This commit is contained in:
mateo-berri 2026-09-03 11:31:35 -07:00
parent cf7abf8136
commit 501f47ba2f
2 changed files with 140 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

84
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

@ -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,

72
tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py
View file

@ -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."""