AudioContextWordTTSService: wait for all requested audio

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.
 

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
+            if context_id:
-            # audio available (i.e. we find None).
+                # Process the audio context until the context doesn't have more
-            await self._handle_audio_context(context_id)
+                # 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