fix(agents): don't close in-use connections when the pool is invalidated - #2437
Open
rosetta-livekit-bot[bot] wants to merge 1 commit into
Open
fix(agents): don't close in-use connections when the pool is invalidated#2437rosetta-livekit-bot[bot] wants to merge 1 commit into
rosetta-livekit-bot[bot] wants to merge 1 commit into
Conversation
🦋 Changeset detectedLatest commit: 25d8ae4 The changes in this PR will be included in the next version bump. This PR includes changesets to release 38 packages
Not sure what this means? Click here to learn what changesets are. Click here if you're a maintainer who wants to add another changeset to this PR |
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
Comment on lines
+119
to
+125
| if (invalidations === this.invalidations) { | ||
| break; | ||
| } | ||
|
|
||
| // The options changed during the handshake, so this connection is stale. | ||
| this.toClose.add(connection); | ||
| await this._drainToClose(); |
Contributor
There was a problem hiding this comment.
🔴 Shutdown resurrects pooled connections
Shutdown during a normal handshake makes _connect create a replacement after close() returns. close does not wait for ordinary acquisitions. The replacement remains open and can serve work after shutdown.
Prompt for agents
Make ConnectionPool shutdown terminal and coordinate it with ordinary get() handshakes. In agents/src/connection_pool.ts, close() currently aborts only prewarm work and does not acquire connectLock or track normal _connect calls. If close() invalidates during get(), _connect observes the generation change, closes the stale result, starts another connectCb attempt, and registers that replacement after close() has already returned. Add lifecycle state and synchronization so close() waits for or cancels all acquisitions, prevents retry or registration once closing begins, and prevents later get()/prewarm() calls from reopening the pool. Add a deterministic test with a deferred normal get() handshake, close() during that handshake, and assertions that shutdown does not return while an owned connection can still appear.
Was this helpful? React with 👍 or 👎 to provide feedback.
Comment on lines
+140
to
+148
| } catch (error) { | ||
| if (error instanceof Error && error.name === 'AbortError') { | ||
| // Keep ownership of a connection whose close was cancelled so it can be retried. | ||
| this.toClose.add(conn); | ||
| throw error; | ||
| } | ||
| if (loggerOptions()) { | ||
| log().warn({ exceptionType: safeErrorType(error) }, 'error closing connection'); | ||
| } |
Contributor
There was a problem hiding this comment.
🔴 Failed closures lose connection ownership
When closeCb rejects without AbortError, _drainToClose logs the failure and drops the connection from toClose. The background remove path does the same. Later drains cannot retry it, so shutdown can leave live sockets untracked.
Prompt for agents
Preserve ownership when closeCb fails in agents/src/connection_pool.ts. Both _drainToClose() and remove() currently retain queued connections only for AbortError and permanently delete them for every other rejection. A failed close does not prove the resource closed. Keep failed connections queued for a later drain and make close() surface incomplete shutdown rather than resolve successfully. Avoid an immediate infinite retry loop by stopping the current drain after restoring the failed connection. Add tests for a close callback that fails once and succeeds on a later get()/close() drain, including the remove() background path.
Was this helpful? React with 👍 or 👎 to provide feedback.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Ports livekit/agents#7139 to the JavaScript connection pool.
Invalidation now retires checked-out connections until they are returned, closes idle connections promptly, retries handshakes invalidated in flight, and preserves queued connections when closing is aborted. Prewarm cancellation is adapted to JavaScript promise semantics so shutdown cannot resurrect the pool and late handshake results are closed.
Source diff coverage
File-by-file classification
livekit-agents/livekit/agents/utils/connection_pool.pyagents/src/connection_pool.ts: ported retired connections, invalidation generation tracking, stale-handshake retry/drain, cancellation-safe close draining, return/remove retirement, idle-only invalidation, and final retired cleanup. Adapted Python task cancellation to an abort race that closes late JavaScript promise results, while preserving prompt background closure for removed Node WebSockets.tests/test_connection_pool.pyagents/src/connection_pool.test.ts: ported all nine added tests to Vitest, using deferred promises forasyncio.EventandAbortErrorrejection for cancellation.No source file was omitted.
Validation
pnpm test agents(151 files passed; 2,499 tests passed, 5 skipped)pnpm build(39/39 workspace packages passed)pnpm --filter @livekit/agents typecheckpnpm lintwas also run, but the workspace baseline currently fails on unrelated existing Prettier violations outside this diff, beginning withplugins/sarvam/src/models.ts; the changed files pass lint.Includes a patch changeset for
@livekit/agents.Source: livekit/agents#7139
Ported from livekit/agents#7139
Original PR description
ConnectionPool.invalidate()is the idiom for "these settings are baked into the handshake, so pooled sockets are now stale". Eleven call sites use it fromupdate_options:tts.py:263tts.py:204,tts_v2.py:235tts.py:305stt.py:555,stt.py:560,stt.py:742stt.py:187stt.py:558tts.py:412tts.py:229Two things went wrong for all of them.
1. It closed connections that were still in use
_connectionsholds every connection, not just idle ones, soinvalidate()swept checked-out sockets into_to_closealong with the rest. The nextget()drains that queue — closing a socket another task was still streaming on. Change a voice mid-utterance and the utterance dies as soon as a second stream acquires a connection.Invalidation now separates the two cases:
put()orremove().put()already refused to re-pool them, sinceinvalidate()clears_connections, so a retired connection was never actually reused — only closed at the wrong moment. Same connections closed, later.aclose()closes retired ones too rather than waiting for holders that may never return them.2. A handshake in flight escaped invalidation entirely
invalidate()can only act on what is already in_connections, and_connect()registers a socket only afterawait connect_cb(...)returns. A handshake that was in flight when the options changed was registered afterwards and stayed eligible for reuse, carrying the old voice, model or language._connect()now records_invalidationsbefore the handshake and, if that count moved while it was in flight, closes the socket it produced and connects again. A caller never receives a connection negotiated with settings that have already changed. The discarded socket is drained inside the retry rather than left queued —get()drains before calling_connect(), so nothing else would reach it until the next acquisition.The loop is deliberately unbounded: per review, a cap would either hand back the stale socket or raise for the caller's own option change, and spinning requires
update_optionsto be called faster than a handshake completes, sustained.Also fixed here: a cancelled close lost the connection
_drain_to_close()pops a connection before awaiting its close callback, andCancelledErroris aBaseException, so it skipped theexcept Exceptionhandler — leaving the connection in neither_to_closenor_connections, owned by nothing, open for good. It is now put back before unwinding, so a later drain oraclose()closes it.This one predates the branch —
get()andaclose()have always drained this way. It is fixed here because the retry path above adds a third call site, so the change depends on that function being cancellation-safe. Happy to split it out if you would rather keep this PR to invalidation alone.Tests
Nine added to
tests/test_connection_pool.py. Mutation-checked in four independent pieces — each mutation fails only its own tests:invalidate()test_invalidate_does_not_close_a_connection_still_in_use,test_invalidate_mid_stream_lets_the_stream_finish_then_reconnectsself._invalidations += 1bumptest_invalidate_during_a_handshake_discards_the_stale_connection,test_prewarm_discards_a_connection_invalidated_mid_handshaketest_invalidate_during_a_handshake_discards_the_stale_connectionexcept BaseExceptionrequeuetest_cancelling_a_drain_leaves_the_connection_queued_for_a_later_closeThe rest pin what the change must not break: idle connections still close promptly, a retired connection closes on
put(), onremove()(the error pathconnection()takes), and viaaclose()if it is never returned. Without those, the easy way to "fix" the first bug is to leak the socket instead of closing it late.16 pass with the change. I also ran the full
-m unitsuite before and after: an identical set of tests fails either way, all from plugins that are not installed in my environment, so nothing here regressed.Related
Found while checking bot feedback on #7132 and #7133, which add
invalidate()calls to two more plugins, and flagged again on #7140. Those PRs follow the existing idiom and are unaffected by this change either way — this makes the idiom correct for all of them at once.