From 10202dc529517a7ada39fd5c0a332a16810e10a3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= Date: Mon, 27 Jan 2025 13:34:51 -0800 Subject: [PATCH] transports(websockets): cancel or wait for tasks to finish --- src/pipecat/transports/network/fastapi_websocket.py | 11 +++++++++++ src/pipecat/transports/network/websocket_server.py | 5 ++--- 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/src/pipecat/transports/network/fastapi_websocket.py b/src/pipecat/transports/network/fastapi_websocket.py index 42f1d3171..bcfbbc8cd 100644 --- a/src/pipecat/transports/network/fastapi_websocket.py +++ b/src/pipecat/transports/network/fastapi_websocket.py @@ -16,6 +16,8 @@ from loguru import logger from pydantic import BaseModel from pipecat.frames.frames import ( + CancelFrame, + EndFrame, Frame, InputAudioRawFrame, OutputAudioRawFrame, @@ -27,6 +29,7 @@ from pipecat.serializers.base_serializer import FrameSerializer, FrameSerializer from pipecat.transports.base_input import BaseInputTransport from pipecat.transports.base_output import BaseOutputTransport from pipecat.transports.base_transport import BaseTransport, TransportParams +from pipecat.utils.asyncio import cancel_task try: from fastapi import WebSocket @@ -72,6 +75,14 @@ class FastAPIWebsocketInputTransport(BaseInputTransport): await self._callbacks.on_client_connected(self._websocket) self._receive_task = self.create_task(self._receive_messages()) + async def stop(self, frame: EndFrame): + await super().stop(frame) + await cancel_task(self._receive_task) + + async def cancel(self, frame: CancelFrame): + await super().cancel(frame) + await cancel_task(self._receive_task) + def _iter_data(self) -> typing.AsyncIterator[bytes | str]: if self._params.serializer.type == FrameSerializerType.BINARY: return self._websocket.iter_bytes() diff --git a/src/pipecat/transports/network/websocket_server.py b/src/pipecat/transports/network/websocket_server.py index 6fe5f91ef..3eb06dc6b 100644 --- a/src/pipecat/transports/network/websocket_server.py +++ b/src/pipecat/transports/network/websocket_server.py @@ -76,12 +76,11 @@ class WebsocketServerInputTransport(BaseInputTransport): async def stop(self, frame: EndFrame): await super().stop(frame) self._stop_server_event.set() - await self._server_task + await self.wait_for_task(self._server_task) async def cancel(self, frame: CancelFrame): await super().cancel(frame) - self._stop_server_event.set() - await self._server_task + await self.cancel_task(self._server_task) async def _server_task_handler(self): logger.info(f"Starting websocket server on {self._host}:{self._port}")