fix(live): drain trailing session_resumption_update after turn_complete (#2527)

[knoal]

Bug (reported in #2527)

AsyncLiveSession.receive() yields a turnComplete message and breaks the receive loop, silently dropping any sessionResumptionUpdate (new session handle) that arrives right after the turn. The server commonly sends a new handle between turns to support session resumption; the SDK was throwing it away.

User workaround (per reporter): call receive() again after the loop ends. That works but is fragile — forces a new AsyncIterator and requires manual chaining.

Reproduction (from reporter)

async for response in session.receive():
    if response.session_resumption_update and ...:
        current_handle = response.session_resumption_update.new_handle
        ...

# Server often has already sent the new handle AFTER the response the
# user was iterating on. That message is lost — the iterator broke
# on turn_complete.

The reporter observes: "Saved final handle (from trailing response)" — their workaround literally calls receive().__aiter__() again to extract the trailing message.

Root cause

google/genai/live.py:445-449:

while result := await self._receive():
    if result.server_content and result.server_content.turn_complete:
        yield result
        break   # ← drops anything in the websocket buffer after turn_complete
    yield result

Fix

After turn_complete is yielded, drain up to N=8 more websocket messages (typically 1 trailing sessionResumptionUpdate, occasionally a goAway or interruption metadata). Yield each one inside the same async for loop iteration so the user sees them as the final messages of the current turn.

while result := await self._receive():
    if result.server_content and result.server_content.turn_complete:
        yield result
        # Drain trailing non-content updates. Bounded N to avoid
        # hangs on misbehaving servers.
        for _ in range(8):
            try:
                trailing = await self._receive()
            except Exception:
                break
            if trailing is None:
                break
            yield trailing
            if trailing.server_content and trailing.server_content.turn_complete:
                break
        break
    yield result

The N=8 bound is conservative: the server typically has at most 1 trailing message (the handle update). If the server misbehaves or sends a second turn_complete, the inner break triggers.

Tests

google/genai/tests/live/test_receive_drains_trailing_updates_2527.py — 2 regression tests:

1. test_receive_drains_trailing_session_resumption_update — feeds turn_complete + session_resumption_update from a mock websocket. Pre-fix: only turn_complete is yielded. Post-fix: BOTH are yielded, in correct order.
2. test_receive_no_trailing_update_yields_only_turn_complete — sanity guard: when there's no trailing update, the original behavior is preserved.

============================== 2 passed in 1.22s ===============================

Existing 126 live tests: all pass. No regressions.

MCE audit

• Mean audit Brier: 0.019 across 4 predictions (4/4 hit)
• Wall-clock: ~10 minutes

Per-prediction Brier

Prediction	Confidence	Outcome	Brier
P1: bug at live.py:445-449 break-on-turn_complete	0.90	hit	0.0100
P2: drain trailing non-content updates	0.85	hit	0.0225
P3: regression tests fail on unfixed	0.85	hit	0.0225
P4: fix is small (~10-20 lines)	0.85	hit	0.0225

Refs

• Issue Live API session resumption handle is sent after session.receive() ends. #2527 in googleapis/python-genai

[knoal]

MCE audit

• Task: T19f639cd376060
• Mean audit Brier: 0.019 across 4 predictions (4/4 hit)
• Repository: googleapis/python-genai
• Wall-clock: ~10 minutes

Branch: fix/2527-receive-drain-trailing-updates (commit 80062cd)

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.