fix(elevenlabs): close TTS websocket on job shutdown

livekit-plugins/livekit-plugins-elevenlabs/livekit/plugins/elevenlabs/tts.py

@@ -34,7 +34,9 @@
     APIError,
     APIStatusError,
     APITimeoutError,
+    JobContext,
     LanguageCode,
+    get_job_context,
     tokenize,
     tts,
     utils,
@@ -211,6 +213,7 @@ def __init__(
 
         self.__current_connection: _Connection | None = None
         self._connection_lock = asyncio.Lock()
+        self._registered_job_ctx_ref: weakref.ReferenceType[JobContext] | None = None
 
     @property
     def model(self) -> str:
@@ -286,6 +289,16 @@ async def _current_connection(self) -> tuple[_Connection, float, bool]:
             Tuple of (connection, acquire_time, connection_reused)
         """
         async with self._connection_lock:
+            job_ctx = get_job_context(required=False)
+            registered_job_ctx = (
+                self._registered_job_ctx_ref() if self._registered_job_ctx_ref is not None else None
+            )
+            if job_ctx is not None and job_ctx is not registered_job_ctx:
+                # AgentSession closes active synthesis streams but does not own the TTS
+                # model. Close our reusable WebSocket before the job-owned HTTP session.
+                job_ctx.add_shutdown_callback(self.aclose)
+                self._registered_job_ctx_ref = weakref.ref(job_ctx)
+
             if (
                 self.__current_connection
                 and self.__current_connection.is_current

tests/test_plugin_elevenlabs_tts.py

@@ -3,16 +3,27 @@
 import asyncio
 import base64
 import json
+import socket
+from collections.abc import Awaitable, Callable
 from types import SimpleNamespace
 
 import aiohttp
 import pytest
+from aiohttp import web
 
 from livekit.plugins.elevenlabs import tts as elevenlabs_tts
 
 pytestmark = pytest.mark.plugin("elevenlabs")
 
 
+class _FakeJobContext:
+    def __init__(self) -> None:
+        self.shutdown_callbacks: list[Callable[[], Awaitable[None]]] = []
+
+    def add_shutdown_callback(self, callback: Callable[[], Awaitable[None]]) -> None:
+        self.shutdown_callbacks.append(callback)
+
+
 class _FakeWebSocket:
     def __init__(self, messages: list[object]) -> None:
         self._messages = messages
@@ -97,6 +108,69 @@ def test_auto_mode_respects_explicit_value_with_chunk_length_schedule() -> None:
     assert tts._opts.auto_mode is True
 
 
+@pytest.mark.asyncio
+async def test_job_shutdown_gracefully_closes_websocket(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    close_codes: asyncio.Queue[int | None] = asyncio.Queue()
+
+    async def websocket_handler(request: web.Request) -> web.WebSocketResponse:
+        websocket = web.WebSocketResponse()
+        await websocket.prepare(request)
+        try:
+            async for _ in websocket:
+                pass
+        finally:
+            close_codes.put_nowait(websocket.close_code)
+        return websocket
+
+    app = web.Application()
+    app.router.add_get("/{path:.*}", websocket_handler)
+    runner = web.AppRunner(app)
+    await runner.setup()
+
+    server_socket = socket.socket()
+    server_socket.bind(("127.0.0.1", 0))
+    port = server_socket.getsockname()[1]
+    site = web.SockSite(runner, server_socket)
+    await site.start()
+
+    job_ctx = _FakeJobContext()
+    monkeypatch.setattr(elevenlabs_tts, "get_job_context", lambda *, required=False: job_ctx)
+
+    try:
+        async with aiohttp.ClientSession() as http_session:
+            tts = elevenlabs_tts.TTS(
+                api_key="test-key",
+                voice_id="test-voice",
+                base_url=f"http://127.0.0.1:{port}",
+                http_session=http_session,
+            )
+            try:
+                await tts._current_connection()  # pyright: ignore[reportPrivateUsage]
+                await tts._current_connection()  # pyright: ignore[reportPrivateUsage]
+
+                assert len(job_ctx.shutdown_callbacks) == 1
+                await job_ctx.shutdown_callbacks[0]()
+
+                assert await asyncio.wait_for(close_codes.get(), timeout=1) == 1000
+                assert not http_session.closed
+
+                job_ctx = _FakeJobContext()
+                await tts._current_connection()  # pyright: ignore[reportPrivateUsage]
+                await tts._current_connection()  # pyright: ignore[reportPrivateUsage]
+
+                assert len(job_ctx.shutdown_callbacks) == 1
+                await job_ctx.shutdown_callbacks[0]()
+
+                assert await asyncio.wait_for(close_codes.get(), timeout=1) == 1000
+                assert not http_session.closed
+            finally:
+                await tts.aclose()
+    finally:
+        await runner.cleanup()
+
+
 def test_build_context_init_packet_includes_generation_config() -> None:
     tts = elevenlabs_tts.TTS(api_key="test-key", chunk_length_schedule=[80, 120], auto_mode=False)
     packet = elevenlabs_tts._build_context_init_packet(  # pyright: ignore[reportPrivateUsage]