[Feature] Flush the speech WebSocket on input.done without closing it (#5517)

Signed-off-by: iancarrasco-b10 <ian.carrasco@baseten.co>
2026-07-28 23:46:19 +00:00
1 parent 0feb92bdfa
commit 23bb4395bb
@@ -33,26 +33,63 @@ tts_ws_server_params = [
 ]


-async def _run_ws_session(host: str, port: int, model: str) -> dict:
-    uri = f"ws://{host}:{port}/v1/audio/speech/stream"
+def _session_config(model: str) -> str:
+    return json.dumps(
+        {
+            "type": "session.config",
+            "model": model,
+            "voice": "vivian",
+            "language": "English",
+            "response_format": "pcm",
+            "stream_audio": True,
+        }
+    )
+
+
+async def _collect_utterance(ws) -> dict:
+    """Read frames until session.done closes out one flushed utterance."""
    starts: list[dict] = []
    dones: list[dict] = []
    chunk_lengths: dict[int, list[int]] = {}
    session_done: dict | None = None

+    while True:
+        message = await asyncio.wait_for(ws.recv(), timeout=180)
+        if isinstance(message, bytes):
+            if not starts:
+                raise AssertionError("Received audio bytes before audio.start")
+            sentence_index = starts[-1]["sentence_index"]
+            chunk_lengths.setdefault(sentence_index, []).append(len(message))
+            continue
+
+        payload = json.loads(message)
+        msg_type = payload.get("type")
+        if msg_type == "audio.start":
+            starts.append(payload)
+            chunk_lengths.setdefault(payload["sentence_index"], [])
+        elif msg_type == "audio.done":
+            dones.append(payload)
+        elif msg_type == "session.done":
+            session_done = payload
+            break
+        elif msg_type == "error":
+            raise AssertionError(f"WebSocket error: {payload['message']}")
+        else:
+            raise AssertionError(f"Unexpected WebSocket message: {payload}")
+
+    return {
+        "starts": starts,
+        "dones": dones,
+        "chunk_lengths": chunk_lengths,
+        "session_done": session_done,
+    }
+
+
+async def _run_ws_session(host: str, port: int, model: str) -> dict:
+    uri = f"ws://{host}:{port}/v1/audio/speech/stream"
+
    async with websockets.connect(uri, max_size=None) as ws:
-        await ws.send(
-            json.dumps(
-                {
-                    "type": "session.config",
-                    "model": model,
-                    "voice": "vivian",
-                    "language": "English",
-                    "response_format": "pcm",
-                    "stream_audio": True,
-                }
-            )
-        )
+        await ws.send(_session_config(model))
        await ws.send(
            json.dumps(
                {
@@ -67,36 +104,23 @@ async def _run_ws_session(host: str, port: int, model: str) -> dict:
        )
        await ws.send(json.dumps({"type": "input.done"}))

-        while True:
-            message = await asyncio.wait_for(ws.recv(), timeout=180)
-            if isinstance(message, bytes):
-                if not starts:
-                    raise AssertionError("Received audio bytes before audio.start")
-                sentence_index = starts[-1]["sentence_index"]
-                chunk_lengths.setdefault(sentence_index, []).append(len(message))
-                continue
+        return await _collect_utterance(ws)

-            payload = json.loads(message)
-            msg_type = payload.get("type")
-            if msg_type == "audio.start":
-                starts.append(payload)
-                chunk_lengths.setdefault(payload["sentence_index"], [])
-            elif msg_type == "audio.done":
-                dones.append(payload)
-            elif msg_type == "session.done":
-                session_done = payload
-                break
-            elif msg_type == "error":
-                raise AssertionError(f"WebSocket error: {payload['message']}")
-            else:
-                raise AssertionError(f"Unexpected WebSocket message: {payload}")

-    return {
-        "starts": starts,
-        "dones": dones,
-        "chunk_lengths": chunk_lengths,
-        "session_done": session_done,
-    }
+async def _run_ws_reused_connection(host: str, port: int, model: str, texts: list[str]) -> list[dict]:
+    """Synthesize several utterances over a single connection."""
+    uri = f"ws://{host}:{port}/v1/audio/speech/stream"
+    results: list[dict] = []
+
+    async with websockets.connect(uri, max_size=None) as ws:
+        await ws.send(_session_config(model))
+        for text in texts:
+            await ws.send(json.dumps({"type": "input.text", "text": text}))
+            await ws.send(json.dumps({"type": "input.done"}))
+            results.append(await _collect_utterance(ws))
+        await ws.send(json.dumps({"type": "session.close"}))
+
+    return results


 class TestQwen3TTSWebSocket:
@@ -113,6 +137,7 @@ class TestQwen3TTSWebSocket:
        session_done = result["session_done"]

        assert session_done is not None
+        assert session_done["utterance_index"] == 0
        assert session_done["total_sentences"] == 1
        assert len(starts) == 1
        assert len(dones) == 1
@@ -131,3 +156,32 @@ class TestQwen3TTSWebSocket:
            assert total_bytes > 0
            assert chunk_lengths[sentence_index], f"Expected binary PCM frames for sentence {sentence_index}"
            assert sum(chunk_lengths[sentence_index]) == total_bytes
+
+    @pytest.mark.advanced_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
+    @pytest.mark.parametrize("omni_server", tts_ws_server_params, indirect=True)
+    def test_input_done_flushes_without_closing_connection(self, omni_server) -> None:
+        # input.done must flush the buffer and leave the connection usable:
+        # the second utterance is synthesized without a second handshake.
+        texts = [
+            "This is the first utterance sent over the websocket connection.",
+            "This is the second utterance, reusing the very same connection.",
+        ]
+        results = asyncio.run(_run_ws_reused_connection(omni_server.host, omni_server.port, omni_server.model, texts))
+
+        assert len(results) == len(texts)
+        for expected_index, result in enumerate(results):
+            assert result["session_done"] == {
+                "type": "session.done",
+                "utterance_index": expected_index,
+                "total_sentences": 1,
+            }
+            assert len(result["starts"]) == 1
+            assert len(result["dones"]) == 1
+            # utterance_index counts the flushes of the connection, while
+            # sentence_index stays within the flush it belongs to.
+            assert result["starts"][0]["utterance_index"] == expected_index
+            assert result["starts"][0]["sentence_index"] == 0
+            assert result["dones"][0]["error"] is False
+            assert result["dones"][0]["total_bytes"] > 0
@@ -72,13 +72,162 @@ class TestStreamingSpeechWebSocket:
                assert audio.startswith(b"RIFF")

                done = ws.receive_json()
-                assert done == {"type": "audio.done", "sentence_index": 0, "total_bytes": len(audio), "error": False}
+                assert done == {
+                    "type": "audio.done",
+                    "utterance_index": 0,
+                    "sentence_index": 0,
+                    "total_bytes": len(audio),
+                    "error": False,
+                }

                session_done = ws.receive_json()
-                assert session_done == {"type": "session.done", "total_sentences": 1}
+                assert session_done == {"type": "session.done", "utterance_index": 0, "total_sentences": 1}

        assert speech_service._generate_audio_bytes.await_count == 1

+    def test_input_done_flushes_and_keeps_connection_open(self, mocker: MockerFixture):
+        app, speech_service = _build_test_app(mocker=mocker)
+
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/audio/speech/stream") as ws:
+                ws.send_json({"type": "session.config", "voice": "Vivian"})
+
+                # The same connection serves several utterances; the config sent
+                # once at the top keeps applying and utterance_index rises.
+                for expected_index, text in enumerate(("First utterance. ", "Second utterance. ")):
+                    ws.send_json({"type": "input.text", "text": text})
+                    ws.send_json({"type": "input.done"})
+
+                    start = ws.receive_json()
+                    assert start["type"] == "audio.start"
+                    assert start["utterance_index"] == expected_index
+                    assert start["sentence_text"] == text.strip()
+                    assert ws.receive_bytes().startswith(b"RIFF")
+                    assert ws.receive_json()["type"] == "audio.done"
+                    assert ws.receive_json() == {
+                        "type": "session.done",
+                        "utterance_index": expected_index,
+                        "total_sentences": 1,
+                    }
+
+        assert speech_service._generate_audio_bytes.await_count == 2
+        assert [call.args[0].voice for call in speech_service._generate_audio_bytes.await_args_list] == [
+            "Vivian",
+            "Vivian",
+        ]
+
+    def test_sentence_index_stays_within_the_flushed_utterance(self, mocker: MockerFixture):
+        # sentence_index counts within one flush and utterance_index counts the
+        # flushes, so a late utterance never reports "sentence 2 of 1".
+        app, _ = _build_test_app(mocker=mocker)
+
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/audio/speech/stream") as ws:
+                ws.send_json({"type": "session.config", "voice": "Vivian"})
+
+                for expected_index in range(3):
+                    ws.send_json({"type": "input.text", "text": "Hello world. "})
+                    ws.send_json({"type": "input.done"})
+
+                    start = ws.receive_json()
+                    assert start["utterance_index"] == expected_index
+                    assert start["sentence_index"] == 0
+                    ws.receive_bytes()
+
+                    done = ws.receive_json()
+                    assert done["utterance_index"] == expected_index
+                    assert done["sentence_index"] == 0
+
+                    session_done = ws.receive_json()
+                    assert session_done["utterance_index"] == expected_index
+                    assert session_done["total_sentences"] == 1
+                    assert start["sentence_index"] < session_done["total_sentences"]
+
+    def test_session_config_between_utterances_replaces_config(self, mocker: MockerFixture):
+        app, speech_service = _build_test_app(mocker=mocker)
+
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/audio/speech/stream") as ws:
+                for expected_index, voice in enumerate(("Vivian", "Serena")):
+                    ws.send_json({"type": "session.config", "voice": voice})
+                    ws.send_json({"type": "input.text", "text": "Hello world. "})
+                    ws.send_json({"type": "input.done"})
+
+                    assert ws.receive_json()["type"] == "audio.start"
+                    ws.receive_bytes()
+                    assert ws.receive_json()["type"] == "audio.done"
+                    assert ws.receive_json() == {
+                        "type": "session.done",
+                        "utterance_index": expected_index,
+                        "total_sentences": 1,
+                    }
+
+        assert [call.args[0].voice for call in speech_service._generate_audio_bytes.await_args_list] == [
+            "Vivian",
+            "Serena",
+        ]
+
+    def test_session_config_rejected_while_input_is_buffered(self, mocker: MockerFixture):
+        app, speech_service = _build_test_app(mocker=mocker)
+
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/audio/speech/stream") as ws:
+                ws.send_json({"type": "session.config", "voice": "Vivian"})
+                ws.send_json({"type": "input.text", "text": "Hello world. "})
+                ws.send_json({"type": "session.config", "voice": "Serena"})
+
+                error = ws.receive_json()
+                assert error["type"] == "error"
+                assert "while input is buffered" in error["message"]
+
+                # The buffered text survives the rejected reconfiguration.
+                ws.send_json({"type": "input.done"})
+                assert ws.receive_json()["sentence_text"] == "Hello world."
+                ws.receive_bytes()
+                assert ws.receive_json()["type"] == "audio.done"
+                assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1}
+
+        assert speech_service._generate_audio_bytes.await_args_list[0].args[0].voice == "Vivian"
+
+    def test_session_close_ends_connection(self, mocker: MockerFixture):
+        app, _ = _build_test_app(mocker=mocker)
+
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/audio/speech/stream") as ws:
+                ws.send_json({"type": "session.config", "voice": "Vivian"})
+                ws.send_json({"type": "input.text", "text": "Hello world. "})
+                ws.send_json({"type": "input.done"})
+
+                assert ws.receive_json()["type"] == "audio.start"
+                ws.receive_bytes()
+                assert ws.receive_json()["type"] == "audio.done"
+                assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1}
+
+                ws.send_json({"type": "session.close"})
+                with pytest.raises(WebSocketDisconnect):
+                    ws.receive_json()
+
+    def test_idle_timeout_closes_reused_connection(self, mocker: MockerFixture):
+        app, _ = _build_test_app(idle_timeout=0.05, mocker=mocker)
+
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/audio/speech/stream") as ws:
+                ws.send_json({"type": "session.config", "voice": "Vivian"})
+                ws.send_json({"type": "input.text", "text": "Hello world. "})
+                ws.send_json({"type": "input.done"})
+
+                assert ws.receive_json()["type"] == "audio.start"
+                ws.receive_bytes()
+                assert ws.receive_json()["type"] == "audio.done"
+                assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1}
+
+                # Holding a connection open is not free: an idle client still
+                # gets timed out after the flush.
+                assert ws.receive_json() == {
+                    "type": "error",
+                    "message": "Idle timeout: no message received",
+                }
+
    def test_streaming_multiple_binary_frames(self, mocker: MockerFixture):
        captured_requests = []

@@ -125,9 +274,15 @@ class TestStreamingSpeechWebSocket:
                assert ws.receive_bytes() == b"\x06"

                done = ws.receive_json()
-                assert done == {"type": "audio.done", "sentence_index": 0, "total_bytes": 6, "error": False}
+                assert done == {
+                    "type": "audio.done",
+                    "utterance_index": 0,
+                    "sentence_index": 0,
+                    "total_bytes": 6,
+                    "error": False,
+                }

-                assert ws.receive_json() == {"type": "session.done", "total_sentences": 1}
+                assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1}

        assert len(captured_requests) == 1
        assert captured_requests[0].stream is True
@@ -210,6 +365,8 @@ class TestStreamingSpeechWebSocket:
                # is fully aligned.
                chunk = ws.receive_json()
                assert chunk["type"] == "audio.chunk"
+                assert chunk["utterance_index"] == 0
+                assert chunk["sentence_index"] == 0
                assert chunk["chunk_id"] == 0
                assert chunk["chunk_start_ms"] == 0
                assert chunk["chunk_end_ms"] == 500
@@ -240,7 +397,13 @@ class TestStreamingSpeechWebSocket:
                ]

                done = ws.receive_json()
-                assert done == {"type": "audio.done", "sentence_index": 0, "total_bytes": 2000, "error": False}
+                assert done == {
+                    "type": "audio.done",
+                    "utterance_index": 0,
+                    "sentence_index": 0,
+                    "total_bytes": 2000,
+                    "error": False,
+                }

        assert captured_requests[0].word_timestamps is True
        assert mock_align.await_count == 1
@@ -309,11 +472,12 @@ class TestStreamingSpeechWebSocket:
                assert ws.receive_bytes()
                assert ws.receive_json() == {
                    "type": "audio.done",
+                    "utterance_index": 0,
                    "sentence_index": 0,
                    "total_bytes": 36,
                    "error": False,
                }
-                assert ws.receive_json() == {"type": "session.done", "total_sentences": 1}
+                assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1}

    def test_invalid_streaming_config(self, mocker: MockerFixture):
        app, _ = _build_test_app(mocker=mocker)
@@ -341,7 +505,7 @@ class TestStreamingSpeechWebSocket:
                ws.send_json({"type": "input.text", "text": ""})
                ws.send_json({"type": "input.done"})

-                assert ws.receive_json() == {"type": "session.done", "total_sentences": 0}
+                assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 0}

        assert speech_service._generate_audio_bytes.await_count == 0

@@ -361,11 +525,12 @@ class TestStreamingSpeechWebSocket:
                ws.receive_bytes()
                assert ws.receive_json() == {
                    "type": "audio.done",
+                    "utterance_index": 0,
                    "sentence_index": 0,
                    "total_bytes": 36,
                    "error": False,
                }
-                assert ws.receive_json() == {"type": "session.done", "total_sentences": 1}
+                assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1}

    def test_unknown_message_type_keeps_session_open(self, mocker: MockerFixture):
        app, _ = _build_test_app(mocker=mocker)
@@ -384,12 +549,13 @@ class TestStreamingSpeechWebSocket:
                ws.receive_bytes()
                assert ws.receive_json() == {
                    "type": "audio.done",
+                    "utterance_index": 0,
                    "sentence_index": 0,
                    "total_bytes": 36,
                    "error": False,
                }

-                assert ws.receive_json() == {"type": "session.done", "total_sentences": 1}
+                assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1}

    def test_config_timeout_closes_session(self, mocker: MockerFixture):
        app, _ = _build_test_app(config_timeout=0.01, mocker=mocker)
@@ -416,10 +582,19 @@ class TestStreamingSpeechWebSocket:
                ws.send_json({"type": "input.done"})

                assert ws.receive_json()["type"] == "audio.start"
-                assert ws.receive_json() == {"type": "error", "message": "Generation failed for sentence 0: boom"}
-                assert ws.receive_json() == {"type": "audio.done", "sentence_index": 0, "total_bytes": 0, "error": True}
+                assert ws.receive_json() == {
+                    "type": "error",
+                    "message": "Generation failed for utterance 0, sentence 0: boom",
+                }
+                assert ws.receive_json() == {
+                    "type": "audio.done",
+                    "utterance_index": 0,
+                    "sentence_index": 0,
+                    "total_bytes": 0,
+                    "error": True,
+                }

-                assert ws.receive_json() == {"type": "session.done", "total_sentences": 1}
+                assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1}

    def test_streaming_generation_error_marks_audio_done(self, mocker: MockerFixture):
        speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech)
@@ -453,16 +628,17 @@ class TestStreamingSpeechWebSocket:
                assert ws.receive_bytes() == b"\x01\x02"
                assert ws.receive_json() == {
                    "type": "error",
-                    "message": "Generation failed for sentence 0: stream boom",
+                    "message": "Generation failed for utterance 0, sentence 0: stream boom",
                }
                assert ws.receive_json() == {
                    "type": "audio.done",
+                    "utterance_index": 0,
                    "sentence_index": 0,
                    "total_bytes": 2,
                    "error": True,
                }

-                assert ws.receive_json() == {"type": "session.done", "total_sentences": 1}
+                assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1}

    def test_invalid_input_text_type_returns_validation_error(self, mocker: MockerFixture):
        app, speech_service = _build_test_app(mocker=mocker)
@@ -478,7 +654,7 @@ class TestStreamingSpeechWebSocket:
                }

                ws.send_json({"type": "input.done"})
-                assert ws.receive_json() == {"type": "session.done", "total_sentences": 0}
+                assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 0}

        assert speech_service._generate_audio_bytes.await_count == 0

@@ -497,7 +673,7 @@ class TestStreamingSpeechWebSocket:
                }

                ws.send_json({"type": "input.done"})
-                assert ws.receive_json() == {"type": "session.done", "total_sentences": 0}
+                assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 0}

        assert speech_service._generate_audio_bytes.await_count == 0

@@ -550,7 +726,15 @@ class TestStreamingSpeechWebSocket:
        config.word_timestamps = False

        with pytest.raises(WebSocketDisconnect):
-            asyncio.run(handler._generate_and_send(websocket, config, "Hello world.", 0))
+            asyncio.run(
+                handler._generate_and_send(
+                    websocket,
+                    config,
+                    "Hello world.",
+                    utterance_index=0,
+                    sentence_index=0,
+                )
+            )

        speech_service.engine_client.abort.assert_awaited_once_with("req-abort")
        assert websocket.send_json.await_count == 2