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

tests/entrypoints/openai_api/test_serving_speech_stream.py

@@ -33,6 +33,12 @@ def _fake_aligner_res(pairs, words):
     )
 
 
+def _fill_audio_format(audio_format: dict | None, sample_rate: int = 24000, channels: int = 1) -> None:
+    """Mirror ``_generate_audio_chunks``: report the format before the first chunk."""
+    if audio_format is not None and not audio_format:
+        audio_format.update(sample_rate=sample_rate, channels=channels)
+
+
 def _build_test_app(
     speech_service=None,
     *,
@@ -57,8 +63,10 @@ def _build_test_app(
             tts_params=None,
             collect=None,
             cumulative_audio=False,
+            audio_format=None,
         ):
             for chunk in (b"\x01\x02", b"\x03\x04\x05"):
+                _fill_audio_format(audio_format)
                 yield (chunk, 24000) if include_sample_rate else chunk
 
         speech_service._generate_pcm_chunks = mock_generate_pcm_chunks
@@ -282,6 +290,7 @@ class TestStreamingSpeechWebSocket:
             request_arrival_ts=None,
             include_sample_rate=False,
             tts_params=None,
+            audio_format=None,
         ):
             assert request_start_s is not None
             assert request_arrival_ts is not None
@@ -289,6 +298,7 @@ class TestStreamingSpeechWebSocket:
             captured_timing["chunk_arrival"] = request_arrival_ts
             captured_tts_params.append(tts_params)
             for chunk in (b"\x01\x02", b"\x03\x04\x05", b"\x06"):
+                _fill_audio_format(audio_format)
                 yield (chunk, 24000) if include_sample_rate else chunk
 
         speech_service._generate_pcm_chunks = mock_generate_pcm_chunks
@@ -312,6 +322,7 @@ class TestStreamingSpeechWebSocket:
                 assert start["type"] == "audio.start"
                 assert start["format"] == "pcm"
                 assert start["sample_rate"] == 24000
+                assert start["channels"] == 1
 
                 assert ws.receive_bytes() == b"\x01\x02"
                 assert ws.receive_bytes() == b"\x03\x04\x05"
@@ -337,6 +348,91 @@ class TestStreamingSpeechWebSocket:
         assert captured_tts_params == [{"_qwen3_tts_effective_max_tokens": [192]}]
         assert speech_service._generate_audio_bytes.await_count == 0
 
+    @staticmethod
+    def _stream_service(mocker: MockerFixture, generate_pcm_chunks, *, forced_aligner: bool = False):
+        speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech)
+        speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav"))
+        speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req", object(), {}))
+        speech_service._generate_pcm_chunks = generate_pcm_chunks
+        speech_service.engine_client = mocker.MagicMock()
+        speech_service.engine_client.abort = mocker.AsyncMock()
+        speech_service.forced_aligner_enabled = forced_aligner
+        return speech_service
+
+    @staticmethod
+    def _send_pcm_session(ws, **config) -> None:
+        ws.send_json(
+            {"type": "session.config", "voice": "Vivian", "stream_audio": True, "response_format": "pcm", **config}
+        )
+        ws.send_json({"type": "input.text", "text": "Hello world. "})
+        ws.send_json({"type": "input.done"})
+
+    def test_streaming_pcm_start_reports_model_native_format(self, mocker: MockerFixture):
+        """audio.start states the real rate/channels, not a nominal 24 kHz."""
+
+        async def mock_generate_pcm_chunks(_generator, _request_id, *, audio_format=None, **_kwargs):
+            _fill_audio_format(audio_format, sample_rate=48000, channels=2)
+            yield b"\x01\x02\x03\x04"
+
+        app, _ = _build_test_app(self._stream_service(mocker, mock_generate_pcm_chunks))
+
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/audio/speech/stream") as ws:
+                self._send_pcm_session(ws)
+                start = ws.receive_json()
+                assert start["type"] == "audio.start"
+                assert start["sample_rate"] == 48000
+                assert start["channels"] == 2
+                assert ws.receive_bytes() == b"\x01\x02\x03\x04"
+                assert ws.receive_json()["type"] == "audio.done"
+
+    def test_streaming_pcm_error_before_audio_keeps_start_done_pairing(self, mocker: MockerFixture):
+        """No audio means no known format: audio.start omits it rather than guessing."""
+
+        async def mock_generate_pcm_chunks(_generator, _request_id, *, audio_format=None, **_kwargs):
+            raise RuntimeError("early boom")
+            yield  # pragma: no cover - generator marker, unreachable
+
+        app, _ = _build_test_app(self._stream_service(mocker, mock_generate_pcm_chunks))
+
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/audio/speech/stream") as ws:
+                self._send_pcm_session(ws)
+                start = ws.receive_json()
+                assert start["type"] == "audio.start"
+                assert "sample_rate" not in start
+                assert "channels" not in start
+                assert ws.receive_json() == {
+                    "type": "error",
+                    "message": "Generation failed for utterance 0, sentence 0: early boom",
+                }
+                done = ws.receive_json()
+                assert done["type"] == "audio.done"
+                assert done["error"] is True
+
+    def test_word_timestamps_chunk_ms_account_for_channels(self, mocker: MockerFixture):
+        """Stereo s16le is 4 bytes per frame; chunk offsets must not double."""
+
+        async def mock_generate_pcm_chunks(
+            _generator, _request_id, *, include_sample_rate=False, collect=None, audio_format=None, **_kwargs
+        ):
+            _fill_audio_format(audio_format, sample_rate=1000, channels=2)
+            yield (b"\x01" * 4000, 1000) if include_sample_rate else b"\x01" * 4000
+
+        speech_service = self._stream_service(mocker, mock_generate_pcm_chunks, forced_aligner=True)
+        app, _ = _build_test_app(speech_service)
+
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/audio/speech/stream") as ws:
+                self._send_pcm_session(ws, word_timestamps=True)
+                start = ws.receive_json()
+                assert (start["sample_rate"], start["channels"]) == (1000, 2)
+                chunk = ws.receive_json()
+                assert (chunk["chunk_start_ms"], chunk["chunk_end_ms"]) == (0, 1000)
+                assert (chunk["sample_rate"], chunk["channels"]) == (1000, 2)
+                final = ws.receive_json()
+                assert final["chunk_end_ms"] == 1000
+
     def test_word_timestamps_requires_configured_aligner(self, mocker: MockerFixture):
         app, _ = _build_test_app(mocker=mocker)
 
@@ -388,11 +484,13 @@ class TestStreamingSpeechWebSocket:
             tts_params=None,
             collect=None,
             cumulative_audio=False,
+            audio_format=None,
         ):
             assert request_start_s is not None
             assert request_arrival_ts is not None
             assert cumulative_audio is True
             for chunk in (first_chunk, second_chunk):
+                _fill_audio_format(audio_format, sample_rate=1000)
                 yield (chunk, 1000) if include_sample_rate else chunk
             if collect is not None:
                 collect["aligner_res"] = _fake_aligner_res([[0, 200], [200, 900]], ["Hello", "world"])
@@ -484,8 +582,10 @@ class TestStreamingSpeechWebSocket:
             tts_params=None,
             collect=None,
             cumulative_audio=False,
+            audio_format=None,
         ):
             chunk = b"\x01" * 1000
+            _fill_audio_format(audio_format, sample_rate=1000)
             yield (chunk, 1000) if include_sample_rate else chunk
             if collect is not None:
                 collect["aligner_res"] = _fake_aligner_res([[0, 1000], [1000, 1200]], ["Hello", "world"])
@@ -673,7 +773,9 @@ class TestStreamingSpeechWebSocket:
             request_arrival_ts=None,
             include_sample_rate=False,
             tts_params=None,
+            audio_format=None,
         ):
+            _fill_audio_format(audio_format)
             yield b"\x01\x02"
             raise RuntimeError("stream boom")
 
@@ -777,7 +879,9 @@ class TestStreamingSpeechWebSocket:
             request_arrival_ts=None,
             include_sample_rate=False,
             tts_params=None,
+            audio_format=None,
         ):
+            _fill_audio_format(audio_format)
             yield b"\x01\x02"
 
         speech_service._generate_pcm_chunks = mock_generate_pcm_chunks