Skip to content

fix(sarvam): close the streams the STT collects - #7058

Open
dhruvladia-sarvam wants to merge 4 commits into
livekit:mainfrom
dhruvladia-sarvam:fix/sarvam-stt-aclose-tracked-streams
Open

fix(sarvam): close the streams the STT collects#7058
dhruvladia-sarvam wants to merge 4 commits into
livekit:mainfrom
dhruvladia-sarvam:fix/sarvam-stt-aclose-tracked-streams

Conversation

@dhruvladia-sarvam

@dhruvladia-sarvam dhruvladia-sarvam commented Aug 31, 2026

Copy link
Copy Markdown
Contributor

The leak

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 into self._streams = weakref.WeakSet[SpeechStream](), but nothing ever reads that set, and the STT class defines no aclose() β€” so STT.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.aclose is documented as "Close the STT, and every stream/requests associated with it", and STT.__aexit__ delegates to it, so async with stt: currently leaks the streams. stt_streaming.py in this same package already implements the intended shape at line 466.

The change

Eleven lines on STT, mirroring stt_streaming.py:

async def aclose(self) -> None:
    for stream in list(self._streams):
        await stream.aclose()
    self._streams.clear()

The STT owns no session of its own (non-streaming paths resolve through utils.http_context.http_session()), so the streams are the whole of it.

Validation

  • Regression tests added (tests/test_plugin_sarvam_stt.py, hermetic β€” the streams' _run is stubbed, no network): per-stream sessions are closed by STT.aclose(), the async with path closes them via __aexit__, and aclose() 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 strict mypy pass on the touched plugin file.
  • Reproduced the pre-fix leak independently: aiohttp's Unclosed client session warnings appear once streams are dropped.

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
@dhruvladia-sarvam
dhruvladia-sarvam requested a review from a team as a code owner August 31, 2026 11:38
devin-ai-integration[bot]

This comment was marked as resolved.

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
devin-ai-integration[bot]

This comment was marked as resolved.

…stt.py

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Rehansanjay added a commit to Rehansanjay/agents that referenced this pull request Aug 31, 2026
`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.
@Rehansanjay

Copy link
Copy Markdown
Contributor

The strong-ownership change is the right call and it is a hole in what I had in #7012 β€” a WeakSet drops a finished stream before aclose() can snapshot it, which is exactly the case that leaks in practice. Same for removing only what the call actually closed rather than clearing the set.

I have rescoped #7012 to simplismart so the two do not overlap. It has the identical bug β€” weakref.WeakSet at stt.py:187, a per-stream aiohttp.ClientSession created in stream(), and no aclose() on the STT β€” and it is not covered here. I followed the approach in this PR there, including the self-discard in SpeechStream.aclose(), so the two plugins stay consistent.

@tinalenguyen tinalenguyen left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants