[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.py

@@ -45,9 +45,15 @@ from vllm_omni.entrypoints.openai.serving_speech import (
     OmniOpenAIServingSpeech,
     _create_wav_header,
 )
-from vllm_omni.entrypoints.openai.tts_adapters.base import DEFAULT_TTS_LANGUAGES, PreparedRequest, SpeechServingContext
+from vllm_omni.entrypoints.openai.tts_adapters.base import (
+    DEFAULT_TTS_LANGUAGES,
+    PreparedRequest,
+    SpeechServingContext,
+    TTSGenerationError,
+)
 from vllm_omni.entrypoints.openai.tts_adapters.capabilities import load_supported_speakers
 from vllm_omni.entrypoints.openai.tts_adapters.ming_tts import MingTTSAdapter
+from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import Qwen3TTSCodecLimitError
 from vllm_omni.entrypoints.openai.tts_adapters.voxtral import VoxtralTTSAdapter
 from vllm_omni.model_executor.models.fish_speech.prompt_utils import (
     FISH_TEXT_ONLY_SYSTEM_PROMPT,
@@ -2928,6 +2934,40 @@ class TestStreamingResponse:
         assert "text/event-stream" not in response.headers["content-type"]
         assert len(response.content) > 0
 
+    def test_raw_stream_passes_tts_params_to_common_guard(self, streaming_app, monkeypatch):
+        """The raw route must not bypass the shared Qwen3-TTS EOS guard."""
+        speech_server = streaming_app.state.speech_server
+        finalized_tts_params = {"_qwen3_tts_effective_max_tokens": [192]}
+        captured: dict = {}
+
+        async def prepare(_request, request_id=None):
+            return request_id, object(), finalized_tts_params
+
+        async def generate_chunks(
+            _generator,
+            _request_id,
+            _response_format="pcm",
+            raw_request=None,
+            request_start_s=None,
+            include_sample_rate=False,
+            usage_acc=None,
+            tts_params=None,
+            collect=None,
+        ):
+            captured["tts_params"] = tts_params
+            yield b"\x00\x00"
+
+        monkeypatch.setattr(speech_server, "_prepare_speech_generation", prepare)
+        monkeypatch.setattr(speech_server, "_generate_audio_chunks", generate_chunks)
+
+        response = TestClient(streaming_app).post(
+            "/v1/audio/speech",
+            json={"input": "Hello", "stream_format": "audio", "response_format": "pcm"},
+        )
+
+        assert response.status_code == 200
+        assert captured["tts_params"] is finalized_tts_params
+
     def test_sse_streaming(self, streaming_app):
         """stream_format=sse without stream=True returns audio deltas as SSE."""
         client = TestClient(streaming_app)
@@ -4848,6 +4888,57 @@ class TestTTSAsyncOffloading:
         assert ("artifact-a", False) not in qwen3_tts_server._ref_audio_model_artifact_ready
         assert "artifact-b" in {entry[3] for entry in qwen3_tts_server._ref_audio_resolve_cache.values()}
 
+    @pytest.mark.asyncio
+    async def test_qwen3_tts_nonstream_retries_codec_limit_once(self, qwen3_tts_server, mocker: MockerFixture) -> None:
+        qwen3_tts_server._check_model = mocker.AsyncMock(return_value=None)
+        qwen3_tts_server._generate_audio_bytes = mocker.AsyncMock(
+            side_effect=[
+                Qwen3TTSCodecLimitError("codec limit"),
+                (b"RIFF" + b"\x00" * 32, "audio/wav"),
+            ]
+        )
+        request = OpenAICreateSpeechRequest(input="Hello", task_type="Base")
+
+        response = await qwen3_tts_server.create_speech(request)
+
+        assert response.status_code == 200
+        assert qwen3_tts_server._generate_audio_bytes.await_count == 2
+        first_call, retry_call = qwen3_tts_server._generate_audio_bytes.await_args_list
+        assert first_call.args[0] is request
+        assert first_call.kwargs["request_id"] != retry_call.kwargs["request_id"]
+        assert retry_call.args[0].seed is not None
+        assert request.seed is None
+
+    @pytest.mark.asyncio
+    async def test_nonstream_does_not_retry_nonretryable_generation_error(
+        self, qwen3_tts_server, mocker: MockerFixture
+    ) -> None:
+        qwen3_tts_server._check_model = mocker.AsyncMock(return_value=None)
+        qwen3_tts_server._generate_audio_bytes = mocker.AsyncMock(side_effect=TTSGenerationError("invalid generation"))
+
+        response = await qwen3_tts_server.create_speech(OpenAICreateSpeechRequest(input="Hello", task_type="Base"))
+
+        assert response.status_code == 500
+        qwen3_tts_server._generate_audio_bytes.assert_awaited_once()
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("request_kwargs", [{"seed": 100423}, {"max_new_tokens": 192}])
+    async def test_qwen3_tts_nonstream_does_not_retry_explicit_sampling_controls(
+        self,
+        qwen3_tts_server,
+        mocker: MockerFixture,
+        request_kwargs: dict[str, int],
+    ) -> None:
+        qwen3_tts_server._check_model = mocker.AsyncMock(return_value=None)
+        qwen3_tts_server._generate_audio_bytes = mocker.AsyncMock(side_effect=Qwen3TTSCodecLimitError("codec limit"))
+
+        response = await qwen3_tts_server.create_speech(
+            OpenAICreateSpeechRequest(input="Hello", task_type="Base", **request_kwargs)
+        )
+
+        assert response.status_code == 500
+        qwen3_tts_server._generate_audio_bytes.assert_awaited_once()
+
     @pytest.mark.asyncio
     async def test_generate_audio_chunks_discards_ref_audio_artifact_warmup_on_error(self, qwen3_tts_server):
         async def failing_generator():
@@ -4862,6 +4953,80 @@ class TestTTSAsyncOffloading:
         assert "req-fail" not in qwen3_tts_server._request_ref_audio_artifact_keys
         assert ("artifact-fail", False) not in qwen3_tts_server._ref_audio_model_artifact_ready
 
+    @pytest.mark.asyncio
+    async def test_generate_audio_chunks_rejects_qwen3_codec_limit_before_success(self, qwen3_tts_server):
+        async def length_limited_generator():
+            yield SimpleNamespace(
+                multimodal_output={"audio": torch.zeros(16, dtype=torch.float32), "sr": 24000},
+                metrics={
+                    "stage_metrics": {
+                        "0": {
+                            "num_tokens_out": 191,
+                            "finish_reason": "length",
+                        }
+                    }
+                },
+            )
+
+        chunks = qwen3_tts_server._generate_audio_chunks(
+            length_limited_generator(),
+            "req-codec-limit",
+            tts_params={"task_type": ["Base"], "_qwen3_tts_effective_max_tokens": [192]},
+        )
+        with pytest.raises(Qwen3TTSCodecLimitError, match="191/192"):
+            async for _ in chunks:
+                pass
+
+    @pytest.mark.asyncio
+    async def test_sse_codec_limit_marks_partial_audio_for_discard(self, qwen3_tts_server):
+        async def length_limited_generator():
+            yield SimpleNamespace(
+                multimodal_output={"audio": torch.zeros(16, dtype=torch.float32), "sr": 24000},
+                metrics={
+                    "stage_metrics": {
+                        "0": {
+                            "num_tokens_out": 191,
+                            "finish_reason": "length",
+                        }
+                    }
+                },
+            )
+
+        events = [
+            event
+            async for event in qwen3_tts_server._generate_audio_sse_events(
+                length_limited_generator(),
+                "req-sse-codec-limit",
+                tts_params={"task_type": ["Base"], "_qwen3_tts_effective_max_tokens": [192]},
+            )
+        ]
+
+        assert any("event: speech.audio.delta" in event for event in events)
+        assert not any("event: speech.audio.done" in event for event in events)
+        error_event = next(event for event in events if "event: speech.audio.error" in event)
+        payload = json.loads(next(line for line in error_event.splitlines() if line.startswith("data: "))[6:])
+        assert payload["error"]["partial_audio"] is True
+        assert payload["error"]["action"] == "discard"
+
+    @pytest.mark.asyncio
+    async def test_sse_error_before_audio_is_not_marked_partial(self, qwen3_tts_server):
+        async def failing_generator():
+            raise RuntimeError("boom")
+            yield  # pragma: no cover
+
+        events = [
+            event
+            async for event in qwen3_tts_server._generate_audio_sse_events(
+                failing_generator(),
+                "req-sse-no-audio",
+            )
+        ]
+
+        error_event = next(event for event in events if "event: speech.audio.error" in event)
+        payload = json.loads(next(line for line in error_event.splitlines() if line.startswith("data: "))[6:])
+        assert "partial_audio" not in payload["error"]
+        assert "action" not in payload["error"]
+
     @pytest.mark.asyncio
     async def test_generate_audio_chunks_discards_ref_audio_artifact_warmup_on_close(self, qwen3_tts_server):
         async def pcm_generator():