fix(google): replay Gemini realtime tool result across session restart

examples/other/gemini_realtime_tool_update_repro.py

@@ -0,0 +1,83 @@
+"""Repro for issue #6479: Gemini Live tool result lost when update_tools() restarts the session.
+
+Scenario
+--------
+1. A function tool (`get_weather`) is registered on a Gemini Live realtime agent.
+2. The user asks something that makes the model call the tool.
+3. While the tool is running, `session.update_tools(...)` is called. On Gemini Live this
+   forces the underlying websocket to restart (`_session_should_close` is set).
+4. The tool returns. Before the fix, its result was sent on the dying socket and never
+   replayed to the reconnected session, so the model hung waiting for a response it would
+   never receive and the turn stalled.
+
+With the fix, the tool result is buffered while the socket is restarting and replayed once the
+new session is established, so the model receives it and continues the turn.
+
+Run (needs GOOGLE_API_KEY):
+    python examples/other/gemini_realtime_tool_update_repro.py console
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+
+from dotenv import load_dotenv
+
+from livekit.agents import Agent, AgentSession, JobContext, WorkerOptions, cli, function_tool
+from livekit.plugins import google
+
+load_dotenv()
+
+logger = logging.getLogger("gemini-tool-update-repro")
+
+
+class WeatherAgent(Agent):
+    def __init__(self) -> None:
+        super().__init__(
+            instructions=(
+                "You are a voice assistant. When asked about the weather, call the "
+                "get_weather tool and then tell the user the result in one sentence."
+            ),
+        )
+
+    @function_tool
+    async def get_weather(self, location: str) -> str:
+        """Get the current weather for a location.
+
+        Args:
+            location: the city to get the weather for
+        """
+        logger.info("get_weather called for %s; triggering update_tools() mid-turn", location)
+
+        # Force a session restart in the middle of the tool call, exactly like a real
+        # update_tools() would while the model is awaiting this tool's result. This is the
+        # window where the result used to be lost.
+        session = self.session
+        assert session._activity is not None
+        realtime_session = session._activity.realtime_llm_session
+        if realtime_session is not None:
+            # change the tool set mid-turn; update_tools() only restarts the socket when the
+            # tools actually differ, so clear them to force the restart this repro needs
+            await realtime_session.update_tools([])
+            logger.info("update_tools() done — session is now restarting")
+
+        # simulate the tool doing a bit of work while the socket reconnects
+        await asyncio.sleep(0.5)
+        return f"It's 22°C and sunny in {location}."
+
+
+async def entrypoint(ctx: JobContext) -> None:
+    session = AgentSession(
+        llm=google.beta.realtime.RealtimeModel(
+            # any Gemini Live model; the fix is model-agnostic (capability-driven)
+            voice="Puck",
+        ),
+    )
+
+    await session.start(agent=WeatherAgent(), room=ctx.room)
+    await session.generate_reply(instructions="Greet the user and ask how you can help.")
+
+
+if __name__ == "__main__":
+    cli.run_app(WorkerOptions(entrypoint_fnc=entrypoint))

livekit-plugins/livekit-plugins-google/livekit/plugins/google/realtime/realtime_api.py

@@ -502,6 +502,20 @@ def __init__(self, realtime_model: RealtimeModel) -> None:
         self._active_session: AsyncSession | None = None
         # indicates if the underlying session should end
         self._session_should_close = asyncio.Event()
+        # a tool result produced while the socket is restarting (e.g. update_tools mid-turn)
+        # is stashed here and replayed once the new session is established, otherwise it would
+        # be sent on the dying session and never reach the model (the turn would hang).
+        self._pending_tool_result: types.LiveClientToolResponse | None = None
+        # whether a websocket session has ever been established. used to tell an in-flight
+        # restart (buffer the tool result for replay) apart from the initial context sync /
+        # agent handoff before the first connect (where the chat context legitimately holds
+        # historical tool outputs that must NOT be resent).
+        self._connected_once = False
+        # tracks whether the current generation has reached its completion signal. it lets us
+        # drop trailing `model_turn` frames that some Live preview models emit after a
+        # generation completed, instead of attaching them to the wrong (finished) generation.
+        # True means "idle / completed"; set False when a new generation starts.
+        self._generation_completed = True
         self._response_created_futures: dict[str, asyncio.Future[llm.GenerationCreatedEvent]] = {}
         self._pending_generation_fut: asyncio.Future[llm.GenerationCreatedEvent] | None = None
         # number of tool calls rejected in the current tool_choice="none" turn; non-zero also
@@ -643,9 +657,7 @@ async def update_chat_ctx(self, chat_ctx: llm.ChatContext) -> None:
             exclude_config_update=True,
         )
         async with self._session_lock:
-            if not self._active_session:
-                self._chat_ctx = chat_ctx
-                return
+            active_session = self._active_session
 
         diff_ops = llm.utils.compute_chat_ctx_diff(self._chat_ctx, chat_ctx)
 
@@ -664,22 +676,56 @@ async def update_chat_ctx(self, chat_ctx: llm.ChatContext) -> None:
                 vertexai=self._opts.vertexai,
                 tool_response_scheduling=self._opts.tool_response_scheduling,
             )
-            if self._realtime_model.capabilities.mutable_chat_context:
-                turns_dict, _ = append_ctx.copy(exclude_function_call=True).to_provider_format(
-                    format="google", inject_dummy_user_message=False
-                )
-                turns = [types.Content.model_validate(turn) for turn in turns_dict]
-                if turns:
-                    self._send_client_event(
-                        types.LiveClientContent(turns=turns, turn_complete=False)
+            # A restart is in flight when the socket is tearing down (e.g. update_tools()
+            # mid-turn, _active_session still set) or an already-established session has been
+            # closed for reconnect (_main_task nulled _active_session). Only in that window can a
+            # tool result correspond to a call the session being replaced issued: the (re)connect
+            # replays self._chat_ctx for the plain turns, but that replay excludes
+            # function_call/function_call_output items, so the tool result must be buffered and
+            # replayed by _main_task or the turn hangs.
+            restarting = self._session_should_close.is_set() or (
+                active_session is None and self._connected_once
+            )
+            if active_session is not None and not self._session_should_close.is_set():
+                # healthy live session: send immediately
+                if self._realtime_model.capabilities.mutable_chat_context:
+                    turns_dict, _ = append_ctx.copy(exclude_function_call=True).to_provider_format(
+                        format="google", inject_dummy_user_message=False
+                    )
+                    turns = [types.Content.model_validate(turn) for turn in turns_dict]
+                    if turns:
+                        self._send_client_event(
+                            types.LiveClientContent(turns=turns, turn_complete=False)
+                        )
+                if tool_results:
+                    self._send_client_event(tool_results)
+            elif restarting:
+                if tool_results:
+                    logger.debug(
+                        "session restarting; buffering tool result to replay after reconnect"
                     )
-            if tool_results:
-                self._send_client_event(tool_results)
+                    self._buffer_pending_tool_result(tool_results)
+            # else: initial context sync / agent handoff before the first connect. The
+            # connect-time replay of self._chat_ctx delivers the turns; historical tool outputs
+            # in that context belong to a prior session and must not be resent (doing so would
+            # make the model reply to stale results), so they are intentionally dropped here.
 
         # since we don't have a view of the history on the server side, we'll assume
         # the current state is accurate. this isn't perfect because removals aren't done.
         self._chat_ctx = chat_ctx
 
+    def _buffer_pending_tool_result(self, tool_results: types.LiveClientToolResponse) -> None:
+        # Accumulate rather than overwrite: more than one tool result can land during a single
+        # reconnect window (update_chat_ctx is called once per tool-execution round), so keep
+        # them all for _main_task to replay once the new session is established.
+        if self._pending_tool_result is None:
+            self._pending_tool_result = tool_results
+        else:
+            self._pending_tool_result.function_responses = [
+                *(self._pending_tool_result.function_responses or []),
+                *(tool_results.function_responses or []),
+            ]
+
     async def update_tools(self, tools: list[llm.Tool]) -> None:
         tool_ctx = llm.ToolContext(tools)
         if self._tools == tool_ctx:
@@ -743,15 +789,6 @@ def generate_reply(
     ) -> asyncio.Future[llm.GenerationCreatedEvent]:
         if is_given(tools):
             logger.warning("per-response tools is not supported by Google Realtime API, ignoring")
-        if not self._realtime_model.capabilities.mutable_chat_context:
-            logger.warning(
-                f"generate_reply is not compatible with '{self._opts.model}' and will be ignored."
-            )
-            fut = asyncio.Future[llm.GenerationCreatedEvent]()
-            fut.set_exception(
-                llm.RealtimeError(f"generate_reply is not compatible with '{self._opts.model}'")
-            )
-            return fut
         if self._pending_generation_fut and not self._pending_generation_fut.done():
             logger.warning(
                 "generate_reply called while another generation is pending, cancelling previous."
@@ -773,13 +810,21 @@ def generate_reply(
             )
             self._in_user_activity = False
 
-        # Gemini requires the last message to end with user's turn
-        # so we need to add a placeholder user turn in order to trigger a new generation
-        turns = []
-        if is_given(instructions):
-            turns.append(types.Content(parts=[types.Part(text=instructions)], role="model"))
-        turns.append(types.Content(parts=[types.Part(text=".")], role="user"))
-        self._send_client_event(types.LiveClientContent(turns=turns, turn_complete=True))
+        # Gemini requires the last message to end with user's turn so we add a placeholder user
+        # turn to trigger a new generation. 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.
+        if self._realtime_model.capabilities.mutable_chat_context:
+            turns = []
+            if is_given(instructions):
+                turns.append(types.Content(parts=[types.Part(text=instructions)], role="model"))
+            turns.append(types.Content(parts=[types.Part(text=".")], role="user"))
+            self._send_client_event(types.LiveClientContent(turns=turns, turn_complete=True))
+        else:
+            if is_given(instructions):
+                self._send_client_event(types.LiveClientRealtimeInput(text=instructions))
+            else:
+                self._send_client_event(types.LiveClientRealtimeInput(text="."))
 
         def _on_timeout() -> None:
             if not fut.done():
@@ -878,6 +923,8 @@ async def _main_task(self) -> None:
             await self._close_active_session()
 
             self._session_should_close.clear()
+            # a fresh session starts idle, with no generation in flight
+            self._generation_completed = True
             config = self._build_connect_config()
             session = None
             try:
@@ -889,6 +936,7 @@ async def _main_task(self) -> None:
                     self._report_connection_acquired(time.perf_counter() - t0)
                     async with self._session_lock:
                         self._active_session = session
+                        self._connected_once = True
 
                         # Check for system/developer messages in initial chat context
                         system_msg_count = sum(
@@ -919,6 +967,24 @@ async def _main_task(self) -> None:
                                 turn_complete=False,
                             )
 
+                    # A tool result produced while the previous session was tearing down was
+                    # buffered instead of being sent on the dying socket. Replay it now so the
+                    # model receives it and can continue the turn.
+                    if self._pending_tool_result is not None and (
+                        function_responses := self._pending_tool_result.function_responses
+                    ):
+                        logger.debug("replaying buffered tool result to the new session")
+                        # Gemini Live auto-generates a reply after a function response
+                        # (auto_tool_reply_generation), so sending the tool result is enough to
+                        # continue the turn. We deliberately do NOT inject a placeholder user
+                        # turn here: it would be a bogus "." message in the transcript and can
+                        # trigger a second, unrelated reply.
+                        await session.send_tool_response(function_responses=function_responses)
+                        # Clear only after the send succeeds. If send_tool_response raises (e.g.
+                        # the fresh socket drops during reconnect churn), the buffer survives so
+                        # the next reconnect can replay it instead of leaving the model hanging.
+                        self._pending_tool_result = None
+
                     # queue up existing chat context
                     send_task = asyncio.create_task(
                         self._send_task(session), name="gemini-realtime-send"
@@ -1093,6 +1159,25 @@ async def _recv_task(self, session: AsyncSession) -> None:
                         self._reject_tool_calls(response.tool_call.function_calls or [])
                         continue
 
+                    if (
+                        (not self._current_generation or self._current_generation._done)
+                        and not self._generation_completed
+                        and response.server_content
+                        and response.server_content.model_turn
+                    ):
+                        # A `model_turn` arrived for a generation that was already torn down but
+                        # never saw a completion signal (`_generation_completed` is still False).
+                        # There is no active, incomplete generation to attach it to, so
+                        # processing it would double-process the content or spin up a spurious
+                        # generation. Strip just the stray model_turn and let the rest of the
+                        # message (turn_complete, usage_metadata, go_away, session resumption)
+                        # flow through the handlers below instead of dropping it wholesale.
+                        if lk_google_debug:
+                            logger.debug(
+                                "dropping trailing model_turn without an active generation"
+                            )
+                        response.server_content.model_turn = None
+
                     if not self._current_generation or self._current_generation._done:
                         if (sc := response.server_content) and sc.interrupted:
                             # two cases an interrupted event is sent without an active generation
@@ -1212,6 +1297,8 @@ def _build_connect_config(self) -> types.LiveConnectConfig:
 
     def _start_new_generation(self) -> None:
         self._rejected_tool_calls = 0
+        # a generation is now in flight; its completion signal will flip this back to True
+        self._generation_completed = False
         if self._current_generation and not self._current_generation._done:
             logger.warning("starting new generation while another is active. Finalizing previous.")
             self._mark_current_generation_done()
@@ -1328,6 +1415,7 @@ def _handle_server_content(self, server_content: types.LiveServerContent) -> Non
                 current_gen.push_text(text)
 
         if server_content.generation_complete or server_content.turn_complete:
+            self._generation_completed = True
             current_gen._completed_timestamp = time.time()
 
         # gemini delays turn_complete until it thinks client-side playback finished, so end
@@ -1454,6 +1542,11 @@ def _handle_tool_calls(self, tool_call: types.LiveServerToolCall) -> None:
                     arguments=arguments,
                 )
             )
+        # A tool call completes the current generation regardless of model family (some Live
+        # preview models don't emit a separate generation_complete afterwards). Recording the
+        # completion here keeps the model's post-tool reply flowing as a fresh generation and
+        # keeps the trailing-model_turn guard in _recv_task from misfiring on it.
+        self._generation_completed = True
         self._mark_current_generation_done()
 
     def _handle_tool_call_cancellation(