Keep a held Gemini filler open until the model carries on, and place inputs as they're sent when nothing is owed

pydantic_ai_slim/pydantic_ai/realtime/_google_lifecycle.py

@@ -100,12 +100,13 @@ class GeminiLifecycle:
             self._audio_since_reply = True
             return
         if self._reply_owed():
+            # It reaches the provider while the model works on a reply, which it follows. (A response that answers
+            # it places it ahead of itself, should that come first.)
             self._held.append(input_id)
-        elif not isinstance(content, (str, ToolResult)):
+        else:
             # Nothing is going on: it joins the conversation as it arrives.
             self._pending.append(InputAdded(input_id=input_id))
         if isinstance(content, (str, ToolResult)):
-            # It asks for a reply: the response that answers it places it, ahead of that response.
             self._unanswered.append(input_id)
 
     def input_failed(self, input_id: InputId) -> None:
@@ -137,16 +138,23 @@ class GeminiLifecycle:
     def _event(self, event: RealtimeCodecEvent) -> tuple[list[LifecycleEvent], list[LifecycleEvent]]:
         before: list[LifecycleEvent] = []
         after: list[LifecycleEvent] = []
+        output = isinstance(event, (AudioDelta, PartStartEvent, ToolCall)) or (
+            isinstance(event, OutputTranscript) and bool(event.text)
+        )
         if self._deferred is not None:
-            if isinstance(event, ToolCall):
-                # The tool call the model was stalling for: it belongs to the filler's response.
+            if isinstance(event, (ToolCall, ResponseDone)):
+                # The tool call the model was stalling for belongs to the filler's response, and so does the next
+                # boundary (handled below).
                 self._deferred = None
-            else:
-                # The stall didn't end in the tool call it was for: the filler was a response of its own.
-                self._end_deferred(before)
-        if isinstance(event, (AudioDelta, PartStartEvent, ToolCall)) or (
-            isinstance(event, OutputTranscript) and event.text
-        ):
+            elif output:
+                # The stall didn't end in the tool call it was for: the filler was a response of its own, which
+                # this one carries on.
+                self._end_deferred(before, continued=True)
+            elif isinstance(event, RealtimeResponseInterruptedEvent):
+                # The user cut in: the exchange is over.
+                self._end_deferred(before, continued=False)
+            # Anything else (the user's transcript, usage) leaves it held: the exchange isn't over.
+        if output:
             self._ensure_response(before)
             if isinstance(event, ToolCall) and event.response_usage_follows:
                 self._calls_awaiting_usage = True
@@ -277,13 +285,14 @@ class GeminiLifecycle:
             self._close_turn(events)
         self._place_held(events)
 
-    def _end_deferred(self, events: list[LifecycleEvent]) -> None:
+    def _end_deferred(self, events: list[LifecycleEvent], *, continued: bool) -> None:
         deferred, self._deferred = self._deferred, None
         assert deferred is not None
-        continued = self._open
+        held = self._open
         self._end(events, status='completed', finish_reason=deferred.finish_reason or 'stop')
-        # The exchange goes on past it, in whatever the model does next.
-        self._continues = continued
+        if continued:
+            # The exchange goes on past it, in the response starting now.
+            self._continues = held
 
     # --- user turns and inputs ----------------------------------------------------------------------
 

tests/realtime/test_session_core.py

@@ -49,6 +49,7 @@ from pydantic_ai.realtime._core import (
     ToolReturned,
     TranscriptOverdue,
 )
+from pydantic_ai.realtime._google_lifecycle import GeminiLifecycle
 from pydantic_ai.realtime._lifecycle import (
     InputAdded,
     InputLost,
@@ -1127,3 +1128,34 @@ def test_a_turn_recorded_after_the_reply_to_it_is_evicted_before_that_reply() ->
     )
     turn, response = session_core._turns['u1'], session_core._responses['r1']  # pyright: ignore[reportPrivateUsage]
     assert list(session_core._recorded_audio) == [turn, response]  # pyright: ignore[reportPrivateUsage]
+
+
+def test_a_held_filler_keeps_the_wait_through_a_transcript_until_the_model_carries_on() -> None:
+    """Gemini's extended-thinking model holds its filler open (`IN_PROGRESS`): the user's words then don't end the wait."""
+    tracker = GeminiLifecycle(transcribes=True)
+    session_core = core()
+
+    def message(*codec: Any) -> None:
+        for event, stale in tracker.message(list(codec)):
+            if not stale:
+                session_core.apply(event)
+
+    session_core.apply(InputSent(input_id=0, request=text_request('Weather?'), solicits=True))
+    tracker.input_sent(0, 'Weather?')
+    message(OutputTranscript('Let me check.'))
+    message(SessionUsage(RequestUsage(input_tokens=1)), ResponseDone(more_expected=True))
+    wait = session_core.wait_tokens()
+    message(InputTranscript('Hm'))
+    assert session_core.still_owed(wait) != frozenset()
+    message(OutputTranscript('Still checking.'))
+    assert session_core.still_owed(wait) != frozenset()
+    message(SessionUsage(RequestUsage(input_tokens=2)), ResponseDone())
+    assert session_core.still_owed(wait) == frozenset()
+    assert summary(session_core.all_messages()) == snapshot(
+        [
+            '{prompt:Weather?}',
+            'None [assistant:Let me check.] complete stop',
+            '{user:Hm}',
+            'None [assistant:Still checking.] complete stop',
+        ]
+    )