dara-core 1.20.1a1__py3-none-any.whl → 1.20.1a3__py3-none-any.whl

dara/core/internal/tasks.py

@@ -15,26 +15,30 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import contextlib
 import inspect
 import math
-from typing import Any, Awaitable, Callable, Dict, List, Optional, Union, overload
+from collections.abc import Awaitable
+from typing import Any, Callable, Dict, List, Optional, Set, Union, overload
 
 from anyio import (
+    BrokenResourceError,
     CancelScope,
     ClosedResourceError,
     create_memory_object_stream,
     create_task_group,
+    get_cancelled_exc_class,
     move_on_after,
 )
 from anyio.abc import TaskGroup
 from anyio.streams.memory import MemoryObjectSendStream
-from exceptiongroup import ExceptionGroup
+from exceptiongroup import ExceptionGroup, catch
 from pydantic import ConfigDict
 
 from dara.core.base_definitions import (
     BaseTask,
-    Cache,
     CachedRegistryEntry,
+    LruCachePolicy,
     PendingTask,
     TaskError,
     TaskMessage,
@@ -45,7 +49,7 @@ from dara.core.internal.cache_store import CacheStore
 from dara.core.internal.devtools import get_error_for_channel
 from dara.core.internal.pandas_utils import remove_index
 from dara.core.internal.pool import TaskPool
-from dara.core.internal.utils import resolve_exception_group, run_user_handler
+from dara.core.internal.utils import exception_group_contains, run_user_handler
 from dara.core.internal.websocket import WebsocketManager
 from dara.core.logging import dev_logger, eng_logger
 from dara.core.metrics import RUNTIME_METRICS_TRACKER
@@ -137,14 +141,12 @@ class Task(BaseTask):
 
             async def on_progress(progress: float, msg: str):
                 if send_stream is not None:
-                    try:
+                    with contextlib.suppress(ClosedResourceError):
                         await send_stream.send(TaskProgressUpdate(task_id=self.task_id, progress=progress, message=msg))
-                    except ClosedResourceError:
-                        pass
 
             async def on_result(result: Any):
                 if send_stream is not None:
-                    try:
+                    with contextlib.suppress(ClosedResourceError):
                         await send_stream.send(
                             TaskResult(
                                 task_id=self.task_id,
@@ -153,19 +155,15 @@ class Task(BaseTask):
                                 reg_entry=self.reg_entry,
                             )
                         )
-                    except ClosedResourceError:
-                        pass
 
             async def on_error(exc: BaseException):
                 if send_stream is not None:
-                    try:
+                    with contextlib.suppress(ClosedResourceError):
                         await send_stream.send(
                             TaskError(
                                 task_id=self.task_id, error=exc, cache_key=self.cache_key, reg_entry=self.reg_entry
                             )
                         )
-                    except ClosedResourceError:
-                        pass
 
             with pool.on_progress(self.task_id, on_progress):
                 pool_task_def = pool.submit(self.task_id, self._func_name, args=tuple(self._args), kwargs=self._kwargs)
@@ -218,7 +216,7 @@ class MetaTask(BaseTask):
         :param notify_channels: If this task is run in a TaskManager instance these channels will also be notified on
                                 completion
         :param process_as_task: Whether to run the process_result function as a task or not, defaults to False
-        :param cache_key: Optional cache key if there is a PendingTask in the store associated with this task
+        :param cache_key: Optional cache key if there is a registry entry to store results for the task in
         :param task_id: Optional task_id to set for the task - otherwise the task generates its id automatically
         """
         self.args = args if args is not None else []
@@ -239,96 +237,122 @@ class MetaTask(BaseTask):
 
         :param send_stream: The stream to send messages to the task manager on
         """
-        tasks: List[BaseTask] = []
-
-        # Collect up the tasks that need to be run and kick them off without awaiting them.
-        tasks.extend(x for x in self.args if isinstance(x, BaseTask))
-        tasks.extend(x for x in self.kwargs.values() if isinstance(x, BaseTask))
-
-        eng_logger.info(f'MetaTask {self.task_id} running sub-tasks', {'task_ids': [x.task_id for x in tasks]})
+        self.cancel_scope = CancelScope()
 
-        # Wait for all tasks to complete
-        results: Dict[str, Any] = {}
-
-        async def _run_and_capture_result(task: BaseTask):
-            """
-            Run a task and capture the result
-            """
-            nonlocal results
-            result = await task.run(send_stream)
-            results[task.task_id] = result
-
-        if len(tasks) > 0:
-            try:
-                async with create_task_group() as tg:
-                    self.cancel_scope = tg.cancel_scope
-                    for task in tasks:
-                        tg.start_soon(_run_and_capture_result, task)
-            except BaseException as e:
-                if send_stream is not None:
-                    await send_stream.send(
-                        TaskError(task_id=self.task_id, error=e, cache_key=self.cache_key, reg_entry=self.reg_entry)
+        try:
+            with self.cancel_scope:
+                tasks: List[BaseTask] = []
+
+                # Collect up the tasks that need to be run and kick them off without awaiting them.
+                tasks.extend(x for x in self.args if isinstance(x, BaseTask))
+                tasks.extend(x for x in self.kwargs.values() if isinstance(x, BaseTask))
+
+                eng_logger.info(f'MetaTask {self.task_id} running sub-tasks', {'task_ids': [x.task_id for x in tasks]})
+
+                # Wait for all tasks to complete
+                results: Dict[str, Any] = {}
+
+                async def _run_and_capture_result(task: BaseTask):
+                    """
+                    Run a task and capture the result
+                    """
+                    nonlocal results
+                    result = await task.run(send_stream)
+                    results[task.task_id] = result
+
+                def handle_exception(err: ExceptionGroup):
+                    with contextlib.suppress(ClosedResourceError, BrokenResourceError):
+                        if send_stream is not None:
+                            send_stream.send_nowait(
+                                TaskError(
+                                    task_id=self.task_id, error=err, cache_key=self.cache_key, reg_entry=self.reg_entry
+                                )
+                            )
+                    raise err
+
+                if len(tasks) > 0:
+                    with catch(
+                        {
+                            BaseException: handle_exception,  # type: ignore
+                            get_cancelled_exc_class(): handle_exception,
+                        }
+                    ):
+                        async with create_task_group() as tg:
+                            for task in tasks:
+                                tg.start_soon(_run_and_capture_result, task)
+
+                        # In testing somehow sometimes cancellations aren't bubbled up, this ensures that's the case
+                        if tg.cancel_scope.cancel_called:
+                            raise get_cancelled_exc_class()('MetaTask caught cancellation')
+
+                eng_logger.debug(f'MetaTask {self.task_id}', 'completed sub-tasks', results)
+
+                # Order the results in the same order as the tasks list
+                result_values = [results[task.task_id] for task in tasks]
+
+                args = []
+                kwargs = {}
+
+                # Rebuild the args and kwargs with the results of the underlying tasks
+                # Here the task results could be DataFrames so make sure we clean the internal __index__ col from them
+                # before passing into the task function
+                for arg in self.args:
+                    if isinstance(arg, BaseTask):
+                        args.append(remove_index(result_values.pop(0)))
+                    else:
+                        args.append(arg)
+
+                for k, val in self.kwargs.items():
+                    if isinstance(val, BaseTask):
+                        kwargs[k] = remove_index(result_values.pop(0))
+                    else:
+                        kwargs[k] = val
+
+                eng_logger.debug(f'MetaTask {self.task_id}', 'processing result', {'args': args, 'kwargs': kwargs})
+
+                # Run the process result function with the completed set of args and kwargs
+                if self.process_as_task:
+                    eng_logger.debug(f'MetaTask {self.task_id}', 'processing result as Task')
+
+                    task = Task(
+                        self.process_result,
+                        args,
+                        kwargs,
+                        # Pass through cache_key so the processing task correctly updates the cache store entry
+                        reg_entry=self.reg_entry,
+                        cache_key=self.cache_key,
+                        task_id=self.task_id,
                     )
-                raise
-            finally:
-                self.cancel_scope = None
-
-        eng_logger.debug(f'MetaTask {self.task_id}', 'completed sub-tasks', results)
+                    res = await task.run(send_stream)
 
-        # Order the results in the same order as the tasks list
-        result_values = [results[task.task_id] for task in tasks]
+                    eng_logger.info(f'MetaTask {self.task_id} returning result', {'result': res})
 
-        args = []
-        kwargs = {}
+                    return res
 
-        # Rebuild the args and kwargs with the results of the underlying tasks
-        # Here the task results could be DataFrames so make sure we clean the internal __index__ col from them
-        # before passing into the task function
-        for arg in self.args:
-            if isinstance(arg, BaseTask):
-                args.append(remove_index(result_values.pop(0)))
-            else:
-                args.append(arg)
-
-        for k, val in self.kwargs.items():
-            if isinstance(val, BaseTask):
-                kwargs[k] = remove_index(result_values.pop(0))
-            else:
-                kwargs[k] = val
-
-        eng_logger.debug(f'MetaTask {self.task_id}', 'processing result', {'args': args, 'kwargs': kwargs})
-
-        # Run the process result function with the completed set of args and kwargs
-        if self.process_as_task:
-            eng_logger.debug(f'MetaTask {self.task_id}', 'processing result as Task')
-            # Pass through cache_key so the processing task correctly updates the cache store entry
-            task = Task(self.process_result, args, kwargs, cache_key=self.cache_key)
-            res = await task.run(send_stream)
+                try:
+                    res = await run_user_handler(self.process_result, args, kwargs)
 
-            eng_logger.info(f'MetaTask {self.task_id} returning result', {'result': res})
+                    # Send MetaTask result - it could be that there is a nested structure
+                    # of MetaTasks so we need to make sure intermediate results are also sent
+                    if send_stream is not None:
+                        await send_stream.send(
+                            TaskResult(
+                                task_id=self.task_id, result=res, cache_key=self.cache_key, reg_entry=self.reg_entry
+                            )
+                        )
+                except BaseException as e:
+                    # Recover from error - update the pending value to prevent subsequent requests getting stuck
+                    if send_stream is not None:
+                        await send_stream.send(
+                            TaskError(task_id=self.task_id, error=e, cache_key=self.cache_key, reg_entry=self.reg_entry)
+                        )
+                    raise
 
-            return res
+                eng_logger.info(f'MetaTask {self.task_id} returning result', {'result': res})
 
-        try:
-            res = await run_user_handler(self.process_result, args, kwargs)
-
-            # Send MetaTask result - it could be that there is a nested structure
-            # of MetaTasks so we need to make sure intermediate results are also sent
-            if send_stream is not None:
-                await send_stream.send(
-                    TaskResult(task_id=self.task_id, result=res, cache_key=self.cache_key, reg_entry=self.reg_entry)
-                )
-        except BaseException as e:
-            # Recover from error - update the pending value to prevent subsequent requests getting stuck
-            if send_stream is not None:
-                await send_stream.send(
-                    TaskError(task_id=self.task_id, error=e, cache_key=self.cache_key, reg_entry=self.reg_entry)
-                )
-            raise
-
-        eng_logger.info(f'MetaTask {self.task_id} returning result', {'result': res})
-
-        return res
+                return res
+        finally:
+            self.cancel_scope = None
 
     async def cancel(self):
         """
@@ -347,8 +371,10 @@ class TaskManager:
     TaskManager is responsible for running tasks and managing their pending state. It is also responsible for
     communicating the state of tasks to the client via the WebsocketManager.
 
-    When a task is run, a PendingTask it is stored in the tasks dict. It is also stored in the store
-    with the key of the task's cache_key. This allows the task to be retrieved by the cache_key from the store.
+    Every task created gets registered with the TaskManager and is tracked by the TaskManager
+    as a PendingTask.
+
+    This allows the task to be retrieved by the cache_key from the store.
 
     When a task is completed, it is removed from the tasks dict and the store entry is updated with the result.
 
@@ -361,13 +387,19 @@ class TaskManager:
         self.ws_manager = ws_manager
         self.store = store
 
+    def register_task(self, task: BaseTask) -> PendingTask:
+        """
+        Register a task. This will ensure the task it tracked and notifications are routed correctly.
+        """
+        pending_task = PendingTask(task.task_id, task)
+        self.tasks[task.task_id] = pending_task
+        return pending_task
+
     @overload
-    async def run_task(self, task: PendingTask, ws_channel: Optional[str] = None) -> Any:
-        ...
+    async def run_task(self, task: PendingTask, ws_channel: Optional[str] = None) -> Any: ...
 
     @overload
-    async def run_task(self, task: BaseTask, ws_channel: Optional[str] = None) -> PendingTask:
-        ...
+    async def run_task(self, task: BaseTask, ws_channel: Optional[str] = None) -> PendingTask: ...
 
     async def run_task(self, task: BaseTask, ws_channel: Optional[str] = None):
         """
@@ -380,7 +412,9 @@ class TaskManager:
         # append the websocket channel to the task
         if isinstance(task, PendingTask):
             if task.task_id in self.tasks:
-                if ws_channel:
+                # Increment subscriber count for this component request
+                self.tasks[task.task_id].add_subscriber()
+                if ws_channel is not None:
                     self.tasks[task.task_id].notify_channels.append(ws_channel)
                 return self.tasks[task.task_id]
 
@@ -394,21 +428,65 @@ class TaskManager:
                 else self.get_result(task.task_id)
             )
 
-        # Create and store the pending task
-        pending_task = PendingTask(task.task_id, task, ws_channel)
-        if task.cache_key is not None and task.reg_entry is not None:
-            await self.store.set(task.reg_entry, key=task.cache_key, value=pending_task)
-
-        self.tasks[task.task_id] = pending_task
+        # Otherwise, we should already have a pending task for this task
+        pending_task = self.tasks[task.task_id]
+        if ws_channel is not None:
+            pending_task.notify_channels.append(ws_channel)
 
         # Run the task in the background
-        self.task_group.start_soon(self._run_task_and_notify, task, ws_channel)
+        self.task_group.start_soon(self._run_task_and_notify, task)
 
         return pending_task
 
+    async def _cancel_tasks(self, task_ids: List[str], notify: bool = True):
+        """
+        Cancel a list of tasks
+
+        :param task_ids: The list of task IDs to cancel
+        :param notify: Whether to send cancellation notifications
+        """
+        with CancelScope(shield=True):
+            # Cancel all tasks in the hierarchy
+            for task_id_to_cancel in task_ids:
+                if task_id_to_cancel in self.tasks:
+                    pending_task = self.tasks[task_id_to_cancel]
+
+                    # Notify channels that this specific task was cancelled
+                    if notify:
+                        await self._send_notification_for_pending_task(
+                            pending_task=pending_task,
+                            messages=[{'status': 'CANCELED', 'task_id': task_id_to_cancel}],
+                        )
+
+                    if not pending_task.event.is_set():
+                        # Cancel the actual task
+                        await pending_task.cancel()
+
+                        # Remove from cache if it has cache settings
+                        if pending_task.task_def.cache_key is not None and pending_task.task_def.reg_entry is not None:
+                            await self.store.delete(
+                                pending_task.task_def.reg_entry, key=pending_task.task_def.cache_key
+                            )
+
+                    # Remove from running tasks
+                    self.tasks.pop(task_id_to_cancel, None)
+
+                    dev_logger.info('Task cancelled', {'task_id': task_id_to_cancel})
+
+    async def _cancel_task_hierarchy(self, task: BaseTask, notify: bool = True):
+        """
+        Recursively cancel all tasks in a task hierarchy
+
+        :param task: The root task to cancel (and its children)
+        :param notify: Whether to send cancellation notifications
+        """
+        all_task_ids = self._collect_all_task_ids_in_hierarchy(task)
+        await self._cancel_tasks(list(all_task_ids), notify)
+
     async def cancel_task(self, task_id: str, notify: bool = True):
         """
-        Cancel a running task by its id
+        Cancel a running task by its id. If the task has child tasks (MetaTask),
+        all child tasks will also be cancelled.
 
         :param task_id: the id of the task
         :param notify: whether to notify, true by default
@@ -423,22 +501,9 @@ class TaskManager:
                 task.remove_subscriber()
                 return
 
-            # Notify any listening channels that the job has been cancelled so that they can handle it correctly
-            if notify:
-                for channel in [*task.notify_channels, *task.task_def.notify_channels]:
-                    await self.ws_manager.send_message(channel, {'status': 'CANCELED', 'task_id': task_id})
-
-            # We're only now cancelling the task to make sure the clients are notified about cancelling
-            # and receive the correct status rather than an error
-            await task.cancel()
-
-            # Then remove the pending task from cache so next requests would recalculate rather than receive
-            # a broken pending task
-            if task.task_def.cache_key is not None and task.task_def.reg_entry is not None:
-                await self.store.set(task.task_def.reg_entry, key=task.task_def.cache_key, value=None)
+            # Cancel the entire task hierarchy (including child tasks)
+            await self._cancel_task_hierarchy(task.task_def, notify)
 
-            # Remove from running tasks
-            self.tasks.pop(task_id, None)
         else:
             raise TaskManagerError('Could not find a task with the passed id to cancel.')
 
@@ -459,12 +524,9 @@ class TaskManager:
 
         :param task_id: the id of the task to fetch
         """
-        result = await self.store.get(TaskResultEntry, key=task_id)
-
-        # Clean up the result afterwards
-        await self.store.delete(TaskResultEntry, key=task_id)
-
-        return result
+        # the result is not deleted, the results are kept in an LRU cache
+        # which will clean up older entries
+        return await self.store.get(TaskResultEntry, key=task_id)
 
     async def set_result(self, task_id: str, value: Any):
         """
@@ -472,7 +534,82 @@ class TaskManager:
         """
         return await self.store.set(TaskResultEntry, key=task_id, value=value)
 
-    async def _run_task_and_notify(self, task: BaseTask, ws_channel: Optional[str]):
+    def _collect_all_task_ids_in_hierarchy(self, task: BaseTask) -> Set[str]:
+        """
+        Recursively collect all task IDs in the task hierarchy
+
+        :param task: The root task to start collecting from
+        :return: Set of all task IDs in the hierarchy
+        """
+        task_ids = {task.task_id}
+
+        if isinstance(task, MetaTask):
+            # Collect from args
+            for arg in task.args:
+                if isinstance(arg, BaseTask):
+                    task_ids.update(self._collect_all_task_ids_in_hierarchy(arg))
+
+            # Collect from kwargs
+            for value in task.kwargs.values():
+                if isinstance(value, BaseTask):
+                    task_ids.update(self._collect_all_task_ids_in_hierarchy(value))
+
+        return task_ids
+
+    async def _multicast_notification(self, task_id: str, messages: List[dict]):
+        """
+        Send notifications to all task IDs that are related to a given task
+
+        :param task: the task the notifications are related to
+        :param messages: List of message dictionaries to send to all related tasks
+        """
+        # prevent cancellation, we need the notifications to be sent
+        with CancelScope(shield=True):
+            # Find all PendingTasks that have the message_task_id in their hierarchy
+            tasks_to_notify = set()
+
+            for pending_task in self.tasks.values():
+                # Check if the message_task_id is in this PendingTask's hierarchy
+                task_ids_in_hierarchy = self._collect_all_task_ids_in_hierarchy(pending_task.task_def)
+                if task_id in task_ids_in_hierarchy:
+                    tasks_to_notify.add(pending_task.task_id)
+
+            # Send notifications for all affected PendingTasks in parallel
+            if tasks_to_notify:
+                async with create_task_group() as task_tg:
+                    for pending_task_id in tasks_to_notify:
+                        if pending_task_id not in self.tasks:
+                            continue
+                        pending_task = self.tasks[pending_task_id]
+                        task_tg.start_soon(self._send_notification_for_pending_task, pending_task, messages)
+
+    async def _send_notification_for_pending_task(self, pending_task: PendingTask, messages: List[dict]):
+        """
+        Send notifications for a specific PendingTask
+
+        :param pending_task: The PendingTask to send notifications for
+        :param messages: The messages to send
+        """
+        # Collect channels for this PendingTask
+        channels_to_notify = set(pending_task.notify_channels)
+        channels_to_notify.update(pending_task.task_def.notify_channels)
+
+        if not channels_to_notify:
+            return
+
+        # Send to all channels for this PendingTask in parallel
+        async def _send_to_channel(channel: str):
+            async with create_task_group() as channel_tg:
+                for message in messages:
+                    # Create message with this PendingTask's task_id (if message has task_id)
+                    message_for_task = {**message, 'task_id': pending_task.task_id} if 'task_id' in message else message
+                    channel_tg.start_soon(self.ws_manager.send_message, channel, message_for_task)
+
+        async with create_task_group() as channel_tg:
+            for channel in channels_to_notify:
+                channel_tg.start_soon(_send_to_channel, channel)
+
+    async def _run_task_and_notify(self, task: BaseTask):
         """
         Run the task to completion and notify the client of progress and completion
 
@@ -481,23 +618,13 @@ class TaskManager:
         """
         cancel_scope = CancelScope()
 
-        self.tasks[task.task_id].cancel_scope = cancel_scope
+        pending_task = self.tasks[task.task_id]
+
+        pending_task.cancel_scope = cancel_scope
 
         with cancel_scope:
             eng_logger.info(f'TaskManager running task {task.task_id}')
 
-            async def notify_channels(*messages: dict):
-                """
-                Notify the channels of the task's progress
-                """
-                channels_to_notify = [*task.notify_channels]
-                if ws_channel:
-                    channels_to_notify.append(ws_channel)
-
-                for channel in channels_to_notify:
-                    for message in messages:
-                        await self.ws_manager.send_message(channel, message)
-
             # Create a memory object stream to capture messages from the tasks
             send_stream, receive_stream = create_memory_object_stream[TaskMessage](math.inf)
 
@@ -505,32 +632,41 @@ class TaskManager:
                 async with receive_stream:
                     async for message in receive_stream:
                         if isinstance(message, TaskProgressUpdate):
-                            # Notify the channels of the task's progress
-                            await notify_channels(
-                                {
-                                    'task_id': task.task_id,
-                                    'status': 'PROGRESS',
-                                    'progress': message.progress,
-                                    'message': message.message,
-                                }
-                            )
+                            # Send progress notifications to related tasks
+                            progress_message = {
+                                'task_id': message.task_id,  # Will be updated per task ID in multicast
+                                'status': 'PROGRESS',
+                                'progress': message.progress,
+                                'message': message.message,
+                            }
+                            await self._multicast_notification(message.task_id, [progress_message])
                             if isinstance(task, Task) and task.on_progress:
                                 await run_user_handler(task.on_progress, args=(message,))
                         elif isinstance(message, TaskResult):
-                            # Resolve the pending task related to the result
+                            # Handle dual coordination patterns:
+                            # 1. Direct-coordinated tasks: resolve via active_tasks registry
                             if message.task_id in self.tasks:
-                                self.tasks[task.task_id].resolve(message.result)
-                            # If the task has a cache key, update the cached value
+                                self.tasks[message.task_id].resolve(message.result)
+
+                            # 2. Cache-coordinated tasks: resolve via cache store (CacheStore.set handles PendingTask resolution)
                             if (
                                 message.cache_key is not None
                                 and message.reg_entry is not None
                                 and message.reg_entry.cache is not None
                             ):
                                 await self.store.set(message.reg_entry, key=message.cache_key, value=message.result)
-                            # Notify the channels of the task's completion
-                            await notify_channels(
-                                {'result': message.result, 'status': 'COMPLETE', 'task_id': message.task_id}
+
+                            # Set final result
+                            await self.set_result(message.task_id, message.result)
+
+                            # Notify all PendingTasks that depend on this specific task
+                            await self._multicast_notification(
+                                task_id=message.task_id,
+                                messages=[{'result': message.result, 'status': 'COMPLETE', 'task_id': message.task_id}],
                             )
+
+                            # Remove the task from the registered tasks - it finished running
+                            self.tasks.pop(message.task_id, None)
                         elif isinstance(message, TaskError):
                             # Fail the pending task related to the error
                             if message.task_id in self.tasks:
@@ -543,46 +679,94 @@ class TaskManager:
                                 and message.reg_entry is not None
                                 and message.reg_entry.cache is not None
                             ):
-                                await self.store.set(message.reg_entry, key=message.cache_key, value=None)
+                                await self.store.delete(message.reg_entry, key=message.cache_key)
+
+                            # Notify all PendingTasks that depend on this specific task
+                            error = get_error_for_channel(message.error)
+                            await self._multicast_notification(
+                                task_id=message.task_id,
+                                messages=[
+                                    {'status': 'ERROR', 'task_id': message.task_id, 'error': error['error']},
+                                    error,
+                                ],
+                            )
 
-            try:
-                async with create_task_group() as tg:
-                    # Handle incoming messages in parallel
-                    tg.start_soon(handle_messages)
+                            # Remove the task from the registered tasks - it finished running
+                            self.tasks.pop(message.task_id, None)
 
-                    # Handle tasks that return other tasks
-                    async with send_stream:
-                        result = task
-                        while isinstance(result, BaseTask):
-                            result = await task.run(send_stream)
+            task_error: Optional[ExceptionGroup] = None
 
-                        # Set final result
-                        await self.set_result(task.task_id, result)
+            # ExceptionGroup handler can't be async so we just mark the task as errored
+            # and run the async handler in the finally block
+            def handle_exception(err: ExceptionGroup):
+                nonlocal task_error
+                task_error = err
 
-                        # Notify any channels that need to be notified about the whole task being completed
-                        await send_stream.send(
-                            TaskResult(
+            try:
+                with catch({BaseException: handle_exception}):  # type: ignore
+                    async with create_task_group() as tg:
+                        # Handle incoming messages in parallel
+                        tg.start_soon(handle_messages)
+
+                        # Handle tasks that return other tasks
+                        async with send_stream:
+                            result = task
+                            while isinstance(result, BaseTask):
+                                result = await task.run(send_stream)
+
+                            # Notify any channels that need to be notified about the whole task being completed
+                            await send_stream.send(
+                                TaskResult(
+                                    task_id=task.task_id,
+                                    result=result,
+                                    cache_key=task.cache_key,
+                                    reg_entry=task.reg_entry,
+                                )
+                            )
+                            eng_logger.info(f'TaskManager finished task {task.task_id}', {'result': result})
+            finally:
+                with CancelScope(shield=True):
+                    # pyright: ignore[reportUnreachable]
+                    if task_error is not None:
+                        err = task_error
+                        # Mark pending task as failed
+                        pending_task.fail(err)
+
+                        await self.set_result(task.task_id, {'error': str(err)})
+
+                        # If the task has a cache key, set cached value to None
+                        # This makes it so that the next request will recalculate the value rather than keep failing
+                        if (
+                            task.cache_key is not None
+                            and task.reg_entry is not None
+                            and task.reg_entry.cache is not None
+                        ):
+                            await self.store.delete(task.reg_entry, key=task.cache_key)
+
+                        # If this is a cancellation, ensure all tasks in the hierarchy are cancelled
+                        if exception_group_contains(err_type=get_cancelled_exc_class(), group=err):
+                            dev_logger.info('Task cancelled', {'task_id': task.task_id})
+
+                            # Cancel any remaining tasks in the hierarchy that might still be running
+                            await self._cancel_task_hierarchy(task, notify=True)
+                        else:
+                            dev_logger.error('Task failed', err, {'task_id': task.task_id})
+
+                            error = get_error_for_channel(err)
+                            message = {'status': 'ERROR', 'task_id': task.task_id, 'error': error['error']}
+                            # Notify about this task failing, and a server broadcast error
+                            await self._send_notification_for_pending_task(
+                                pending_task=pending_task,
+                                messages=[message, error],
+                            )
+                            # notify related tasks
+                            await self._multicast_notification(
                                 task_id=task.task_id,
-                                result=result,
-                                cache_key=task.cache_key,
-                                reg_entry=task.reg_entry,
+                                messages=[message],
                             )
-                        )
-                        eng_logger.info(f'TaskManager finished task {task.task_id}', {'result': result})
-            except (Exception, ExceptionGroup) as err:
-                err = resolve_exception_group(err)
-
-                # Mark pending task as failed
-                self.tasks[task.task_id].fail(err)
 
-                dev_logger.error('Task failed', err, {'task_id': task.task_id})
-                await self.set_result(task.task_id, {'error': str(err)})
-
-                # Notify any channels that need to be notified
-                await notify_channels({'status': 'ERROR', 'task_id': task.task_id}, get_error_for_channel())
-            finally:
-                # Remove the task from the running tasks
-                self.tasks.pop(task.task_id, None)
+                    # Remove the task from the running tasks
+                    self.tasks.pop(task.task_id, None)
 
                 # Make sure streams are closed
                 with move_on_after(3, shield=True):
@@ -590,7 +774,7 @@ class TaskManager:
                     await receive_stream.aclose()
 
 
-TaskResultEntry = CachedRegistryEntry(uid='task-results', cache=Cache.Policy.KeepAll())
+TaskResultEntry = CachedRegistryEntry(uid='task-results', cache=LruCachePolicy(max_size=256))
 """
 Global registry entry for task results.
 This is global because task ids are unique and accessed one time only so it's effectively a one-time use random key.