fix(speech): await the TTS bridge under aspeech

Under aspeech the Gemini and Vertex TTS bridge awaits acompletion instead of running
sync completion in an executor thread, so a slow upstream no longer holds a worker
thread. This completes 176808e268 now that c5fbd1c3c3 adds the internal sub-call scope.
The inner chat call shares the speech call's logging object, so the bridge keeps the
speech call's acompletion flag, and sync success callbacks still run once with the
speech response

litellm/endpoints/speech/speech_to_completion_bridge/handler.py

@@ -3,7 +3,7 @@ Handler for transforming /chat/completions api requests to litellm.responses req
 """
 
 import json
-from collections.abc import Mapping
+from collections.abc import Coroutine, Mapping
 from typing import TYPE_CHECKING, Final
 
 from typing_extensions import TypedDict
@@ -32,7 +32,7 @@ class SpeechToCompletionBridgeHandlerInputKwargs(TypedDict):
     model: str
     input: str
     voice: str | dict | None
-    optional_params: dict
+    optional_params: dict[str, object]
     litellm_params: dict
     logging_obj: "LiteLLMLoggingObj"
     headers: dict
@@ -106,10 +106,10 @@ class SpeechToCompletionBridgeHandler:
         extra_headers: Mapping[str, object] | None,
         timeout: "float | httpx.Timeout | None",
         cache_controls: "CacheControlsHint",
-    ) -> "HttpxBinaryResponseContent":
+        aspeech: bool,
+    ) -> "HttpxBinaryResponseContent | Coroutine[object, object, HttpxBinaryResponseContent]":
         received_args: Final = locals()
-        from litellm import completion
-        from litellm.types.utils import ModelResponse
+        from litellm import acompletion, completion
 
         validated_kwargs: Final = self.validate_input_kwargs(received_args)
         model = validated_kwargs["model"]
@@ -138,19 +138,36 @@ class SpeechToCompletionBridgeHandler:
             voice=voice,
         )
 
+        requested_response_format: Final = validated_kwargs["optional_params"].get("response_format")
+        response_format: Final = requested_response_format if isinstance(requested_response_format, str) else None
+        if aspeech:
+            return self._aspeech(acompletion(**request_data), response_format, logging_obj)
+
         with internal_sub_call():
             result: Final = completion(
                 **request_data,
             )
+        return self._to_speech_response(result, response_format)
 
-        requested_response_format: Final = optional_params.get("response_format")
-        if isinstance(result, ModelResponse):
-            return self.transformation_handler.transform_response(
-                model_response=result,
-                response_format=requested_response_format if isinstance(requested_response_format, str) else None,
-            )
-        else:
+    async def _aspeech(
+        self,
+        pending_completion: Coroutine[object, object, object],
+        response_format: str | None,
+        logging_obj: "LiteLLMLoggingObj",
+    ) -> "HttpxBinaryResponseContent":
+        with logging_obj.preserve_entrypoint_acompletion(), internal_sub_call():
+            result: Final = await pending_completion
+        return self._to_speech_response(result, response_format)
+
+    def _to_speech_response(self, result: object, response_format: str | None) -> "HttpxBinaryResponseContent":
+        from litellm.types.utils import ModelResponse
+
+        if not isinstance(result, ModelResponse):
             raise Exception(f"Unmapped response type. Got type: {type(result)}")
+        return self.transformation_handler.transform_response(
+            model_response=result,
+            response_format=response_format,
+        )
 
 
 speech_to_completion_bridge_handler = SpeechToCompletionBridgeHandler()

litellm/litellm_core_utils/litellm_logging.py

@@ -12,7 +12,8 @@ import subprocess
 import sys
 import time
 import traceback
-from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence
+from collections.abc import Awaitable, Callable, Generator, Iterable, Iterator, Mapping, Sequence
+from contextlib import contextmanager
 from datetime import datetime as dt_object
 from functools import lru_cache
 from types import MappingProxyType, TracebackType
@@ -73,6 +74,7 @@ from litellm.litellm_core_utils.classifier_logging import (
     is_classifier_call,
 )
 from litellm.litellm_core_utils.core_helpers import (
+    as_str_mapping,
     get_provider_response_headers_from_hidden_params,
     is_expected_client_error,
     proxy_stamped_used_client_oauth_token,
@@ -1047,6 +1049,16 @@ class Logging(LiteLLMLoggingBaseClass):
         if "custom_llm_provider" in self.model_call_details:
             self.custom_llm_provider = self.model_call_details["custom_llm_provider"]
 
+    @contextmanager
+    def preserve_entrypoint_acompletion(self) -> Generator[None]:
+        acompletion: Final = (as_str_mapping(self._model_call_details_view.get("litellm_params")) or {}).get(
+            CallTypes.acompletion.value
+        )
+        try:
+            yield
+        finally:
+            self.litellm_params[CallTypes.acompletion.value] = acompletion
+
     def update_from_kwargs(
         self,
         kwargs: dict,

litellm/main.py

@@ -8753,6 +8753,7 @@ def speech(
                 extra_headers=_DEPLOYMENT_HEADERS_HINT.validate_python(kwargs).get("extra_headers"),
                 timeout=timeout,
                 cache_controls=_CACHE_CONTROLS_HINT.validate_python(kwargs),
+                aspeech=aspeech or False,
             )
 
         # Vertex AI Text-to-Speech (Google Cloud TTS)
@@ -8808,6 +8809,7 @@ def speech(
             extra_headers=_DEPLOYMENT_HEADERS_HINT.validate_python(kwargs).get("extra_headers"),
             timeout=timeout,
             cache_controls=_CACHE_CONTROLS_HINT.validate_python(kwargs),
+            aspeech=aspeech or False,
         )
     elif custom_llm_provider == "runwayml":
         from litellm.llms.runwayml.text_to_speech.transformation import (

tests/unit/test_main.py

@@ -4547,6 +4547,39 @@ async def test_aspeech_gemini_bridge_logs_only_the_speech_call_under_its_call_id
     _assert_logged_once_as_speech(events)
 
 
+
+@pytest.mark.asyncio
+async def test_aspeech_gemini_bridge_runs_sync_function_callbacks_once_with_the_speech_response(
+    respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    monkeypatch.delenv("GEMINI_API_KEY", raising=False)
+    monkeypatch.delenv("GOOGLE_API_KEY", raising=False)
+    logged: Final[list[tuple[type, object]]] = []  # mutable-ok: records what the sync callback receives
+
+    def sync_callback(kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object) -> None:
+        logged.append((type(response_obj), kwargs.get("litellm_call_id")))
+
+    monkeypatch.setattr(litellm, "success_callback", [sync_callback])
+    respx_mock.post(
+        url__regex=r"https://generativelanguage\.googleapis\.com/v1beta/models/gemini-2\.5-flash-preview-tts:generateContent.*"
+    ).mock(return_value=httpx.Response(200, json=_gemini_tts_generate_content_response()))
+
+    await litellm.aspeech(
+        model="gemini/gemini-2.5-flash-preview-tts",
+        input="sync callback check",
+        voice="Kore",
+        api_key="fake-gemini-key",
+        litellm_call_id="speech-call-id",
+    )
+    for _ in range(100):
+        if logged:
+            break
+        await asyncio.sleep(0.05)
+    await asyncio.sleep(0.3)
+
+    assert logged == [(HttpxBinaryResponseContent, "speech-call-id")]
+
 @pytest.mark.asyncio
 async def test_aspeech_gemini_bridge_logs_a_cached_reply_only_as_the_speech_call(
     respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
@@ -5063,6 +5096,36 @@ def test_speech_gemini_tts_keeps_cached_audio_only_for_the_request_ttl(
     assert route.call_count == 2
 
 
+async def test_aspeech_gemini_tts_waits_for_a_stalled_upstream_without_holding_a_worker_thread(
+    respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
+):
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    upstream_reached: Final = asyncio.Event()
+    release_upstream: Final = asyncio.Event()
+
+    async def stalled_upstream(request: httpx.Request) -> httpx.Response:
+        upstream_reached.set()
+        await release_upstream.wait()
+        return httpx.Response(200, json=_gemini_tts_generate_content_response())
+
+    respx_mock.post(url__regex=GEMINI_DEFAULT_GENERATE_CONTENT_URL).mock(side_effect=stalled_upstream)
+    loop: Final = asyncio.get_running_loop()
+    with ThreadPoolExecutor(max_workers=1) as single_worker:
+        loop.set_default_executor(single_worker)
+        speech_task: Final = asyncio.create_task(
+            litellm.aspeech(model=f"gemini/{GEMINI_TTS_MODEL}", input="hello", voice="Kore", api_key="fake-gemini-key")
+        )
+        await asyncio.wait_for(upstream_reached.wait(), timeout=10)
+
+        worker_answer: Final = await asyncio.wait_for(loop.run_in_executor(None, lambda: "free"), timeout=10)
+
+        release_upstream.set()
+        response: Final = await asyncio.wait_for(speech_task, timeout=10)
+
+    assert worker_answer == "free"
+    assert b"pcm-audio-bytes" in response.content
+
+
 def test_speech_gemini_tts_sends_number_and_boolean_deployment_headers_as_their_config_text(
     respx_mock: respx.MockRouter,
 ):