[Bugfix][Qwen3-TTS] Handle codec generations that exhaust their token budget (#6728)

Signed-off-by: Sy03 <1370724210@qq.com>

tests/entrypoints/openai_api/test_serving_speech_stream.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import asyncio
 import base64
 from types import SimpleNamespace
@@ -44,7 +47,9 @@ def _build_test_app(
         speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-1", object(), {}))
         speech_service.forced_aligner_enabled = False
 
-        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None):
+        async def mock_generate_pcm_chunks(
+            _generator, _request_id, *, include_sample_rate=False, tts_params=None, collect=None
+        ):
             for chunk in (b"\x01\x02", b"\x03\x04\x05"):
                 yield (chunk, 24000) if include_sample_rate else chunk
 
@@ -244,6 +249,7 @@ class TestStreamingSpeechWebSocket:
 
     def test_streaming_multiple_binary_frames(self, mocker: MockerFixture):
         captured_requests = []
+        captured_tts_params = []
 
         speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech)
         speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav"))
@@ -253,11 +259,12 @@ class TestStreamingSpeechWebSocket:
 
         async def mock_prepare_speech_generation(request):
             captured_requests.append(request)
-            return "req-stream", object(), {}
+            return "req-stream", object(), {"_qwen3_tts_effective_max_tokens": [192]}
 
         speech_service._prepare_speech_generation = mock_prepare_speech_generation
 
-        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
+        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, tts_params=None):
+            captured_tts_params.append(tts_params)
             for chunk in (b"\x01\x02", b"\x03\x04\x05", b"\x06"):
                 yield (chunk, 24000) if include_sample_rate else chunk
 
@@ -302,6 +309,7 @@ class TestStreamingSpeechWebSocket:
         assert captured_requests[0].stream is True
         assert captured_requests[0].response_format == "pcm"
         assert captured_requests[0].initial_codec_chunk_frames == 12
+        assert captured_tts_params == [{"_qwen3_tts_effective_max_tokens": [192]}]
         assert speech_service._generate_audio_bytes.await_count == 0
 
     def test_word_timestamps_requires_configured_aligner(self, mocker: MockerFixture):
@@ -344,7 +352,9 @@ class TestStreamingSpeechWebSocket:
 
         # The forced-aligner stage rides the same generator: its pooling output
         # is surfaced via the ``collect`` channel once the audio has streamed.
-        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None):
+        async def mock_generate_pcm_chunks(
+            _generator, _request_id, *, include_sample_rate=False, tts_params=None, collect=None
+        ):
             for chunk in (first_chunk, second_chunk):
                 yield (chunk, 1000) if include_sample_rate else chunk
             if collect is not None:
@@ -427,7 +437,9 @@ class TestStreamingSpeechWebSocket:
         speech_service.forced_aligner_enabled = True
         speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req", object(), {}))
 
-        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None):
+        async def mock_generate_pcm_chunks(
+            _generator, _request_id, *, include_sample_rate=False, tts_params=None, collect=None
+        ):
             chunk = b"\x01" * 1000
             yield (chunk, 1000) if include_sample_rate else chunk
             if collect is not None:
@@ -608,7 +620,7 @@ class TestStreamingSpeechWebSocket:
         speech_service.engine_client.abort = mocker.AsyncMock()
         speech_service.forced_aligner_enabled = False
 
-        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
+        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, tts_params=None):
             yield b"\x01\x02"
             raise RuntimeError("stream boom")
 
@@ -633,6 +645,8 @@ class TestStreamingSpeechWebSocket:
                 assert ws.receive_json() == {
                     "type": "error",
                     "message": "Generation failed for utterance 0, sentence 0: stream boom",
+                    "partial_audio": True,
+                    "action": "discard",
                 }
                 assert ws.receive_json() == {
                     "type": "audio.done",
@@ -702,7 +716,7 @@ class TestStreamingSpeechWebSocket:
         speech_service.engine_client.abort = mocker.AsyncMock()
         speech_service.forced_aligner_enabled = False
 
-        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
+        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, tts_params=None):
             yield b"\x01\x02"
 
         speech_service._generate_pcm_chunks = mock_generate_pcm_chunks