From 689cf7c17232ba74c0f215bc72be72ee0d6c5d49 Mon Sep 17 00:00:00 2001 From: Aleksandar Mitov <140423361+aleks-mitov@users.noreply.github.com> Date: Sun, 27 Sep 2026 12:10:08 +0000 Subject: [PATCH 1/2] fix(streaming): hold TLS reads until the opening handshake request is written websockets' sync client starts its reader thread, blocked in recv(), before the calling thread writes the opening handshake request on the same socket. Against servers that send TLS 1.3 post-handshake messages, such as an AWS Application Load Balancer, a small fraction of connections opened this way never receive the HTTP 101 response. RealTimeTranscriber.connect() then succeeds only after its 1 s open timeout and a retry, about 1.5 s late. In 600 sequential connects to such an endpoint, 14 hit this without the change and none with it, and the median connect time was the same (25 ms). The asyncio client is unaffected. For wss:// URIs the sync client now passes a TLS context whose sockets hold reads until the first send() or sendall() has returned, or until shutdown() is called, with a 30 s bound. Certificate and hostname verification are those of ssl.create_default_context(), which websockets uses by default. The context is passed as ssl on websockets 13 and later and as ssl_context on websockets 11 and 12. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_019UNrfUwp3uuUFEEbNvY3zH --- assemblyai/__version__.py | 2 +- assemblyai/streaming/v3/_tls.py | 91 +++++++++ assemblyai/streaming/v3/client.py | 23 ++- tests/unit/test_streaming_tls.py | 324 ++++++++++++++++++++++++++++++ 4 files changed, 438 insertions(+), 2 deletions(-) create mode 100644 assemblyai/streaming/v3/_tls.py create mode 100644 tests/unit/test_streaming_tls.py diff --git a/assemblyai/__version__.py b/assemblyai/__version__.py index f49459c7..51bbb3f2 100644 --- a/assemblyai/__version__.py +++ b/assemblyai/__version__.py @@ -1 +1 @@ -__version__ = "1.6.1" +__version__ = "1.6.2" diff --git a/assemblyai/streaming/v3/_tls.py b/assemblyai/streaming/v3/_tls.py new file mode 100644 index 00000000..c2730d4b --- /dev/null +++ b/assemblyai/streaming/v3/_tls.py @@ -0,0 +1,91 @@ +"""TLS context for the thread-based streaming client. + +websockets' sync client starts the thread that reads the connection before the +calling thread writes the opening handshake request. With servers that send +TLS 1.3 post-handshake messages, such as session tickets, a read already +blocked on the socket while that request is written occasionally leaves the +connection without a handshake response, and the connect times out. Sockets +created from ``_create_tls_context()`` hold reads until the first write has +returned, so the request is written before anything is read. +""" + +from __future__ import annotations + +import ssl +import threading +from typing import TYPE_CHECKING, Optional + +if TYPE_CHECKING: + from _typeshed import ReadableBuffer, WriteableBuffer + +# Upper bound on how long a read is held, so a socket that is never written to +# can still be read. +_FIRST_WRITE_WAIT_SEC = 30.0 + +# Serializes the creation of each socket's first-write event. +_FIRST_WRITE_LOCK = threading.Lock() + + +class _WriteFirstSSLSocket(ssl.SSLSocket): + """``SSLSocket`` whose reads wait until the first write has returned. + + The first ``send()`` or ``sendall()`` releases held reads whether or not it + succeeds, and so does ``shutdown()``, which lets closing the socket + interrupt a held read. Once released, reads and writes proceed + concurrently as with a plain ``SSLSocket``. + """ + + _first_write_event: Optional[threading.Event] = None + + def _first_write_done(self) -> threading.Event: + # SSLContext.wrap_socket() builds instances without calling __init__, + # so the event is created on first use. + done = self._first_write_event + if done is None: + with _FIRST_WRITE_LOCK: + done = self._first_write_event + if done is None: + done = self._first_write_event = threading.Event() + return done + + def send(self, data: ReadableBuffer, flags: int = 0) -> int: + try: + return super().send(data, flags) + finally: + self._first_write_done().set() + + def sendall(self, data: ReadableBuffer, flags: int = 0) -> None: + try: + super().sendall(data, flags) + finally: + self._first_write_done().set() + + def recv(self, buflen: int = 1024, flags: int = 0) -> bytes: + self._first_write_done().wait(_FIRST_WRITE_WAIT_SEC) + return super().recv(buflen, flags) + + def recv_into( + self, + buffer: WriteableBuffer, + nbytes: Optional[int] = None, + flags: int = 0, + ) -> int: + self._first_write_done().wait(_FIRST_WRITE_WAIT_SEC) + return super().recv_into(buffer, nbytes, flags) + + def shutdown(self, how: int) -> None: + self._first_write_done().set() + super().shutdown(how) + + +def _create_tls_context() -> ssl.SSLContext: + """Return a default client context whose sockets are + ``_WriteFirstSSLSocket``. + + Certificate and hostname verification are those of + ``ssl.create_default_context()``, which is also what websockets uses when + no context is given. + """ + context = ssl.create_default_context() + context.sslsocket_class = _WriteFirstSSLSocket + return context diff --git a/assemblyai/streaming/v3/client.py b/assemblyai/streaming/v3/client.py index 29e4c9d0..d0c510fe 100644 --- a/assemblyai/streaming/v3/client.py +++ b/assemblyai/streaming/v3/client.py @@ -1,14 +1,17 @@ +import inspect import json import logging import queue import threading import time from typing import Any, Dict, Generator, Iterable, Optional, Union +from urllib.parse import urlsplit import httpx import websockets from pydantic import BaseModel -from websockets.sync.client import connect as websocket_connect +from websockets.sync.client import ClientConnection +from websockets.sync.client import connect as _ws_sync_connect from ._base import ( _BaseStreamingClient, @@ -21,6 +24,7 @@ _resolve_options, _user_agent, ) +from ._tls import _create_tls_context from .models import ( ErrorEvent, EventMessage, @@ -40,6 +44,23 @@ logger = logging.getLogger(__name__) +# The sync ``connect()`` TLS parameter is ``ssl`` from websockets 13 on and +# ``ssl_context`` before (``setup.py`` allows ``websockets>=11.0``). +_WS_SSL_KW = ( + "ssl" if "ssl" in inspect.signature(_ws_sync_connect).parameters else "ssl_context" +) + + +def websocket_connect(uri: str, **kwargs: Any) -> ClientConnection: + """Open a sync websocket connection with ``websockets.sync.client.connect``. + + ``wss://`` URIs use the context from ``_create_tls_context()``. Module-level + indirection so tests can patch a single attribute. + """ + if urlsplit(uri).scheme == "wss": + kwargs[_WS_SSL_KW] = _create_tls_context() + return _ws_sync_connect(uri, **kwargs) + class RealTimeTranscriber(_BaseStreamingClient): def __init__( diff --git a/tests/unit/test_streaming_tls.py b/tests/unit/test_streaming_tls.py new file mode 100644 index 00000000..dee94713 --- /dev/null +++ b/tests/unit/test_streaming_tls.py @@ -0,0 +1,324 @@ +import contextlib +import inspect +import shutil +import socket +import ssl +import subprocess +import threading +import warnings +from typing import Iterator, List, Optional, Tuple + +import pytest +import websockets +from pytest_mock import MockFixture +from websockets.sync.server import serve as websocket_serve + +from assemblyai.streaming.v3 import ( + SpeechModel, + StreamingClient, + StreamingClientOptions, + StreamingParameters, + _tls, +) +from assemblyai.streaming.v3.client import websocket_connect + +# websockets 13 renamed the sync ``connect()`` / ``serve()`` TLS parameter +# from ``ssl_context`` to ``ssl``. +_WEBSOCKETS_MAJOR = int(websockets.__version__.split(".")[0]) +_EXPECTED_SSL_KW = "ssl" if _WEBSOCKETS_MAJOR >= 13 else "ssl_context" +_SERVE_SSL_KW = ( + "ssl" if "ssl" in inspect.signature(websocket_serve).parameters else "ssl_context" +) + +_HOST = "127.0.0.1" + + +@pytest.fixture(scope="module") +def tls_cert(tmp_path_factory: pytest.TempPathFactory) -> Tuple[str, str]: + """A throwaway self-signed certificate and key for ``127.0.0.1``.""" + openssl = shutil.which("openssl") + if openssl is None: + pytest.skip("the openssl CLI is needed to generate a test certificate") + directory = tmp_path_factory.mktemp("tls") + cert = str(directory / "cert.pem") + key = str(directory / "key.pem") + subprocess.run( + [ + openssl, + "req", + "-x509", + "-newkey", + "ec", + "-pkeyopt", + "ec_paramgen_curve:prime256v1", + "-nodes", + "-days", + "1", + "-subj", + f"/CN={_HOST}", + "-addext", + f"subjectAltName=IP:{_HOST}", + "-keyout", + key, + "-out", + cert, + ], + check=True, + capture_output=True, + ) + return cert, key + + +@pytest.fixture +def server_context(tls_cert: Tuple[str, str]) -> ssl.SSLContext: + """A server context that only speaks TLS 1.3.""" + cert, key = tls_cert + context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + context.minimum_version = ssl.TLSVersion.TLSv1_3 + context.load_cert_chain(cert, key) + return context + + +@pytest.fixture +def client_context(tls_cert: Tuple[str, str]) -> ssl.SSLContext: + """The SDK context, additionally trusting the test certificate.""" + cert, _ = tls_cert + context = _tls._create_tls_context() + context.load_verify_locations(cert) + return context + + +@contextlib.contextmanager +def _tls_server(context: ssl.SSLContext, greeting: bytes) -> Iterator[int]: + """Accept one TLS connection, send ``greeting`` as soon as the handshake + completes, then read until the client goes away. Yields the port.""" + listener = socket.create_server((_HOST, 0)) + listener.settimeout(10) + + def serve() -> None: + with contextlib.suppress(OSError): + conn, _ = listener.accept() + with context.wrap_socket(conn, server_side=True) as tls: + tls.sendall(greeting) + while tls.recv(1024): + pass + + thread = threading.Thread(target=serve, daemon=True) + thread.start() + try: + yield listener.getsockname()[1] + finally: + listener.close() + thread.join(timeout=10) + + +@contextlib.contextmanager +def _websocket_server(context: ssl.SSLContext) -> Iterator[int]: + """Run a websockets echo server over TLS. Yields the port.""" + + def echo(connection) -> None: + for message in connection: + connection.send(message) + + server = websocket_serve(echo, _HOST, 0, **{_SERVE_SSL_KW: context}) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield server.socket.getsockname()[1] + finally: + server.shutdown() + thread.join(timeout=10) + + +class _Read: + """``recv()`` on a background thread, as websockets' reader thread does.""" + + def __init__(self, sock: ssl.SSLSocket): + self.done = threading.Event() + self.result: Optional[bytes] = None + self.error: Optional[BaseException] = None + threading.Thread(target=self._run, args=(sock,), daemon=True).start() + + def _run(self, sock: ssl.SSLSocket) -> None: + try: + self.result = sock.recv(1024) + except (OSError, ValueError) as exc: + self.error = exc + finally: + self.done.set() + + +def _connect_tls(context: ssl.SSLContext, port: int) -> ssl.SSLSocket: + return context.wrap_socket( + socket.create_connection((_HOST, port), timeout=10), + server_hostname=_HOST, + ) + + +def test_tls_context_keeps_verification_on(): + context = _tls._create_tls_context() + + assert context.sslsocket_class is _tls._WriteFirstSSLSocket + assert context.verify_mode == ssl.CERT_REQUIRED + assert context.check_hostname is True + + +def test_tls_context_rejects_untrusted_certificate( + server_context: ssl.SSLContext, +): + # Given: a server whose self-signed certificate the default trust store + # does not contain + with _tls_server(server_context, greeting=b"hello") as port: + # When / Then: the handshake fails verification + with pytest.raises(ssl.SSLCertVerificationError): + _connect_tls(_tls._create_tls_context(), port) + + +def test_read_waits_for_first_write( + client_context: ssl.SSLContext, server_context: ssl.SSLContext +): + # Given: a TLS 1.3 connection whose server writes as soon as the handshake + # completes + with _tls_server(server_context, greeting=b"hello") as port: + with _connect_tls(client_context, port) as sock: + assert isinstance(sock, _tls._WriteFirstSSLSocket) + assert sock.version() == "TLSv1.3" + + # When: a read starts before anything is written + read = _Read(sock) + + # Then: it is held although the server's data is available + assert not read.done.wait(0.5) + + # When: the first write returns + sock.sendall(b"ping") + + # Then: the read proceeds and returns the server's data + assert read.done.wait(10) + assert read.error is None + assert read.result == b"hello" + + +def test_shutdown_releases_held_read( + client_context: ssl.SSLContext, server_context: ssl.SSLContext +): + # Given: a read held because nothing has been written + with _tls_server(server_context, greeting=b"hello") as port: + with _connect_tls(client_context, port) as sock: + read = _Read(sock) + assert not read.done.wait(0.2) + + # When: the socket is shut down without ever being written to + sock.shutdown(socket.SHUT_RDWR) + + # Then: the read returns well before the wait bound + assert read.done.wait(5) + + +def test_read_proceeds_after_wait_bound( + monkeypatch: pytest.MonkeyPatch, + client_context: ssl.SSLContext, + server_context: ssl.SSLContext, +): + # Given: a short wait bound + monkeypatch.setattr(_tls, "_FIRST_WRITE_WAIT_SEC", 0.2) + + with _tls_server(server_context, greeting=b"hello") as port: + with _connect_tls(client_context, port) as sock: + # When: a read starts and nothing is ever written + read = _Read(sock) + + # Then: the read returns the server's data once the bound elapses + assert read.done.wait(10) + assert read.error is None + assert read.result == b"hello" + + +def test_websocket_handshake_over_wss( + monkeypatch: pytest.MonkeyPatch, + tls_cert: Tuple[str, str], + server_context: ssl.SSLContext, +): + # Given: the SDK context, trusting the test certificate, and a TLS 1.3 + # websockets echo server + cert, _ = tls_cert + create_tls_context = _tls._create_tls_context + + def trusting_tls_context() -> ssl.SSLContext: + context = create_tls_context() + context.load_verify_locations(cert) + return context + + monkeypatch.setattr( + "assemblyai.streaming.v3.client._create_tls_context", trusting_tls_context + ) + + with _websocket_server(server_context) as port: + # When: connecting over wss:// with the sync client's connect + with warnings.catch_warnings(): + # The TLS parameter name matches the installed websockets, so no + # rename DeprecationWarning is raised. + warnings.simplefilter("error", DeprecationWarning) + connection = websocket_connect(f"wss://{_HOST}:{port}/", open_timeout=10) + + # Then: the opening handshake completes over the write-first socket + with connection: + assert isinstance(connection.socket, _tls._WriteFirstSSLSocket) + assert connection.socket.version() == "TLSv1.3" + connection.send("ping") + assert connection.recv(timeout=10) == "ping" + + +def _record_connect(mocker: MockFixture) -> List[Tuple[str, dict]]: + calls: List[Tuple[str, dict]] = [] + + def fake_connect(uri: str, **kwargs) -> None: + calls.append((uri, kwargs)) + + mocker.patch("assemblyai.streaming.v3.client._ws_sync_connect", new=fake_connect) + mocker.patch("threading.Thread.start", return_value=None) + return calls + + +def _connect_client(api_host: str) -> None: + client = StreamingClient(StreamingClientOptions(api_key="test", api_host=api_host)) + client.connect( + StreamingParameters( + sample_rate=16000, + speech_model=SpeechModel.universal_streaming_english, + ) + ) + + +def test_client_connect_passes_tls_context_for_wss(mocker: MockFixture): + # Given: a client pointed at a wss:// host + calls = _record_connect(mocker) + + # When: connecting + _connect_client("api.example.com") + + # Then: websockets gets the SDK context under the TLS parameter name of + # the installed version, alongside the usual arguments + [(uri, kwargs)] = calls + assert uri.startswith("wss://api.example.com/v3/ws?") + assert {"ssl", "ssl_context"} & kwargs.keys() == {_EXPECTED_SSL_KW} + context = kwargs[_EXPECTED_SSL_KW] + assert isinstance(context, ssl.SSLContext) + assert context.sslsocket_class is _tls._WriteFirstSSLSocket + assert context.verify_mode == ssl.CERT_REQUIRED + assert context.check_hostname is True + assert kwargs["additional_headers"]["Authorization"] == "test" + assert kwargs["open_timeout"] == 1.0 + + +def test_client_connect_passes_no_tls_context_for_ws(mocker: MockFixture): + # Given: a client pointed at a plain ws:// host + calls = _record_connect(mocker) + + # When: connecting + _connect_client("ws://localhost:8080") + + # Then: no TLS argument is passed, which websockets rejects for ws:// + [(uri, kwargs)] = calls + assert uri.startswith("ws://localhost:8080/v3/ws?") + assert not {"ssl", "ssl_context"} & kwargs.keys() From 781cfa4d1c143527dbd50cdea3d346b6f40c829c Mon Sep 17 00:00:00 2001 From: Aleksandar Mitov <140423361+aleks-mitov@users.noreply.github.com> Date: Sun, 27 Sep 2026 12:40:31 +0000 Subject: [PATCH 2/2] fix(streaming): leave foreign TLS socket classes in place _WriteFirstSSLSocket extends the ssl.SSLSocket present when the SDK is imported. If gevent monkey-patches ssl after that import, ssl.create_default_context() returns gevent's context, whose wrap_socket() calls sslsocket_class(...) directly, and the stdlib class refuses public construction. Every wss:// connect from the sync client then raised TypeError, which the connect retry loop does not catch. _create_tls_context() now installs _WriteFirstSSLSocket only when it extends the context's sslsocket_class, and otherwise leaves the context as ssl.create_default_context() returned it. SSLContext replacements that do not build sockets from sslsocket_class also keep their own sockets; the module docstring now says so. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_019UNrfUwp3uuUFEEbNvY3zH --- assemblyai/streaming/v3/_tls.py | 14 +++++++++++++- tests/unit/test_streaming_tls.py | 19 +++++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/assemblyai/streaming/v3/_tls.py b/assemblyai/streaming/v3/_tls.py index c2730d4b..252fcdfd 100644 --- a/assemblyai/streaming/v3/_tls.py +++ b/assemblyai/streaming/v3/_tls.py @@ -7,6 +7,11 @@ connection without a handshake response, and the connect times out. Sockets created from ``_create_tls_context()`` hold reads until the first write has returned, so the request is written before anything is read. + +Reads are held only on sockets the context builds from its +``sslsocket_class``. An ``SSLContext`` replacement whose ``wrap_socket()`` +builds sockets some other way keeps its own sockets, and reads on them are not +held. """ from __future__ import annotations @@ -85,7 +90,14 @@ def _create_tls_context() -> ssl.SSLContext: Certificate and hostname verification are those of ``ssl.create_default_context()``, which is also what websockets uses when no context is given. + + ``_WriteFirstSSLSocket`` extends the ``ssl.SSLSocket`` present when this + module was imported. If the default context uses a socket class that + ``_WriteFirstSSLSocket`` does not extend, for example because gevent + patched ``ssl`` afterwards, the context keeps that class: its + ``wrap_socket()`` may not be able to construct ``_WriteFirstSSLSocket``. """ context = ssl.create_default_context() - context.sslsocket_class = _WriteFirstSSLSocket + if issubclass(_WriteFirstSSLSocket, context.sslsocket_class): + context.sslsocket_class = _WriteFirstSSLSocket return context diff --git a/tests/unit/test_streaming_tls.py b/tests/unit/test_streaming_tls.py index dee94713..6a5b452d 100644 --- a/tests/unit/test_streaming_tls.py +++ b/tests/unit/test_streaming_tls.py @@ -155,6 +155,10 @@ def _connect_tls(context: ssl.SSLContext, port: int) -> ssl.SSLSocket: ) +class _OtherSSLSocket(ssl.SSLSocket): + """A socket class that ``_WriteFirstSSLSocket`` does not extend.""" + + def test_tls_context_keeps_verification_on(): context = _tls._create_tls_context() @@ -163,6 +167,21 @@ def test_tls_context_keeps_verification_on(): assert context.check_hostname is True +def test_tls_context_keeps_socket_class_it_does_not_extend( + monkeypatch: pytest.MonkeyPatch, +): + # Given: default contexts that build sockets of another class + monkeypatch.setattr(ssl.SSLContext, "sslsocket_class", _OtherSSLSocket) + + # When: creating the SDK context + context = _tls._create_tls_context() + + # Then: the context keeps that class and its verification settings + assert context.sslsocket_class is _OtherSSLSocket + assert context.verify_mode == ssl.CERT_REQUIRED + assert context.check_hostname is True + + def test_tls_context_rejects_untrusted_certificate( server_context: ssl.SSLContext, ):