Merge pull request #12 from pipecat-ai/aleix/flush-before-close

WhiskerObserver: flush data before closing and finish task

Author: aconchillo

CHANGELOG.md

@@ -21,6 +21,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Fixed an issue where the buffer data was not being sent to the client when
+  closing.
+
+- Fixed an issue that was causing data to be written to the output file after
+  the file was already closed.
+
 - Fixed an issue that could cause `WhiskerObserver` to crash if the given
   pipeline didn't have a previous processor.
 

pipecat/src/pipecat_whisker/observer.py

@@ -146,7 +146,7 @@ def __init__(
         self._id = 0
         self._client = None
         self._server_future = asyncio.get_running_loop().create_future()
-        self._server_task = asyncio.create_task(self._start_task_handler())
+        self._server_task = asyncio.create_task(self._server_task_handler())
         self._send_task = asyncio.create_task(self._send_task_handler())
         self._send_queue = asyncio.Queue()
         self._batch = []
@@ -158,16 +158,13 @@ async def cleanup(self):
         """Clean up resources and close the Whisker server."""
         await super().cleanup()
 
-        await self._maybe_close_file()
+        await self._stop_send_task()
 
-        if self._client:
-            await self._client.close(reason="Whisker shutting down")
+        await self._close_client()
 
-        if not self._server_future.done():
-            self._server_future.set_result(None)
+        await self._stop_server()
 
-        if self._server_task:
-            await self._server_task
+        await self._maybe_close_file()
 
     async def on_process_frame(self, data: FrameProcessed):
         """Handle frame processing events.
@@ -187,7 +184,7 @@ async def on_push_frame(self, data: FramePushed):
         if not isinstance(data.frame, self._exclude_frames):
             await self._send_push_frame(data)
 
-    async def _start_task_handler(self):
+    async def _server_task_handler(self):
         """Start the Whisker server and handle incoming connections.
 
         This method runs in a separate task and manages the websocket server lifecycle.
@@ -198,11 +195,18 @@ async def _start_task_handler(self):
         # Queue initial pipeline structure
         await self._send_pipeline()
 
-        async with serve(self._server_handler, self._host, self._port):
+        async with serve(self._client_handler, self._host, self._port):
             logger.debug(f"ᓚᘏᗢ Whisker running at ws://{self._host}:{self._port}")
             await self._server_future
 
-    async def _server_handler(self, client):
+    async def _stop_server(self):
+        if not self._server_future.done():
+            self._server_future.set_result(None)
+
+        if self._server_task:
+            await self._server_task
+
+    async def _client_handler(self, client):
         """Handle a new Whisker client connection.
 
         Args:
@@ -227,6 +231,10 @@ async def _server_handler(self, client):
             logger.debug("ᓚᘏᗢ Whisker: client disconnected")
             await self._reset_client()
 
+    async def _close_client(self):
+        if self._client:
+            await self._client.close(reason="Whisker shutting down")
+
     async def _reset_client(self):
         self._client = None
 
@@ -243,18 +251,27 @@ async def _maybe_close_file(self):
 
     async def _send_task_handler(self):
         """Handle sending batched messages to the client."""
-        while True:
+        running = True
+        while running:
             try:
                 data, flush = await asyncio.wait_for(self._send_queue.get(), timeout=0.5)
 
-                self._batch.append(data)
+                if data:
+                    self._batch.append(data)
 
                 await self._maybe_send_batch(flush=flush)
 
                 self._send_queue.task_done()
+
+                running = data is not None
             except asyncio.TimeoutError:
                 await self._maybe_send_batch(flush=True)
 
+    async def _stop_send_task(self):
+        await self._queue_data(None, True)
+        if self._send_task:
+            await self._send_task
+
     async def _maybe_send_batch(self, *, flush: bool = False):
         """Send batched messages to the client.
 
@@ -418,9 +435,9 @@ async def _send_push_frame(self, data: FramePushed):
 
         await self._queue_data(msg_packed)
 
-    async def _queue_data(self, msg: bytes, flush: bool = False):
+    async def _queue_data(self, msg: Optional[bytes], flush: bool = False):
         await self._send_queue.put((msg, flush))
-        if self._file:
+        if self._file and msg:
             await self._file.write(msg)
 
     async def _send(self, msg: bytes):