transports(websocket): close connection from last transport

CHANGELOG.md

@@ -98,6 +98,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Fixed a `FastAPIWebsocketTransport` and `WebsocketClientTransport` issue that
+  would cause the transport to be closed prematurely, preventing the internally
+  queued audio to be sent. The same issue could also cause an infinite loop
+  while using an output mixer and when sending an `EndFrame`, preventing the bot
+  to finish.
+
 - Fixed an issue that could cause the `TranscriptionUpdateFrame` being pushed
   because of an interruption to be discarded.
 

src/pipecat/transports/network/fastapi_websocket.py

@@ -61,6 +61,10 @@ class FastAPIWebsocketClient:
         self._closing = False
         self._is_binary = is_binary
         self._callbacks = callbacks
+        self._leave_counter = 0
+
+    async def setup(self, _: StartFrame):
+        self._leave_counter += 1
 
     def receive(self) -> typing.AsyncIterator[bytes | str]:
         return self._websocket.iter_bytes() if self._is_binary else self._websocket.iter_text()
@@ -73,6 +77,10 @@ class FastAPIWebsocketClient:
                 await self._websocket.send_text(data)
 
     async def disconnect(self):
+        self._leave_counter -= 1
+        if self._leave_counter > 0:
+            return
+
         if self.is_connected and not self.is_closing:
             self._closing = True
             await self._websocket.close()
@@ -116,6 +124,7 @@ class FastAPIWebsocketInputTransport(BaseInputTransport):
 
     async def start(self, frame: StartFrame):
         await super().start(frame)
+        await self._client.setup(frame)
         await self._params.serializer.setup(frame)
         if not self._monitor_websocket_task and self._params.session_timeout:
             self._monitor_websocket_task = self.create_task(self._monitor_websocket())
@@ -192,6 +201,7 @@ class FastAPIWebsocketOutputTransport(BaseOutputTransport):
 
     async def start(self, frame: StartFrame):
         await super().start(frame)
+        await self._client.setup(frame)
         await self._params.serializer.setup(frame)
         self._send_interval = (self._audio_chunk_size / self.sample_rate) / 2
 

src/pipecat/transports/network/websocket_client.py

@@ -56,8 +56,8 @@ class WebsocketClientSession:
         self._callbacks = callbacks
         self._transport_name = transport_name
 
+        self._leave_counter = 0
         self._task_manager: Optional[BaseTaskManager] = None
-
         self._websocket: Optional[websockets.WebSocketClientProtocol] = None
 
     @property
@@ -69,6 +69,7 @@ class WebsocketClientSession:
         return self._task_manager
 
     async def setup(self, frame: StartFrame):
+        self._leave_counter += 1
         if not self._task_manager:
             self._task_manager = frame.task_manager
 
@@ -87,7 +88,8 @@ class WebsocketClientSession:
             logger.error(f"Timeout connecting to {self._uri}")
 
     async def disconnect(self):
-        if not self._websocket:
+        self._leave_counter -= 1
+        if not self._websocket or self._leave_counter > 0:
             return
 
         await self.task_manager.cancel_task(self._client_task)