Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ByteMaster-1 wants to merge 1 commit into
base: main
Choose a base branch
from
+285 −20
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
83 changes: 83 additions & 0 deletions examples/other/gemini_realtime_tool_update_repro.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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))
104 changes: 87 additions & 17 deletions livekit-plugins/livekit-plugins-google/livekit/plugins/google/realtime/realtime_api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -500,6 +500,15 @@ 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
# 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
Expand Down Expand Up @@ -672,7 +681,17 @@ async def update_chat_ctx(self, chat_ctx: llm.ChatContext) -> None:
types.LiveClientContent(turns=turns, turn_complete=False)
)
if tool_results:
self._send_client_event(tool_results)
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)

# 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.
Expand Down Expand Up @@ -741,15 +760,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."
Expand All @@ -771,13 +781,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():
Expand Down Expand Up @@ -867,6 +885,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:
Expand Down Expand Up @@ -908,6 +928,30 @@ 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")
await session.send_tool_response(function_responses=function_responses)
Comment thread
devin-ai-integration[bot] marked this conversation as resolved.
# 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
# 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=".")

# queue up existing chat context
send_task = asyncio.create_task(
self._send_task(session), name="gemini-realtime-send"
Expand Down Expand Up @@ -1082,6 +1126,24 @@ 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. Drop this frame; a genuine next turn will start its own
# generation once a completion signal has been recorded.
if lk_google_debug:
logger.debug(
"dropping trailing model_turn without an active generation"
)
continue

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
Expand Down Expand Up @@ -1201,6 +1263,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()
Expand Down Expand Up @@ -1317,6 +1381,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
Expand Down Expand Up @@ -1443,6 +1508,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(
Expand Down
Loading