[Bugfix] Report sample rate and channels for streamed speech PCM (#8567)

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>

tests/entrypoints/openai_api/test_serving_speech.py

@@ -3468,88 +3468,105 @@ class TestStreamingProtocolValidation:
         assert req.is_raw_audio_stream() is False
 
 
+def _build_streaming_speech_app(
+    mocker: MockerFixture,
+    chunk: torch.Tensor | None = None,
+    sample_rate: int | None = None,
+) -> FastAPI:
+    """Speech app whose mock engine yields one intermediate then one final chunk.
+
+    ``sample_rate`` is reported as the model's ``sr`` when given; omitted, the
+    serving layer's 24 kHz default applies (as for most TTS models).
+    """
+    mock_engine_client = mocker.MagicMock()
+    mock_engine_client.errored = False
+    if chunk is None:
+        chunk = torch.sin(torch.linspace(0, 440 * 2 * torch.pi, 24000))
+
+    def _make_output(finished: bool) -> OmniRequestOutput:
+        mm_output: dict = {"audio": chunk}
+        if sample_rate is not None:
+            mm_output["sr"] = torch.tensor(sample_rate)
+
+        class MockCompletionOutput:
+            def __init__(self, index: int = 0):
+                self.index = index
+                self.text = ""
+                self.token_ids: list[int] = []
+                self.finish_reason = "stop"
+                self.stop_reason = None
+                self.logprobs = None
+
+        class MockRequestOutput:
+            def __init__(self, multimodal_output: dict):
+                self.request_id = "speech-stream-test"
+                self.outputs = [MockCompletionOutput(index=0)]
+                self.multimodal_output = multimodal_output
+                self.finished = finished
+                self.prompt_token_ids = None
+                self.encoder_prompt_token_ids = None
+                self.num_cached_tokens = None
+                self.prompt_logprobs = None
+                self.kv_transfer_params = None
+
+        return OmniRequestOutput.from_stage_output(
+            MockRequestOutput(mm_output),
+            stage_id=0,
+            final_output_type="audio",
+            finished=finished,
+        )
+
+    async def mock_generate_streaming(*args, **kwargs):
+        yield _make_output(finished=False)
+        yield _make_output(finished=True)
+
+    mock_engine_client.generate = mocker.MagicMock(side_effect=mock_generate_streaming)
+    mock_engine_client.default_sampling_params_list = [{}]
+    mock_models = mocker.MagicMock()
+    mock_models.is_base_model.return_value = True
+
+    speech_server = OmniOpenAIServingSpeech(
+        engine_client=mock_engine_client,
+        models=mock_models,
+        request_logger=mocker.MagicMock(),
+    )
+    speech_server._tts_model_type = None
+    speech_server._adapter = None
+    speech_server._diffusion_mode = False
+    mocker.patch.object(speech_server, "_uses_native_speed_control", return_value=False)
+    mocker.patch.object(
+        speech_server,
+        "create_error_response",
+        side_effect=lambda message, **_: JSONResponse(
+            status_code=400,
+            content={"error": {"message": message}},
+        ),
+    )
+
+    original_create_speech = speech_server.create_speech
+    sig = signature(original_create_speech)
+    new_parameters = [p for name, p in sig.parameters.items() if name != "raw_request"]
+    new_sig = Signature(parameters=new_parameters, return_annotation=sig.return_annotation)
+
+    async def awaitable_create_speech(*args, **kwargs):
+        return await original_create_speech(*args, **kwargs)
+
+    awaitable_create_speech.__signature__ = new_sig  # type: ignore[attr-defined]
+    speech_server.create_speech = awaitable_create_speech
+
+    app = FastAPI()
+    app.add_api_route("/v1/audio/speech", speech_server.create_speech, methods=["POST"], response_model=None)
+    app.state.speech_server = speech_server
+    return app
+
+
 class TestStreamingResponse:
     """Integration tests for the streaming audio response path."""
 
     @pytest.fixture
     def streaming_app(self, mocker: MockerFixture):
         """Test app whose mock engine yields one intermediate chunk then a final chunk."""
-        mock_engine_client = mocker.MagicMock()
-        mock_engine_client.errored = False
-
-        def _make_output(finished: bool) -> OmniRequestOutput:
-            chunk = torch.sin(torch.linspace(0, 440 * 2 * torch.pi, 24000))
-
-            class MockCompletionOutput:
-                def __init__(self, index: int = 0):
-                    self.index = index
-                    self.text = ""
-                    self.token_ids: list[int] = []
-                    self.finish_reason = "stop"
-                    self.stop_reason = None
-                    self.logprobs = None
-
-            class MockRequestOutput:
-                def __init__(self, audio_tensor: torch.Tensor):
-                    self.request_id = "speech-stream-test"
-                    self.outputs = [MockCompletionOutput(index=0)]
-                    self.multimodal_output = {"audio": audio_tensor}
-                    self.finished = finished
-                    self.prompt_token_ids = None
-                    self.encoder_prompt_token_ids = None
-                    self.num_cached_tokens = None
-                    self.prompt_logprobs = None
-                    self.kv_transfer_params = None
-
-            return OmniRequestOutput.from_stage_output(
-                MockRequestOutput(audio_tensor=chunk),
-                stage_id=0,
-                final_output_type="audio",
-                finished=finished,
-            )
-
-        async def mock_generate_streaming(*args, **kwargs):
-            yield _make_output(finished=False)
-            yield _make_output(finished=True)
-
-        mock_engine_client.generate = mocker.MagicMock(side_effect=mock_generate_streaming)
-        mock_engine_client.default_sampling_params_list = [{}]
-        mock_models = mocker.MagicMock()
-        mock_models.is_base_model.return_value = True
-
-        speech_server = OmniOpenAIServingSpeech(
-            engine_client=mock_engine_client,
-            models=mock_models,
-            request_logger=mocker.MagicMock(),
-        )
-        speech_server._tts_model_type = None
-        speech_server._adapter = None
-        speech_server._diffusion_mode = False
-        mocker.patch.object(speech_server, "_uses_native_speed_control", return_value=False)
-        mocker.patch.object(
-            speech_server,
-            "create_error_response",
-            side_effect=lambda message, **_: JSONResponse(
-                status_code=400,
-                content={"error": {"message": message}},
-            ),
-        )
-
-        original_create_speech = speech_server.create_speech
-        sig = signature(original_create_speech)
-        new_parameters = [p for name, p in sig.parameters.items() if name != "raw_request"]
-        new_sig = Signature(parameters=new_parameters, return_annotation=sig.return_annotation)
-
-        async def awaitable_create_speech(*args, **kwargs):
-            return await original_create_speech(*args, **kwargs)
-
-        awaitable_create_speech.__signature__ = new_sig  # type: ignore[attr-defined]
-        speech_server.create_speech = awaitable_create_speech
-
-        app = FastAPI()
-        app.add_api_route("/v1/audio/speech", speech_server.create_speech, methods=["POST"], response_model=None)
-        app.state.speech_server = speech_server
-        return app
+        return _build_streaming_speech_app(mocker)
 
     @staticmethod
     def _assert_sse_audio_response(response, response_format: str = "pcm"):
@@ -3603,6 +3620,7 @@ class TestStreamingResponse:
             tts_params=None,
             collect=None,
             target_sample_rate=None,
+            audio_format=None,
         ):
             captured["tts_params"] = tts_params
             captured["target_sample_rate"] = target_sample_rate
@@ -3790,6 +3808,146 @@ class TestStreamingResponse:
         assert "audio/wav" in response.headers["content-type"]
 
 
+class TestSpeechAudioFormatMetadata:
+    """pcm carries no header, so responses must state the model-native format."""
+
+    # 24 kHz mono (most TTS models) and 48 kHz [channels, samples] stereo
+    # (e.g. MOSS-TTS-Local v1.5); the serving layer must not care which.
+    _MONO = torch.sin(torch.linspace(0, 440 * 2 * torch.pi, 2400))
+    _STEREO = torch.stack([_MONO.repeat(2) * 0.5, -_MONO.repeat(2) * 0.25])
+
+    _CASES = [
+        pytest.param(None, None, 24000, 1, id="default-24k-mono"),
+        pytest.param(_MONO, 24000, 24000, 1, id="explicit-24k-mono"),
+        pytest.param(_STEREO, 48000, 48000, 2, id="48k-stereo"),
+    ]
+
+    @staticmethod
+    def _expected_pcm(chunk: torch.Tensor | None) -> bytes:
+        """Two engine outputs, each the chunk as interleaved s16le."""
+        # ``None`` is the builder's default chunk.
+        audio = torch.sin(torch.linspace(0, 440 * 2 * torch.pi, 24000)) if chunk is None else chunk
+        samples = audio.numpy().T if audio.ndim == 2 else audio.numpy()
+        # libsndfile-compatible PCM_16 conversion used by ``create_audio``.
+        one = np.clip(np.floor(samples.astype(np.float64) * 32768.0), -32768, 32767).astype("<i2").tobytes()
+        return one * 2
+
+    @pytest.mark.parametrize(("chunk", "model_sr", "rate", "channels"), _CASES)
+    def test_raw_pcm_stream_headers(self, mocker: MockerFixture, chunk, model_sr, rate, channels):
+        app = _build_streaming_speech_app(mocker, chunk=chunk, sample_rate=model_sr)
+        response = TestClient(app).post(
+            "/v1/audio/speech",
+            json={"input": "Hello", "stream_format": "audio", "response_format": "pcm"},
+        )
+
+        assert response.status_code == 200
+        assert response.headers["content-type"].startswith("audio/pcm")
+        assert response.headers["x-audio-sample-rate"] == str(rate)
+        assert response.headers["x-audio-channels"] == str(channels)
+        exposed = {h.strip().lower() for h in response.headers["access-control-expose-headers"].split(",")}
+        assert {"x-audio-sample-rate", "x-audio-channels"} <= exposed
+        # The body stays the model-native interleaved s16le stream, byte for byte.
+        assert response.content == self._expected_pcm(chunk)
+
+    # WAV streaming requires the model to report ``sr`` on its first chunk.
+    @pytest.mark.parametrize(("chunk", "model_sr", "rate", "channels"), _CASES[1:])
+    def test_raw_wav_stream_header_matches_metadata(self, mocker: MockerFixture, chunk, model_sr, rate, channels):
+        app = _build_streaming_speech_app(mocker, chunk=chunk, sample_rate=model_sr)
+        response = TestClient(app).post(
+            "/v1/audio/speech",
+            json={"input": "Hello", "stream_format": "audio", "response_format": "wav"},
+        )
+
+        assert response.status_code == 200
+        assert response.headers["x-audio-sample-rate"] == str(rate)
+        assert response.headers["x-audio-channels"] == str(channels)
+        num_channels, sample_rate = struct.unpack("<HI", response.content[22:28])
+        assert (sample_rate, num_channels) == (rate, channels)
+
+    @pytest.mark.parametrize(("chunk", "model_sr", "rate", "channels"), _CASES)
+    def test_sse_deltas_carry_format(self, mocker: MockerFixture, chunk, model_sr, rate, channels):
+        app = _build_streaming_speech_app(mocker, chunk=chunk, sample_rate=model_sr)
+        response = TestClient(app).post(
+            "/v1/audio/speech",
+            json={"input": "Hello", "stream": True, "response_format": "pcm"},
+        )
+
+        deltas = [
+            json.loads(line.removeprefix("data: "))
+            for line in response.text.splitlines()
+            if line.startswith("data: ") and '"speech.audio.delta"' in line
+        ]
+        assert deltas
+        assert all((d["sample_rate"], d["channels"]) == (rate, channels) for d in deltas)
+
+    @pytest.mark.parametrize(("chunk", "model_sr", "rate", "channels"), _CASES)
+    def test_non_streaming_pcm_headers(self, mocker: MockerFixture, chunk, model_sr, rate, channels):
+        app = _build_streaming_speech_app(mocker, chunk=chunk, sample_rate=model_sr)
+        response = TestClient(app).post("/v1/audio/speech", json={"input": "Hello", "response_format": "pcm"})
+
+        assert response.status_code == 200
+        assert response.headers["x-audio-sample-rate"] == str(rate)
+        assert response.headers["x-audio-channels"] == str(channels)
+        assert len(response.content) % (2 * channels) == 0
+
+    def test_requested_sample_rate_is_reported(self, mocker: MockerFixture):
+        """The headers describe the bytes sent, i.e. after request.sample_rate resampling."""
+        app = _build_streaming_speech_app(mocker)
+        speech_server = app.state.speech_server
+        mocker.patch.object(speech_server, "_validate_speech_sample_rate", return_value=None)
+        response = TestClient(app).post(
+            "/v1/audio/speech",
+            json={"input": "Hello", "stream_format": "audio", "response_format": "pcm", "sample_rate": 8000},
+        )
+
+        assert response.status_code == 200
+        assert response.headers["x-audio-sample-rate"] == "8000"
+        assert response.headers["x-audio-channels"] == "1"
+        # Two 1 s chunks at 8 kHz mono s16le, within resampler edge tolerance.
+        assert abs(len(response.content) - 2 * 8000 * 2) <= 2 * 64
+
+    def test_raw_stream_error_before_audio_is_an_http_error(self, mocker: MockerFixture):
+        """Priming the first chunk lets a pre-audio failure return a real status code."""
+        app = _build_streaming_speech_app(mocker)
+        speech_server = app.state.speech_server
+
+        async def failing_generate(*args, **kwargs):
+            raise RuntimeError("boom before audio")
+            yield  # pragma: no cover - generator marker, unreachable
+
+        speech_server.engine_client.generate = mocker.MagicMock(side_effect=failing_generate)
+        response = TestClient(app).post(
+            "/v1/audio/speech",
+            json={"input": "Hello", "stream_format": "audio", "response_format": "pcm"},
+        )
+
+        assert response.status_code >= 400
+        assert "audio/pcm" not in response.headers.get("content-type", "")
+        assert "x-audio-sample-rate" not in response.headers
+
+    def test_primed_stream_closes_inner_generator(self):
+        """Closing the response body before iteration still closes the engine stream."""
+        from vllm_omni.entrypoints.openai.serving_speech import _PrimedAudioStream
+
+        closed: list[bool] = []
+
+        async def chunks():
+            try:
+                yield b"\x01\x00"
+                yield b"\x02\x00"
+            finally:
+                closed.append(True)
+
+        async def run() -> list[bytes]:
+            body = await _PrimedAudioStream.start(chunks())
+            await body.aclose()
+            assert closed == [True]
+            body = await _PrimedAudioStream.start(chunks())
+            return [c async for c in body]
+
+        assert asyncio.run(run()) == [b"\x01\x00", b"\x02\x00"]
+
+
 class TestSpeechBatchAPI:
     """Tests for the /v1/audio/speech/batch endpoint."""