Skip to content

fix: close AsyncStream client on session end#457

Merged
aliev merged 10 commits intomainfrom
fix/close-stream-client-on-session-end
Mar 29, 2026
Merged

fix: close AsyncStream client on session end#457
aliev merged 10 commits intomainfrom
fix/close-stream-client-on-session-end

Conversation

@aliev
Copy link
Copy Markdown
Member

@aliev aliev commented Mar 27, 2026

Why

After each agent session, HTTP clients and WebSocket connections from STT, TTS, and Edge plugins were never closed. Each session left behind orphaned TCP connections and asyncio tasks. On production GKE, pods accumulated up to 3.3 GiB of non-evictable memory and OOMKilled.

Problem

Agent._close() calls close() on each plugin via _apply("close"), but:

  1. StreamEdge.close() only set self._call = None without closing the AsyncStream httpx client or calling leave() on the ConnectionManager
  2. ElevenLabs TTS had no close() method at all
  3. Deepgram STT close() closed the WebSocket but not the underlying httpx client
  4. Both ElevenLabs and Deepgram referenced self.client._client which doesn't exist in current SDK versions. The hasattr guard silently skipped the close.
  5. Race condition: _real_connection was stored only after __aenter__(), so close() during join() couldn't clean up WebSocket tasks

Changes

  • StreamEdge: close AsyncStream client, call leave() on ConnectionManager, store _real_connection immediately after rtc.join()
  • ElevenLabs TTS: add close() with correct httpx client path (client._client_wrapper.httpx_client.httpx_client)
  • Deepgram STT: close httpx client with correct path
  • Tests: unit test for StreamEdge.close() verifying client is closed, connection is None, and leave() is called. Red-green tests for ElevenLabs and Deepgram verifying httpx client is_closed.

Benchmark

Tested in a local k8s cluster (OrbStack) with 50 create+close sessions (no audio):

Metric Before After Improvement
RSS delta after 10min +150.5 MiB +0 MiB (returns to baseline) -100%
Leaked TCP connections 87 2 -98%
Leaked asyncio Tasks 65 14 (8 from aioice, 5 system) -78%
StreamAPIWS reader/heartbeat 30 0 -100%

Memory now fully returns to baseline after sessions end. Previously it accumulated ~3 MiB/session permanently.

Remaining 2 TCP connections and 8 aioice tasks (TurnClientMixin.refresh, Connection.check_start) are from aiortc's ICE layer, outside our control.

Companion PR

Summary by CodeRabbit

  • New Features

    • Added explicit shutdown handling to TTS to ensure underlying network resources are closed.
  • Bug Fixes

    • Improved resource cleanup for Deepgram, ElevenLabs and GetStream integrations to reliably close HTTP clients and live connections during shutdown.
  • Tests

    • Added tests verifying proper cleanup of underlying HTTP clients and connection release.

StreamEdge.close() only set self._call = None, leaving the
AsyncStream httpx client and its connection pool open. Each
session leaked ~1-2 TCP connections and associated asyncio tasks
(heartbeat, reader), accumulating ~1.7 MiB/session.

Call self.client.aclose() to release HTTP connection pools,
WebSocket connections, and background tasks.
@coderabbitai
Copy link
Copy Markdown

coderabbitai bot commented Mar 27, 2026

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: e0fea4a5-80e7-4481-b453-1e55a417fad7

📥 Commits

Reviewing files that changed from the base of the PR and between 228f65b and 1bfa9e2.

📒 Files selected for processing (2)
  • plugins/deepgram/vision_agents/plugins/deepgram/deepgram_stt.py
  • plugins/elevenlabs/vision_agents/plugins/elevenlabs/tts.py
🚧 Files skipped from review as they are similar to previous changes (1)
  • plugins/elevenlabs/vision_agents/plugins/elevenlabs/tts.py

📝 Walkthrough

Walkthrough

Adds explicit resource cleanup across Deepgram STT, ElevenLabs TTS, and GetStream StreamEdge by closing underlying httpx clients and ensuring real connection teardown; adds unit tests that assert client is_closed state changes and connection leave behavior.

Changes

Cohort / File(s) Summary
Deepgram STT Cleanup
plugins/deepgram/vision_agents/plugins/deepgram/deepgram_stt.py, plugins/deepgram/tests/test_deepgram_stt_close.py
close() now attempts to find and aclose() the nested httpx client (_client_wrapper.httpx_client); new async test verifies httpx_client.is_closed flips from False to True after await stt.close().
ElevenLabs TTS Cleanup
plugins/elevenlabs/vision_agents/plugins/elevenlabs/tts.py, plugins/elevenlabs/tests/test_tts_close.py
Adds async def close(self) to TTS that looks up and aclose()s a nested httpx client (_client_wrapper.httpx_client.httpx_client) and calls super().close() in finally; test verifies underlying client's is_closed state transitions.
GetStream Edge Transport Cleanup
plugins/getstream/vision_agents/plugins/getstream/stream_edge_transport.py, plugins/getstream/tests/test_stream_edge_transport.py
join() assigns self._real_connection immediately after join completes; close() calls leave() on _real_connection, clears it, and awaits self.client.aclose(); tests updated to use AsyncMock and assert client closure, leave() call, and _real_connection cleared.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Poem

The small machines unlace themselves at dusk,
A darkened thread of sockets pulled through my hands—
I watch the quiet death of humming pipes,
The slow, exacting click that seals the bands,
And name the silence: everything is closed.

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 55.56% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (2 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title accurately describes the primary change: fixing resource leaks by properly closing AsyncStream client and HTTP clients across multiple plugins when sessions end.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch fix/close-stream-client-on-session-end

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

aliev added 8 commits March 27, 2026 21:28
ElevenLabs TTS had no close() method at all - the AsyncElevenLabs
httpx client was never closed, leaking TCP connections per session.

Deepgram STT close() closed the WebSocket but not the underlying
AsyncDeepgramClient httpx client, also leaking connections.

Together with the StreamEdge fix, this addresses all three sources
of TCP connection leaks found during memory investigation.
…se()

rtc.join() creates a ConnectionManager with WebSocket tasks (heartbeat,
reader) immediately. If the session is closed before join() completes,
_real_connection was not yet set, so close() couldn't clean up these tasks.

Store the connection reference right after rtc.join() returns, and call
leave() explicitly in StreamEdge.close() to ensure WebSocket tasks are
cancelled even if the Agent's join flow was interrupted.
Checks that after close():
- httpx client is closed (is_closed == True)
- _real_connection is set to None
- leave() is called on the ConnectionManager
… close()

Both plugins referenced self.client._client which doesn't exist in
current SDK versions. The actual path is
client._client_wrapper.httpx_client.httpx_client. The hasattr guard
silently skipped the close, leaving HTTP connections open.

Found by writing tests that check httpx_client.is_closed after close().
Copy link
Copy Markdown

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

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

Actionable comments posted: 4

🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@plugins/deepgram/vision_agents/plugins/deepgram/deepgram_stt.py`:
- Around line 316-325: The nested reflective access to Deepgram SDK internals
(self.client -> _client_wrapper -> httpx_client -> httpx_client) in the cleanup
block should be removed or replaced: do not access private attributes like
"_client_wrapper" or "httpx_client"; instead either delete the manual httpx
client close and rely on the SDK/user lifecycle, or if you must attempt a
workaround, wrap any access in a clear SDK-internal comment and a defensive
try/except that skips cleanup if attributes are missing or change, referencing
the cleanup site where self.client is inspected and the nested
_client_wrapper/httpx_client chain is currently used; also add a short TODO
pointing to checking Deepgram SDK docs for an official shutdown API.

In `@plugins/elevenlabs/vision_agents/plugins/elevenlabs/tts.py`:
- Around line 66-76: TTS.close currently bypasses the base class cleanup and
uses a nested getattr chain to reach an internal httpx client; replace this with
a call to await super().close() and remove the nested
getattr(getattr(getattr(...))) logic so the SDK/base class handles closing the
internal httpx client (update the TTS.close method to simply await
super().close() and remove manual httpx_client access to comply with repo
guidelines and ElevenLabs SDK recommendations).

In `@plugins/getstream/tests/test_stream_edge_transport.py`:
- Line 3: The test currently uses unittest.mock.AsyncMock to stub
_real_connection and assert calls; replace this with a real ConnectionManager
instance from the existing connection_manager fixture or a small hand-rolled
test double class that implements the same interface (e.g., open/close/send
methods) and records invocation counts; update the test to inject that real
ConnectionManager or test double into the StreamEdgeTransport (or the class
under test that sets _real_connection) and replace assert_called_once() with
assertions on the test double's recorded call counters or the real
ConnectionManager's observable state, removing all uses of AsyncMock and Mock.

In `@plugins/getstream/vision_agents/plugins/getstream/stream_edge_transport.py`:
- Around line 500-508: The code currently uses a bare except Exception: when
awaiting self._real_connection.leave(); change this to catch the specific
exceptions handled elsewhere in this module (match the pattern in
StreamConnection.close()), e.g., use except (asyncio.TimeoutError,
RuntimeError): and log an appropriate message with logger.exception("Error
during connection leave") for those cases, and let any other exceptions
propagate (do not catch them broadly) so unexpected errors aren't swallowed;
update the block around await self._real_connection.leave() and keep resetting
self._real_connection = None afterward.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: ee82fea1-cc1a-419f-8d86-d2c3d506d7aa

📥 Commits

Reviewing files that changed from the base of the PR and between edff829 and 228f65b.

📒 Files selected for processing (6)
  • plugins/deepgram/tests/test_deepgram_stt_close.py
  • plugins/deepgram/vision_agents/plugins/deepgram/deepgram_stt.py
  • plugins/elevenlabs/tests/test_tts_close.py
  • plugins/elevenlabs/vision_agents/plugins/elevenlabs/tts.py
  • plugins/getstream/tests/test_stream_edge_transport.py
  • plugins/getstream/vision_agents/plugins/getstream/stream_edge_transport.py

- Add super().close() call in finally block
- Break nested getattr into readable steps in both ElevenLabs and Deepgram
- Add SDK workaround comments
@aliev aliev merged commit 37c479c into main Mar 29, 2026
6 checks passed
@aliev aliev deleted the fix/close-stream-client-on-session-end branch March 29, 2026 12:44
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants