fix(google): preserve non-blocking realtime generation

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

@@ -47,6 +47,12 @@
 # stop rejecting tool calls after this many in a row to avoid a loop (tool_choice="none")
 MAX_TOOL_CALL_REJECTIONS = 3
 
+# A NON_BLOCKING tool call can be followed by output in the same Gemini turn. Keep
+# that output open until it goes quiet, but do not wait forever for a completion
+# event that some models omit while waiting for the tool response.
+NON_BLOCKING_TOOL_DRAIN_QUIESCENCE_SECONDS = 0.25
+NON_BLOCKING_TOOL_DRAIN_TIMEOUT_SECONDS = 5.0
+
 # Known VertexAI models for the Live API
 # See: https://docs.cloud.google.com/vertex-ai/generative-ai/docs/live-api
 KNOWN_VERTEXAI_MODELS: frozenset[str] = frozenset(
@@ -180,6 +186,10 @@ class _ResponseGeneration:
     """Whether the generation is done (set when the turn is complete)"""
     _extra_content_warned: bool = False
     """Whether we've warned about audio/text arriving after generation completed"""
+    _last_output_activity_at: float | None = None
+    """Last server output observed after a NON_BLOCKING tool call"""
+    _tool_output_drain_atask: asyncio.Task[None] | None = None
+    """Task that bounds how long a NON_BLOCKING generation remains open"""
 
     def push_text(self, text: str) -> None:
         if self.text_ch.closed:
@@ -1272,6 +1282,11 @@ def _handle_server_content(self, server_content: types.LiveServerContent) -> Non
                 logger.warning("received server content but no active generation.")
             return
 
+        if server_content.model_turn or (
+            server_content.output_transcription and server_content.output_transcription.text
+        ):
+            current_gen._last_output_activity_at = time.monotonic()
+
         if model_turn := server_content.model_turn:
             for part in model_turn.parts or []:
                 if part.thought:
@@ -1346,11 +1361,15 @@ def _mark_current_generation_done(self) -> None:
         if not self._current_generation or self._current_generation._done:
             return
 
+        self._mark_generation_done(self._current_generation)
+
+    def _mark_generation_done(self, gen: _ResponseGeneration) -> None:
+        if gen._done:
+            return
+
         # emit input_speech_stopped event after the generation is done
         self._handle_input_speech_stopped()
 
-        gen = self._current_generation
-
         # The only way we'd know that the transcription is complete is by when they are
         # done with generation
         if gen.input_transcription:
@@ -1380,6 +1399,10 @@ def _mark_current_generation_done(self) -> None:
 
         self._close_output_streams(gen)
 
+        drain_task = gen._tool_output_drain_atask
+        if drain_task and not drain_task.done() and drain_task is not asyncio.current_task():
+            drain_task.cancel()
+
         gen.function_ch.close()
         gen.message_ch.close()
         gen._done = True
@@ -1444,7 +1467,8 @@ def _handle_tool_calls(self, tool_call: types.LiveServerToolCall) -> None:
             return
 
         gen = self._current_generation
-        for fnc_call in tool_call.function_calls or []:
+        function_calls = tool_call.function_calls or []
+        for fnc_call in function_calls:
             arguments = json.dumps(fnc_call.args)
 
             gen.function_ch.send_nowait(
@@ -1454,8 +1478,53 @@ def _handle_tool_calls(self, tool_call: types.LiveServerToolCall) -> None:
                     arguments=arguments,
                 )
             )
+
+        # NON_BLOCKING calls may be followed by more output in the same turn.
+        # Keep the generation open briefly for that output, with a bounded fallback
+        # for models that wait for the tool response instead of completing the turn.
+        if self._opts.tool_behavior == types.Behavior.NON_BLOCKING and function_calls:
+            gen._last_output_activity_at = time.monotonic()
+            self._schedule_non_blocking_tool_output_drain(gen)
+            return
+
         self._mark_current_generation_done()
 
+    def _schedule_non_blocking_tool_output_drain(self, gen: _ResponseGeneration) -> None:
+        drain_task = gen._tool_output_drain_atask
+        if drain_task and not drain_task.done():
+            return
+
+        gen._tool_output_drain_atask = asyncio.create_task(
+            self._finalize_non_blocking_generation_after_output_drain(gen),
+            name="gemini-realtime-tool-output-drain",
+        )
+
+    async def _finalize_non_blocking_generation_after_output_drain(
+        self,
+        gen: _ResponseGeneration,
+    ) -> None:
+        started_at = time.monotonic()
+        last_activity_at = gen._last_output_activity_at
+
+        while not gen._done:
+            elapsed = time.monotonic() - started_at
+            if elapsed >= NON_BLOCKING_TOOL_DRAIN_TIMEOUT_SECONDS:
+                break
+
+            await asyncio.sleep(
+                min(
+                    NON_BLOCKING_TOOL_DRAIN_QUIESCENCE_SECONDS,
+                    NON_BLOCKING_TOOL_DRAIN_TIMEOUT_SECONDS - elapsed,
+                )
+            )
+
+            current_activity_at = gen._last_output_activity_at
+            if current_activity_at == last_activity_at:
+                break
+            last_activity_at = current_activity_at
+
+        self._mark_generation_done(gen)
+
     def _handle_tool_call_cancellation(
         self, tool_call_cancellation: types.LiveServerToolCallCancellation
     ) -> None:

tests/test_plugin_google_realtime.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import logging
 from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
@@ -8,6 +9,7 @@
 from google.genai import types
 
 from livekit.agents import utils
+from livekit.plugins.google.realtime import realtime_api
 from livekit.plugins.google.realtime.realtime_api import RealtimeModel, RealtimeSession
 
 pytestmark = pytest.mark.unit
@@ -17,15 +19,22 @@
 
 
 @asynccontextmanager
-async def _make_session(monkeypatch: pytest.MonkeyPatch) -> AsyncIterator[RealtimeSession]:
+async def _make_session(
+    monkeypatch: pytest.MonkeyPatch,
+    *,
+    tool_behavior: types.Behavior | None = None,
+) -> AsyncIterator[RealtimeSession]:
     """A session whose background connect loop is stopped before it hits the network.
 
     Closed on exit so the genai http clients are released here instead of by
     ``AsyncClient.__del__``, which schedules ``aclose()`` on whatever event loop
     is running when the collector happens to reach them.
     """
     monkeypatch.setenv("GOOGLE_API_KEY", "fake-key")
-    session = RealtimeModel().session()
+    model = (
+        RealtimeModel(tool_behavior=tool_behavior) if tool_behavior is not None else RealtimeModel()
+    )
+    session = model.session()
     # cancel the connect loop before the event loop ever schedules it, so no
     # websocket connection is attempted
     session._msg_ch.close()
@@ -130,3 +139,108 @@ async def _spy() -> None:
         monkeypatch.setattr(session._client.aio, "aclose", _spy)
 
     assert closed
+
+
+async def test_blocking_tool_call_finalizes_generation(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    async with _make_session(monkeypatch) as session:
+        session._start_new_generation()
+        gen = session._current_generation
+        assert gen is not None
+
+        session._handle_tool_calls(
+            types.LiveServerToolCall(
+                function_calls=[types.FunctionCall(id="call-1", name="get_weather", args={})]
+            )
+        )
+
+        function_call = gen.function_ch.recv_nowait()
+        assert function_call.call_id == "call-1"
+        assert function_call.name == "get_weather"
+        assert gen._done
+        assert gen.message_ch.closed
+        assert gen.audio_ch.closed
+        assert gen.text_ch.closed
+
+
+async def test_non_blocking_tool_call_keeps_generation_open(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    async with _make_session(
+        monkeypatch,
+        tool_behavior=types.Behavior.NON_BLOCKING,
+    ) as session:
+        session._start_new_generation()
+        gen = session._current_generation
+        assert gen is not None
+
+        session._handle_tool_calls(
+            types.LiveServerToolCall(
+                function_calls=[types.FunctionCall(id="call-1", name="get_weather", args={})]
+            )
+        )
+
+        function_call = gen.function_ch.recv_nowait()
+        assert function_call.call_id == "call-1"
+        assert function_call.name == "get_weather"
+        assert not gen._done
+        assert not gen.message_ch.closed
+        assert not gen.audio_ch.closed
+        assert not gen.text_ch.closed
+
+        session._handle_server_content(
+            _audio_content(
+                output_transcription=types.Transcription(text="still speaking"),
+                generation_complete=True,
+            )
+        )
+
+        assert gen.output_text == "still speaking"
+        assert gen.audio_ch.qsize() == 1
+        assert gen.audio_ch.closed
+        assert gen.text_ch.closed
+        assert not gen._done
+
+        session._handle_server_content(types.LiveServerContent(turn_complete=True))
+
+        assert gen._done
+        assert gen.message_ch.closed
+
+
+async def test_non_blocking_tool_call_finalizes_without_completion_event(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(
+        realtime_api,
+        "NON_BLOCKING_TOOL_DRAIN_QUIESCENCE_SECONDS",
+        0.01,
+    )
+    monkeypatch.setattr(
+        realtime_api,
+        "NON_BLOCKING_TOOL_DRAIN_TIMEOUT_SECONDS",
+        0.1,
+    )
+
+    async with _make_session(
+        monkeypatch,
+        tool_behavior=types.Behavior.NON_BLOCKING,
+    ) as session:
+        session._start_new_generation()
+        gen = session._current_generation
+        assert gen is not None
+
+        session._handle_tool_calls(
+            types.LiveServerToolCall(
+                function_calls=[types.FunctionCall(id="call-1", name="get_weather", args={})]
+            )
+        )
+
+        assert not gen._done
+        await asyncio.sleep(0.02)
+
+        assert gen._done
+        assert gen.message_ch.closed
+        assert gen.function_ch.closed
+        assert gen.audio_ch.closed
+        assert gen.text_ch.closed