diff --git a/pyproject.toml b/pyproject.toml index 27a93dc..8d03250 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ dependencies = [ "httpx>=0.27.2", "ormsgpack>=1.5.0", "pydantic>=2.9.1", - "httpx-ws>=0.6.2", + "httpx-ws>=0.7.2", "typing-extensions>=4.15.0", ] requires-python = ">=3.9" diff --git a/src/fish_audio_sdk/websocket.py b/src/fish_audio_sdk/websocket.py index 10f972c..37dbd12 100644 --- a/src/fish_audio_sdk/websocket.py +++ b/src/fish_audio_sdk/websocket.py @@ -6,6 +6,8 @@ import ormsgpack from httpx_ws import WebSocketDisconnect, aconnect_ws, connect_ws +from fishaudio.core._ws_utils import drain_reader + from .exceptions import WebSocketErr from .schemas import Backends, CloseEvent, StartEvent, TextEvent, TTSRequest @@ -70,22 +72,25 @@ def sender(): sender_future = self._executor.submit(sender) - while True: - try: - message = ws.receive_bytes() - data = ormsgpack.unpackb(message) - event = data["event"] - if event == "audio": - yield data["audio"] - elif event == "finish": - if data["reason"] == "error": - raise WebSocketErr - elif data["reason"] == "stop": - break - except WebSocketDisconnect: - raise WebSocketErr - - sender_future.result() + try: + while True: + try: + message = ws.receive_bytes() + data = ormsgpack.unpackb(message) + event = data["event"] + if event == "audio": + yield data["audio"] + elif event == "finish": + if data["reason"] == "error": + raise WebSocketErr + elif data["reason"] == "stop": + break + except WebSocketDisconnect: + raise WebSocketErr + + sender_future.result() + finally: + drain_reader(ws) class AsyncWebSocketSession: diff --git a/src/fishaudio/core/_ws_utils.py b/src/fishaudio/core/_ws_utils.py new file mode 100644 index 0000000..9bc0574 --- /dev/null +++ b/src/fishaudio/core/_ws_utils.py @@ -0,0 +1,66 @@ +"""Internal helpers for synchronous httpx-ws sessions.""" + +from __future__ import annotations + +import contextlib +import logging +import socket +from typing import Any + +import wsproto.events + +logger = logging.getLogger(__name__) + + +def drain_reader(ws: Any) -> None: + """Send the close frame, shut the socket down, and wait for httpx-ws's reader + thread to exit before the session's own ``close()`` frees the fd. + + Why this is needed (Linux only in practice): + + httpx-ws's sync ``WebSocketSession`` reads on a background thread and closes the + socket from the calling thread. A read that has *not started yet* is harmless: + CPython's ``SSLSocket._real_close`` sets ``_sslobj = None`` before closing the + fd, so a later read falls through to ``socket.recv`` on fd ``-1`` and fails with + ``EBADF``. The problem is the read that is *already in flight*. On Linux, + ``close()`` does not interrupt a thread blocked in ``recv()``. OpenSSL (with + read_ahead off) fetches a TLS record header and body with separate ``read(fd)`` + calls on the fd number cached in its BIO, so when the server's close-frame reply + arrives the second ``read()`` can land on a brand-new connection that has since + reused the same fd number, and that new session fails with + ``[SSL: WRONG_VERSION_NUMBER]`` or a similar TLS error. macOS wakes the blocked + reader with ``EBADF`` on ``close()``, so it is unaffected. + + The fix is to end the in-flight read *before* the fd is freed: send the close + frame (so the server still sees a clean 1000 close rather than 1006), then + ``shutdown(SHUT_RDWR)`` the socket, which wakes the blocked reader with EOF, and + join the reader thread. The later ``close()`` sees ``LOCAL_CLOSING`` and skips + sending a second frame. + + Residual risk: if httpx-ws has already closed the socket on one of its own error + paths (keepalive ping timeout, or a sender ``WriteError`` triggering ``close()``), + the fd is already freed, ``shutdown()`` fails with ``EBADF``, and the in-flight + read can still race. That can only be fixed inside httpx-ws's ``close()``. + """ + try: + stream = getattr(ws, "stream", None) + if stream is None: + response = getattr(ws, "response", None) + stream = response.extensions.get("network_stream") if response else None + sock = stream.get_extra_info("socket") if stream is not None else None + + # Moves the wsproto state to LOCAL_CLOSING; later close() calls skip the frame. + with contextlib.suppress(Exception): + ws.send(wsproto.events.CloseConnection(1000)) + + if sock is not None: + with contextlib.suppress(OSError): + sock.shutdown(socket.SHUT_RDWR) + + thread = getattr(ws, "_background_receive_task", None) + if thread is not None: + # 5 s is the worst case only when the reader is stuck; httpx-ws's own + # untimed join in __exit__ would hang afterwards anyway. + thread.join(timeout=5) + except Exception: + logger.debug("drain_reader failed; falling back to plain close", exc_info=True) diff --git a/src/fishaudio/resources/tts.py b/src/fishaudio/resources/tts.py index 2a5114a..62307a4 100644 --- a/src/fishaudio/resources/tts.py +++ b/src/fishaudio/resources/tts.py @@ -14,6 +14,7 @@ RequestOptions, WebSocketOptions, ) +from fishaudio.core._ws_utils import drain_reader from fishaudio.core.iterators import AsyncAudioStream, AudioStream from fishaudio.types import ( AudioFormat, @@ -362,10 +363,15 @@ def sender(): sender_future = executor.submit(sender) - # Process incoming audio messages - yield from iter_websocket_audio(ws) + try: + # Process incoming audio messages + yield from iter_websocket_audio(ws) - sender_future.result() + sender_future.result() + finally: + # Stop httpx-ws's reader thread before the socket is closed so it + # cannot read from a reused fd belonging to the next connection. + drain_reader(ws) finally: executor.shutdown(wait=False) diff --git a/tests/unit/test_ws_utils.py b/tests/unit/test_ws_utils.py new file mode 100644 index 0000000..b72a3c0 --- /dev/null +++ b/tests/unit/test_ws_utils.py @@ -0,0 +1,198 @@ +"""Tests for the sync WebSocket reader-draining helper.""" + +import socket +import threading +from unittest.mock import Mock, patch + +import httpx_ws +import pytest +import wsproto +from httpcore._backends.sync import SyncStream + +from fishaudio.core._ws_utils import drain_reader + + +def _make_ws(sock=None, thread=None): + ws = Mock() + stream = Mock() + stream.get_extra_info.return_value = sock + ws.stream = stream + ws._background_receive_task = thread + return ws + + +class TestDrainReader: + def test_sends_close_frame_then_shuts_down_socket_and_joins_reader(self): + sock = Mock() + thread = Mock() + ws = _make_ws(sock, thread) + parent = Mock() + parent.attach_mock(ws.send, "send") + parent.attach_mock(sock.shutdown, "shutdown") + parent.attach_mock(thread.join, "join") + + drain_reader(ws) + + assert [c[0] for c in parent.mock_calls] == ["send", "shutdown", "join"] + (event,) = ws.send.call_args.args + assert isinstance(event, wsproto.events.CloseConnection) + assert event.code == 1000 + sock.shutdown.assert_called_once_with(socket.SHUT_RDWR) + thread.join.assert_called_once_with(timeout=5) + + def test_falls_back_to_response_network_stream(self): + sock = Mock() + ws = Mock() + ws.stream = None + stream = Mock() + stream.get_extra_info.return_value = sock + ws.response.extensions = {"network_stream": stream} + ws._background_receive_task = Mock() + drain_reader(ws) + sock.shutdown.assert_called_once_with(socket.SHUT_RDWR) + + def test_shuts_down_socket_when_close_frame_send_fails(self): + sock = Mock() + ws = _make_ws(sock, Mock()) + ws.send.side_effect = httpx_ws.WebSocketNetworkError() + drain_reader(ws) + sock.shutdown.assert_called_once_with(socket.SHUT_RDWR) + + def test_joins_reader_when_socket_shutdown_fails(self): + sock = Mock() + sock.shutdown.side_effect = OSError("already closed") + thread = Mock() + drain_reader(_make_ws(sock, thread)) + thread.join.assert_called_once_with(timeout=5) + + def test_tolerates_missing_socket_and_thread(self): + ws = Mock() + ws.stream = None + ws.response = None + del ws._background_receive_task + drain_reader(ws) # must not raise + + def test_swallows_unexpected_errors(self, caplog): + ws = Mock() + ws.stream = Mock() + ws.stream.get_extra_info.side_effect = RuntimeError("boom") + with caplog.at_level("DEBUG", logger="fishaudio.core._ws_utils"): + drain_reader(ws) # must not raise + assert "drain_reader failed" in caplog.text + + +class TestDrainReaderRealSession: + def test_real_httpx_ws_session_over_socketpair(self): + client_sock, server_sock = socket.socketpair() + server_sock.settimeout(1) + hook_calls = [] + old_hook = threading.excepthook + threading.excepthook = lambda args: hook_calls.append(args) + try: + ws = httpx_ws.WebSocketSession( + SyncStream(client_sock), keepalive_ping_interval_seconds=None + ) + with ws: + drain_reader(ws) + assert not ws._background_receive_task.is_alive() + assert ( + ws.connection.state + is wsproto.connection.ConnectionState.LOCAL_CLOSING + ) + + # The peer must have received a clean close frame before FIN. + received = b"" + while True: + chunk = server_sock.recv(4096) + if not chunk: + break + received += chunk + server = wsproto.connection.Connection(wsproto.ConnectionType.SERVER) + server.receive_data(received) + events = list(server.events()) + assert len(events) == 1 + assert isinstance(events[0], wsproto.events.CloseConnection) + assert events[0].code == 1000 + assert hook_calls == [] + finally: + threading.excepthook = old_hook + server_sock.close() + client_sock.close() + + +class TestDrainReaderIsCalledFromSyncPaths: + @patch("fishaudio.resources.tts.drain_reader") + @patch("fishaudio.resources.tts.connect_ws") + @patch("fishaudio.resources.tts.ThreadPoolExecutor") + def test_new_client_calls_drain_reader_on_success( + self, mock_executor, mock_connect_ws, mock_drain + ): + from fishaudio.core import ClientWrapper + from fishaudio.resources.tts import TTSClient + + ws = Mock() + ws.__enter__ = Mock(return_value=ws) + ws.__exit__ = Mock(return_value=None) + mock_connect_ws.return_value = ws + mock_executor.return_value.submit.return_value.result.return_value = None + parent = Mock() + parent.attach_mock(mock_drain, "drain") + parent.attach_mock(ws.__exit__, "exit") + + client = TTSClient(ClientWrapper(api_key="k", base_url="https://x")) + with patch("fishaudio.resources.tts.iter_websocket_audio") as recv: + recv.return_value = iter([b"a"]) + assert list(client.stream_websocket(iter(["hi"]))) == [b"a"] + + mock_drain.assert_called_once_with(ws) + # drain must happen before the session's __exit__ closes the socket + assert [c[0] for c in parent.mock_calls] == ["drain", "exit"] + + @patch("fishaudio.resources.tts.drain_reader") + @patch("fishaudio.resources.tts.connect_ws") + @patch("fishaudio.resources.tts.ThreadPoolExecutor") + def test_new_client_calls_drain_reader_on_error( + self, mock_executor, mock_connect_ws, mock_drain + ): + from fishaudio.core import ClientWrapper + from fishaudio.resources.tts import TTSClient + + ws = Mock() + ws.__enter__ = Mock(return_value=ws) + ws.__exit__ = Mock(return_value=None) + mock_connect_ws.return_value = ws + + client = TTSClient(ClientWrapper(api_key="k", base_url="https://x")) + with patch("fishaudio.resources.tts.iter_websocket_audio") as recv: + recv.side_effect = RuntimeError("boom") + with pytest.raises(RuntimeError): + list(client.stream_websocket(iter(["hi"]))) + + mock_drain.assert_called_once_with(ws) + + def test_legacy_session_calls_drain_reader(self): + import ormsgpack + + try: + from fish_audio_sdk import WebSocketSession + except TypeError: # legacy schemas use `X | None`, which fails on 3.9 + pytest.skip("fish_audio_sdk does not import on this Python version") + + ws = Mock() + ws.__enter__ = Mock(return_value=ws) + ws.__exit__ = Mock(return_value=None) + ws.receive_bytes.side_effect = [ + ormsgpack.packb({"event": "audio", "audio": b"a"}), + ormsgpack.packb({"event": "finish", "reason": "stop"}), + ] + + connect = patch("fish_audio_sdk.websocket.connect_ws", return_value=ws) + drain = patch("fish_audio_sdk.websocket.drain_reader") + with connect, drain as mock_drain, WebSocketSession("k") as session: + # Bypass the real sender thread; only the receive path is under test. + session._executor = Mock() + session._executor.submit.return_value.result.return_value = None + chunks = list(session.tts(Mock(), iter(["hi"]))) + + assert chunks == [b"a"] + mock_drain.assert_called_once_with(ws) diff --git a/uv.lock b/uv.lock index 45bc7ad..e32d642 100644 --- a/uv.lock +++ b/uv.lock @@ -292,7 +292,7 @@ resolution-markers = [ "python_full_version < '3.10'", ] dependencies = [ - { name = "colorama", marker = "python_full_version < '3.10' and sys_platform == 'win32'" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/b9/2e/0090cbf739cee7d23781ad4b89a9894a41538e4fcf4c31dcdd705b78eb8b/click-8.1.8.tar.gz", hash = "sha256:ed53c9d8990d83c2a27deae68e4ee337473f6330c040a31d4225c9574d16096a", size = 226593, upload-time = "2024-12-21T18:38:44.339Z" } wheels = [ @@ -308,7 +308,7 @@ resolution-markers = [ "python_full_version == '3.10.*'", ] dependencies = [ - { name = "colorama", marker = "python_full_version >= '3.10' and sys_platform == 'win32'" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/46/61/de6cd827efad202d7057d93e0fed9294b96952e188f7384832791c7b2254/click-8.3.0.tar.gz", hash = "sha256:e7b8232224eba16f4ebe410c25ced9f7875cb5f3263ffc93cc3e8da705e229c4", size = 276943, upload-time = "2025-09-18T17:32:23.696Z" } wheels = [ @@ -440,7 +440,7 @@ wheels = [ [package.optional-dependencies] toml = [ - { name = "tomli", marker = "python_full_version < '3.10'" }, + { name = "tomli" }, ] [[package]] @@ -548,7 +548,7 @@ wheels = [ [package.optional-dependencies] toml = [ - { name = "tomli", marker = "python_full_version >= '3.10' and python_full_version <= '3.11'" }, + { name = "tomli", marker = "python_full_version <= '3.11'" }, ] [[package]] @@ -643,7 +643,7 @@ name = "exceptiongroup" version = "1.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "python_full_version < '3.11'" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/0b/9f/a65090624ecf468cdca03533906e7c69ed7588582240cfe7cc9e770b50eb/exceptiongroup-1.3.0.tar.gz", hash = "sha256:b241f5885f560bc56a59ee63ca4c6a8bfa46ae4ad651af316d4e81817bb9fd88", size = 29749, upload-time = "2025-05-10T17:42:51.123Z" } wheels = [ @@ -652,7 +652,7 @@ wheels = [ [[package]] name = "fish-audio-sdk" -version = "1.2.0" +version = "1.3.0" source = { editable = "." } dependencies = [ { name = "httpx" }, @@ -689,7 +689,7 @@ dev = [ [package.metadata] requires-dist = [ { name = "httpx", specifier = ">=0.27.2" }, - { name = "httpx-ws", specifier = ">=0.6.2" }, + { name = "httpx-ws", specifier = ">=0.7.2" }, { name = "ormsgpack", specifier = ">=1.5.0" }, { name = "pydantic", specifier = ">=2.9.1" }, { name = "sounddevice", marker = "extra == 'utils'", specifier = ">=0.4.6" }, @@ -757,10 +757,10 @@ resolution-markers = [ "python_full_version < '3.10'", ] dependencies = [ - { name = "anyio", marker = "python_full_version < '3.10'" }, - { name = "httpcore", marker = "python_full_version < '3.10'" }, - { name = "httpx", marker = "python_full_version < '3.10'" }, - { name = "wsproto", marker = "python_full_version < '3.10'" }, + { name = "anyio" }, + { name = "httpcore" }, + { name = "httpx" }, + { name = "wsproto" }, ] sdist = { url = "https://files.pythonhosted.org/packages/56/ba/e310ccdb8f18a2b894bfacd085ef390cf6cc70bb10ff9f109d58d94f6b47/httpx_ws-0.7.2.tar.gz", hash = "sha256:93edea6c8fc313464fc287bff7d2ad20e6196b7754c76f946f73b4af79886d4e", size = 24513, upload-time = "2025-03-28T13:20:03.039Z" } wheels = [ @@ -776,10 +776,10 @@ resolution-markers = [ "python_full_version == '3.10.*'", ] dependencies = [ - { name = "anyio", marker = "python_full_version >= '3.10'" }, - { name = "httpcore", marker = "python_full_version >= '3.10'" }, - { name = "httpx", marker = "python_full_version >= '3.10'" }, - { name = "wsproto", marker = "python_full_version >= '3.10'" }, + { name = "anyio" }, + { name = "httpcore" }, + { name = "httpx" }, + { name = "wsproto" }, ] sdist = { url = "https://files.pythonhosted.org/packages/bb/3d/93a17bc7d769f8b47facc9c8e2cb5435b04185c1ebfdf49f2c7485f892b6/httpx_ws-0.8.1.tar.gz", hash = "sha256:fdc8f471980d572f371d113a50710bbf397245048e9a542606af0bba2b956ab5", size = 105819, upload-time = "2025-10-28T07:31:42.298Z" } wheels = [ @@ -1366,9 +1366,9 @@ resolution-markers = [ "python_full_version < '3.10'", ] dependencies = [ - { name = "pastel", marker = "python_full_version < '3.10'" }, - { name = "pyyaml", marker = "python_full_version < '3.10'" }, - { name = "tomli", marker = "python_full_version < '3.10'" }, + { name = "pastel" }, + { name = "pyyaml" }, + { name = "tomli" }, ] sdist = { url = "https://files.pythonhosted.org/packages/a5/f2/273fe54a78dc5c6c8dd63db71f5a6ceb95e4648516b5aeaeff4bde804e44/poethepoet-0.37.0.tar.gz", hash = "sha256:73edf458707c674a079baa46802e21455bda3a7f82a408e58c31b9f4fe8e933d", size = 68570, upload-time = "2025-08-11T18:00:29.103Z" } wheels = [ @@ -1384,9 +1384,9 @@ resolution-markers = [ "python_full_version == '3.10.*'", ] dependencies = [ - { name = "pastel", marker = "python_full_version >= '3.10'" }, - { name = "pyyaml", marker = "python_full_version >= '3.10'" }, - { name = "tomli", marker = "python_full_version == '3.10.*'" }, + { name = "pastel" }, + { name = "pyyaml" }, + { name = "tomli", marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/a8/b9/fa92286560f70eaa40d473ea48376d20c6c21f63627d33c6bb1c5e385175/poethepoet-0.41.0.tar.gz", hash = "sha256:dcaad621dc061f6a90b17d091bebb9ca043d67bfe9bd6aa4185aea3ebf7ff3e6", size = 87780, upload-time = "2026-02-08T20:45:36.061Z" } wheels = [ @@ -1639,8 +1639,8 @@ resolution-markers = [ "python_full_version < '3.10'", ] dependencies = [ - { name = "packaging", marker = "python_full_version < '3.10'" }, - { name = "pytest", marker = "python_full_version < '3.10'" }, + { name = "packaging" }, + { name = "pytest" }, ] sdist = { url = "https://files.pythonhosted.org/packages/26/53/a543a76f922a5337d10df22441af8bf68f1b421cadf9aedf8a77943b81f6/pytest_rerunfailures-16.0.1.tar.gz", hash = "sha256:ed4b3a6e7badb0a720ddd93f9de1e124ba99a0cb13bc88561b3c168c16062559", size = 27612, upload-time = "2025-09-02T06:48:25.193Z" } wheels = [ @@ -1656,8 +1656,8 @@ resolution-markers = [ "python_full_version == '3.10.*'", ] dependencies = [ - { name = "packaging", marker = "python_full_version >= '3.10'" }, - { name = "pytest", marker = "python_full_version >= '3.10'" }, + { name = "packaging" }, + { name = "pytest" }, ] sdist = { url = "https://files.pythonhosted.org/packages/de/04/71e9520551fc8fe2cf5c1a1842e4e600265b0815f2016b7c27ec85688682/pytest_rerunfailures-16.1.tar.gz", hash = "sha256:c38b266db8a808953ebd71ac25c381cb1981a78ff9340a14bcb9f1b9bff1899e", size = 30889, upload-time = "2025-10-10T07:06:01.238Z" } wheels = [