Drive task end/cancel shutdown from BaseTask by default
``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.
This commit is contained in:
@@ -1153,31 +1153,43 @@ class BaseTask(BaseObject, BusSubscriber):
|
|||||||
await self._call_event_handler("on_deactivated")
|
await self._call_event_handler("on_deactivated")
|
||||||
|
|
||||||
async def _handle_task_end(self, message: BusEndTaskMessage) -> None:
|
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
|
Subclasses with their own runtime (e.g. `PipelineTask`) call
|
||||||
down their own runtime after children have finished.
|
:meth:`_propagate_end_to_children` directly and drive their own
|
||||||
|
shutdown so ``_finished_event`` fires at the right moment.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
message: The ``BusEndTaskMessage`` requesting a graceful end.
|
message: The ``BusEndTaskMessage`` requesting a graceful end.
|
||||||
"""
|
"""
|
||||||
logger.debug(f"Task '{self}': received 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:
|
for child in self._children:
|
||||||
await self.send_message(
|
await self.send_message(
|
||||||
BusEndTaskMessage(source=self.name, target=child.name, reason=message.reason)
|
BusEndTaskMessage(source=self.name, target=child.name, reason=message.reason)
|
||||||
)
|
)
|
||||||
await asyncio.gather(*(child.wait() for child in self._children))
|
await asyncio.gather(*(child.wait() for child in self._children))
|
||||||
|
|
||||||
async def _handle_task_cancel(self, message: BusCancelTaskMessage) -> None:
|
async def _propagate_cancel_to_children(self, message: BusCancelTaskMessage) -> None:
|
||||||
"""Propagate cancel to children.
|
"""Forward a ``BusCancelTaskMessage`` to each child."""
|
||||||
|
|
||||||
Subclasses (e.g. `PipelineTask`) override to additionally cancel
|
|
||||||
their own runtime.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
message: The ``BusCancelTaskMessage`` requesting cancellation.
|
|
||||||
"""
|
|
||||||
logger.debug(f"Task '{self}': received cancel")
|
|
||||||
for child in self._children:
|
for child in self._children:
|
||||||
await self.send_message(
|
await self.send_message(
|
||||||
BusCancelTaskMessage(source=self.name, target=child.name, reason=message.reason)
|
BusCancelTaskMessage(source=self.name, target=child.name, reason=message.reason)
|
||||||
|
|||||||
@@ -860,13 +860,25 @@ class PipelineTask(BaseTask):
|
|||||||
await self._pipeline.cleanup()
|
await self._pipeline.cleanup()
|
||||||
|
|
||||||
async def _handle_task_end(self, message: BusEndTaskMessage) -> None:
|
async def _handle_task_end(self, message: BusEndTaskMessage) -> None:
|
||||||
"""End the pipeline after propagating end to children."""
|
"""End the pipeline after propagating end to children.
|
||||||
await super()._handle_task_end(message)
|
|
||||||
|
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))
|
await self.queue_frame(EndFrame(reason=message.reason))
|
||||||
|
|
||||||
async def _handle_task_cancel(self, message: BusCancelTaskMessage) -> None:
|
async def _handle_task_cancel(self, message: BusCancelTaskMessage) -> None:
|
||||||
"""Cancel the pipeline after propagating cancel to children."""
|
"""Cancel the pipeline after propagating cancel to children.
|
||||||
await super()._handle_task_cancel(message)
|
|
||||||
|
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)
|
await self.cancel(reason=message.reason)
|
||||||
|
|
||||||
async def _process_push_queue(self):
|
async def _process_push_queue(self):
|
||||||
|
|||||||
@@ -10,12 +10,7 @@ import asyncio
|
|||||||
|
|
||||||
from loguru import logger
|
from loguru import logger
|
||||||
|
|
||||||
from pipecat.bus import (
|
from pipecat.bus import BusMessage, BusTaskRegistryMessage
|
||||||
BusCancelTaskMessage,
|
|
||||||
BusEndTaskMessage,
|
|
||||||
BusMessage,
|
|
||||||
BusTaskRegistryMessage,
|
|
||||||
)
|
|
||||||
from pipecat.bus.messages import BusLocalMessage
|
from pipecat.bus.messages import BusLocalMessage
|
||||||
from pipecat.bus.serializers import JSONMessageSerializer
|
from pipecat.bus.serializers import JSONMessageSerializer
|
||||||
from pipecat.bus.serializers.base import MessageSerializer
|
from pipecat.bus.serializers.base import MessageSerializer
|
||||||
@@ -157,16 +152,6 @@ class WebSocketProxyClientTask(BaseTask):
|
|||||||
if message.source == self._local_task_name:
|
if message.source == self._local_task_name:
|
||||||
await self._send_ws(message)
|
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:
|
async def _send_ws(self, message: BusMessage) -> None:
|
||||||
"""Serialize and send a message over the WebSocket."""
|
"""Serialize and send a message over the WebSocket."""
|
||||||
if not self._ws:
|
if not self._ws:
|
||||||
|
|||||||
@@ -10,12 +10,7 @@ import asyncio
|
|||||||
|
|
||||||
from loguru import logger
|
from loguru import logger
|
||||||
|
|
||||||
from pipecat.bus import (
|
from pipecat.bus import BusMessage, BusTaskRegistryMessage
|
||||||
BusCancelTaskMessage,
|
|
||||||
BusEndTaskMessage,
|
|
||||||
BusMessage,
|
|
||||||
BusTaskRegistryMessage,
|
|
||||||
)
|
|
||||||
from pipecat.bus.messages import BusLocalMessage
|
from pipecat.bus.messages import BusLocalMessage
|
||||||
from pipecat.bus.serializers import JSONMessageSerializer
|
from pipecat.bus.serializers import JSONMessageSerializer
|
||||||
from pipecat.bus.serializers.base import MessageSerializer
|
from pipecat.bus.serializers.base import MessageSerializer
|
||||||
@@ -174,16 +169,6 @@ class WebSocketProxyServerTask(BaseTask):
|
|||||||
elif isinstance(message, self._forward_messages):
|
elif isinstance(message, self._forward_messages):
|
||||||
await self._send_ws(message)
|
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:
|
async def _send_ws(self, message: BusMessage) -> None:
|
||||||
"""Serialize and send a message over the WebSocket."""
|
"""Serialize and send a message over the WebSocket."""
|
||||||
if not self._ws:
|
if not self._ws:
|
||||||
|
|||||||
Reference in New Issue
Block a user