From 7dc2b414122a0e50cc138ed74d9314faeb4b0435 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= Date: Thu, 14 May 2026 23:15:33 -0700 Subject: [PATCH] Drive task end/cancel shutdown from BaseTask by default MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ``BaseTask._handle_task_end`` and ``_handle_task_cancel`` now call ``stop()`` after propagating to children, so bus-only subclasses (``WebSocketProxyClientTask``, ``WebSocketProxyServerTask``, custom worker tasks like ``CodeWorker``) don't need to override these handlers just to set ``_finished_event``. Children-propagation is extracted into ``_propagate_end_to_children`` and ``_propagate_cancel_to_children`` so ``PipelineTask`` can call them directly without invoking ``stop()`` prematurely — the pipeline still drives its own shutdown through the ``EndFrame`` / ``CancelFrame`` path, which triggers ``on_pipeline_finished`` and ``stop()`` after the pipeline drains. Drop the now-redundant overrides from the WebSocket proxy tasks. --- src/pipecat/pipeline/base_task.py | 38 ++++++++++++++------- src/pipecat/pipeline/task.py | 20 ++++++++--- src/pipecat/tasks/proxy/websocket/client.py | 17 +-------- src/pipecat/tasks/proxy/websocket/server.py | 17 +-------- 4 files changed, 43 insertions(+), 49 deletions(-) diff --git a/src/pipecat/pipeline/base_task.py b/src/pipecat/pipeline/base_task.py index ef7daeee3..04acce5c1 100644 --- a/src/pipecat/pipeline/base_task.py +++ b/src/pipecat/pipeline/base_task.py @@ -1153,31 +1153,43 @@ class BaseTask(BaseObject, BusSubscriber): await self._call_event_handler("on_deactivated") async def _handle_task_end(self, message: BusEndTaskMessage) -> None: - """Propagate end to children and wait for them. + """Propagate end to children, wait for them, then stop this task. - Subclasses (e.g. `PipelineTask`) override to additionally shut - down their own runtime after children have finished. + Subclasses with their own runtime (e.g. `PipelineTask`) call + :meth:`_propagate_end_to_children` directly and drive their own + shutdown so ``_finished_event`` fires at the right moment. Args: message: The ``BusEndTaskMessage`` requesting a graceful end. """ logger.debug(f"Task '{self}': received end") + await self._propagate_end_to_children(message) + await self.stop() + + async def _handle_task_cancel(self, message: BusCancelTaskMessage) -> None: + """Propagate cancel to children, then stop this task. + + Subclasses with their own runtime (e.g. `PipelineTask`) call + :meth:`_propagate_cancel_to_children` directly and drive their own + shutdown so ``_finished_event`` fires at the right moment. + + Args: + message: The ``BusCancelTaskMessage`` requesting cancellation. + """ + logger.debug(f"Task '{self}': received cancel") + await self._propagate_cancel_to_children(message) + await self.stop() + + async def _propagate_end_to_children(self, message: BusEndTaskMessage) -> None: + """Forward a ``BusEndTaskMessage`` to each child and wait for them.""" for child in self._children: await self.send_message( BusEndTaskMessage(source=self.name, target=child.name, reason=message.reason) ) await asyncio.gather(*(child.wait() for child in self._children)) - async def _handle_task_cancel(self, message: BusCancelTaskMessage) -> None: - """Propagate cancel to children. - - Subclasses (e.g. `PipelineTask`) override to additionally cancel - their own runtime. - - Args: - message: The ``BusCancelTaskMessage`` requesting cancellation. - """ - logger.debug(f"Task '{self}': received cancel") + async def _propagate_cancel_to_children(self, message: BusCancelTaskMessage) -> None: + """Forward a ``BusCancelTaskMessage`` to each child.""" for child in self._children: await self.send_message( BusCancelTaskMessage(source=self.name, target=child.name, reason=message.reason) diff --git a/src/pipecat/pipeline/task.py b/src/pipecat/pipeline/task.py index c9136c6af..3634f1dc3 100644 --- a/src/pipecat/pipeline/task.py +++ b/src/pipecat/pipeline/task.py @@ -860,13 +860,25 @@ class PipelineTask(BaseTask): await self._pipeline.cleanup() async def _handle_task_end(self, message: BusEndTaskMessage) -> None: - """End the pipeline after propagating end to children.""" - await super()._handle_task_end(message) + """End the pipeline after propagating end to children. + + Drives shutdown through the pipeline (``EndFrame``) so + ``_finished_event`` fires once the frame drains through the + sink, rather than calling ``stop()`` directly. + """ + logger.debug(f"Task '{self}': received end") + await self._propagate_end_to_children(message) await self.queue_frame(EndFrame(reason=message.reason)) async def _handle_task_cancel(self, message: BusCancelTaskMessage) -> None: - """Cancel the pipeline after propagating cancel to children.""" - await super()._handle_task_cancel(message) + """Cancel the pipeline after propagating cancel to children. + + Drives shutdown through the pipeline (``CancelFrame``) so + ``_finished_event`` fires once the frame drains, rather than + calling ``stop()`` directly. + """ + logger.debug(f"Task '{self}': received cancel") + await self._propagate_cancel_to_children(message) await self.cancel(reason=message.reason) async def _process_push_queue(self): diff --git a/src/pipecat/tasks/proxy/websocket/client.py b/src/pipecat/tasks/proxy/websocket/client.py index ab7770588..99bb1c04d 100644 --- a/src/pipecat/tasks/proxy/websocket/client.py +++ b/src/pipecat/tasks/proxy/websocket/client.py @@ -10,12 +10,7 @@ import asyncio from loguru import logger -from pipecat.bus import ( - BusCancelTaskMessage, - BusEndTaskMessage, - BusMessage, - BusTaskRegistryMessage, -) +from pipecat.bus import BusMessage, BusTaskRegistryMessage from pipecat.bus.messages import BusLocalMessage from pipecat.bus.serializers import JSONMessageSerializer from pipecat.bus.serializers.base import MessageSerializer @@ -157,16 +152,6 @@ class WebSocketProxyClientTask(BaseTask): if message.source == self._local_task_name: await self._send_ws(message) - async def _handle_task_end(self, message: BusEndTaskMessage) -> None: - """Signal the run loop to finish on a graceful end.""" - await super()._handle_task_end(message) - self._finished_event.set() - - async def _handle_task_cancel(self, message: BusCancelTaskMessage) -> None: - """Signal the run loop to finish on cancellation.""" - await super()._handle_task_cancel(message) - self._finished_event.set() - async def _send_ws(self, message: BusMessage) -> None: """Serialize and send a message over the WebSocket.""" if not self._ws: diff --git a/src/pipecat/tasks/proxy/websocket/server.py b/src/pipecat/tasks/proxy/websocket/server.py index e67ddfd8d..9f6b76fc3 100644 --- a/src/pipecat/tasks/proxy/websocket/server.py +++ b/src/pipecat/tasks/proxy/websocket/server.py @@ -10,12 +10,7 @@ import asyncio from loguru import logger -from pipecat.bus import ( - BusCancelTaskMessage, - BusEndTaskMessage, - BusMessage, - BusTaskRegistryMessage, -) +from pipecat.bus import BusMessage, BusTaskRegistryMessage from pipecat.bus.messages import BusLocalMessage from pipecat.bus.serializers import JSONMessageSerializer from pipecat.bus.serializers.base import MessageSerializer @@ -174,16 +169,6 @@ class WebSocketProxyServerTask(BaseTask): elif isinstance(message, self._forward_messages): await self._send_ws(message) - async def _handle_task_end(self, message: BusEndTaskMessage) -> None: - """Signal the run loop to finish on a graceful end.""" - await super()._handle_task_end(message) - self._finished_event.set() - - async def _handle_task_cancel(self, message: BusCancelTaskMessage) -> None: - """Signal the run loop to finish on cancellation.""" - await super()._handle_task_cancel(message) - self._finished_event.set() - async def _send_ws(self, message: BusMessage) -> None: """Serialize and send a message over the WebSocket.""" if not self._ws: