fix(sarvam): close the streams the STT collects - #7058
Conversation
stream() gives each SpeechStream its own aiohttp.ClientSession, and only SpeechStream.aclose() closes it. The STT collected streams into a WeakSet but never read it, and STT.aclose() resolved to the base-class no-op, so the per-stream sessions outlived the STT and aiohttp reported unclosed client sessions once the streams were dropped. Add STT.aclose() to close every tracked stream and clear the set, mirroring the existing implementation in this package's stt_streaming.py and the base contract (\"Close the STT, and every stream/requests associated with it\"). The STT owns no session of its own, so the streams are the whole of it. πͺ· Generated with Sarvam Code
Add three unit tests for the STT.aclose() fix: tracked per-stream aiohttp sessions are closed by STT.aclose(), the async-with path closes them via __aexit__, and aclose tolerates already-closed streams. The streams' _run is stubbed so the tests stay hermetic (no network). Also raise StopAsyncIteration from None in _FakeWS to satisfy the repo-wide ruff bugbear rule now that this file is tracked. πͺ· Generated with Sarvam Code
Review follow-up to the STT.aclose() fix: with weakref.WeakSet tracking, a completed SpeechStream the caller dropped could be garbage-collected before STT.aclose() snapshotted the set. Its per-stream aiohttp.ClientSession then became unreachable and stayed open. Track owned streams in a plain set instead. Streams discard themselves from the set in SpeechStream.aclose() once cleanup completes, so long- lived STT use does not retain closed streams, and STT.aclose() removes only the streams it actually closed (snapshot + difference_update), so one cancelled cleanup cannot strand another stream untracked. The new regression test finishes a stream task, drops the caller's reference, forces gc.collect(), and verifies STT.aclose() still closes the retained session. It fails against the WeakSet implementation with aiohttp's 'Unclosed client session' warning. πͺ· Generated with Sarvam Code
β¦stt.py Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
`STT.stream()` gives every `SpeechStream` its own `aiohttp.ClientSession`
("a fresh session for this stream to avoid conflicts"), and only
`SpeechStream.aclose()` closes it. The STT collects its streams but
defines no `aclose()`, so `STT.aclose()` resolves to the base class
no-op: the streams are never closed, their sessions stay open, and
aiohttp reports unclosed client sessions once the streams are dropped.
The base contract asks for this β `STT.aclose` is documented as "Close
the STT, and every stream/requests associated with it", and
`__aexit__` delegates to it, so `async with stt:` leaks today.
Streams are now owned strongly rather than through a `WeakSet`: a stream
that finishes and is garbage collected would otherwise take its
still-open session with it, before `aclose()` ever sees it. Closed
streams discard themselves from the STT, so the set does not grow
without bound, and `aclose()` removes only what it actually closed so a
concurrently created stream stays tracked.
Regression tests are hermetic β `_run` is stubbed, no network. Three of
the four fail with the fix reverted, with aiohttp's "Unclosed client
session" appearing exactly as the leak predicts.
The sarvam half of this change is superseded by livekit#7058, which fixes the
identical bug in that plugin; this keeps the simplismart half, which
livekit#7058 does not cover. The strong-ownership and self-discard approach here
follows livekit#7058 so the two plugins stay consistent.
|
The strong-ownership change is the right call and it is a hole in what I had in #7012 β a I have rescoped #7012 to simplismart so the two do not overlap. It has the identical bug β |
tinalenguyen
left a comment
There was a problem hiding this comment.
hi, thanks for the pr! could you address this finding:
Nothing in livekit-agents/ calls STT.aclose(), so the new strong set is filled on every stream() and drained only by a detached, cancellable path β pinning dead streams and their open sessions for the process lifetime. That's a worse leak than the WeakSet it replaces.
Fix before merge: gather(..., return_exceptions=True) in aclose(), discard in a finally (not gated on session.closed), a weakref.finalize to close the session on GC, and repair the test branch that tolerates the regression. No API-level breaking changes.
The leak
STT.stream()gives everySpeechStreamits ownaiohttp.ClientSession("a fresh session for this stream to avoid conflicts"), and onlySpeechStream.aclose()closes it. TheSTTcollects its streams intoself._streams = weakref.WeakSet[SpeechStream](), but nothing ever reads that set, and theSTTclass defines noaclose()β soSTT.aclose()resolves to the abstract no-op on the base class. The streams are never closed, their sessions stay open, and aiohttp reports unclosed client sessions once the streams are dropped.The base contract asks for this:
STT.acloseis documented as "Close the STT, and every stream/requests associated with it", andSTT.__aexit__delegates to it, soasync with stt:currently leaks the streams.stt_streaming.pyin this same package already implements the intended shape at line 466.The change
Eleven lines on
STT, mirroringstt_streaming.py:The
STTowns no session of its own (non-streaming paths resolve throughutils.http_context.http_session()), so the streams are the whole of it.Validation
tests/test_plugin_sarvam_stt.py, hermetic β the streams'_runis stubbed, no network): per-stream sessions are closed bySTT.aclose(), theasync withpath closes them via__aexit__, andaclose()tolerates already-closed streams. With the fix reverted, the first two tests fail exactly as the leak predicts (plus the suite's leaked-task detection fires); with the fix, all 16 tests in the file pass.ruff check,ruff format --check, and strictmypypass on the touched plugin file.Unclosed client sessionwarnings appear once streams are dropped.