From 9a50f33e36a8b6b542c22ff8cc208e9b099e2a11 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= <aleix@daily.co>
Date: Tue, 25 Feb 2025 15:35:47 -0800
Subject: [PATCH] AudioContextWordTTSService: wait for all requested audio

---
 CHANGELOG.md                        |  4 ++++
 src/pipecat/services/ai_services.py | 36 +++++++++++++++++++----------
 2 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index e5d3ae6d1..ade408c14 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -72,6 +72,10 @@ stt = DeepgramSTTService(..., live_options=LiveOptions(model="nova-2-general"))
 
 ### Fixed
 
+- Fixed an `AudioContextWordTTSService` issue that would cause an `EndFrame` to
+  disconnect from the TTS service before audio from all the contexts was
+  received. This affected services like Cartesia and Rime.
+
 - Fixed an issue that was not allowing to pass an `OpenAILLMContext` to create
   `GoogleLLMService`'s context aggregators.
 
diff --git a/src/pipecat/services/ai_services.py b/src/pipecat/services/ai_services.py
index a47320b7a..136afb47a 100644
--- a/src/pipecat/services/ai_services.py
+++ b/src/pipecat/services/ai_services.py
@@ -652,7 +652,12 @@ class AudioContextWordTTSService(WebsocketWordTTSService):
 
     async def stop(self, frame: EndFrame):
         await super().stop(frame)
-        await self._stop_audio_context_task()
+        if self._audio_context_task:
+            # Indicate no more audio contexts are available. this will end the
+            # task cleanly after all contexts have been processed.
+            await self._contexts_queue.put(None)
+            await self.wait_for_task(self._audio_context_task)
+            self._audio_context_task = None
 
     async def cancel(self, frame: CancelFrame):
         await super().cancel(frame)
@@ -676,22 +681,29 @@ class AudioContextWordTTSService(WebsocketWordTTSService):
 
     async def _audio_context_task_handler(self):
         """In this task we process audio contexts in order."""
-        while True:
+        running = True
+        while running:
             context_id = await self._contexts_queue.get()
 
-            # Process the audio context until the context doesn't have more
-            # audio available (i.e. we find None).
-            await self._handle_audio_context(context_id)
+            if context_id:
+                # Process the audio context until the context doesn't have more
+                # audio available (i.e. we find None).
+                await self._handle_audio_context(context_id)
+
+                # We just finished processing the context, so we can safely remove it.
+                del self._contexts[context_id]
+
+                # Append some silence between sentences.
+                silence = b"\x00" * self.sample_rate
+                frame = TTSAudioRawFrame(
+                    audio=silence, sample_rate=self.sample_rate, num_channels=1
+                )
+                await self.push_frame(frame)
+            else:
+                running = False
 
-            # We just finished processing the context, so we can safely remove it.
-            del self._contexts[context_id]
             self._contexts_queue.task_done()
 
-            # Append some silence between sentences.
-            silence = b"\x00" * self.sample_rate
-            frame = TTSAudioRawFrame(audio=silence, sample_rate=self.sample_rate, num_channels=1)
-            await self.push_frame(frame)
-
     async def _handle_audio_context(self, context_id: str):
         # If we don't receive any audio during this time, we consider the context finished.
         AUDIO_CONTEXT_TIMEOUT = 3.0