Skip to content

fix(llminternal): close the live connection when SendHistory fails - #1261

Merged
baptmont merged 2 commits into
google:mainfrom
baptmont:fix/runlive-sendhistory-cleanup
Aug 7, 2026
Merged

baptmont merged 2 commits into
google:mainfrom
baptmont:fix/runlive-sendhistory-cleanup

Conversation

@baptmont

@baptmont baptmont commented Aug 6, 2026

Copy link
Copy Markdown
Contributor

Please ensure you have read the contribution guide before creating a pull request.

Link to Issue or Description of Change

1. Link to an existing issue (if applicable):

Problem:

Flow.RunLive's SendHistory failure path is the only return in the function's own body that releases neither connCtx nor liveConn:

if err := liveConn.SendHistory(ctx, nreq.Contents); err != nil {
    log.Printf("failed to send history: %v\n", err)
    sess.pushError(err)
    return          // no cleanup()
}

Auditing all 11 returns in the body (excluding the two producer goroutines, where returning without cleanup is correct), this is the only one that releases nothing.

Cancelling the context would not be enough. genai dials the live socket with websocket.DefaultDialer.Dial (genai/live.go:123), which takes no context.Context, and nothing in the SDK watches one — Session.Close() closing the underlying conn is the only thing that releases it. So this early return strands the connection for the life of the process, not merely until the invocation ends. go vet's lostcancel misses it because cancelConn is reached through the cleanup closure.

Solution:

Call cleanup() before returning, matching every other return in the body. One line, plus a comment recording why cancelling the context is not a substitute.

Testing Plan

Unit Tests:

  • I have added or updated unit tests for my change.
  • All unit tests pass locally.

TestRunLiveClosesConnectionOnSendHistoryFailure drives a real websocket through the fake Live server already in base_flow_live_test.go and asserts the server observes the client's close.

One design note: the send has to fail client-side. A broken socket would also fail it, but would leave no connection to observe. The test feeds SendHistory a turn carrying NaN, which json.Marshal rejects inside the SDK's deepMarshal before it touches the wire — so the send fails while the websocket stays healthy, and the test can distinguish "client closed it" from "client abandoned it" via a server-side read deadline.

Written before the fix and confirmed to fail on unmodified main:

--- FAIL: TestRunLiveClosesConnectionOnSendHistoryFailure (3.03s)
    base_flow_live_test.go:480: RunLive returned without closing the live connection: the socket
        is leaked for the life of the process (missing cleanup on the SendHistory path)

With the fix, and a re-audit showing 0 remaining unreleased returns (was 1 of 11):

go build -mod=readonly ./...                                              # pass
go test -race -mod=readonly -count=10 -run TestRunLive ./internal/llminternal/   # pass
go test -race -mod=readonly -count=1 -shuffle=on ./internal/llminternal/         # pass
go test -race -mod=readonly -count=1 ./...                                # pass, full repo
golangci-lint run ./internal/llminternal/...                              # 0 issues
gofmt -l internal/llminternal/                                            # clean
go mod tidy -diff                                                         # clean, no new deps

Manual End-to-End (E2E) Tests:

Not run against a live Gemini endpoint — the failure is a local connection-lifecycle bug and the regression test exercises it over a real websocket, including the pre-fix failure. Reproducing it against the real service would require inducing a server-side history-send rejection, which the fake server models directly.

Checklist

  • I have read the CONTRIBUTING.md document.
  • I have performed a self-review of my own code.
  • I have commented my code, particularly in hard-to-understand areas.
  • I have added tests that prove my fix is effective or that my feature works.
  • New and existing unit tests pass locally with my changes.
  • I have manually tested my changes end-to-end. — see note above
  • Any dependent changes have been merged and published in downstream modules. — branched off main after fix(llminternal): guard RunLive errChan sends so abandoned producers exit on reconnect #1167 landed

Additional context

Credit to @hanorik, who spotted this reviewing #1167 and flagged it as outside that diff.

RunLive's SendHistory failure path was the one return in the function's own
body that released neither connCtx nor liveConn, so the websocket outlived
the invocation.

Cancelling the context would not have been enough: genai dials the live
socket with websocket.DefaultDialer.Dial, which takes no context, and
nothing in the SDK watches one -- Session.Close closing the underlying conn
is the only thing that releases it. So an early return here stranded the
connection for the life of the process, not merely until the invocation
ended. go vet's lostcancel does not catch it because cancelConn is reached
through the cleanup closure.

Call cleanup() before returning, matching every other return in the body.

The regression test drives a real websocket through the existing fake Live
server and asserts the server observes the client's close. To keep the
observation meaningful the send has to fail client-side -- a broken socket
would fail it too, but leave nothing to observe -- so the test feeds
SendHistory a turn carrying NaN, which json.Marshal rejects inside the SDK
while the connection stays healthy. On the unfixed code the server sits
until its read deadline expires and the test fails.
@baptmont
baptmont requested a review from hanorik August 6, 2026 12:29
@baptmont
baptmont merged commit 73c71de into google:main Aug 7, 2026
13 checks passed
wolo-lab added a commit that referenced this pull request Aug 10, 2026
…1172, #1167, #1261) (#1279)

* fix(memory): hold read lock while scanning sessions in SearchMemory (#1139)

InMemoryService.SearchMemory read the per-user session map under the read
lock but released it before iterating, while AddSessionToMemory writes into
that same map under the write lock. Concurrent Search and Add on the same
app/user therefore race and can panic with "concurrent map iteration and
map write", even though the service is documented as thread-safe.

Hold the read lock across the whole scan. AddSessionToMemory replaces session
slices wholesale rather than mutating them in place, so a read lock over the
iteration is sufficient and still lets searches run concurrently.

Co-authored-by: João Westerberg <[email protected]>

(cherry picked from commit 0420029)

* fix(gcsartifact): prevent artifact version race in concurrent Save (#1172)

* fix(gcsartifact): prevent artifact version race in concurrent Save

Save assigned versions non-atomically: it listed the existing versions,
computed max+1, and wrote that blob. Two concurrent Save calls for the same
artifact would read the same version set, pick the same next version, and the
second write would silently overwrite the first — duplicate returned versions
and a lost object. The code carried a TODO acknowledging this.

Assign versions with optimistic concurrency instead: write the new blob with a
"does not exist" generation precondition (x-goog-if-generation-match: 0). GCS
has no transaction spanning the list and the write, so a colliding writer now
fails the precondition (HTTP 412); Save re-reads the versions and retries with
a fresh number, using a jittered backoff between attempts and bounded by
maxSaveAttempts. It returns errVersionConflict if the budget is exhausted
rather than overwriting an existing version. N concurrent saves therefore yield
exactly versions 1..N with no lost writes.

Testing Plan:
- go build -mod=readonly ./...
- go test -race -mod=readonly -count=1 -shuffle=on ./artifact/...
- golangci-lint run ./artifact/...
- go mod tidy -diff
- TestGCSArtifactServiceConcurrentSave: concurrent savers assert distinct,
  contiguous versions and that nothing is overwritten (saver count kept within
  the retry budget). Verified it fails against the pre-fix code.
- TestSaveReturnsWriteErrorWithoutRetry: a non-precondition write error is
  surfaced as-is and not retried.
- TestSaveExhaustsRetriesOnPersistentConflict: persistent 412s give up with
  errVersionConflict after maxSaveAttempts.
- TestBackoffDelayBounds: jittered backoff stays within bounds.

* test(gcsartifact): cover the real precondition and zero-byte artifacts

The does-not-exist precondition on the production ObjectHandle wrapper had
no test: replacing it with a no-op left the whole suite green, so a
regression there would silently restore last-writer-wins. Add an
httptest-backed test that drives the real wrapper offline and asserts the
upload carries ifGenerationMatch=0 (the client sends the condition as a
query parameter, not a header).

The fake also used data != nil as its existence test, so a zero-byte
artifact -- reachable through the public API, since SaveRequest.Validate
only requires Part.InlineData -- was modelled as not existing: saving it
twice returned version 1 both times. Track existence in an explicit
exists flag instead, which also subsumes the deleted flag.

Also wrap the context error from the retry backoff, drop the unreachable
nil check on the sleep hook, and note the adk-python parity gap and the
per-attempt listing cost in the Save doc.

* docs(gcsartifact): drop the cross-implementation claim from Save

The caveat about other writers is a property of the scheme, so state it in
those terms rather than pinning it to another implementation's current
behaviour, which this repo cannot keep in sync.

(cherry picked from commit a90b910)

* fix(llminternal): guard RunLive errChan sends so abandoned producers exit on reconnect (#1167)

* fix(llminternal): guard RunLive errChan sends so abandoned producers exit on reconnect

RunLive's reader and sender goroutines did bare sends on the unbuffered,
per-attempt errChan. Once the consumer loop stops reading that channel --
after a resumable-error reconnect, a parent context cancellation, or any
other teardown -- cleanup() closes the connection, the producer's in-flight
call fails, and its send blocks forever, leaking one goroutine (plus its
captured connection state) per occurrence. Long-lived /run_live servers
accumulate these across sessions and reconnects.

Guard all three sends with the same select-on-connCtx idiom the function
already uses for eventsChan. cleanup() always cancels connCtx before the
next attempt (and before returning), so a guarded send is guaranteed to
unblock. connCtx, not ctx: ctx outlives reconnects and would not unblock
an abandoned producer.

The regression test fakes the Gemini Live wire protocol over a local
websocket server and fails deterministically on the unfixed code with
goroutines parked in 'chan send' at the bare-send sites.

Fixes #1152

* test(llminternal): cover the SendRealtime guard and stop the leak detector failing open

Addresses review feedback on the RunLive errChan leak fix.

runLiveStacks sized its buffer once at 1 MiB and never checked for
truncation. runtime.Stack truncates silently, so a dump larger than the
buffer would under-count parked goroutines and let a real leak pass --
the wrong failure direction for a leak detector. Grow until the whole
dump fits.

Add a case driving RealtimeInput across a connection loss. Reverting only
the SendRealtime guard left every existing case green, so that third
errChan send site was uncovered; nothing in the suite populated
LiveRequest.RealtimeInput. Each of the three guards is now caught
individually when reverted.

Rename "non-resumable error terminates" and record why it passes with or
without the guards: the sender exits through the existing connCtx.Done()
arm, so nothing is stranded. It pins terminal-path behaviour rather than
guarding the leak regression.

---------

Co-authored-by: westerberg <[email protected]>

(cherry picked from commit c108489)

* fix(llminternal): close the live connection when SendHistory fails (#1261)

RunLive's SendHistory failure path was the one return in the function's own
body that released neither connCtx nor liveConn, so the websocket outlived
the invocation.

Cancelling the context would not have been enough: genai dials the live
socket with websocket.DefaultDialer.Dial, which takes no context, and
nothing in the SDK watches one -- Session.Close closing the underlying conn
is the only thing that releases it. So an early return here stranded the
connection for the life of the process, not merely until the invocation
ended. go vet's lostcancel does not catch it because cancelConn is reached
through the cleanup closure.

Call cleanup() before returning, matching every other return in the body.

The regression test drives a real websocket through the existing fake Live
server and asserts the server observes the client's close. To keep the
observation meaningful the send has to fail client-side -- a broken socket
would fail it too, but leave nothing to observe -- so the test feeds
SendHistory a turn carrying NaN, which json.Marshal rejects inside the SDK
while the connection stays healthy. On the unfixed code the server sits
until its read deadline expires and the test fails.

(cherry picked from commit 73c71de)

---------

Co-authored-by: Madan kumar <[email protected]>
Co-authored-by: chethanuk <[email protected]>
Co-authored-by: João Westerberg <[email protected]>
kdroste-google pushed a commit to kdroste-google/adk-go that referenced this pull request Sep 29, 2026
…oogle#1261)

RunLive's SendHistory failure path was the one return in the function's own
body that released neither connCtx nor liveConn, so the websocket outlived
the invocation.

Cancelling the context would not have been enough: genai dials the live
socket with websocket.DefaultDialer.Dial, which takes no context, and
nothing in the SDK watches one -- Session.Close closing the underlying conn
is the only thing that releases it. So an early return here stranded the
connection for the life of the process, not merely until the invocation
ended. go vet's lostcancel does not catch it because cancelConn is reached
through the cleanup closure.

Call cleanup() before returning, matching every other return in the body.

The regression test drives a real websocket through the existing fake Live
server and asserts the server observes the client's close. To keep the
observation meaningful the send has to fail client-side -- a broken socket
would fail it too, but leave nothing to observe -- so the test feeds
SendHistory a turn carrying NaN, which json.Marshal rejects inside the SDK
while the connection stays healthy. On the unfixed code the server sits
until its read deadline expires and the test fails.
li-xuankun312 added a commit to li-xuankun312/einsteinOS that referenced this pull request Sep 29, 2026
…oogle#1261)

RunLive's SendHistory failure path was the one return in the function's own
body that released neither connCtx nor liveConn, so the websocket outlived
the invocation.

Cancelling the context would not have been enough: genai dials the live
socket with websocket.DefaultDialer.Dial, which takes no context, and
nothing in the SDK watches one -- Session.Close closing the underlying conn
is the only thing that releases it. So an early return here stranded the
connection for the life of the process, not merely until the invocation
ended. go vet's lostcancel does not catch it because cancelConn is reached
through the cleanup closure.

Call cleanup() before returning, matching every other return in the body.

The regression test drives a real websocket through the existing fake Live
server and asserts the server observes the client's close. To keep the
observation meaningful the send has to fail client-side -- a broken socket
would fail it too, but leave nothing to observe -- so the test feeds
SendHistory a turn carrying NaN, which json.Marshal rejects inside the SDK
while the connection stays healthy. On the unfixed code the server sits
until its read deadline expires and the test fails.
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.

Flow.RunLive leaks the live websocket when SendHistory fails: early return skips cleanup

2 participants