From 981269d594baadefb8d42ad7529f4f3ff50d1dde Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= <aleix@daily.co>
Date: Thu, 15 Aug 2024 15:22:40 -0700
Subject: [PATCH] pipeline(task): process ErrorFrame in same task and stop
 pipeline task

---
 src/pipecat/frames/frames.py |  8 ++++----
 src/pipecat/pipeline/task.py | 26 ++++++++++++++++++++------
 2 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/src/pipecat/frames/frames.py b/src/pipecat/frames/frames.py
index 320d25e13..68ec1ec38 100644
--- a/src/pipecat/frames/frames.py
+++ b/src/pipecat/frames/frames.py
@@ -239,7 +239,7 @@ class CancelFrame(SystemFrame):
 class ErrorFrame(SystemFrame):
     """This is used notify upstream that an error has occurred downstream the
     pipeline."""
-    error: str | None
+    error: str
 
     def __str__(self):
         return f"{self.name}(error: {self.error})"
@@ -247,9 +247,9 @@ class ErrorFrame(SystemFrame):
 
 @dataclass
 class StopTaskFrame(SystemFrame):
-    """Indicates that a pipeline task should be stopped. This should inform the
-    pipeline processors that they should stop pushing frames but that they
-    should be kept in a running state.
+    """Indicates that a pipeline task should be stopped but that the pipeline
+    processors should be kept in a running state. This is normally queued from
+    the pipeline task.
 
     """
     pass
diff --git a/src/pipecat/pipeline/task.py b/src/pipecat/pipeline/task.py
index c9481271f..1f09ad233 100644
--- a/src/pipecat/pipeline/task.py
+++ b/src/pipecat/pipeline/task.py
@@ -10,7 +10,14 @@ from typing import AsyncIterable, Iterable
 
 from pydantic import BaseModel
 
-from pipecat.frames.frames import CancelFrame, EndFrame, ErrorFrame, Frame, MetricsFrame, StartFrame, StopTaskFrame
+from pipecat.frames.frames import (
+    CancelFrame,
+    EndFrame,
+    ErrorFrame,
+    Frame,
+    MetricsFrame,
+    StartFrame,
+    StopTaskFrame)
 from pipecat.pipeline.base_pipeline import BasePipeline
 from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
 from pipecat.utils.utils import obj_count, obj_id
@@ -37,10 +44,18 @@ class Source(FrameProcessor):
 
         match direction:
             case FrameDirection.UPSTREAM:
-                await self._up_queue.put(frame)
+                await self._handle_upstream_frame(frame)
             case FrameDirection.DOWNSTREAM:
                 await self.push_frame(frame, direction)
 
+    async def _handle_upstream_frame(self, frame: Frame):
+        if isinstance(frame, ErrorFrame):
+            logger.error(f"Error running app: {frame.error}")
+            # Cancel all tasks downstream.
+            await self.push_frame(CancelFrame())
+            # Tell the task we should stop.
+            await self._up_queue.put(StopTaskFrame())
+
 
 class PipelineTask:
 
@@ -70,7 +85,7 @@ class PipelineTask:
         # 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)
+        await self._source.push_frame(CancelFrame())
         self._process_down_task.cancel()
         self._process_up_task.cancel()
         await self._process_down_task
@@ -134,9 +149,8 @@ class PipelineTask:
         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())
+                if isinstance(frame, StopTaskFrame):
+                    await self.queue_frame(StopTaskFrame())
                 self._up_queue.task_done()
             except asyncio.CancelledError:
                 break