fix: fix speech to speech model transitions (#545)

* fix: fix transition logic for realtime providers

* chore: run formatter

* chore: generate SDK and fix other realtime providers

* fix: fix ultravox node transitions

api/services/pipecat/realtime/gemini_live.py

@@ -9,8 +9,8 @@ Layers Dograh engine integration quirks onto upstream-pristine
 - **Reconnect on node transitions.** Gemini Live cannot update
   ``system_instruction`` mid-session, so a setting change triggers a
   reconnect (deferred until the bot turn ends if currently responding).
-- **Function-call deferral.** Tool calls emitted mid-turn are queued and run
-  when the bot stops speaking, to avoid racing the turn's audio.
+- **Node-transition deferral.** Node-transition calls emitted mid-turn are
+  queued and run when the bot stops speaking, to avoid cutting off its audio.
 - **User-mute audio gating.** ``UserMuteStarted/StoppedFrame`` from the
   user aggregator gates whether incoming audio is forwarded to Gemini.
 - **TTSSpeakFrame as greeting trigger.** The engine queues a TTSSpeakFrame
@@ -18,6 +18,7 @@ Layers Dograh engine integration quirks onto upstream-pristine
   it and runs the initial-context path.
 """
 
+import asyncio
 from typing import Any
 
 from google.genai.types import Content, Part
@@ -44,6 +45,11 @@ from pipecat.utils.tracing.service_decorators import traced_gemini_live
 class DograhGeminiLiveLLMService(GeminiLiveLLMService):
     """Gemini Live with Dograh engine integration quirks. See module docstring."""
 
+    # Gemini input transcription is delivered independently from tool calls.
+    # Give late transcription messages a small window to arrive before running
+    # a node-transition function and tearing down the current Live connection.
+    _NODE_TRANSITION_TRANSCRIPTION_GRACE_SECONDS = 0.5
+
     # Route tool schemas through Gemini's ``parameters_json_schema`` field so
     # MCP/imported tools that use JSON Schema keywords (``const``, ``not``,
     # nested ``anyOf``) rejected by the strict ``Schema`` model are accepted.
@@ -59,15 +65,19 @@ class DograhGeminiLiveLLMService(GeminiLiveLLMService):
         # Guards initial-response triggering against double-firing across the
         # initial TTSSpeakFrame and any LLMContextFrame that may arrive.
         self._handled_initial_context: bool = False
-        # When a system_instruction change arrives mid-bot-turn, the reconnect
-        # is queued and drained when the turn ends.
-        self._reconnect_pending: bool = False
-        # Function calls emitted by Gemini mid-bot-turn are deferred here and
-        # invoked when the turn ends, so they don't race the turn's audio.
-        self._pending_function_calls: list[FunctionCallFromLLM] = []
+        # Node-transition calls emitted mid-bot-turn are deferred here so the
+        # transition does not tear down Gemini while it is still producing audio.
+        self._pending_node_transition_function_calls: list[FunctionCallFromLLM] = []
         # Text greeting captured from the first TTSSpeakFrame while the Gemini
         # session is still connecting.
         self._pending_initial_greeting_text: str | None = None
+        self._transition_function_call_task: asyncio.Task | None = None
+        # Intentional node changes use a fresh, context-seeded connection rather
+        # than a potentially stale session-resumption handle. The new connection
+        # remains gated until the function-call result has landed in LLMContext.
+        self._awaiting_node_transition_context: bool = False
+        self._node_transition_context_received: bool = False
+        self._node_transition_context_seed_started: bool = False
 
     # ------------------------------------------------------------------
     # Hooks from upstream GeminiLiveLLMService
@@ -78,33 +88,109 @@ class DograhGeminiLiveLLMService(GeminiLiveLLMService):
         # lets pre-call fetch populate template variables first.
         return bool(self._settings.system_instruction)
 
+    def _requires_node_transition_context_aggregation(self) -> bool:
+        # A node transition replaces the current Gemini Live connection and
+        # seeds the new one from our local LLMContext. Wait for the upstream
+        # user aggregator to commit any final TranscriptionFrame before
+        # set_node() changes the prompt and starts that reconnect.
+        return True
+
+    async def cleanup(self) -> None:
+        """Cancel a delayed transition before tearing down the Live session."""
+        if self._transition_function_call_task:
+            await self.cancel_task(self._transition_function_call_task)
+            self._transition_function_call_task = None
+        await super().cleanup()
+
     async def _handle_changed_settings(self, changed: dict[str, Any]) -> set[str]:
         if "system_instruction" not in changed:
             return set()
+
+        # PipecatEngine updates system_instruction only from set_node(). The
+        # first set_node happens before a Live session exists; every later one
+        # is a node transition whose tool call has already been deferred until
+        # the current bot turn finishes.
         if not self._session:
             # First-time setting after deferred-connect.
             await self._connect()
-        elif self._bot_is_responding:
-            # Bot is mid-turn — drain the reconnect when it ends so we don't
-            # cut the bot off mid-utterance.
-            self._reconnect_pending = True
         else:
-            await self._reconnect()
+            await self._reconnect_for_node_transition()
         return {"system_instruction"}
 
     async def _run_or_defer_function_calls(
         self, function_calls_llm: list[FunctionCallFromLLM]
     ):
+        if not self._contains_node_transition(function_calls_llm):
+            await super()._run_or_defer_function_calls(function_calls_llm)
+            return
+
+        # Keep a provider tool-call batch together. Splitting a mixed batch here
+        # would discard Pipecat's shared function-call group and could trigger an
+        # LLM run before every result from the original batch has arrived.
         if self._bot_is_responding:
             # Latest batch wins; Gemini emits tool calls as one batch per
             # tool_call message, so this overwrite is intentional.
-            self._pending_function_calls = function_calls_llm
+            self._pending_node_transition_function_calls = function_calls_llm
             logger.debug(
-                f"{self}: deferring {len(function_calls_llm)} function call(s) "
+                f"{self}: deferring {len(function_calls_llm)} node-transition "
+                "function call(s) "
                 "until bot turn ends"
             )
             return
-        await super()._run_or_defer_function_calls(function_calls_llm)
+
+        self._schedule_node_transition_function_calls(function_calls_llm)
+
+    def _contains_node_transition(
+        self, function_calls_llm: list[FunctionCallFromLLM]
+    ) -> bool:
+        return any(self._is_node_transition(fc) for fc in function_calls_llm)
+
+    def _is_node_transition(self, function_call: FunctionCallFromLLM) -> bool:
+        return self._function_is_node_transition(function_call.function_name)
+
+    def _schedule_node_transition_function_calls(
+        self, function_calls_llm: list[FunctionCallFromLLM]
+    ) -> None:
+        """Run transition calls after late input transcription has settled."""
+        if (
+            self._transition_function_call_task
+            and not self._transition_function_call_task.done()
+        ):
+            logger.warning(
+                f"{self}: node-transition function call already pending; "
+                "ignoring duplicate batch"
+            )
+            return
+
+        async def _run_after_transcription_grace() -> None:
+            try:
+                await asyncio.sleep(self._NODE_TRANSITION_TRANSCRIPTION_GRACE_SECONDS)
+                await self._flush_pending_user_transcription()
+                await self.run_function_calls(function_calls_llm)
+            finally:
+                self._transition_function_call_task = None
+
+        self._transition_function_call_task = self.create_task(
+            _run_after_transcription_grace(),
+            name=f"{self}::node-transition-function-calls",
+        )
+
+    async def _flush_pending_user_transcription(self) -> None:
+        """Publish any punctuationless user transcript before a node handoff."""
+        if self._transcription_timeout_task:
+            if not self._transcription_timeout_task.done():
+                await self.cancel_task(self._transcription_timeout_task)
+            self._transcription_timeout_task = None
+
+        if not self._user_transcription_buffer:
+            return
+
+        text = self._user_transcription_buffer
+        self._user_transcription_buffer = ""
+        logger.debug(
+            f"{self}: flushing pending user transcription before node transition"
+        )
+        await self._push_user_transcription(text, result=None)
 
     # ------------------------------------------------------------------
     # State-transition side effects
@@ -114,22 +200,35 @@ class DograhGeminiLiveLLMService(GeminiLiveLLMService):
         was_responding = self._bot_is_responding
         await super()._set_bot_is_responding(responding)
         if was_responding and not responding:
-            await self._run_pending_function_calls()
-            if self._reconnect_pending:
-                self._reconnect_pending = False
-                await self._reconnect()
+            await self._run_pending_node_transition_function_calls()
 
-    async def _run_pending_function_calls(self):
-        """Run any function calls deferred during the bot's last turn."""
-        if not self._pending_function_calls:
+    async def _run_pending_node_transition_function_calls(self):
+        """Run any node-transition calls deferred during the bot's last turn."""
+        if not self._pending_node_transition_function_calls:
             return
-        fcs = self._pending_function_calls
-        self._pending_function_calls = []
+        fcs = self._pending_node_transition_function_calls
+        self._pending_node_transition_function_calls = []
         logger.debug(
-            f"{self}: executing {len(fcs)} deferred function call(s) "
+            f"{self}: executing {len(fcs)} deferred node-transition call(s) "
             "after bot turn ended"
         )
-        await self.run_function_calls(fcs)
+        self._schedule_node_transition_function_calls(fcs)
+
+    async def _reconnect_for_node_transition(self) -> None:
+        """Start a fresh connection and wait to seed the completed context.
+
+        Gemini can report ``resumable=False`` while generating or executing a
+        function call. A workflow transition happens at exactly that boundary,
+        so using the last (older) resumption handle can omit the triggering user
+        turn. Use the local LLMContext as the source of truth for this intentional
+        handoff instead.
+        """
+        self._awaiting_node_transition_context = True
+        self._node_transition_context_received = False
+        self._node_transition_context_seed_started = False
+        self._session_resumption_handle = None
+        await self._disconnect()
+        await self._connect(session_resumption_handle=None)
 
     # ------------------------------------------------------------------
     # Frame handling: mute, TTSSpeakFrame, BotStoppedSpeakingFrame flush
@@ -165,9 +264,9 @@ class DograhGeminiLiveLLMService(GeminiLiveLLMService):
         if isinstance(frame, BotStoppedSpeakingFrame):
             # Belt-and-suspenders: the main drain happens in
             # _set_bot_is_responding(False), but if Gemini delays turn_complete
-            # past the audible end of the turn, flushing here ensures pending
-            # function calls fire promptly.
-            await self._run_pending_function_calls()
+            # past the audible end of the turn, flushing here ensures a pending
+            # node transition fires promptly.
+            await self._run_pending_node_transition_function_calls()
             # Fall through to super for the actual push.
         await super().process_frame(frame, direction)
 
@@ -185,6 +284,11 @@ class DograhGeminiLiveLLMService(GeminiLiveLLMService):
     # ------------------------------------------------------------------
 
     async def _handle_context(self, context: LLMContext):
+        if self._awaiting_node_transition_context:
+            self._context = context
+            self._node_transition_context_received = True
+            await self._maybe_seed_node_transition_context()
+            return
         if not self._handled_initial_context:
             self._handled_initial_context = True
             self._context = context
@@ -249,6 +353,12 @@ class DograhGeminiLiveLLMService(GeminiLiveLLMService):
             f"In _handle_session_ready self._run_llm_when_session_ready: {self._run_llm_when_session_ready}"
         )
         self._session = session
+        if self._awaiting_node_transition_context:
+            # Do not accept realtime input until the function-call result frame
+            # has updated the shared context and that complete history is seeded.
+            self._ready_for_realtime_input = False
+            await self._maybe_seed_node_transition_context()
+            return
         self._ready_for_realtime_input = True
         if self._run_llm_when_session_ready:
             # Context arrived before session was ready — fulfil the queued
@@ -266,3 +376,25 @@ class DograhGeminiLiveLLMService(GeminiLiveLLMService):
         # a handle (e.g. node transitions before any handle was issued) are
         # followed by a function-call-result LLMContextFrame which feeds the
         # updated-context branch in _handle_context.
+
+    async def _maybe_seed_node_transition_context(self) -> None:
+        if (
+            not self._awaiting_node_transition_context
+            or not self._node_transition_context_received
+            or not self._session
+            or self._node_transition_context_seed_started
+        ):
+            return
+
+        self._node_transition_context_seed_started = True
+        try:
+            # The complete tool result is already present in the history being
+            # seeded, so mark it delivered locally instead of sending a provider
+            # tool response for a call that the fresh session never issued.
+            await self._process_completed_function_calls(send_new_results=False)
+            await self._create_initial_response()
+            self._awaiting_node_transition_context = False
+            self._node_transition_context_received = False
+            await self._drain_pending_tool_results()
+        finally:
+            self._node_transition_context_seed_started = False