fix(simplismart): close the streams the STT collects

livekit-plugins/livekit-plugins-simplismart/livekit/plugins/simplismart/stt.py

@@ -23,7 +23,6 @@
 import json
 import os
 import traceback
-import weakref
 from typing import Any, Literal
 from urllib.parse import urlparse
 
@@ -184,7 +183,10 @@ def __init__(
         )
         self._base_url = base_url
         self._session = http_session
-        self._streams = weakref.WeakSet[SpeechStream]()
+        # Strong ownership: a stream that finishes and gets garbage collected would
+        # otherwise take its still-open per-stream aiohttp.ClientSession with it.
+        # Streams discard themselves in SpeechStream.aclose() once closed.
+        self._streams: set[SpeechStream] = set()
 
     @property
     def provider(self) -> str:
@@ -304,6 +306,22 @@ def stream(
         self._streams.add(stream)
         return stream
 
+    async def aclose(self) -> None:
+        """Close every stream this instance created.
+
+        ``stream()`` gives each ``SpeechStream`` its own ``aiohttp.ClientSession``,
+        and only ``SpeechStream.aclose()`` closes it. Streams are owned strongly
+        (see ``self._streams``) so a finished, garbage-collected stream cannot take
+        its open session with it; closed streams discard themselves, so the set
+        never grows without bound.
+        """
+        closed = list(self._streams)
+        for stream in closed:
+            await stream.aclose()
+        # Remove only what this call closed: a stream created concurrently (or one
+        # whose aclose was cancelled mid-cleanup) stays tracked for a later close.
+        self._streams.difference_update(closed)
+
 
 class SpeechStream(stt.SpeechStream):
     """Simplismart streaming speech-to-text implementation."""
@@ -439,6 +457,10 @@ async def aclose(self) -> None:
         await super().aclose()
         if self._session and not self._session.closed:
             await self._session.close()
+        # Release the STT's strong ownership only after session cleanup succeeds.
+        stt = self._stt
+        if self._session.closed and isinstance(stt, STT):
+            stt._streams.discard(self)
 
     async def _send_initial_config(self, ws: aiohttp.ClientWebSocketResponse) -> None:
         """Send initial configuration message with language for Simplismart models."""

tests/test_plugin_simplismart_stt.py

@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+import asyncio
+import gc
+
+import pytest
+
+from livekit.plugins.simplismart import stt as simplismart_stt
+
+pytestmark = pytest.mark.unit
+
+
+async def _idle_run(self: object) -> None:
+    del self
+    await asyncio.Event().wait()  # cancelled by aclose()
+
+
+@pytest.mark.asyncio
+async def test_simplismart_stt_aclose_closes_tracked_stream_sessions(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(simplismart_stt.SpeechStream, "_run", _idle_run)
+
+    stt = simplismart_stt.STT(api_key="sk_test")
+    stream_a = stt.stream()
+    stream_b = stt.stream()
+    sessions = [stream_a._session, stream_b._session]
+
+    assert all(not s.closed for s in sessions)
+    assert len(stt._streams) == 2
+
+    await stt.aclose()
+
+    assert all(s.closed for s in sessions), (
+        "STT.aclose() must close every per-stream aiohttp session"
+    )
+    assert len(stt._streams) == 0
+
+
+@pytest.mark.asyncio
+async def test_simplismart_stt_async_context_closes_stream_sessions(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(simplismart_stt.SpeechStream, "_run", _idle_run)
+
+    async with simplismart_stt.STT(api_key="sk_test") as stt:
+        stream = stt.stream()
+        session = stream._session
+        assert not session.closed
+
+    assert session.closed, "exiting `async with` must close per-stream sessions"
+
+
+@pytest.mark.asyncio
+async def test_simplismart_stt_aclose_tolerates_already_closed_streams(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(simplismart_stt.SpeechStream, "_run", _idle_run)
+
+    stt = simplismart_stt.STT(api_key="sk_test")
+    stream = stt.stream()
+    session = stream._session
+
+    await stream.aclose()
+    await stt.aclose()  # must not raise on an already-closed stream
+
+    assert session.closed
+
+
+@pytest.mark.asyncio
+async def test_simplismart_stt_aclose_closes_session_of_dropped_stream(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A finished stream the caller dropped must not take its session to the GC.
+
+    With weak tracking, a completed SpeechStream can be collected before
+    STT.aclose() snapshots the set, leaving its per-stream ClientSession
+    unreachable and unclosed. Strong ownership keeps it reachable until
+    aclose closes it.
+    """
+
+    async def _immediate_run(self: object) -> None:
+        del self  # return immediately: the stream task finishes on its own
+
+    monkeypatch.setattr(simplismart_stt.SpeechStream, "_run", _immediate_run)
+
+    stt = simplismart_stt.STT(api_key="sk_test")
+    stream = stt.stream()
+    session = stream._session
+
+    del stream
+    gc.collect()
+
+    assert len(stt._streams) == 1, "a dropped stream must stay tracked until aclose"
+
+    await stt.aclose()
+
+    assert session.closed, "aclose() must close the session of a dropped stream"
+    assert len(stt._streams) == 0