From 0935d773f5a1264e6388b077a183c5c92cd7fbff Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= <aleix@daily.co>
Date: Wed, 4 Dec 2024 15:15:43 -0800
Subject: [PATCH] transport(websockets): fix initial busy loop when using audio
 mixers

---
 CHANGELOG.md                                  |  3 +++
 .../transports/network/fastapi_websocket.py   | 23 +++++++++++++------
 .../transports/network/websocket_server.py    | 20 ++++++++++------
 3 files changed, 32 insertions(+), 14 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 69f1d1fa2..de29d0c74 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -61,6 +61,9 @@ async def on_audio_data(processor, audio, sample_rate, num_channels):
 
 ### Fixed
 
+- Fixed an issue in `WebsocketServerTransport` and `FastAPIWebsocketTransport`
+  that would cause a busy loop when using audio mixer.
+
 - Fixed a `DailyTransport` and `LiveKitTransport` issue where connections were
   being closed in the input transport prematurely. This was causing frames
   queued inside the pipeline being discarded.
diff --git a/src/pipecat/transports/network/fastapi_websocket.py b/src/pipecat/transports/network/fastapi_websocket.py
index db1dcccdd..0ca969463 100644
--- a/src/pipecat/transports/network/fastapi_websocket.py
+++ b/src/pipecat/transports/network/fastapi_websocket.py
@@ -105,6 +105,11 @@ class FastAPIWebsocketOutputTransport(BaseOutputTransport):
             self._next_send_time = 0
 
     async def write_raw_audio_frames(self, frames: bytes):
+        if self._websocket.client_state != WebSocketState.CONNECTED:
+            # Simulate audio playback with a sleep.
+            await self._write_audio_sleep()
+            return
+
         frame = AudioRawFrame(
             audio=frames,
             sample_rate=self._params.audio_out_sample_rate,
@@ -125,10 +130,21 @@ class FastAPIWebsocketOutputTransport(BaseOutputTransport):
             )
             frame = wav_frame
 
+        payload = self._params.serializer.serialize(frame)
+        if payload:
+            await self._websocket.send_text(payload)
+
+        self._websocket_audio_buffer = bytes()
+
+        # Simulate audio playback with a sleep.
+        await self._write_audio_sleep()
+
+    async def _write_frame(self, frame: Frame):
         payload = self._params.serializer.serialize(frame)
         if payload and self._websocket.client_state == WebSocketState.CONNECTED:
             await self._websocket.send_text(payload)
 
+    async def _write_audio_sleep(self):
         # Simulate a clock.
         current_time = time.monotonic()
         sleep_duration = max(0, self._next_send_time - current_time)
@@ -138,13 +154,6 @@ class FastAPIWebsocketOutputTransport(BaseOutputTransport):
         else:
             self._next_send_time += self._send_interval
 
-        self._websocket_audio_buffer = bytes()
-
-    async def _write_frame(self, frame: Frame):
-        payload = self._params.serializer.serialize(frame)
-        if payload and self._websocket.client_state == WebSocketState.CONNECTED:
-            await self._websocket.send_text(payload)
-
 
 class FastAPIWebsocketTransport(BaseTransport):
     def __init__(
diff --git a/src/pipecat/transports/network/websocket_server.py b/src/pipecat/transports/network/websocket_server.py
index 10a42828b..2b88a7334 100644
--- a/src/pipecat/transports/network/websocket_server.py
+++ b/src/pipecat/transports/network/websocket_server.py
@@ -153,6 +153,8 @@ class WebsocketServerOutputTransport(BaseOutputTransport):
 
     async def write_raw_audio_frames(self, frames: bytes):
         if not self._websocket:
+            # Simulate audio playback with a sleep.
+            await self._write_audio_sleep()
             return
 
         frame = AudioRawFrame(
@@ -179,6 +181,17 @@ class WebsocketServerOutputTransport(BaseOutputTransport):
         if proto:
             await self._websocket.send(proto)
 
+        self._websocket_audio_buffer = bytes()
+
+        # Simulate audio playback with a sleep.
+        await self._write_audio_sleep()
+
+    async def _write_frame(self, frame: Frame):
+        payload = self._params.serializer.serialize(frame)
+        if payload and self._websocket:
+            await self._websocket.send(payload)
+
+    async def _write_audio_sleep(self):
         # Simulate a clock.
         current_time = time.monotonic()
         sleep_duration = max(0, self._next_send_time - current_time)
@@ -188,13 +201,6 @@ class WebsocketServerOutputTransport(BaseOutputTransport):
         else:
             self._next_send_time += self._send_interval
 
-        self._websocket_audio_buffer = bytes()
-
-    async def _write_frame(self, frame: Frame):
-        payload = self._params.serializer.serialize(frame)
-        if payload and self._websocket:
-            await self._websocket.send(payload)
-
 
 class WebsocketServerTransport(BaseTransport):
     def __init__(