fix(google): replay Gemini realtime tool result across session restart - #6481
fix(google): replay Gemini realtime tool result across session restart#6481ByteMaster-1 wants to merge 1 commit into
Conversation
b16c719 to
30dea9c
Compare
Routine + handled on our deployment: post-completion re-emissions are dropped by design (livekit#6431/livekit#6481) and the worker's capture wrap folds the text into the stored transcript before this handler runs. The WARNING paged Rollbar on essentially every Gemini call (cai items livekit#1543/livekit#1544, livekit#1606/livekit#1607).
When update_tools() restarts the Gemini Live websocket mid-turn, a function-tool result produced in that window was sent on the closing session and never replayed, so the model never received it and the turn hung. Buffer the tool result while the socket is restarting and replay it via send_tool_response once the new session is established, then nudge the model to continue the turn. Also decouple generation-completion handling from the model name so newer Live preview models work: track generation state explicitly via _generation_completed, drop trailing model_turn frames that arrive with no active generation, and allow generate_reply() for models without mutable_chat_context (gated on the capability, not a model-name literal). Fixes livekit#6479 Co-Authored-By: heer <heer434@users.noreply.github.com>
30dea9c to
87e3efe
Compare
| if self._session_should_close.is_set(): | ||
| # The socket is tearing down (e.g. update_tools() mid-turn). Sending the | ||
| # tool result now would deliver it to the dying session and it would never | ||
| # reach the model, hanging the turn. Stash it so _main_task can replay it | ||
| # once the new session is established. | ||
| logger.debug( | ||
| "session restarting; buffering tool result to replay after reconnect" | ||
| ) | ||
| self._pending_tool_result = tool_results | ||
| else: | ||
| self._send_client_event(tool_results) |
There was a problem hiding this comment.
π΄ Tool results are still lost when they finish slightly later during a session restart
The tool result is only saved for replay when the old connection object still exists (self._session_should_close.is_set() check at livekit-plugins/livekit-plugins-google/livekit/plugins/google/realtime/realtime_api.py:684), but the whole update is abandoned earlier whenever the old connection is already gone, so results that finish while reconnecting are still thrown away and the turn still hangs.
Impact: In the exact scenario this change targets (a tool finishing while the connection is being restarted), the assistant can still get stuck waiting forever and never answer the user.
Early return in update_chat_ctx bypasses the new buffering window
update_chat_ctx returns early when self._active_session is falsy (livekit-plugins/livekit-plugins-google/livekit/plugins/google/realtime/realtime_api.py:652-655): it just assigns self._chat_ctx = chat_ctx and returns, so tool_results is never computed nor buffered.
During a restart, _main_task exits the async with block and its finally calls _close_active_session(), which sets self._active_session = None (realtime_api.py:531-540), and it stays None for the whole duration of _close_active_session() + live.connect(). Any tool result completing in that window (the repro in examples/other/gemini_realtime_tool_update_repro.py sleeps 0.5s after update_tools(), so it almost certainly lands there) hits the early return.
Because self._chat_ctx is nonetheless updated, the next diff will not re-emit the item, and the connect-time replay of self._chat_ctx uses exclude_function_call=True, which drops both function_call and function_call_output items (livekit-agents/livekit/agents/llm/chat_context.py:497-501). The result is therefore never delivered to any session.
The buffering should also happen when there is no active session (i.e. compute the tool results and stash them before/instead of the early return).
Prompt for agents
In RealtimeSession.update_chat_ctx (livekit-plugins/livekit-plugins-google/livekit/plugins/google/realtime/realtime_api.py), the new buffering of tool results for replay after a reconnect is only reachable when self._active_session is still set. However, during a restart _main_task closes the active session (_close_active_session sets _active_session = None) and keeps it None while it reconnects, and update_chat_ctx early-returns in that state (it only assigns self._chat_ctx and returns). Since self._chat_ctx is updated, the item will not appear in a later diff, and the connect-time replay of chat context excludes function_call/function_call_output items, so the tool result is dropped permanently and the turn hangs β the very scenario this PR fixes. Rework the flow so that when a restart is pending or no active session exists, the tool responses for newly created items are still computed (get_tool_results_for_realtime) and stashed into _pending_tool_result for replay, rather than returning early. Consider also accumulating (rather than overwriting) multiple buffered tool responses.
Was this helpful? React with π or π to provide feedback.
| # Gemini only generates after a user turn. Mutable-context models accept | ||
| # an appended client-content turn; the live-preview family ignores | ||
| # appended turns until the next session, so nudge those with a realtime | ||
| # text input instead (mirrors how generate_reply() nudges generation). | ||
| if self._realtime_model.capabilities.mutable_chat_context: | ||
| await session.send_client_content( | ||
| turns=[types.Content(parts=[types.Part(text=".")], role="user")], | ||
| turn_complete=True, | ||
| ) | ||
| else: | ||
| await session.send_realtime_input(text=".") |
There was a problem hiding this comment.
π‘ Replayed tool result is followed by a stray dummy user message that can trigger a duplicate reply
After the saved tool result is re-sent, an artificial user message containing "." is also sent and marked as a completed user turn (send_client_content(..., turn_complete=True) at livekit-plugins/livekit-plugins-google/livekit/plugins/google/realtime/realtime_api.py:948-951), even though the model already answers on its own after a tool result, so the user can hear a second, unrelated reply.
Impact: After a reconnect mid tool call, the assistant may speak twice or cut off its own answer, and a bogus "." turn is injected into the conversation.
Gemini Live auto-generates after a function response
The plugin itself declares auto_tool_reply_generation=True (realtime_api.py:307), i.e. Gemini Live starts generating a reply as soon as it receives a function response β the normal (non-restart) path at realtime_api.py:694 sends only send_tool_response(...) with no nudge and the model replies. Adding a turn_complete=True user turn (or send_realtime_input(text=".") for the non-mutable branch at realtime_api.py:953) therefore adds a second trigger, which the server treats as a new user turn β it can interrupt the in-flight post-tool reply and/or produce an extra response to ".".
Prompt for agents
In _main_task's replay of _pending_tool_result (livekit-plugins/livekit-plugins-google/livekit/plugins/google/realtime/realtime_api.py), after send_tool_response the code additionally sends a dummy "." user turn with turn_complete=True (or a realtime text input for models without mutable_chat_context). Gemini Live already generates a reply automatically upon receiving a function response β the plugin advertises auto_tool_reply_generation=True and the normal non-restart path sends only the tool response. The extra nudge risks interrupting the in-flight post-tool reply or producing a second, spurious reply to ".". Consider dropping the nudge (relying on the auto tool reply), or only nudging when the model is verified not to auto-continue after a replayed tool response.
Was this helpful? React with π or π to provide feedback.
| logger.debug( | ||
| "session restarting; buffering tool result to replay after reconnect" | ||
| ) | ||
| self._pending_tool_result = tool_results |
There was a problem hiding this comment.
π‘ Only the most recent tool result survives a reconnect when several are produced
The saved-for-later tool result is stored in a single slot that is overwritten by the next one (self._pending_tool_result = tool_results at livekit-plugins/livekit-plugins-google/livekit/plugins/google/realtime/realtime_api.py:692), so if more than one result is produced while the connection is restarting, the earlier ones are silently discarded.
Impact: With several tool calls completing during a reconnect, some answers never reach the model and the conversation can stall or answer with missing information.
Overwrite mechanism
update_chat_ctx is invoked once per tool-execution round (livekit-agents/livekit/agents/voice/agent_activity.py:4041-4044). Because self._chat_ctx is updated at the end of each call (realtime_api.py:698), a subsequent call during the same restart window produces a new, disjoint tool_results that replaces the previously stashed one instead of being appended. Accumulating the function_responses lists (or keeping a list of pending responses) would preserve all of them.
Prompt for agents
_pending_tool_result in the Google realtime plugin is a single-slot buffer that is overwritten on each buffered tool result while the socket is restarting (realtime_api.py, update_chat_ctx). If two tool-result batches are produced during the same restart window, the earlier batch is lost and never replayed, leaving the model waiting for those responses. Change the buffer to accumulate function_responses (e.g. extend an existing LiveClientToolResponse or keep a list) and replay all of them in _main_task after reconnect.
Was this helpful? React with π or π to provide feedback.
When update_tools() restarts the Gemini Live websocket mid-turn, a function-tool result produced in that window was sent on the closing session and never replayed, so the model never received it and the turn hung. Buffer the tool result while the socket is restarting and replay it via send_tool_response once the new session is established, then nudge the model to continue the turn.
Also decouple generation-completion handling from the model name so newer Live preview models work: track generation state explicitly via _generation_completed, drop trailing model_turn frames that arrive with no active generation, and allow generate_reply() for models without mutable_chat_context (gated on the capability, not a model-name literal). Fixes #6479