fix(google): close realtime output streams on generation_complete

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

@@ -178,8 +178,18 @@ class _ResponseGeneration:
     """The timestamp when the generation is completed"""
     _done: bool = False
     """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"""
 
     def push_text(self, text: str) -> None:
+        if self.text_ch.closed:
+            # generation_complete already finalized the output; a turn should not emit
+            # more text, so drop it (see _handle_server_content)
+            if not self._extra_content_warned:
+                self._extra_content_warned = True
+                logger.warning("Gemini sent text after generation completed; dropping it")
+            return
+
         if self.output_text:
             self.output_text += text
         else:
@@ -1259,6 +1269,15 @@ def _handle_server_content(self, server_content: types.LiveServerContent) -> Non
                 if part.text:
                     current_gen.push_text(part.text)
                 if part.inline_data:
+                    if current_gen.audio_ch.closed:
+                        # generation_complete already closed the audio stream; a turn
+                        # should not emit more audio, so drop any late frame
+                        if not current_gen._extra_content_warned:
+                            current_gen._extra_content_warned = True
+                            logger.warning(
+                                "Gemini sent audio after generation completed; dropping it"
+                            )
+                        continue
                     if not current_gen._first_token_timestamp:
                         current_gen._first_token_timestamp = time.time()
                     frame_data = part.inline_data.data
@@ -1300,6 +1319,11 @@ def _handle_server_content(self, server_content: types.LiveServerContent) -> Non
         if server_content.generation_complete or server_content.turn_complete:
             current_gen._completed_timestamp = time.time()
 
+        # gemini delays turn_complete until it thinks client-side playback finished, so end
+        # the output streams on generation_complete instead
+        if server_content.generation_complete:
+            self._close_output_streams(current_gen)
+
         if server_content.interrupted and not self._pending_generation_fut:
             # interrupt agent if there is no pending user initiated generation
             self._handle_input_speech_started()
@@ -1343,6 +1367,17 @@ def _mark_current_generation_done(self) -> None:
                 id=gen.response_id,
             )
 
+        self._close_output_streams(gen)
+
+        gen.function_ch.close()
+        gen.message_ch.close()
+        gen._done = True
+        if lk_google_debug:
+            logger.debug(f"generation done {gen}")
+
+    def _close_output_streams(self, gen: _ResponseGeneration) -> None:
+        # ends the audio segment and finalizes the output transcript. called on
+        # generation_complete (audio/text are done by then) and again at final teardown.
         if not gen.text_ch.closed:
             if self._opts.output_audio_transcription is None:
                 # close the text data of transcription synchronizer
@@ -1351,12 +1386,6 @@ def _mark_current_generation_done(self) -> None:
         if not gen.audio_ch.closed:
             gen.audio_ch.close()
 
-        gen.function_ch.close()
-        gen.message_ch.close()
-        gen._done = True
-        if lk_google_debug:
-            logger.debug(f"generation done {gen}")
-
     def _handle_input_speech_started(self) -> None:
         self.emit("input_speech_started", llm.InputSpeechStartedEvent())
 

tests/test_plugin_google_realtime.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+import logging
+
+import pytest
+from google.genai import types
+
+from livekit.agents import utils
+from livekit.plugins.google.realtime.realtime_api import RealtimeModel, RealtimeSession
+
+pytestmark = pytest.mark.unit
+
+# 10ms of silence at the output sample rate (24kHz mono, 16-bit)
+_PCM_FRAME = b"\x00\x01" * 240
+
+
+async def _make_session(monkeypatch: pytest.MonkeyPatch) -> RealtimeSession:
+    """A session whose background connect loop is stopped before it hits the network."""
+    monkeypatch.setenv("GOOGLE_API_KEY", "fake-key")
+    session = RealtimeModel().session()
+    # cancel the connect loop before the event loop ever schedules it, so no
+    # websocket connection is attempted
+    session._msg_ch.close()
+    await utils.aio.cancel_and_wait(session._main_atask)
+    return session
+
+
+def _audio_content(**kwargs: object) -> types.LiveServerContent:
+    return types.LiveServerContent(
+        model_turn=types.Content(
+            parts=[types.Part(inline_data=types.Blob(data=_PCM_FRAME, mime_type="audio/pcm"))]
+        ),
+        **kwargs,  # type: ignore[arg-type]
+    )
+
+
+async def test_output_streams_close_on_generation_complete(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """generation_complete ends the audio/text segment; finalization waits for turn_complete.
+
+    Gemini delays turn_complete until it estimates client-side playback has finished, so
+    keying the stream close off turn_complete makes AudioSegmentEnd (and the finalized
+    transcript) arrive seconds late (issue #6421). Both streams must close on
+    generation_complete, while the generation stays open until turn_complete for input
+    transcription and metrics.
+    """
+    session = await _make_session(monkeypatch)
+    session._start_new_generation()
+    gen = session._current_generation
+    assert gen is not None
+
+    session._handle_server_content(
+        _audio_content(
+            output_transcription=types.Transcription(text="hello"),
+            generation_complete=True,
+        )
+    )
+
+    # audio and text were consumed and both segments ended immediately
+    assert gen._first_token_timestamp is not None
+    assert gen.output_text == "hello"
+    assert gen.audio_ch.closed
+    assert gen.text_ch.closed
+    # but the generation is still open for trailing input transcription until turn_complete
+    assert not gen._done
+    assert not gen.message_ch.closed
+
+    session._handle_server_content(types.LiveServerContent(turn_complete=True))
+
+    assert gen._done
+    assert gen.message_ch.closed
+
+
+async def test_late_content_after_generation_complete_is_dropped(
+    monkeypatch: pytest.MonkeyPatch,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Stray audio/text after generation_complete is dropped (not pushed to a closed stream)."""
+    session = await _make_session(monkeypatch)
+    session._start_new_generation()
+    gen = session._current_generation
+    assert gen is not None
+
+    session._handle_server_content(_audio_content(generation_complete=True))
+    assert gen.audio_ch.closed and gen.text_ch.closed
+
+    with caplog.at_level(logging.WARNING):
+        # must not raise ChanClosed, must not append to the transcript, and must warn
+        session._handle_server_content(
+            _audio_content(output_transcription=types.Transcription(text="late"))
+        )
+
+    assert gen.audio_ch.closed and gen.text_ch.closed
+    assert gen.output_text == ""
+    assert not gen._done
+    assert any("after generation completed" in r.message for r in caplog.records)