From 5b6d9a1050dd74bb1c007c612de269ffa64950dc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= <aleix@daily.co>
Date: Wed, 4 Dec 2024 10:56:57 -0800
Subject: [PATCH] transports: fix premature input transport closing

---
 CHANGELOG.md                               |  4 ++++
 src/pipecat/transports/services/daily.py   | 21 +++++++++++----------
 src/pipecat/transports/services/livekit.py | 11 ++++++++++-
 3 files changed, 25 insertions(+), 11 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c449d450c..b7cb1f80d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -43,6 +43,10 @@ async def on_audio_data(processor, audio, sample_rate, num_channels):
 
 ### Fixed
 
+- 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.
+
 - Fixed an issue in `DailyTransport` that would cause some internal callbacks to
   not be executed.
 
diff --git a/src/pipecat/transports/services/daily.py b/src/pipecat/transports/services/daily.py
index 07c023233..c4c50f7a6 100644
--- a/src/pipecat/transports/services/daily.py
+++ b/src/pipecat/transports/services/daily.py
@@ -200,8 +200,7 @@ class DailyTransportClient(EventHandler):
         self._other_participant_has_joined = False
 
         self._joined = False
-        self._joining = False
-        self._leaving = False
+        self._leave_counter = 0
 
         self._executor = ThreadPoolExecutor(max_workers=5)
 
@@ -263,7 +262,7 @@ class DailyTransportClient(EventHandler):
         self._callbacks = callbacks
 
     async def send_message(self, frame: TransportMessageFrame | TransportMessageUrgentFrame):
-        if not self._joined or self._leaving:
+        if not self._joined:
             return
 
         participant_id = None
@@ -315,13 +314,13 @@ class DailyTransportClient(EventHandler):
 
     async def join(self):
         # Transport already joined, ignore.
-        if self._joined or self._joining:
+        if self._joined:
+            # Increment leave counter if we already joined.
+            self._leave_counter += 1
             return
 
         logger.info(f"Joining {self._room_url}")
 
-        self._joining = True
-
         # For performance reasons, never subscribe to video streams (unless a
         # video renderer is registered).
         self._client.update_subscription_profiles(
@@ -335,7 +334,8 @@ class DailyTransportClient(EventHandler):
 
             if not error:
                 self._joined = True
-                self._joining = False
+                # Increment leave counter if we successfully joined.
+                self._leave_counter += 1
 
                 logger.info(f"Joined {self._room_url}")
 
@@ -423,12 +423,14 @@ class DailyTransportClient(EventHandler):
         return await asyncio.wait_for(future, timeout=10)
 
     async def leave(self):
+        # Decrement leave counter when leaving.
+        self._leave_counter -= 1
+
         # Transport not joined, ignore.
-        if not self._joined or self._leaving:
+        if not self._joined or self._leave_counter > 0:
             return
 
         self._joined = False
-        self._leaving = True
 
         logger.info(f"Leaving {self._room_url}")
 
@@ -438,7 +440,6 @@ class DailyTransportClient(EventHandler):
         try:
             error = await self._leave()
             if not error:
-                self._leaving = False
                 logger.info(f"Left {self._room_url}")
                 await self._callbacks.on_left()
             else:
diff --git a/src/pipecat/transports/services/livekit.py b/src/pipecat/transports/services/livekit.py
index 6c99f6a66..f53c8332f 100644
--- a/src/pipecat/transports/services/livekit.py
+++ b/src/pipecat/transports/services/livekit.py
@@ -83,6 +83,7 @@ class LiveKitTransportClient:
         self._room = rtc.Room(loop=loop)
         self._participant_id: str = ""
         self._connected = False
+        self._disconnect_counter = 0
         self._audio_source: rtc.AudioSource | None = None
         self._audio_track: rtc.LocalAudioTrack | None = None
         self._audio_tracks = {}
@@ -105,6 +106,8 @@ class LiveKitTransportClient:
     @retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
     async def connect(self):
         if self._connected:
+            # Increment disconnect counter if already connected.
+            self._disconnect_counter += 1
             return
 
         logger.info(f"Connecting to {self._room_name}")
@@ -116,6 +119,9 @@ class LiveKitTransportClient:
                 options=rtc.RoomOptions(auto_subscribe=True),
             )
             self._connected = True
+            # Increment disconnect counter if we successfully connected.
+            self._disconnect_counter += 1
+
             self._participant_id = self._room.local_participant.sid
             logger.info(f"Connected to {self._room_name}")
 
@@ -142,7 +148,10 @@ class LiveKitTransportClient:
             raise
 
     async def disconnect(self):
-        if not self._connected:
+        # Decrement leave counter when leaving.
+        self._disconnect_counter -= 1
+
+        if not self._connected or self._disconnect_counter > 0:
             return
 
         logger.info(f"Disconnecting from {self._room_name}")