task: handle StopFrame and StopTaskFrame gracefully

src/pipecat/pipeline/task.py

@@ -22,6 +22,7 @@ from pipecat.frames.frames import (
     HeartbeatFrame,
     MetricsFrame,
     StartFrame,
+    StopFrame,
     StopTaskFrame,
 )
 from pipecat.metrics.metrics import ProcessingMetricsData, TTFBMetricsData
@@ -163,9 +164,9 @@ class PipelineTask(BaseTask):
         # This is the heartbeat queue. When a heartbeat frame is received in the
         # down queue we add it to the heartbeat queue for processing.
         self._heartbeat_queue = asyncio.Queue()
-        # This event is used to indicate an EndFrame has been received in the
-        # down queue.
-        self._endframe_event = asyncio.Event()
+        # This event is used to indicate a finalize frame (e.g. EndFrame,
+        # StopFrame) has been received in the down queue.
+        self._pipeline_end_event = asyncio.Event()
 
         self._source = PipelineTaskSource(self._up_queue)
         self._source.link(pipeline)
@@ -224,9 +225,12 @@ class PipelineTask(BaseTask):
         """Starts and manages the pipeline execution until completion or cancellation."""
         if self.has_finished():
             return
+        cleanup_pipeline = True
         try:
             push_task = await self._create_tasks()
             await self._task_manager.wait_for_task(push_task)
+            # We have already cleaned up the pipeline inside the task.
+            cleanup_pipeline = False
         except asyncio.CancelledError:
             # We are awaiting on the push task and it might be cancelled
             # (e.g. Ctrl-C). This means we will get a CancelledError here as
@@ -234,7 +238,7 @@ class PipelineTask(BaseTask):
             # awaiting a task.
             pass
         await self._cancel_tasks()
-        await self._cleanup()
+        await self._cleanup(cleanup_pipeline)
         if self._check_dangling_tasks:
             self._print_dangling_tasks()
         self._finished = True
@@ -305,13 +309,14 @@ class PipelineTask(BaseTask):
             data.append(ProcessingMetricsData(processor=p.name, value=0.0))
         return MetricsFrame(data=data)
 
-    async def _wait_for_endframe(self):
-        await self._endframe_event.wait()
-        self._endframe_event.clear()
+    async def _wait_for_pipeline_end(self):
+        await self._pipeline_end_event.wait()
+        self._pipeline_end_event.clear()
 
-    async def _cleanup(self):
+    async def _cleanup(self, cleanup_pipeline: bool):
         await self._source.cleanup()
-        await self._pipeline.cleanup()
+        if cleanup_pipeline:
+            await self._pipeline.cleanup()
         await self._sink.cleanup()
 
     async def _process_push_queue(self):
@@ -342,18 +347,16 @@ class PipelineTask(BaseTask):
             await self._source.queue_frame(self._initial_metrics_frame(), FrameDirection.DOWNSTREAM)
 
         running = True
-        should_cleanup = True
+        cleanup_pipeline = True
         while running:
             frame = await self._push_queue.get()
             await self._source.queue_frame(frame, FrameDirection.DOWNSTREAM)
-            if isinstance(frame, EndFrame):
-                await self._wait_for_endframe()
-            running = not isinstance(frame, (CancelFrame, EndFrame, StopTaskFrame))
-            should_cleanup = not isinstance(frame, StopTaskFrame)
+            if isinstance(frame, (EndFrame, StopFrame)):
+                await self._wait_for_pipeline_end()
+            running = not isinstance(frame, (CancelFrame, EndFrame, StopFrame))
+            cleanup_pipeline = not isinstance(frame, StopFrame)
             self._push_queue.task_done()
-        # Cleanup only if we need to.
-        if should_cleanup:
-            await self._cleanup()
+        await self._cleanup(cleanup_pipeline)
 
     async def _process_up_queue(self):
         """This is the task that processes frames coming upstream from the
@@ -371,7 +374,8 @@ class PipelineTask(BaseTask):
                 # Tell the task we should end right away.
                 await self.queue_frame(CancelFrame())
             elif isinstance(frame, StopTaskFrame):
-                await self.queue_frame(StopTaskFrame())
+                # Tell the task we should stop nicely.
+                await self.queue_frame(StopFrame())
             elif isinstance(frame, ErrorFrame):
                 logger.error(f"Error running app: {frame}")
                 if frame.fatal:
@@ -390,8 +394,8 @@ class PipelineTask(BaseTask):
         """
         while True:
             frame = await self._down_queue.get()
-            if isinstance(frame, EndFrame):
-                self._endframe_event.set()
+            if isinstance(frame, (EndFrame, StopFrame)):
+                self._pipeline_end_event.set()
             elif isinstance(frame, HeartbeatFrame):
                 await self._heartbeat_queue.put(frame)
             self._down_queue.task_done()