From 0836723dcb6a7e19adcebf3b3216c294625c31ff Mon Sep 17 00:00:00 2001 From: Cale Shapera Date: Fri, 11 Sep 2026 14:00:01 -0700 Subject: [PATCH 1/2] fix(websocket): drain httpx-ws reader thread before closing sync sessions --- src/fish_audio_sdk/websocket.py | 37 +++++----- src/fishaudio/core/_ws_utils.py | 31 +++++++++ src/fishaudio/resources/tts.py | 12 +++- tests/unit/test_ws_utils.py | 120 ++++++++++++++++++++++++++++++++ 4 files changed, 181 insertions(+), 19 deletions(-) create mode 100644 src/fishaudio/core/_ws_utils.py create mode 100644 tests/unit/test_ws_utils.py 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..d17aed6 --- /dev/null +++ b/src/fishaudio/core/_ws_utils.py @@ -0,0 +1,31 @@ +"""Internal helpers for synchronous httpx-ws sessions.""" + +from __future__ import annotations + +import contextlib +import socket +from typing import Any + + +def drain_reader(ws: Any) -> None: + """Wake httpx-ws's background reader and wait for it to exit before the socket is + closed, so it can never read() a reused fd that belongs to the next connection. + + httpx-ws's sync ``WebSocketSession`` reads on a background thread. When the + session is closed, the socket is closed while that reader may be about to call + ``read()``. If the OS hands the same file descriptor to the next connection, the + stale reader consumes the new connection's bytes through the old SSL object and + the new session fails with ``[SSL: WRONG_VERSION_NUMBER]`` or similar errors. + """ + try: + 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 + 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: + thread.join(timeout=5) + except Exception: + pass 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..81bc5f0 --- /dev/null +++ b/tests/unit/test_ws_utils.py @@ -0,0 +1,120 @@ +"""Tests for the sync WebSocket reader-draining helper.""" + +import socket +from unittest.mock import Mock, patch + +import pytest + +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.response.extensions = {"network_stream": stream} + ws._background_receive_task = thread + return ws + + +class TestDrainReader: + def test_shuts_down_socket_and_joins_reader(self): + sock = Mock() + thread = Mock() + drain_reader(_make_ws(sock, thread)) + sock.shutdown.assert_called_once_with(socket.SHUT_RDWR) + thread.join.assert_called_once_with(timeout=5) + + 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.response = None + del ws._background_receive_task + drain_reader(ws) # must not raise + + def test_swallows_unexpected_errors(self): + ws = Mock() + ws.response.extensions = Mock() + ws.response.extensions.get.side_effect = RuntimeError("boom") + drain_reader(ws) # must not raise + + +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 + + 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 ws.__exit__.called + + @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) From 090c91b45324ba651b35622cdce33738baf74e8a Mon Sep 17 00:00:00 2001 From: Cale Shapera Date: Fri, 11 Sep 2026 14:23:03 -0700 Subject: [PATCH 2/2] fix(websocket): send close frame before draining reader; correct mechanism notes; real-socket test --- src/fishaudio/core/_ws_utils.py | 57 ++++++++++++++++---- tests/unit/test_ws_utils.py | 94 ++++++++++++++++++++++++++++++--- 2 files changed, 132 insertions(+), 19 deletions(-) diff --git a/src/fishaudio/core/_ws_utils.py b/src/fishaudio/core/_ws_utils.py index d17aed6..9bc0574 100644 --- a/src/fishaudio/core/_ws_utils.py +++ b/src/fishaudio/core/_ws_utils.py @@ -3,29 +3,64 @@ 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: - """Wake httpx-ws's background reader and wait for it to exit before the socket is - closed, so it can never read() a reused fd that belongs to the next connection. - - httpx-ws's sync ``WebSocketSession`` reads on a background thread. When the - session is closed, the socket is closed while that reader may be about to call - ``read()``. If the OS hands the same file descriptor to the next connection, the - stale reader consumes the new connection's bytes through the old SSL object and - the new session fails with ``[SSL: WRONG_VERSION_NUMBER]`` or similar errors. + """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: - response = getattr(ws, "response", None) - stream = response.extensions.get("network_stream") if response else None + 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: - pass + logger.debug("drain_reader failed; falling back to plain close", exc_info=True) diff --git a/tests/unit/test_ws_utils.py b/tests/unit/test_ws_utils.py index 81bc5f0..b72a3c0 100644 --- a/tests/unit/test_ws_utils.py +++ b/tests/unit/test_ws_utils.py @@ -1,9 +1,13 @@ """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 @@ -12,19 +16,48 @@ def _make_ws(sock=None, thread=None): ws = Mock() stream = Mock() stream.get_extra_info.return_value = sock - ws.response.extensions = {"network_stream": stream} + ws.stream = stream ws._background_receive_task = thread return ws class TestDrainReader: - def test_shuts_down_socket_and_joins_reader(self): + def test_sends_close_frame_then_shuts_down_socket_and_joins_reader(self): sock = Mock() thread = Mock() - drain_reader(_make_ws(sock, thread)) + 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") @@ -34,15 +67,57 @@ def test_joins_reader_when_socket_shutdown_fails(self): 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): + def test_swallows_unexpected_errors(self, caplog): ws = Mock() - ws.response.extensions = Mock() - ws.response.extensions.get.side_effect = RuntimeError("boom") - drain_reader(ws) # must not raise + 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: @@ -60,6 +135,9 @@ def test_new_client_calls_drain_reader_on_success( 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: @@ -68,7 +146,7 @@ def test_new_client_calls_drain_reader_on_success( mock_drain.assert_called_once_with(ws) # drain must happen before the session's __exit__ closes the socket - assert ws.__exit__.called + assert [c[0] for c in parent.mock_calls] == ["drain", "exit"] @patch("fishaudio.resources.tts.drain_reader") @patch("fishaudio.resources.tts.connect_ws")