From f62fe059b1d04aab93cb0ac7a80228a12be20345 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= <aleix@daily.co>
Date: Fri, 17 May 2024 10:43:52 -0700
Subject: [PATCH] fix issues with Ctrl-C tasks cancellation

---
 CHANGELOG.md                         |  4 +++
 src/pipecat/pipeline/task.py         | 46 ++++++++++++++++------------
 src/pipecat/transports/base_input.py | 13 ++++----
 3 files changed, 37 insertions(+), 26 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5f5049a57..b1c9876ee 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,6 +11,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Added initial interruptions support.
 
+### Fixed
+
+- Fixed issues with Ctrl-C program termination.
+
 ## [0.0.16] - 2024-05-16
 
 ### Fixed
diff --git a/src/pipecat/pipeline/task.py b/src/pipecat/pipeline/task.py
index b693c0d58..221d588f5 100644
--- a/src/pipecat/pipeline/task.py
+++ b/src/pipecat/pipeline/task.py
@@ -38,7 +38,7 @@ class PipelineTask:
         self._pipeline = pipeline
         self._allow_interruptions = allow_interruptions
 
-        self._task_queue = asyncio.Queue()
+        self._down_queue = asyncio.Queue()
         self._up_queue = asyncio.Queue()
 
         self._source = Source(self._up_queue)
@@ -50,15 +50,22 @@ class PipelineTask:
 
     async def cancel(self):
         logger.debug(f"Canceling pipeline task {self}")
-        await self.queue_frame(CancelFrame())
+        # Make sure everything is cleaned up downstream. This is sent
+        # out-of-band from the main streaming task which is what we want since
+        # we want to cancel right away.
+        await self._source.process_frame(CancelFrame(), FrameDirection.DOWNSTREAM)
+        self._process_down_task.cancel()
+        self._process_up_task.cancel()
 
     async def run(self):
-        await asyncio.gather(self._process_task_queue(), self._process_up_queue())
+        self._process_up_task = asyncio.create_task(self._process_up_queue())
+        self._process_down_task = asyncio.create_task(self._process_down_queue())
+        await asyncio.gather(self._process_up_task, self._process_down_task)
         await self._source.cleanup()
         await self._pipeline.cleanup()
 
     async def queue_frame(self, frame: Frame):
-        await self._task_queue.put(frame)
+        await self._down_queue.put(frame)
 
     async def queue_frames(self, frames: Iterable[Frame] | AsyncIterable[Frame]):
         if isinstance(frames, AsyncIterable):
@@ -70,30 +77,31 @@ class PipelineTask:
         else:
             raise Exception("Frames must be an iterable or async iterable")
 
-    async def _process_task_queue(self):
+    async def _process_down_queue(self):
         await self._source.process_frame(
             StartFrame(allow_interruptions=self._allow_interruptions), FrameDirection.DOWNSTREAM)
         running = True
         while running:
-            frame = await self._task_queue.get()
-            await self._source.process_frame(frame, FrameDirection.DOWNSTREAM)
-            self._task_queue.task_done()
-            running = not (isinstance(frame, StopTaskFrame) or
-                           isinstance(frame, CancelFrame) or
-                           isinstance(frame, EndFrame))
-        # We just enqueue None to terminate the task.
-        await self._up_queue.put(None)
+            try:
+                frame = await self._down_queue.get()
+                await self._source.process_frame(frame, FrameDirection.DOWNSTREAM)
+                self._down_queue.task_done()
+                running = not (isinstance(frame, StopTaskFrame) or isinstance(frame, EndFrame))
+            except asyncio.CancelledError:
+                break
+        # We just enqueue None to terminate the task gracefully.
+        self._process_up_task.cancel()
 
     async def _process_up_queue(self):
-        running = True
-        while running:
-            frame = await self._up_queue.get()
-            if frame:
+        while True:
+            try:
+                frame = await self._up_queue.get()
                 if isinstance(frame, ErrorFrame):
                     logger.error(f"Error running app: {frame.error}")
                     await self.queue_frame(CancelFrame())
-            self._up_queue.task_done()
-            running = frame is not None
+                self._up_queue.task_done()
+            except asyncio.CancelledError:
+                break
 
     def __str__(self):
         return self.name
diff --git a/src/pipecat/transports/base_input.py b/src/pipecat/transports/base_input.py
index 6000127b4..92473c769 100644
--- a/src/pipecat/transports/base_input.py
+++ b/src/pipecat/transports/base_input.py
@@ -65,8 +65,7 @@ class BaseInputTransport(FrameProcessor):
             await self._audio_in_thread
             await self._audio_out_thread
 
-        await self._internal_push_frame(None, None)
-        await self._push_frame_task
+        self._push_frame_task.cancel()
 
     def vad_analyze(self, audio_frames: bytes) -> VADState:
         pass
@@ -112,12 +111,12 @@ class BaseInputTransport(FrameProcessor):
         await self._push_queue.put((frame, direction))
 
     async def _push_frame_task_handler(self):
-        running = True
-        while running:
-            (frame, direction) = await self._push_queue.get()
-            if frame:
+        while True:
+            try:
+                (frame, direction) = await self._push_queue.get()
                 await self.push_frame(frame, direction)
-            running = frame is not None
+            except asyncio.CancelledError:
+                break
 
     #
     # Handle interruptions