From 94fd515f9183ff43ef3255e8f98e26938ab26bf0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 4 Nov 2026 13:08:18 -0800 Subject: [PATCH] fix(proxy): close streamed speech when the route fails before responding A response headers hook, or any other step between aspeech returning an unread audio stream and the route returning its StreamingResponse, could raise and leave that stream open. The upstream connection and the deployment's max_parallel_requests slot then stayed held, so later clips got 429s. The route now closes the audio on any error before it returns --- litellm/proxy/proxy_server.py | 110 +++++++++--------- .../proxy/proxy_server/test_routes_audio.py | 65 ++++++++++- 2 files changed, 119 insertions(+), 56 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d23396422d1..e6c6722e21b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -13740,64 +13740,68 @@ async def audio_speech( ### ALERTING ### asyncio.create_task(proxy_logging_obj.update_request_status(litellm_call_id=litellm_call_id, status="success")) - ### RESPONSE HEADERS ### - hidden_params: Final = getattr(response, "_hidden_params", {}) or {} - model_id: Final = hidden_params.get("model_id", None) or "" - cache_key: Final = hidden_params.get("cache_key", None) or "" - api_base: Final = hidden_params.get("api_base", None) or "" - response_cost: Final = hidden_params.get("response_cost", None) or "" - response_call_id: Final = hidden_params.get("litellm_call_id", None) or "" + try: + ### RESPONSE HEADERS ### + hidden_params: Final = getattr(response, "_hidden_params", {}) or {} + model_id: Final = hidden_params.get("model_id", None) or "" + cache_key: Final = hidden_params.get("cache_key", None) or "" + api_base: Final = hidden_params.get("api_base", None) or "" + response_cost: Final = hidden_params.get("response_cost", None) or "" + response_call_id: Final = hidden_params.get("litellm_call_id", None) or "" - custom_headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers( - user_api_key_dict=user_api_key_dict, - model_id=model_id, - cache_key=cache_key, - api_base=api_base, - version=version, - response_cost=response_cost, - model_region=getattr(user_api_key_dict, "allowed_model_region", ""), - fastest_response_batch_completion=None, - call_id=response_call_id, - request_data=data, - hidden_params=hidden_params, - ) + custom_headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=user_api_key_dict, + model_id=model_id, + cache_key=cache_key, + api_base=api_base, + version=version, + response_cost=response_cost, + model_region=getattr(user_api_key_dict, "allowed_model_region", ""), + fastest_response_batch_completion=None, + call_id=response_call_id, + request_data=data, + hidden_params=hidden_params, + ) - # Call response headers hook (matches audio_transcription behavior) - callback_headers: Final = await proxy_logging_obj.post_call_response_headers_hook( - data=data, - user_api_key_dict=user_api_key_dict, - response=response, - request_headers=dict(request.headers), - ) - if callback_headers: - custom_headers.update(callback_headers) + # Call response headers hook (matches audio_transcription behavior) + callback_headers: Final = await proxy_logging_obj.post_call_response_headers_hook( + data=data, + user_api_key_dict=user_api_key_dict, + response=response, + request_headers=dict(request.headers), + ) + if callback_headers: + custom_headers.update(callback_headers) - requested_format: Final = data.get("response_format") - upstream_content_type: Final = ( - response.content_type if isinstance(response, HttpxBinaryResponseContent) else None - ) - media_type: Final = resolve_speech_media_type( - upstream_content_type=upstream_content_type, - response_format=requested_format if isinstance(requested_format, str) else None, - ) + requested_format: Final = data.get("response_format") + upstream_content_type: Final = ( + response.content_type if isinstance(response, HttpxBinaryResponseContent) else None + ) + media_type: Final = resolve_speech_media_type( + upstream_content_type=upstream_content_type, + response_format=requested_format if isinstance(requested_format, str) else None, + ) - audio: Final = cast( # cast-ok: the speech routes return this, and resolve_speech_media_type above already reads it as such - HttpxBinaryResponseContent, response - ) - return StreamingResponse( - _audio_speech_chunk_generator( - audio, - on_upstream_error=partial( - _log_audio_speech_stream_failure, - request_data=data, - user_api_key_dict=user_api_key_dict, - response_cost=float(hidden_params.get("response_cost") or 0.0), + audio: Final = cast( # cast-ok: the speech routes return this, and resolve_speech_media_type above already reads it as such + HttpxBinaryResponseContent, response + ) + return StreamingResponse( + _audio_speech_chunk_generator( + audio, + on_upstream_error=partial( + _log_audio_speech_stream_failure, + request_data=data, + user_api_key_dict=user_api_key_dict, + response_cost=float(hidden_params.get("response_cost") or 0.0), + ), ), - ), - media_type=media_type, - headers=custom_headers, - background=BackgroundTask(_close_audio_speech_response, audio), - ) + media_type=media_type, + headers=custom_headers, + background=BackgroundTask(_close_audio_speech_response, audio), + ) + except BaseException: + await _close_audio_speech_response(response) + raise except Exception as e: await proxy_logging_obj.post_call_failure_hook( diff --git a/tests/unit/proxy/proxy_server/test_routes_audio.py b/tests/unit/proxy/proxy_server/test_routes_audio.py index b3ed1235500..9e160399c78 100644 --- a/tests/unit/proxy/proxy_server/test_routes_audio.py +++ b/tests/unit/proxy/proxy_server/test_routes_audio.py @@ -16,7 +16,7 @@ import io import json import math import wave -from collections.abc import AsyncIterator, Callable, Mapping +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping from contextlib import AbstractContextManager from typing import Final from unittest.mock import AsyncMock, MagicMock @@ -498,7 +498,11 @@ async def _serve_until_disconnect(response: StreamingResponse) -> None: await response({"type": "http", "asgi": {"spec_version": "2.4"}}, _receive, _send) -def _route_speech_to(monkeypatch: pytest.MonkeyPatch, response: HttpxBinaryResponseContent) -> AsyncMock: +def _route_speech_to( + monkeypatch: pytest.MonkeyPatch, + response: HttpxBinaryResponseContent, + response_headers_hook: Callable[..., Awaitable[dict[str, str]]] | None = None, +) -> AsyncMock: failure_hook: Final = AsyncMock() monkeypatch.setattr(proxy_server, "llm_router", MagicMock()) monkeypatch.setattr( @@ -507,7 +511,7 @@ def _route_speech_to(monkeypatch: pytest.MonkeyPatch, response: HttpxBinaryRespo MagicMock( pre_call_hook=AsyncMock(side_effect=lambda **kw: kw["data"]), post_call_failure_hook=failure_hook, - post_call_response_headers_hook=AsyncMock(return_value={}), + post_call_response_headers_hook=response_headers_hook or AsyncMock(return_value={}), update_request_status=AsyncMock(), ), ) @@ -586,6 +590,61 @@ async def test_audio_speech_route_releases_an_upstream_whose_body_is_never_read( assert stream.closed is True +def _elevenlabs_client(stream: httpx.AsyncByteStream) -> litellm.AsyncHTTPHandler: + return litellm.AsyncHTTPHandler( + transport=httpx.MockTransport( + lambda _: httpx.Response(200, headers={"content-type": "audio/mpeg"}, stream=stream) + ) + ) + + +async def _headers_from_the_audio_body(*, response: HttpxBinaryResponseContent, **_: object) -> dict[str, str]: + return {"x-audio-bytes": str(len(response.content))} + + +@pytest.mark.asyncio +async def test_audio_speech_route_frees_the_deployment_slot_when_a_response_hook_fails( + monkeypatch: pytest.MonkeyPatch, +): + """A response hook that fails before any audio is sent must not keep the deployment's only + slot, or every later clip on max_parallel_requests: 1 is turned away.""" + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = Router( + model_list=[ + { + "model_name": "tts", + "litellm_params": { + "model": "elevenlabs/eleven_multilingual_v2", + "api_key": "fake-key", + "max_parallel_requests": 1, + }, + } + ] + ) + first_upstream: Final = _RecordingAudioStream([b"audio"], []) + first_clip: Final = await router.aspeech( + model="tts", input="hello", voice="voice-id", client=_elevenlabs_client(first_upstream), stream=True + ) + _route_speech_to(monkeypatch, first_clip, response_headers_hook=_headers_from_the_audio_body) + + with pytest.raises(ProxyException): + await proxy_server.audio_speech( + request=_speech_request({"model": "tts", "input": "hello", "voice": "voice-id", "stream": True}), + fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + next_clip: Final = await router.aspeech( + model="tts", + input="hello", + voice="voice-id", + client=_elevenlabs_client(_RecordingAudioStream([b"audio"], [])), + stream=True, + ) + await next_clip.aclose() + assert first_upstream.closed is True + + class _BreaksMidClipAudioStream(httpx.AsyncByteStream): async def __aiter__(self): yield b"a" * AUDIO_SPEECH_CHUNK_SIZE