keble-task 2.22.0__py3-none-any.whl

keble_task/main.py

@@ -0,0 +1,2708 @@
+import asyncio
+import logging
+import secrets
+import string
+import traceback
+import re
+from datetime import datetime, timedelta
+from enum import Enum
+from typing import (
+    Any,
+    Awaitable,
+    Callable,
+    List,
+    Literal,
+    Mapping,
+    Optional,
+    Union,
+    cast,
+    overload,
+)
+
+import keble_exceptions
+import tenacity
+from keble_db import AgentDbDeps, ExtendedAsyncRedis, QueryBase
+from keble_helpers import (
+    AgenticActionEventStatus,
+    AgenticEventEmitter,
+    Currency,
+    ExchangeRateInUsd,
+    Language,
+    Money,
+    ObjectId,
+    PydanticModelConfig,
+    SharingScope,
+    UsageAccountingRecorderProtocol,
+    UsageAccountingSource,
+    utc_now,
+)
+from motor.motor_asyncio import AsyncIOMotorClient
+from neo4j import AsyncDriver as Neo4jAsyncDriver
+from pydantic import BaseModel, ConfigDict, SkipValidation
+from pydantic_ai.usage import RunUsage
+from pymongo import ASCENDING, DESCENDING
+from pymongo.errors import DuplicateKeyError
+from qdrant_client import AsyncQdrantClient
+
+from keble_task.exceptions import (
+    TaskException,
+    TaskExceptionType,
+    TaskFailedToStartException,
+)
+
+from .actions import (
+    CreateRelatedTaskAction,
+    CreateRelatedTaskActionedResult,
+    TaskAction,
+    TaskActionCreatedRelation,
+    TaskActionCreatedTask,
+    TaskActionedResults,
+    TaskActionEvent,
+    TaskActions,
+    TaskActionType,
+    TaskEventType,
+    TaskLifecycleEvent,
+    TaskLifecycleEventPayload,
+    TaskUxContext,
+)
+from .crud import CRUDTask, CRUDTaskCost, CRUDTaskRelation
+from .schemas import (
+    TaskBase,
+    TaskCostAggregateRequest,
+    TaskCostAggregateResponse,
+    TaskCostCreate,
+    TaskCostListRequest,
+    TaskCostListResponse,
+    TaskCostMetadata,
+    TaskCostMongoObject,
+    TaskCostTokenRates,
+    TaskMetadata,
+    TaskMongoObject,
+    TaskMongoObjectExtended,
+    TaskPublicRef,
+    TaskRoomGraphContext,
+    TaskRoomResolution,
+    TaskRelationCreate,
+    TaskRelationMongoObject,
+    TaskStage,
+    TaskUpdate,
+)
+from .task_tree import build_task_tree, find_task_in_tree
+
+
+class TokenConsumptionType(str, Enum):
+    CONSUME = "CONSUME"
+    RECOVER = "RECOVER"
+
+
+class TokenConsumptionPayload(BaseModel):
+    """Host-owned token consumption/recovery request for one task.
+
+    Step by step:
+    1. carry the consumption type, owner, token amount, and task id;
+    2. carry the minimal DB clients the host token handler needs to price and
+       persist the consumption (replacing the retired `TaskResources` bag);
+    3. carry optional handler metadata for pricing context.
+
+    Side effects:
+    - consumed by the host-provided `_token_consumption_handler` in keble.backend;
+      that handler now reads `payload.amongo`/`extended_aredis`/`aneo4j` directly
+      instead of `payload.resources.*`.
+    """
+
+    model_config = ConfigDict(
+        **PydanticModelConfig.default_dict(),
+        arbitrary_types_allowed=True,  # for database clients
+    )
+    consumption_type: TokenConsumptionType
+    owner: str
+    token: int
+    task_id: Optional[ObjectId] = None
+    amongo: AsyncIOMotorClient
+    extended_aredis: ExtendedAsyncRedis
+    aneo4j: Optional[Neo4jAsyncDriver] = None
+    metadata: Optional[TaskMetadata] = None
+
+
+class TaskHandlerResponse(BaseModel):
+    model_config = PydanticModelConfig.default()
+    task: TaskMongoObject
+    success: bool
+    consuming_token: int
+    exception_type: Optional[TaskExceptionType] = None
+    error: Optional[str] = None
+
+
+class TaskHandlerRequest(AgentDbDeps):
+    """One task execution request: the DB-rooted agent deps plus the task itself.
+
+    Step by step:
+    1. inherits ALL DB clients + cross-cutting context (owner, owner_type,
+       owner_scope, event_emitter, usage_recorder, marketplace, language) from
+       `AgentDbDeps` — so a handler uses `request` directly as its deps object;
+    2. carries the persisted task row and optional handler metadata;
+    3. carries NO completion callback and NO resources bag — handlers RETURN a
+       `TaskHandlerResponse` (runtime finalizes) or `None` (the handler delegated
+       completion to another process; the task stays PROCESSING).
+
+    Side effects:
+    - every `atask_handler` (positioning, amz-product-report, backend handlers)
+      now reads clients off `request` directly and returns a response or None.
+    """
+
+    model_config = PydanticModelConfig.default(arbitrary_types_allowed=True)
+    task: TaskMongoObject
+    metadata: Optional[TaskMetadata] = None
+
+
+logger = logging.getLogger(__name__)
+
+
+def _merge_metadata(
+    base: Optional[Mapping[str, Any]], override: Optional[Mapping[str, Any]]
+) -> Optional[dict[str, Any]]:
+    if base is None and override is None:
+        return None
+    if base is None:
+        assert override is not None
+        return dict(override)
+    if override is None:
+        return dict(base)
+    return {**base, **override}
+
+
+def _task_lifecycle_event_status(*, stage: TaskStage) -> AgenticActionEventStatus:
+    """Map persisted task stage to the shared agentic event status.
+
+    Step by step:
+    1. processing is a lifecycle start/progress signal for room UIs;
+    2. success and failure map to terminal event statuses;
+    3. pending is retained as a non-terminal progress state for direct calls.
+    """
+
+    if stage is TaskStage.PROCESSING:
+        return AgenticActionEventStatus.STARTED
+    if stage is TaskStage.SUCCESS:
+        return AgenticActionEventStatus.SUCCEEDED
+    if stage is TaskStage.FAILURE:
+        return AgenticActionEventStatus.FAILED
+    return AgenticActionEventStatus.PROGRESSED
+
+
+class TaskClient:
+    def __init__(
+        self,
+        *,
+        # Token consumption handler can be synchronous or asynchronous
+        # If sync, it should return a bool
+        # If async, it should return a coroutine that resolves to a bool
+        token_consumption_handler: Callable[
+            [TokenConsumptionPayload], Union[bool, Awaitable[bool]]
+        ],
+        # Run one task. The handler RETURNS a `TaskHandlerResponse` (the runtime
+        # finalizes success/failure and emits the terminal lifecycle event) OR
+        # `None` when it deliberately delegated completion to another process
+        # (e.g. fire-and-forget bootstrap, multi-invocation pipeline) — the task
+        # then stays PROCESSING until that process finalizes it.
+        task_handler: Callable[
+            [TaskHandlerRequest], Awaitable[Optional[TaskHandlerResponse]]
+        ],
+        # Database config
+        mongo_database: str = "__keble_task__",
+        task_collection: str = "__keble_task__task__",
+        task_relation_collection: str = "__keble_task__task_relation__",
+        task_cost_collection: str = "__keble_task__task_cost__",
+        # Public short id config (for shared URLs)
+        public_id_prefix: str = "t",
+        public_id_length: int = 10,
+    ):
+        """Initialize the TaskClient.
+
+        Step by step:
+        1. store runtime handlers without opening database connections;
+        2. configure the task, relation, and cost collections used by later calls;
+        3. validate public-id settings before any public task is created.
+
+        Note: Resources are not initialized here but passed to each method when needed.
+        """
+        self._token_consumption_handler = token_consumption_handler
+        self._mongo_database = mongo_database
+        self._task_collection = task_collection
+        self._task_relation_collection = task_relation_collection
+        self._task_cost_collection = task_cost_collection
+        self._task_handler = task_handler
+        self._public_id_prefix = public_id_prefix
+        self._public_id_length = public_id_length
+        self._public_id_alphabet = string.ascii_letters + string.digits
+
+        self._validate_public_id_config(
+            prefix=self._public_id_prefix, length=self._public_id_length
+        )
+
+        self.crud_task = CRUDTask(
+            model=TaskMongoObject,
+            collection=self._task_collection,
+            database=mongo_database,
+        )
+        self.crud_task_relation = CRUDTaskRelation(
+            model=TaskRelationMongoObject,
+            collection=self._task_relation_collection,
+            database=mongo_database,
+        )
+        self.crud_task_cost = CRUDTaskCost(
+            model=TaskCostMongoObject,
+            collection=self._task_cost_collection,
+            database=mongo_database,
+        )
+
+    def _validate_public_id_config(self, *, prefix: str, length: int) -> None:
+        if not isinstance(prefix, str):
+            raise keble_exceptions.ServerSideInvalidParams(
+                admin_note={"prefix": prefix},
+                alert_admin=True,
+                but_got="prefix is not str",
+                expected="prefix is str",
+                invalid_params="prefix",
+            )
+        if not isinstance(length, int):
+            raise keble_exceptions.ServerSideInvalidParams(
+                admin_note={"length": length},
+                alert_admin=True,
+                but_got="length is not int",
+                expected="length is int",
+                invalid_params="length",
+            )
+        if length <= len(prefix):
+            raise keble_exceptions.ServerSideInvalidParams(
+                admin_note={"prefix": prefix, "length": length},
+                alert_admin=True,
+                but_got="length <= len(prefix)",
+                expected="length > len(prefix)",
+                invalid_params="length,prefix",
+            )
+
+    async def _ensure_public_id_index(self, *, amongo: AsyncIOMotorClient) -> None:
+        await self.crud_task.aensure_public_id_index(amongo)
+
+    async def _ensure_task_indexes(self, *, amongo: AsyncIOMotorClient) -> None:
+        """Ensure the task collection has the tree/stage read indexes.
+
+        Step by step:
+        1. delegate the index definitions to the task CRUD helper;
+        2. keep index creation out of request handlers and worker loops;
+        3. share one idempotent setup between startup and package methods.
+        """
+
+        await self.crud_task.aensure_task_indexes(amongo)
+
+    def _generate_public_id_candidate(self, *, prefix: str, length: int) -> str:
+        suffix_len = length - len(prefix)
+        suffix = "".join(
+            secrets.choice(self._public_id_alphabet) for _ in range(suffix_len)
+        )
+        return f"{prefix}{suffix}"
+
+    async def _ensure_task_public_id(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        task_id: ObjectId,
+        public_id_prefix: Optional[str],
+        public_id_length: Optional[int],
+    ) -> str:
+        prefix = (
+            public_id_prefix if public_id_prefix is not None else self._public_id_prefix
+        )
+        length = (
+            public_id_length if public_id_length is not None else self._public_id_length
+        )
+        self._validate_public_id_config(prefix=prefix, length=length)
+        await self._ensure_public_id_index(amongo=amongo)
+
+        collection = amongo[self._mongo_database][self._task_collection]
+        existing = await collection.find_one({"_id": task_id}, {"public_id": 1})
+        if existing is None:
+            raise keble_exceptions.ServerSideInvalidParams(
+                admin_note={"task_id": task_id},
+                alert_admin=True,
+                but_got="task not found",
+                expected="task exists",
+                invalid_params="task_id",
+            )
+        if isinstance(existing.get("public_id"), str) and existing.get("public_id"):
+            return str(existing["public_id"])
+
+        for _ in range(20):
+            candidate = self._generate_public_id_candidate(prefix=prefix, length=length)
+            try:
+                result = await collection.update_one(
+                    {"_id": task_id, "public_id": None},
+                    {"$set": {"public_id": candidate}},
+                )
+            except DuplicateKeyError:
+                continue
+
+            if result.matched_count == 0:
+                refreshed = await collection.find_one(
+                    {"_id": task_id}, {"public_id": 1}
+                )
+                if refreshed and refreshed.get("public_id"):
+                    return str(refreshed["public_id"])
+                raise keble_exceptions.ServerSideInvalidParams(
+                    admin_note={"task_id": task_id, "candidate": candidate},
+                    alert_admin=True,
+                    but_got="task public_id not updated",
+                    expected="task public_id updated",
+                    invalid_params="task_id,public_id",
+                )
+            return candidate
+
+        raise keble_exceptions.ServerSideInvalidParams(
+            admin_note={"task_id": task_id, "prefix": prefix, "length": length},
+            alert_admin=True,
+            but_got="exceeded retries to generate unique public_id",
+            expected="unique public_id generated",
+            invalid_params="public_id",
+        )
+
+    async def _ensure_task_relation_indexes(
+        self, *, amongo: AsyncIOMotorClient
+    ) -> None:
+        """Ensure the relation collection has the required query indexes.
+
+        Step by step:
+        1. delegate index ownership to the relation CRUD helper;
+        2. keep this method private so callers use relation APIs, not storage details.
+        """
+
+        await self.crud_task_relation.aensure_relation_indexes(amongo)
+
+    async def _ensure_task_cost_indexes(self, *, amongo: AsyncIOMotorClient) -> None:
+        """Ensure the cost collection has the required query indexes.
+
+        Step by step:
+        1. delegate index definitions to the task-cost CRUD helper;
+        2. keep index creation outside request handlers and worker loops;
+        3. let startup and package methods share the same idempotent setup.
+        """
+
+        await self.crud_task_cost.aensure_cost_indexes(amongo)
+
+    async def aensure_indexes(self, *, amongo: AsyncIOMotorClient) -> None:
+        """Create all Mongo indexes owned by the task package.
+
+        Step by step:
+        1. create the public task id index used by shared task URLs;
+        2. create task tree/stage indexes for root, child, and retry reads;
+        3. create relation indexes for root-scoped graph lookups;
+        4. create cost indexes for admin list and aggregate reports.
+        """
+
+        await self._ensure_public_id_index(amongo=amongo)
+        await self._ensure_task_indexes(amongo=amongo)
+        await self._ensure_task_relation_indexes(amongo=amongo)
+        await self._ensure_task_cost_indexes(amongo=amongo)
+
+    async def ainit(self, *, amongo: AsyncIOMotorClient) -> None:
+        """Initialize task Mongo indexes for backend startup.
+
+        Step by step:
+        1. expose one public startup hook for backend lifespan wiring;
+        2. delegate to the existing package-owned index initializer.
+        """
+
+        await self.aensure_indexes(amongo=amongo)
+
+    def _task_workspace_root_id(self, task: TaskMongoObject) -> ObjectId:
+        """Resolve the root-task id used to scope workspace-local relation rows.
+
+        Step by step:
+        1. prefer the persisted `root_task` field for normal modern rows;
+        2. fall back to the task id for old root rows that predate root backfill.
+        """
+
+        return task.root_task or task.id
+
+    async def _aget_relation_task(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        task_id: ObjectId,
+        field_name: str,
+    ) -> TaskMongoObject:
+        """Load one task referenced by a relation row and raise a typed error.
+
+        Step by step:
+        1. load the task without a task-type filter;
+        2. return the typed task row when it exists;
+        3. raise a package-consistent invalid-params error when the reference is stale.
+        """
+
+        task = await self.aget(amongo=amongo, task_id=task_id, task_type=None)
+        if task is None:
+            raise keble_exceptions.ServerSideInvalidParams(
+                admin_note={field_name: task_id},
+                alert_admin=True,
+                but_got=f"{field_name} task not found",
+                expected=f"{field_name} task exists",
+                invalid_params=field_name,
+            )
+        return task
+
+    def _validate_relation_task_root(
+        self,
+        *,
+        relation: TaskRelationCreate,
+        task: TaskMongoObject,
+        field_name: str,
+    ) -> None:
+        """Validate that one relation endpoint belongs to the relation root.
+
+        Step by step:
+        1. resolve the endpoint task's canonical workspace root;
+        2. compare it with the relation root;
+        3. raise a typed error when a caller attempts cross-workspace lineage.
+        """
+
+        task_root = self._task_workspace_root_id(task)
+        if task_root != relation.root_task:
+            raise keble_exceptions.ServerSideInvalidParams(
+                admin_note={
+                    "root_task": relation.root_task,
+                    field_name: task.id,
+                    f"{field_name}_root_task": task_root,
+                },
+                alert_admin=True,
+                but_got=f"{field_name} belongs to a different root task",
+                expected=f"{field_name} belongs to relation.root_task",
+                invalid_params=f"root_task,{field_name}",
+            )
+
+    @staticmethod
+    def _validate_unique_task_relation_edges(
+        *,
+        relations: list[TaskRelationCreate],
+        invalid_params: str,
+    ) -> None:
+        """Reject duplicate relation edge payloads before any write happens.
+
+        Step by step:
+        1. derive the same edge identity protected by the Mongo unique index;
+        2. collect duplicates without mutating task or relation collections;
+        3. raise a caller-facing validation error instead of relying on a partial
+           insert followed by a database duplicate-key failure.
+        """
+
+        seen_edge_keys: set[tuple[ObjectId, ObjectId, ObjectId, str]] = set()
+        duplicate_edges: list[str] = []
+        for relation in relations:
+            edge_key = (
+                relation.root_task,
+                relation.from_task_id,
+                relation.to_task_id,
+                relation.relation_type.value,
+            )
+            if edge_key in seen_edge_keys:
+                duplicate_edges.append(
+                    "root_task="
+                    f"{relation.root_task};from_task_id={relation.from_task_id};"
+                    f"to_task_id={relation.to_task_id};"
+                    f"relation_type={relation.relation_type.value}"
+                )
+                continue
+            seen_edge_keys.add(edge_key)
+        if duplicate_edges:
+            raise keble_exceptions.ClientSideInvalidParams(
+                invalid_params=invalid_params,
+                expected="unique relation edge keys",
+                but_got="; ".join(duplicate_edges),
+                alert_admin=False,
+            )
+
+    @staticmethod
+    def _validate_unique_create_related_source_ids(
+        *,
+        from_task_ids: list[ObjectId],
+    ) -> None:
+        """Reject duplicate action source ids before creating the child task.
+
+        Step by step:
+        1. check the source id list that will become relation edges;
+        2. stop duplicate source ids before the child task is inserted;
+        3. return normally when each source would create one unique edge.
+        """
+
+        seen_task_ids: set[ObjectId] = set()
+        duplicate_task_ids: list[str] = []
+        for from_task_id in from_task_ids:
+            if from_task_id in seen_task_ids:
+                duplicate_task_ids.append(str(from_task_id))
+                continue
+            seen_task_ids.add(from_task_id)
+        if duplicate_task_ids:
+            raise keble_exceptions.ClientSideInvalidParams(
+                invalid_params="actions[].from_task_ids",
+                expected="unique source task ids per created relation edge",
+                but_got=",".join(duplicate_task_ids),
+                alert_admin=False,
+            )
+
+    async def _avalidate_task_relation_create(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        relation: TaskRelationCreate,
+    ) -> None:
+        """Validate one task-relation create payload before insertion.
+
+        Step by step:
+        1. load the declared root, source task, and derived task;
+        2. prove the declared root is the canonical workspace root;
+        3. prove both endpoints belong to the same root workspace.
+        """
+
+        root_task = await self._aget_relation_task(
+            amongo=amongo,
+            task_id=relation.root_task,
+            field_name="root_task",
+        )
+        self._validate_relation_task_root(
+            relation=relation,
+            task=root_task,
+            field_name="root_task",
+        )
+        from_task = await self._aget_relation_task(
+            amongo=amongo,
+            task_id=relation.from_task_id,
+            field_name="from_task_id",
+        )
+        self._validate_relation_task_root(
+            relation=relation,
+            task=from_task,
+            field_name="from_task_id",
+        )
+        to_task = await self._aget_relation_task(
+            amongo=amongo,
+            task_id=relation.to_task_id,
+            field_name="to_task_id",
+        )
+        self._validate_relation_task_root(
+            relation=relation,
+            task=to_task,
+            field_name="to_task_id",
+        )
+
+    async def _execute_token_consumption_handler(
+        self, payload: TokenConsumptionPayload
+    ) -> bool:
+        """Execute token consumption handler supporting both sync and async implementations.
+
+        Args:
+            payload: The token consumption payload
+
+        Returns:
+            bool: True if successful, False otherwise
+        """
+        result = self._token_consumption_handler(payload)
+        if asyncio.iscoroutine(result):
+            # If it's a coroutine, await it
+            return await cast(Awaitable[bool], result)
+        # If it's a boolean, return it directly
+        return cast(bool, result)
+
+    #
+    #
+    #
+    #
+    # Core running Task API
+    #
+    #
+    #
+    #
+    # Create
+    async def acreate(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        expected_token: int,
+        owner: str,
+        progress_key: Optional[str] = None,
+        image: Optional[str] = None,
+        title: Optional[str] = None,
+        subtitle: Optional[str] = None,
+        metadata: Optional[TaskMetadata] = None,
+        language: Optional[Language] = None,
+        task_type: str,
+        timeout_mins: int = 120,
+        sharing_scope: SharingScope = SharingScope.PRIVATE,
+        stage: TaskStage = TaskStage.PENDING,
+        attempts: int = 0,
+        consumed_token: int = 0,
+        # task tree fields
+        root_task: Optional[ObjectId] = None,
+        parent_task: Optional[ObjectId] = None,
+        public_id_prefix: Optional[str] = None,
+        public_id_length: Optional[int] = None,
+    ) -> TaskMongoObject:
+        """Create a new task asynchronously."""
+        if parent_task is None and root_task is not None:
+            raise keble_exceptions.ServerSideInvalidParams(
+                admin_note={"root_task": root_task, "parent_task": parent_task},
+                alert_admin=True,
+                but_got="root_task provided for root task creation",
+                expected="root_task is None when parent_task is None",
+                invalid_params="root_task,parent_task",
+            )
+        if parent_task is not None and root_task is None:
+            parent = await self.aget(amongo=amongo, task_id=parent_task, task_type=None)
+            if parent is None:
+                raise keble_exceptions.ServerSideInvalidParams(
+                    admin_note={"parent_task": parent_task},
+                    alert_admin=True,
+                    but_got="parent task not found",
+                    expected="parent task exists",
+                    invalid_params="parent_task",
+                )
+            root_task = parent.root_task or parent.id
+
+        if sharing_scope == SharingScope.PUBLIC:
+            prefix = (
+                public_id_prefix
+                if public_id_prefix is not None
+                else self._public_id_prefix
+            )
+            length = (
+                public_id_length
+                if public_id_length is not None
+                else self._public_id_length
+            )
+            self._validate_public_id_config(prefix=prefix, length=length)
+
+        result = await self.crud_task.acreate(
+            amongo,
+            obj_in=cast(
+                TaskMongoObject,
+                TaskBase(
+                    progress_key=progress_key,
+                    image=image,
+                    title=title,
+                    subtitle=subtitle,
+                    metadata=metadata,
+                    language=language,
+                    owner=str(owner),
+                    consumed_token=consumed_token,
+                    expected_token=expected_token,
+                    attempts=attempts,
+                    sharing_scope=sharing_scope,
+                    stage=stage,
+                    task_type=task_type,
+                    timeout_mins=timeout_mins,
+                    root_task=root_task,
+                    parent_task=parent_task,
+                ),
+            ),
+        )
+        _id = result.inserted_id
+        if parent_task is None and root_task is None:
+            await self.crud_task.aupdate(amongo, _id=_id, obj_in={"root_task": _id})
+        if sharing_scope == SharingScope.PUBLIC:
+            await self._ensure_task_public_id(
+                amongo=amongo,
+                task_id=_id,
+                public_id_prefix=public_id_prefix,
+                public_id_length=public_id_length,
+            )
+        task = await self.crud_task.afirst_by_id(amongo, _id=_id)
+        assert task is not None
+        return task
+
+    async def acreate_task_relation(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        obj_in: TaskRelationCreate,
+    ) -> TaskRelationMongoObject:
+        """Create one extra lineage relation without changing task parentage.
+
+        Step by step:
+        1. validate the source and derived tasks belong to the declared root;
+        2. insert a relation row into the separate relation collection;
+        3. return the persisted relation object for downstream API responses.
+        """
+
+        results = await self.acreate_task_relations(amongo=amongo, objs_in=[obj_in])
+        return results[0]
+
+    async def acreate_task_relations(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        objs_in: list[TaskRelationCreate],
+    ) -> list[TaskRelationMongoObject]:
+        """Create many extra lineage relations as independent durable rows.
+
+        Step by step:
+        1. ensure lookup and duplicate-protection indexes exist;
+        2. validate every relation before inserting any new rows;
+        3. insert each relation without touching `root_task` or `parent_task`.
+        """
+
+        if not objs_in:
+            return []
+
+        self._validate_unique_task_relation_edges(
+            relations=objs_in,
+            invalid_params="objs_in",
+        )
+        await self._ensure_task_relation_indexes(amongo=amongo)
+        for relation in objs_in:
+            await self._avalidate_task_relation_create(
+                amongo=amongo,
+                relation=relation,
+            )
+
+        created: list[TaskRelationMongoObject] = []
+        for relation in objs_in:
+            result = await self.crud_task_relation.acreate_relation(
+                amongo,
+                obj_in=relation,
+            )
+            db_obj = await self.crud_task_relation.afirst_by_id(
+                amongo,
+                _id=result.inserted_id,
+            )
+            assert db_obj is not None
+            created.append(db_obj)
+        return created
+
+    async def acreate_task_cost(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        task_id: ObjectId,
+        tags: list[str] | None = None,
+        run_usage: RunUsage,
+        source: UsageAccountingSource | None = None,
+        token_rates_per_million: TaskCostTokenRates | None = None,
+        additional_cost: Money | None = None,
+        seconds: float = 0,
+        retry: int = 0,
+        occurred_at: datetime | None = None,
+        metadata: TaskCostMetadata | None = None,
+    ) -> TaskCostMongoObject:
+        """Create one task-cost row denormalized from the current task row.
+
+        Step by step:
+        1. load the task so owner, type, and root scope cannot drift from storage;
+        2. build a cost create schema with raw `RunUsage`, rates, timing, and metadata;
+        3. persist the row in the separate cost collection and reload it as a typed object.
+        """
+
+        task = await self.aget(amongo=amongo, task_id=task_id, task_type=None)
+        if task is None:
+            raise keble_exceptions.ServerSideInvalidParams(
+                admin_note={"task_id": task_id},
+                alert_admin=True,
+                but_got="task not found",
+                expected="task exists",
+                invalid_params="task_id",
+            )
+        result = await self.crud_task_cost.acreate_cost(
+            amongo,
+            obj_in=TaskCostCreate.from_task_usage(
+                task=task,
+                root_task=self._task_workspace_root_id(task),
+                tags=tags,
+                run_usage=run_usage,
+                source=source,
+                token_rates_per_million=token_rates_per_million,
+                additional_cost=additional_cost,
+                seconds=seconds,
+                retry=retry,
+                occurred_at=occurred_at,
+                metadata=metadata,
+            ),
+        )
+        db_obj = await self.crud_task_cost.afirst_by_id(
+            amongo,
+            _id=result.inserted_id,
+        )
+        assert db_obj is not None
+        return db_obj
+
+    async def alist_task_costs(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        query: TaskCostListRequest,
+    ) -> TaskCostListResponse:
+        """List task-cost rows through the package-owned indexed query shape.
+
+        Step by step:
+        1. list the requested page using `CRUDTaskCost`;
+        2. count the same filter for a compact paginated response;
+        3. rely on startup/`ainit(...)` for index creation.
+        """
+
+        costs = await self.crud_task_cost.alist_costs(amongo, query=query)
+        total = await self.crud_task_cost.acount_costs(amongo, query=query)
+        return TaskCostListResponse(
+            costs=costs,
+            total=total,
+            skip=query.skip,
+            limit=query.limit,
+        )
+
+    async def aaggregate_task_costs(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        query: TaskCostAggregateRequest,
+        exchange_rates: list[ExchangeRateInUsd],
+    ) -> TaskCostAggregateResponse:
+        """Aggregate task costs with explicit currency conversion.
+
+        Step by step:
+        1. load rows through the same indexed filters used by list reads;
+        2. calculate each row's precise Decimal cost in the requested currency;
+        3. group by requested time bucket and optionally fan rows into tag buckets.
+        """
+
+        costs = await self.crud_task_cost.alist_costs_for_aggregate(
+            amongo,
+            query=query,
+        )
+        return TaskCostAggregateResponse.build(
+            costs=costs,
+            query=query,
+            exchange_rates=exchange_rates,
+        )
+
+    def _validate_task_action_root_context(
+        self,
+        *,
+        current_task: TaskMongoObject,
+        ux_context: TaskUxContext | None,
+    ) -> None:
+        """Validate optional UI root context before running task actions.
+
+        Step by step:
+        1. skip validation when no UI root was provided;
+        2. resolve the current task's canonical root workspace id;
+        3. reject mismatched UI roots before creating any child task or relation.
+        """
+
+        if ux_context is None or ux_context.selected_root_task_id is None:
+            return
+        current_root_id = self._task_workspace_root_id(current_task)
+        if ux_context.selected_root_task_id != current_root_id:
+            raise keble_exceptions.ClientSideInvalidParams(
+                invalid_params="ux_context.selected_root_task_id",
+                expected=str(current_root_id),
+                but_got=str(ux_context.selected_root_task_id),
+                alert_admin=False,
+            )
+
+    async def _aprepare_create_related_task_action(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        current_task: TaskMongoObject,
+        action: CreateRelatedTaskAction,
+    ) -> TaskMongoObject:
+        """Load the parent task for a generic create-related-task action.
+
+        Step by step:
+        1. resolve `parent_task_id` to the current task when omitted;
+        2. load the parent task so `acreate(...)` can preserve root linkage;
+        3. return the parent task without any feature-specific side effects.
+        """
+
+        parent_task_id = action.parent_task_id or current_task.id
+        return await self._aget_relation_task(
+            amongo=amongo,
+            task_id=parent_task_id,
+            field_name="parent_task_id",
+        )
+
+    async def _acomplete_created_task_if_needed(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        extended_aredis: ExtendedAsyncRedis | None,
+        created_task: TaskMongoObject,
+        complete_created_task: bool,
+        consuming_token: int,
+    ) -> TaskMongoObject:
+        """Mark a lightweight created task successful when the action requests it.
+
+        Step by step:
+        1. return the task unchanged for normal queued child work;
+        2. require Redis resources for task finalization;
+        3. delegate token/status accounting to the existing task success API.
+        """
+
+        if not complete_created_task:
+            return created_task
+        if extended_aredis is None:
+            raise keble_exceptions.ServerSideMissingParams(
+                missing_params="extended_aredis",
+                alert_admin=True,
+            )
+        return await self.aon_task_success(
+            amongo=amongo,
+            extended_aredis=extended_aredis,
+            task=created_task,
+            consuming_token=consuming_token,
+        )
+
+    async def _arollback_created_related_task(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        created_task_id: ObjectId,
+    ) -> None:
+        """Compensate a failed create-related-task action.
+
+        Step by step:
+        1. delete any relation rows already inserted that point at the orphaned child
+           (`to_task_id == created_task_id`);
+        2. delete the orphaned child task itself;
+        3. leave the original failure to propagate from the caller.
+
+        Used only on the relation-write failure path because standalone Mongo cannot run
+        the child + relation inserts in one atomic transaction.
+        """
+
+        await self.crud_task_relation.adelete_multi(
+            amongo,
+            query=QueryBase(filters={"to_task_id": created_task_id}),
+        )
+        await self.crud_task.adelete(amongo, _id=created_task_id)
+
+    async def _aapply_create_related_task_action(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        current_task: TaskMongoObject,
+        action: CreateRelatedTaskAction,
+        extended_aredis: ExtendedAsyncRedis | None,
+    ) -> CreateRelatedTaskActionedResult:
+        """Apply one generic create-related-task action.
+
+        Step by step:
+        1. resolve parent/root context;
+        2. validate every source task belongs to the same root room;
+        3. create the child task from generic task fields only;
+        4. create relation rows from source tasks to the created task;
+        5. optionally mark the child task successful for lightweight generic work.
+        """
+
+        parent_task = await self._aprepare_create_related_task_action(
+            amongo=amongo,
+            current_task=current_task,
+            action=action,
+        )
+        from_task_ids = action.from_task_ids or [current_task.id]
+        self._validate_unique_create_related_source_ids(from_task_ids=from_task_ids)
+        parent_root_id = self._task_workspace_root_id(parent_task)
+        for from_task_id in from_task_ids:
+            from_task = await self._aget_relation_task(
+                amongo=amongo,
+                task_id=from_task_id,
+                field_name="from_task_ids",
+            )
+            from_root_id = self._task_workspace_root_id(from_task)
+            if from_root_id != parent_root_id:
+                raise keble_exceptions.ClientSideInvalidParams(
+                    invalid_params="from_task_ids",
+                    expected=str(parent_root_id),
+                    but_got=f"from_task_id={from_task_id}; from_task_root_id={from_root_id}",
+                    alert_admin=False,
+                )
+        if action.task_type is None:
+            raise keble_exceptions.ClientSideMissingParams(
+                missing_params="actions[].task_type",
+                alert_admin=False,
+            )
+        created_task = await self.acreate(
+            amongo=amongo,
+            expected_token=action.expected_token,
+            owner=current_task.owner,
+            progress_key=action.progress_key,
+            image=action.image,
+            title=action.title,
+            subtitle=action.subtitle,
+            metadata=action.metadata,
+            task_type=action.task_type,
+            timeout_mins=action.timeout_mins,
+            sharing_scope=action.sharing_scope or current_task.sharing_scope,
+            parent_task=parent_task.id,
+        )
+        # Saga compensation: standalone Mongo has no multi-document transactions, so the
+        # child task and its relation rows are written as separate operations. If any
+        # relation write fails (a concurrent duplicate hitting the unique edge index, or
+        # an infra error mid-loop), roll back the just-created child and any partial
+        # relation rows pointing at it, then re-raise unchanged so the workspace never
+        # keeps an orphaned child without its intended relations.
+        try:
+            relations = await self.acreate_task_relations(
+                amongo=amongo,
+                objs_in=[
+                    TaskRelationCreate(
+                        root_task=parent_root_id,
+                        from_task_id=from_task_id,
+                        to_task_id=created_task.id,
+                        relation_type=action.relation_type,
+                        metadata=action.metadata,
+                    )
+                    for from_task_id in from_task_ids
+                ],
+            )
+        except Exception:
+            await self._arollback_created_related_task(
+                amongo=amongo, created_task_id=created_task.id
+            )
+            raise
+        finalized_task = await self._acomplete_created_task_if_needed(
+            amongo=amongo,
+            extended_aredis=extended_aredis,
+            created_task=created_task,
+            complete_created_task=action.complete_created_task,
+            consuming_token=action.consuming_token,
+        )
+        return CreateRelatedTaskActionedResult(
+            created_task=TaskActionCreatedTask.build(task=finalized_task),
+            created_relations=[
+                TaskActionCreatedRelation.build(relation=relation)
+                for relation in relations
+            ],
+        )
+
+    async def _aapply_one_task_action(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        current_task: TaskMongoObject,
+        action: TaskAction,
+        extended_aredis: ExtendedAsyncRedis | None,
+    ) -> tuple[CreateRelatedTaskActionedResult, TaskActionEvent]:
+        """Execute exactly one canonical task action and build its event.
+
+        The single canonical executor shared by the one-action agent surface
+        (`aapply_action`) and the REST/programmatic batch (`aapply_actions`), so
+        action semantics and event shape never fork between the two callers.
+        """
+
+        # CREATE_RELATED_TASK is the only action type today; keep the explicit
+        # branch so a future action type cannot silently no-op.
+        if action.action_type is not TaskActionType.CREATE_RELATED_TASK:
+            raise keble_exceptions.ServerSideInvalidParams(
+                admin_note={"action_type": str(action.action_type)},
+                alert_admin=True,
+                but_got=f"unsupported task action type {action.action_type}",
+                expected="a supported TaskActionType (CREATE_RELATED_TASK)",
+                invalid_params="action_type",
+            )
+        result = await self._aapply_create_related_task_action(
+            amongo=amongo,
+            current_task=current_task,
+            action=action,
+            extended_aredis=extended_aredis,
+        )
+        event = TaskActionEvent(
+            action_type=result.action_type.value,
+            payload=result,
+            root_id=str(
+                result.created_task.root_task_id or result.created_task.task_id
+            ),
+            object_id=str(result.created_task.task_id),
+            correlation_id=result.created_task.progress_key,
+        )
+        return result, event
+
+    async def aapply_action(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        current_task: TaskMongoObject,
+        action: TaskAction,
+        ux_context: TaskUxContext | None = None,
+        event_emitter: AgenticEventEmitter | None = None,
+        extended_aredis: ExtendedAsyncRedis | None = None,
+    ) -> CreateRelatedTaskActionedResult:
+        """Apply ONE canonical task action (the agent one-action-per-call surface).
+
+        Multiplicity is expressed by the agent emitting several
+        `mutate_task_workspace` tool calls; each becomes its own deferred
+        approval card. This method intentionally takes a single
+        `AgenticActionPayload` action — never a list wrapper.
+
+        Step by step:
+        1. validate optional UI root context against the current task workspace;
+        2. execute the single action through the shared canonical executor;
+        3. emit the action event when an emitter is provided;
+        4. return the single action result.
+        """
+
+        self._validate_task_action_root_context(
+            current_task=current_task,
+            ux_context=ux_context,
+        )
+        result, event = await self._aapply_one_task_action(
+            amongo=amongo,
+            current_task=current_task,
+            action=action,
+            extended_aredis=extended_aredis,
+        )
+        if event_emitter is not None:
+            await event_emitter.aemit(event)
+            # Drain detached event callbacks before returning (same boundary
+            # discipline as the REST batch path), so the agent tool observes a
+            # flushed, error-surfaced event stream.
+            await event_emitter.adrain()
+        return result
+
+    async def aapply_actions(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        current_task: TaskMongoObject,
+        payload: TaskActions,
+        ux_context: TaskUxContext | None = None,
+        event_emitter: AgenticEventEmitter | None = None,
+        extended_aredis: ExtendedAsyncRedis | None = None,
+    ) -> TaskActionedResults:
+        """Apply a plain REST/programmatic batch of canonical task actions in order.
+
+        This is the non-agent seam (REST endpoint, internal callers). The agent
+        chat surface uses `aapply_action` (one action per tool call). Both share
+        `_aapply_one_task_action`, so behavior never diverges.
+
+        Step by step:
+        1. validate optional UI root context against the current task workspace;
+        2. execute each action in caller-provided order via the shared executor;
+        3. emit each event when an emitter is provided;
+        4. return ordered action results plus the emitted event payloads.
+        """
+
+        self._validate_task_action_root_context(
+            current_task=current_task,
+            ux_context=ux_context,
+        )
+        actioned_results: list[CreateRelatedTaskActionedResult] = []
+        events: list[TaskActionEvent] = []
+        for action in payload.actions:
+            result, event = await self._aapply_one_task_action(
+                amongo=amongo,
+                current_task=current_task,
+                action=action,
+                extended_aredis=extended_aredis,
+            )
+            actioned_results.append(result)
+            events.append(event)
+            if event_emitter is not None:
+                await event_emitter.aemit(event)
+        # Drain at the action-batch boundary: each action event's callbacks
+        # (publish/persist) are detached by the helpers emitter, so join them
+        # before returning so the caller observes a flushed, error-surfaced event
+        # stream rather than racing in-flight callbacks.
+        if event_emitter is not None:
+            await event_emitter.adrain()
+        return TaskActionedResults(
+            actioned_results=actioned_results,
+            events=events,
+        )
+
+    #
+    #
+    #
+    #
+    # Read
+    #
+    #
+    #
+    #
+    @overload
+    async def aget(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        task_id: ObjectId,
+        task_type: Optional[str],
+        include_childs: Literal[False] = False,
+    ) -> TaskMongoObject | None: ...
+
+    @overload
+    async def aget(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        task_id: ObjectId,
+        task_type: Optional[str],
+        include_childs: Literal[True],
+    ) -> TaskMongoObjectExtended | None: ...
+
+    async def aget(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        task_id: ObjectId,
+        task_type: Optional[str],
+        include_childs: bool = False,
+    ) -> TaskMongoObject | TaskMongoObjectExtended | None:
+        """Get a task by ID asynchronously."""
+        _filter: dict = {
+            "_id": task_id,
+        }
+        if task_type is not None:
+            _filter["task_type"] = task_type
+        task = await self.crud_task.afirst(amongo, query=QueryBase(filters=_filter))
+        if task is None or not include_childs:
+            return task
+
+        base_filter: dict = {}
+        if task_type is not None:
+            base_filter["task_type"] = task_type
+        return await self._aget_task_tree_node(
+            amongo=amongo,
+            task=task,
+            target_task_id=task_id,
+            base_filter=base_filter,
+        )
+
+    @overload
+    async def aowner_get(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        owner: str,
+        task_id: ObjectId,
+        task_type: Optional[str],
+        sharing_scope: Optional[SharingScope],
+        include_childs: Literal[False] = False,
+    ) -> TaskMongoObject | None: ...
+
+    @overload
+    async def aowner_get(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        owner: str,
+        task_id: ObjectId,
+        task_type: Optional[str],
+        sharing_scope: Optional[SharingScope],
+        include_childs: Literal[True],
+    ) -> TaskMongoObjectExtended | None: ...
+
+    async def aowner_get(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        owner: str,
+        task_id: ObjectId,
+        task_type: Optional[str],
+        sharing_scope: Optional[SharingScope],
+        include_childs: bool = False,
+    ) -> TaskMongoObject | TaskMongoObjectExtended | None:
+        """Get a task by owner and ID asynchronously."""
+        _filter: dict = {"_id": task_id, "owner": owner}
+        if task_type is not None:
+            _filter["task_type"] = task_type
+        if sharing_scope is not None:
+            _filter["sharing_scope"] = sharing_scope
+        task = await self.crud_task.afirst(amongo, query=QueryBase(filters=_filter))
+        if task is None or not include_childs:
+            return task
+
+        base_filter: dict = {"owner": owner}
+        if task_type is not None:
+            base_filter["task_type"] = task_type
+        if sharing_scope is not None:
+            base_filter["sharing_scope"] = sharing_scope
+        return await self._aget_task_tree_node(
+            amongo=amongo,
+            task=task,
+            target_task_id=task_id,
+            base_filter=base_filter,
+        )
+
+    async def aresolve_task_room(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        owner: str,
+        task_id: ObjectId,
+    ) -> TaskRoomResolution:
+        """Resolve one owner-visible task into its canonical root-room identity.
+
+        Step by step:
+        1. load the requested task under owner permissions;
+        2. resolve the durable root-task id from the stored task row;
+        3. return only the canonical room identity without prescribing room-context shape.
+        """
+
+        requested_task = await self.aowner_get(
+            amongo=amongo,
+            owner=owner,
+            task_id=task_id,
+            task_type=None,
+            sharing_scope=None,
+        )
+        if requested_task is None:
+            raise keble_exceptions.NoObjectPermission(
+                object_id=str(task_id),
+                object_type="task",
+            )
+        return TaskRoomResolution(
+            root_task_id=self._task_workspace_root_id(requested_task),
+            requested_task_id=requested_task.id,
+        )
+
+    async def aget_task_room_graph_context(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        owner: str,
+        root_task_id: ObjectId,
+        focused_task_id: ObjectId | None = None,
+    ) -> TaskRoomGraphContext:
+        """Load generic task graph context for one owner-visible root task room.
+
+        Step by step:
+        1. load the full root task tree under owner permissions;
+        2. validate the focused task belongs to the same root task room;
+        3. attach root-scoped relation rows without adding feature-specific fields.
+        """
+
+        root_task = await self.aowner_get(
+            amongo=amongo,
+            owner=owner,
+            task_id=root_task_id,
+            task_type=None,
+            sharing_scope=None,
+            include_childs=True,
+        )
+        if root_task is None:
+            raise keble_exceptions.NoObjectPermission(
+                object_id=str(root_task_id),
+                object_type="task",
+            )
+        if self._task_workspace_root_id(root_task) != root_task.id:
+            raise keble_exceptions.ServerSideInvalidParams(
+                admin_note={
+                    "root_task_id": root_task_id,
+                    "resolved_root_task_id": self._task_workspace_root_id(root_task),
+                },
+                alert_admin=True,
+                but_got="requested task is not a root task",
+                expected="root_task_id points to the root task room",
+                invalid_params="root_task_id",
+            )
+
+        resolved_focused_task_id = focused_task_id or root_task.id
+        if resolved_focused_task_id != root_task.id:
+            focused_task = await self.aowner_get(
+                amongo=amongo,
+                owner=owner,
+                task_id=resolved_focused_task_id,
+                task_type=None,
+                sharing_scope=None,
+            )
+            if focused_task is None:
+                raise keble_exceptions.NoObjectPermission(
+                    object_id=str(resolved_focused_task_id),
+                    object_type="task",
+                )
+            if self._task_workspace_root_id(focused_task) != root_task.id:
+                raise keble_exceptions.ServerSideInvalidParams(
+                    admin_note={
+                        "root_task_id": root_task.id,
+                        "focused_task_id": resolved_focused_task_id,
+                        "focused_root_task_id": self._task_workspace_root_id(
+                            focused_task
+                        ),
+                    },
+                    alert_admin=True,
+                    but_got="focused task belongs to another root task room",
+                    expected="focused task belongs to root_task_id",
+                    invalid_params="focused_task_id",
+                )
+
+        relations = await self.alist_task_relations_by_root(
+            amongo=amongo,
+            root_task=root_task.id,
+        )
+        return TaskRoomGraphContext.build(
+            root_task=root_task,
+            focused_task_id=resolved_focused_task_id,
+            relations=relations,
+        )
+
+    @overload
+    async def apublic_get(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        task_id: ObjectId,
+        task_type: Optional[str],
+        include_childs: Literal[False] = False,
+    ) -> TaskMongoObject | None: ...
+
+    @overload
+    async def apublic_get(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        task_id: ObjectId,
+        task_type: Optional[str],
+        include_childs: Literal[True],
+    ) -> TaskMongoObjectExtended | None: ...
+
+    async def apublic_get(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        task_id: ObjectId,
+        task_type: Optional[str],
+        include_childs: bool = False,
+    ) -> TaskMongoObject | TaskMongoObjectExtended | None:
+        """Get a public task by ID asynchronously."""
+        _filter = {"_id": task_id, "sharing_scope": SharingScope.PUBLIC}
+        if task_type is not None:
+            _filter["task_type"] = task_type
+        task = await self.crud_task.afirst(amongo, query=QueryBase(filters=_filter))
+        if task is None or not include_childs:
+            return task
+
+        base_filter: dict = {"sharing_scope": SharingScope.PUBLIC}
+        if task_type is not None:
+            base_filter["task_type"] = task_type
+        return await self._aget_task_tree_node(
+            amongo=amongo,
+            task=task,
+            target_task_id=task_id,
+            base_filter=base_filter,
+        )
+
+    @overload
+    async def apublic_get_by_public_id(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        public_id: str,
+        task_type: Optional[str],
+        include_childs: Literal[False] = False,
+    ) -> TaskMongoObject | None: ...
+
+    @overload
+    async def apublic_get_by_public_id(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        public_id: str,
+        task_type: Optional[str],
+        include_childs: Literal[True],
+    ) -> TaskMongoObjectExtended | None: ...
+
+    async def apublic_get_by_public_id(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        public_id: str,
+        task_type: Optional[str],
+        include_childs: bool = False,
+    ) -> TaskMongoObject | TaskMongoObjectExtended | None:
+        """Get a public task by public_id asynchronously."""
+        _filter: dict = {"public_id": public_id, "sharing_scope": SharingScope.PUBLIC}
+        if task_type is not None:
+            _filter["task_type"] = task_type
+        task = await self.crud_task.afirst(amongo, query=QueryBase(filters=_filter))
+        if task is None or not include_childs:
+            return task
+
+        base_filter: dict = {"sharing_scope": SharingScope.PUBLIC}
+        if task_type is not None:
+            base_filter["task_type"] = task_type
+        return await self._aget_task_tree_node(
+            amongo=amongo,
+            task=task,
+            target_task_id=task.id,
+            base_filter=base_filter,
+        )
+
+    async def _aget_multi(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        extended_aredis: ExtendedAsyncRedis,
+        skip: Optional[int],
+        limit: Optional[int],
+        filter_: dict,
+    ) -> List[TaskMongoObject]:
+        """Get multiple tasks by filter asynchronously."""
+        tasks = await self.crud_task.aget_multi(
+            amongo,
+            query=QueryBase(
+                filters=filter_,
+                order_by=[("created", DESCENDING)],  # by created time
+                offset=skip,
+                limit=limit,
+            ),
+        )
+        for i, task in enumerate(tasks):
+            if task.unfinshed_timeout:
+                # mark it as failure
+                logger.warning(
+                    f"[Task] Found an unfinished timeout task: {task.id}. Task stage = {task.stage}, created = {task.created}, task type = {task.task_type}, attempts = {task.attempts}"
+                )
+                if task.allow_to_retry:
+                    # try to restart the task
+                    logger.warning(
+                        f"[Task] Found an unfinished timeout task: {task.id}. Start to retry."
+                    )
+                    await self.astart(
+                        amongo=amongo,
+                        extended_aredis=extended_aredis,
+                        task=task,
+                    )
+                else:
+                    # mark it as timeout failure
+                    logger.warning(
+                        f"[Task] Found an unfinished timeout task: {task.id}. Make it as Timeout."
+                    )
+                    await self.aon_task_timeout(
+                        amongo=amongo,
+                        task=task,
+                    )
+
+                updated_task = await self.crud_task.afirst_by_id(amongo, _id=task.id)
+                assert updated_task is not None, (
+                    f"[Task] Internal error, expected not none object for id {task.id}, but got none."
+                )
+                tasks[i] = updated_task
+
+        return tasks
+
+    async def _aget_root_tasks(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        extended_aredis: ExtendedAsyncRedis,
+        skip: int,
+        limit: int,
+        base_filter: dict,
+    ) -> List[TaskMongoObject]:
+        filter_ = dict(base_filter)
+        filter_["parent_task"] = None
+        return await self._aget_multi(
+            amongo=amongo,
+            extended_aredis=extended_aredis,
+            skip=skip,
+            limit=limit,
+            filter_=filter_,
+        )
+
+    async def _aget_descendant_tasks(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        extended_aredis: ExtendedAsyncRedis,
+        root_ids: List[ObjectId],
+        base_filter: dict,
+    ) -> List[TaskMongoObject]:
+        if len(root_ids) == 0:
+            return []
+        filter_ = dict(base_filter)
+        filter_["root_task"] = {"$in": root_ids}
+        filter_["_id"] = {"$nin": root_ids}
+        return await self._aget_multi(
+            amongo=amongo,
+            extended_aredis=extended_aredis,
+            skip=None,
+            limit=None,
+            filter_=filter_,
+        )
+
+    async def _aget_tasks_in_root_tree(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        root_id: ObjectId,
+        base_filter: dict,
+    ) -> List[TaskMongoObject]:
+        return await self.crud_task.aget_tasks_in_root_tree(
+            amongo, root_id=root_id, base_filter=base_filter
+        )
+
+    async def alist_task_relations_by_root(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        root_task: ObjectId,
+    ) -> list[TaskRelationMongoObject]:
+        """List extra lineage relations scoped to one root-task workspace.
+
+        Step by step:
+        1. query relation rows by the durable `root_task` workspace id;
+        2. order rows by creation time for stable canvas rendering;
+        3. return only relation rows, leaving tree composition to callers.
+        """
+
+        await self._ensure_task_relation_indexes(amongo=amongo)
+        return await self.crud_task_relation.aget_multi(
+            amongo,
+            query=QueryBase(
+                filters={"root_task": root_task},
+                order_by=[("created", ASCENDING)],
+            ),
+        )
+
+    async def _aget_task_tree_node(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        task: TaskMongoObject,
+        target_task_id: ObjectId,
+        base_filter: dict,
+    ) -> TaskMongoObjectExtended:
+        root_id = task.root_task or task.id
+        tasks = await self._aget_tasks_in_root_tree(
+            amongo=amongo, root_id=root_id, base_filter=base_filter
+        )
+        subtree_roots = self._build_task_tree(roots=[task], tasks=tasks)
+        found = self._find_task_in_tree(nodes=subtree_roots, task_id=target_task_id)
+        if found is not None:
+            return found
+        return TaskMongoObjectExtended(**task.model_dump(), childs=[])
+
+    def _build_task_tree(
+        self, *, roots: List[TaskMongoObject], tasks: List[TaskMongoObject]
+    ) -> List[TaskMongoObjectExtended]:
+        return build_task_tree(roots=roots, tasks=tasks)
+
+    def _find_task_in_tree(
+        self, *, nodes: List[TaskMongoObjectExtended], task_id: ObjectId
+    ) -> Optional[TaskMongoObjectExtended]:
+        return find_task_in_tree(nodes=nodes, task_id=task_id)
+
+    async def _aget_multi_with_optional_childs(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        extended_aredis: ExtendedAsyncRedis,
+        skip: int,
+        limit: int,
+        base_filter: dict,
+        include_childs: bool,
+        root_task: Optional[ObjectId],
+        parent_task: Optional[ObjectId],
+    ) -> List[TaskMongoObject] | List[TaskMongoObjectExtended]:
+        """List either flat root tasks or task trees from one shared filter.
+
+        Step by step:
+        1. flat calls list root tasks by default so normal task lists stay workspace-level;
+        2. explicit `parent_task` flat calls still list that parent's direct children;
+        3. tree calls preserve the existing paginated-root plus descendant assembly.
+        """
+
+        if not include_childs:
+            flat_filter = dict(base_filter)
+            if parent_task is None:
+                flat_filter["parent_task"] = None
+            return await self._aget_multi(
+                amongo=amongo,
+                extended_aredis=extended_aredis,
+                skip=skip,
+                limit=limit,
+                filter_=flat_filter,
+            )
+
+        if parent_task is not None:
+            page_roots = await self._aget_multi(
+                amongo=amongo,
+                extended_aredis=extended_aredis,
+                skip=skip,
+                limit=limit,
+                filter_=base_filter,
+            )
+            if len(page_roots) == 0:
+                return []
+            top_root_ids = list(
+                {
+                    task.root_task if task.root_task is not None else task.id
+                    for task in page_roots
+                }
+            )
+            tree_filter = {
+                k: v
+                for k, v in base_filter.items()
+                if k not in {"parent_task", "root_task"}
+            }
+            if root_task is not None:
+                tree_filter["root_task"] = root_task
+            else:
+                tree_filter["root_task"] = {"$in": top_root_ids}
+            tasks_in_trees = await self._aget_multi(
+                amongo=amongo,
+                extended_aredis=extended_aredis,
+                skip=None,
+                limit=None,
+                filter_=tree_filter,
+            )
+            return self._build_task_tree(roots=page_roots, tasks=tasks_in_trees)
+
+        roots = await self._aget_root_tasks(
+            amongo=amongo,
+            extended_aredis=extended_aredis,
+            skip=skip,
+            limit=limit,
+            base_filter=base_filter,
+        )
+        children = await self._aget_descendant_tasks(
+            amongo=amongo,
+            extended_aredis=extended_aredis,
+            root_ids=[t.id for t in roots],
+            base_filter=base_filter,
+        )
+        return self._build_task_tree(roots=roots, tasks=roots + children)
+
+    @overload
+    async def aget_multi(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        extended_aredis: ExtendedAsyncRedis,
+        skip: int,
+        limit: int,
+        task_types: Optional[List[str]],
+        include_childs: Literal[False] = False,
+        root_task: Optional[ObjectId] = None,
+        stages: Optional[List[str]] = None,
+        title_contains: Optional[str] = None,
+        parent_task: Optional[ObjectId] = None,
+    ) -> List[TaskMongoObject]: ...
+
+    @overload
+    async def aget_multi(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        extended_aredis: ExtendedAsyncRedis,
+        skip: int,
+        limit: int,
+        task_types: Optional[List[str]],
+        include_childs: Literal[True],
+        root_task: Optional[ObjectId] = None,
+        stages: Optional[List[str]] = None,
+        title_contains: Optional[str] = None,
+        parent_task: Optional[ObjectId] = None,
+    ) -> List[TaskMongoObjectExtended]: ...
+
+    async def aget_multi(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        extended_aredis: ExtendedAsyncRedis,
+        skip: int,
+        limit: int,
+        task_types: Optional[List[str]],
+        include_childs: bool = False,
+        root_task: Optional[ObjectId] = None,
+        stages: Optional[List[str]] = None,
+        title_contains: Optional[str] = None,
+        parent_task: Optional[ObjectId] = None,
+    ) -> List[TaskMongoObject] | List[TaskMongoObjectExtended]:
+        """Get multiple tasks asynchronously.
+
+        Side effect if changes:
+        - `aowner_get_multi` and agent query registry tests mirror this generic
+          read shape; keep stage/title filters aligned across owner/public
+          variants.
+        - Tree/list endpoints rely on `_aget_multi_with_optional_childs` for
+          the `include_childs` return-type split.
+        """
+        base_filter: dict = {}
+        if task_types is not None:
+            base_filter["task_type"] = {"$in": task_types}
+        if root_task is not None:
+            base_filter["root_task"] = root_task
+        if parent_task is not None:
+            base_filter["parent_task"] = parent_task
+        if stages is not None:
+            base_filter["stage"] = {"$in": stages}
+        if title_contains is not None and title_contains.strip():
+            base_filter["title"] = {
+                "$regex": re.escape(title_contains.strip()),
+                "$options": "i",
+            }
+        return await self._aget_multi_with_optional_childs(
+            amongo=amongo,
+            extended_aredis=extended_aredis,
+            skip=skip,
+            limit=limit,
+            base_filter=base_filter,
+            include_childs=include_childs,
+            root_task=root_task,
+            parent_task=parent_task,
+        )
+
+    @overload
+    async def aowner_get_multi(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        extended_aredis: ExtendedAsyncRedis,
+        owner: str,
+        skip: int,
+        limit: int,
+        task_types: Optional[List[str]],
+        sharing_scopes: Optional[List[SharingScope]],
+        include_childs: Literal[False] = False,
+        root_task: Optional[ObjectId] = None,
+        stages: Optional[List[str]] = None,
+        title_contains: Optional[str] = None,
+        parent_task: Optional[ObjectId] = None,
+    ) -> List[TaskMongoObject]: ...
+
+    @overload
+    async def aowner_get_multi(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        extended_aredis: ExtendedAsyncRedis,
+        owner: str,
+        skip: int,
+        limit: int,
+        task_types: Optional[List[str]],
+        sharing_scopes: Optional[List[SharingScope]],
+        include_childs: Literal[True],
+        root_task: Optional[ObjectId] = None,
+        stages: Optional[List[str]] = None,
+        title_contains: Optional[str] = None,
+        parent_task: Optional[ObjectId] = None,
+    ) -> List[TaskMongoObjectExtended]: ...
+
+    async def aowner_get_multi(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        extended_aredis: ExtendedAsyncRedis,
+        owner: str,
+        skip: int,
+        limit: int,
+        task_types: Optional[List[str]],
+        sharing_scopes: Optional[List[SharingScope]],
+        include_childs: bool = False,
+        root_task: Optional[ObjectId] = None,
+        stages: Optional[List[str]] = None,
+        title_contains: Optional[str] = None,
+        parent_task: Optional[ObjectId] = None,
+    ) -> List[TaskMongoObject] | List[TaskMongoObjectExtended]:
+        """Get multiple tasks by owner asynchronously.
+
+        Collection: the task collection. Filter shape: `owner` equality rides
+        the `owner + parent_task + created` index prefix; `task_type`,
+        `sharing_scope`, `root_task`, `parent_task` are indexed-or-residual
+        equality filters as before. Additive (2.10.0): `stages` (`$in`
+        equality) and `title_contains` (escaped case-insensitive regex) are
+        RESIDUAL filters inside the owner-bounded shape — intended for
+        bounded agent/list reads (limit <= 50), never unbounded sweeps.
+        """
+        base_filter: dict = {"owner": owner}
+        if task_types is not None:
+            base_filter["task_type"] = {"$in": task_types}
+        if sharing_scopes is not None:
+            base_filter["sharing_scope"] = {"$in": sharing_scopes}
+        if root_task is not None:
+            base_filter["root_task"] = root_task
+        if parent_task is not None:
+            base_filter["parent_task"] = parent_task
+        if stages is not None:
+            base_filter["stage"] = {"$in": stages}
+        if title_contains is not None and title_contains.strip():
+            base_filter["title"] = {
+                "$regex": re.escape(title_contains.strip()),
+                "$options": "i",
+            }
+        return await self._aget_multi_with_optional_childs(
+            amongo=amongo,
+            extended_aredis=extended_aredis,
+            skip=skip,
+            limit=limit,
+            base_filter=base_filter,
+            include_childs=include_childs,
+            root_task=root_task,
+            parent_task=parent_task,
+        )
+
+    @overload
+    async def apublic_get_multi(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        extended_aredis: ExtendedAsyncRedis,
+        skip: int,
+        limit: int,
+        task_types: Optional[List[str]],
+        include_childs: Literal[False] = False,
+        root_task: Optional[ObjectId] = None,
+        stages: Optional[List[str]] = None,
+        title_contains: Optional[str] = None,
+        parent_task: Optional[ObjectId] = None,
+    ) -> List[TaskMongoObject]: ...
+
+    @overload
+    async def apublic_get_multi(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        extended_aredis: ExtendedAsyncRedis,
+        skip: int,
+        limit: int,
+        task_types: Optional[List[str]],
+        include_childs: Literal[True],
+        root_task: Optional[ObjectId] = None,
+        stages: Optional[List[str]] = None,
+        title_contains: Optional[str] = None,
+        parent_task: Optional[ObjectId] = None,
+    ) -> List[TaskMongoObjectExtended]: ...
+
+    async def apublic_get_multi(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        extended_aredis: ExtendedAsyncRedis,
+        skip: int,
+        limit: int,
+        task_types: Optional[List[str]],
+        include_childs: bool = False,
+        root_task: Optional[ObjectId] = None,
+        stages: Optional[List[str]] = None,
+        title_contains: Optional[str] = None,
+        parent_task: Optional[ObjectId] = None,
+    ) -> List[TaskMongoObject] | List[TaskMongoObjectExtended]:
+        """Get multiple public tasks asynchronously.
+
+        `stages` (stage equality) and `title_contains` (escaped case-insensitive
+        regex) are RESIDUAL filters inside the PUBLIC-bounded shape — mirroring
+        `aowner_get_multi`. The overloads always promised them; the implementation
+        now honors them instead of silently dropping them.
+        """
+        base_filter: dict = {"sharing_scope": SharingScope.PUBLIC}
+        if task_types is not None:
+            base_filter["task_type"] = {"$in": task_types}
+        if root_task is not None:
+            base_filter["root_task"] = root_task
+        if parent_task is not None:
+            base_filter["parent_task"] = parent_task
+        if stages is not None:
+            base_filter["stage"] = {"$in": stages}
+        if title_contains is not None and title_contains.strip():
+            base_filter["title"] = {
+                "$regex": re.escape(title_contains.strip()),
+                "$options": "i",
+            }
+        return await self._aget_multi_with_optional_childs(
+            amongo=amongo,
+            extended_aredis=extended_aredis,
+            skip=skip,
+            limit=limit,
+            base_filter=base_filter,
+            include_childs=include_childs,
+            root_task=root_task,
+            parent_task=parent_task,
+        )
+
+    async def apublic_list_indexable(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        task_types: List[str],
+        stages: List[str],
+        limit: int,
+    ) -> List[TaskPublicRef]:
+        """List PUBLIC tasks of given types+stages, newest-first, id+updated only.
+
+        Objective: the lean read backing the dynamic-page sitemap. Unlike
+        `apublic_get_multi` this is PROJECTED (only `_id` + `updated`), needs no
+        `extended_aredis`, never loads childs, and returns the tiny `TaskPublicRef`
+        instead of the full `TaskMongoObject` — so a crawl-time listing of many
+        public reports stays cheap.
+
+        Filter shape (served by `sharing_scope + task_type + stage + updated DESC`):
+        equality on `sharing_scope=PUBLIC`, `task_type $in`, `stage $in`, sorted by
+        `updated` DESC and capped by `limit`. Converts raw projection docs ->
+        `TaskPublicRef`.
+        """
+        docs = await self.crud_task.aget_multi_docs(
+            amongo,
+            query=QueryBase(
+                filters={
+                    "sharing_scope": SharingScope.PUBLIC,
+                    "task_type": {"$in": task_types},
+                    "stage": {"$in": stages},
+                },
+                order_by=[("updated", DESCENDING)],
+                limit=limit,
+            ),
+            project={"_id": 1, "updated": 1},
+        )
+        return [TaskPublicRef(**doc) for doc in docs]
+
+    #
+    #
+    #
+    #
+    # Update
+    #
+    #
+    #
+    #
+    async def aupdate(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        extended_aredis: ExtendedAsyncRedis,
+        task_id: ObjectId,
+        obj_in: TaskUpdate,
+        public_id_prefix: Optional[str] = None,
+        public_id_length: Optional[int] = None,
+    ):
+        """Update a task asynchronously."""
+        task = await self.aget(amongo=amongo, task_id=task_id, task_type=None)
+        if task is None:
+            raise keble_exceptions.ServerSideInvalidParams(
+                admin_note={"task_id": task_id},
+                alert_admin=True,
+                but_got="task is None",
+                expected="task exists",
+                invalid_params=f"task_id={task_id}",
+            )
+        if obj_in.sharing_scope == SharingScope.PUBLIC:
+            await self._ensure_task_public_id(
+                amongo=amongo,
+                task_id=task_id,
+                public_id_prefix=public_id_prefix,
+                public_id_length=public_id_length,
+            )
+        await self.crud_task.aupdate(
+            amongo,
+            _id=task_id,
+            obj_in=obj_in.model_dump(exclude_unset=True),
+        )
+
+    async def aowner_update(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        owner: str,
+        task_id: ObjectId,
+        obj_in: TaskUpdate,
+        public_id_prefix: Optional[str] = None,
+        public_id_length: Optional[int] = None,
+    ):
+        """Update a task by owner asynchronously."""
+        r = await self.aowner_get(
+            amongo=amongo,
+            owner=owner,
+            sharing_scope=None,
+            task_id=task_id,
+            task_type=None,
+        )
+        if r is None:
+            raise keble_exceptions.ServerSideInvalidParams(
+                admin_note={"task_id": task_id, "owner": owner},
+                alert_admin=True,
+                but_got="task is None",
+                expected="task exists",
+                invalid_params=f"task_id={task_id}",
+            )
+        if obj_in.sharing_scope == SharingScope.PUBLIC:
+            await self._ensure_task_public_id(
+                amongo=amongo,
+                task_id=task_id,
+                public_id_prefix=public_id_prefix,
+                public_id_length=public_id_length,
+            )
+        await self.crud_task.aupdate(
+            amongo,
+            _id=task_id,
+            obj_in=obj_in.model_dump(exclude_unset=True),
+        )
+
+    #
+    #
+    #
+    #
+    # Event Handler
+    #
+    #
+    #
+    #
+    def _get_task_start_lock_key(self, task: TaskMongoObject) -> str:
+        """Get the Redis key for task start lock."""
+        return f"just_started:{task.id}"
+
+    async def _acheck_start_lock(
+        self, extended_aredis: ExtendedAsyncRedis, task: TaskMongoObject
+    ) -> bool:
+        """Check if task is locked. Returns True if not locked."""
+        return await extended_aredis.get(self._get_task_start_lock_key(task)) is None
+
+    async def _alock_start(
+        self, extended_aredis: ExtendedAsyncRedis, task: TaskMongoObject
+    ):
+        """Lock a task for starting."""
+        await extended_aredis.set(
+            self._get_task_start_lock_key(task),
+            "1",
+            ex=60,  # expire after 1 min
+        )
+
+    async def _aclear_lock(
+        self, extended_aredis: ExtendedAsyncRedis, task: TaskMongoObject
+    ):
+        """Clear the task start lock."""
+        await extended_aredis.delete(self._get_task_start_lock_key(task))
+
+    # `before_sleep` logs EVERY failed attempt with its full traceback.
+    # Without it, `reraise=True` surfaces only the LAST attempt's exception —
+    # in a live incident the real attempt-1 failure (a crashed spawn tool) was
+    # silently discarded and only a misleading downstream error reached the
+    # task row, costing the diagnosis.
+    @tenacity.retry(
+        stop=tenacity.stop_after_attempt(3),
+        reraise=True,
+        before_sleep=tenacity.before_sleep_log(logger, logging.ERROR, exc_info=True),
+    )
+    async def _ahandle_task(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        extended_aredis: ExtendedAsyncRedis,
+        aneo4j: Optional[Neo4jAsyncDriver] = None,
+        qdrant_client: Optional[AsyncQdrantClient] = None,
+        task: TaskMongoObject,
+        task_handler_metadata: Optional[TaskMetadata] = None,
+        event_emitter: AgenticEventEmitter,
+    ):
+        """Run one task handler and route its returned response.
+
+        Step by step:
+        1. increment attempts and merge metadata;
+        2. build the request AS an `AgentDbDeps` (DB clients + owner + the shared
+           room `event_emitter` + `usage_recorder`) — the handler uses it directly;
+        3. await the handler's `TaskHandlerResponse | None`;
+        4. `None` => the handler delegated completion to another process; leave the
+           task PROCESSING (a separate worker finalizes it later);
+        5. otherwise route success/failure through the runtime finalizers, which
+           emit the terminal `KEBLE_TASK` lifecycle event via `event_emitter`.
+
+        Side effects:
+        - replaces the old `aon_task_completed` callback inversion-of-control: the
+          runtime (not the handler) owns persistence, token accounting, and the
+          terminal event.
+        """
+        # increment attempts for each retry
+        task = await self.crud_task.aincrement_attempts(amongo, _id=task.id)
+        merged_metadata = _merge_metadata(task.metadata, task_handler_metadata)
+
+        # `ExtendedAsyncRedis` IS an `AsyncRedis` subclass, so it satisfies both the
+        # base `aredis` field and the extended `extended_aredis` field.
+        request = TaskHandlerRequest(
+            amongo=amongo,
+            aredis=extended_aredis,
+            extended_aredis=extended_aredis,
+            aneo4j=aneo4j,
+            aqdrant=qdrant_client,
+            owner=task.owner,
+            language=task.language or Language.ENGLISH,
+            event_emitter=event_emitter,
+            task=task,
+            metadata=merged_metadata,
+        )
+        response = await self._task_handler(request)
+
+        # None => the handler delegated completion elsewhere; do NOT finalize.
+        if response is None:
+            return
+
+        if response.success:
+            await self.aon_task_success(
+                amongo=amongo,
+                task=response.task,
+                consuming_token=response.consuming_token,
+                extended_aredis=extended_aredis,
+                aneo4j=aneo4j,
+                metadata=merged_metadata,
+                event_emitter=event_emitter,
+            )
+        else:
+            await self.aon_task_failure(
+                amongo=amongo,
+                extended_aredis=extended_aredis,
+                error=response.error
+                or f"Task handled response has a status = {response.success}",
+                exception_type=response.exception_type or TaskExceptionType.UNKNOWN,
+                task=response.task,
+                aneo4j=aneo4j,
+                metadata=merged_metadata,
+                event_emitter=event_emitter,
+            )
+
+    async def astart(
+        self,
+        *,
+        get_amongo: Optional[Callable[[], AsyncIOMotorClient]] = None,
+        get_extended_aredis: Optional[Callable[[], ExtendedAsyncRedis]] = None,
+        get_aneo4j: Optional[Callable[[], Optional[Neo4jAsyncDriver]]] = None,
+        get_qdrant_client: Optional[Callable[[], Optional[AsyncQdrantClient]]] = None,
+        amongo: Optional[AsyncIOMotorClient] = None,
+        extended_aredis: Optional[ExtendedAsyncRedis] = None,
+        aneo4j: Optional[Neo4jAsyncDriver] = None,
+        qdrant_client: Optional[AsyncQdrantClient] = None,
+        task_id: Optional[ObjectId] = None,
+        task: Optional[TaskMongoObject] = None,
+        # metadata to pass into handler request
+        task_handler_metadata: Optional[TaskMetadata] = None,
+        task_lifecycle_event_emitter: AgenticEventEmitter | None = None,
+    ):
+        """Start one task and always finalize uncaught handler failures into task state.
+
+        Step by step this entrypoint:
+        1. resolves the required runtime resources;
+        2. checks the single-start Redis lock;
+        3. moves the task into `PROCESSING`;
+        4. runs the configured task handler with retry support;
+        5. if the handler still raises, marks the task as `FAILURE` unless it was
+           already finalized by the handler itself;
+        6. emits lifecycle events after persisted stage transitions when provided;
+        7. always clears the single-start Redis lock before the method exits, even
+           when the processing transition itself fails.
+        """
+        if amongo is None:
+            if get_amongo is None:
+                raise keble_exceptions.ServerSideMissingParams(
+                    admin_note={"get_amongo": get_amongo},
+                    alert_admin=True,
+                    missing_params="get_amongo",
+                )
+            amongo = get_amongo()
+        if extended_aredis is None:
+            if get_extended_aredis is None:
+                raise keble_exceptions.ServerSideMissingParams(
+                    admin_note={"get_extended_aredis": get_extended_aredis},
+                    alert_admin=True,
+                    missing_params="get_extended_aredis",
+                )
+            extended_aredis = get_extended_aredis()
+        # neo4j + getter can be none
+        if aneo4j is None and get_aneo4j is not None:
+            aneo4j = get_aneo4j()
+        # qdrant + getter can be none
+        if qdrant_client is None and get_qdrant_client is not None:
+            qdrant_client = get_qdrant_client()
+
+        db_obj = await self._aget_task_for_event(
+            amongo=amongo, task_id=task_id, task=task
+        )
+        # Try to start the task
+        allow_proceed = await self._acheck_start_lock(
+            extended_aredis=extended_aredis, task=db_obj
+        )
+        if not allow_proceed:
+            logger.warning(f"[Task] Task {db_obj.id} already started, skip to start.")
+            return db_obj
+
+        if not db_obj.stage.allow_to_start:
+            logger.warning(
+                f"[Task] Can not start task: {db_obj.id}, current stage = {db_obj.stage}"
+            )
+            raise TaskFailedToStartException(
+                f"Can not start task: {db_obj.id}, current stage = {db_obj.stage}"
+            )
+
+        try:
+            await self._alock_start(extended_aredis=extended_aredis, task=db_obj)
+            await self.aon_task_processing(
+                amongo=amongo,
+                task=db_obj,
+                event_emitter=task_lifecycle_event_emitter,
+            )
+            try:
+                await self._ahandle_task(
+                    amongo=amongo,
+                    extended_aredis=extended_aredis,
+                    task=db_obj,
+                    aneo4j=aneo4j,
+                    qdrant_client=qdrant_client,
+                    task_handler_metadata=task_handler_metadata,
+                    event_emitter=task_lifecycle_event_emitter
+                    or AgenticEventEmitter(),
+                )
+            except Exception as exc:
+                logger.critical(
+                    f"[Task] Task failed during execution: {db_obj.id}, with error: {str(exc)}\n{traceback.format_exc()}"
+                )
+                latest_task = await self.crud_task.afirst_by_id(amongo, _id=db_obj.id)
+                assert latest_task is not None, (
+                    f"Task with ID {db_obj.id} not found after handler failure"
+                )
+                if latest_task.stage not in {TaskStage.SUCCESS, TaskStage.FAILURE}:
+                    failure_error = str(exc)
+                    failure_exception_type = TaskExceptionType.UNKNOWN
+                    if isinstance(exc, TaskException):
+                        failure_error = exc.error or failure_error
+                        failure_exception_type = exc.exception_type
+                    await self.aon_task_failure(
+                        amongo=amongo,
+                        extended_aredis=extended_aredis,
+                        aneo4j=aneo4j,
+                        task=latest_task,
+                        error=failure_error,
+                        exception_type=failure_exception_type,
+                        metadata=task_handler_metadata,
+                        event_emitter=task_lifecycle_event_emitter,
+                    )
+        finally:
+            await self._aclear_lock(extended_aredis=extended_aredis, task=db_obj)
+            # Backstop drain: every aon_task_* lifecycle method already self-drains,
+            # but join any residual detached emit here so the worker envelope never
+            # returns (and releases the Celery lease) with spawn-on-terminal or
+            # terminal-subtree-settle callbacks still in flight.
+            if task_lifecycle_event_emitter is not None:
+                await task_lifecycle_event_emitter.adrain()
+
+        result = await self.crud_task.afirst_by_id(amongo, _id=db_obj.id)
+        assert result is not None, f"Task with ID {db_obj.id} not found after start"
+        return result
+
+    async def aretry_all_undone(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        extended_aredis: ExtendedAsyncRedis,
+        aneo4j: Optional[Neo4jAsyncDriver] = None,
+        get_aneo4j: Optional[Callable[[], Optional[Neo4jAsyncDriver]]] = None,
+        task_handler_metadata: Optional[TaskMetadata] = None,
+        include_tasks_created_within_minutes: int = 30,
+    ):
+        """Retry all undone tasks asynchronously."""
+        # fetch all undone
+        after = utc_now() - timedelta(
+            minutes=include_tasks_created_within_minutes
+        )
+        db_objs = await self.crud_task.aget_multi(
+            amongo,
+            query=QueryBase(
+                filters={
+                    "stage": {"$in": [TaskStage.PENDING, TaskStage.PROCESSING]},
+                    "created": {"$gt": after},
+                },
+                order_by=[("created", DESCENDING)],  # by created time
+            ),
+        )
+        for db_obj in db_objs:
+            if db_obj.allow_to_retry:
+                # restart it
+                await self.astart(
+                    amongo=amongo,
+                    extended_aredis=extended_aredis,
+                    aneo4j=aneo4j,
+                    get_aneo4j=get_aneo4j,
+                    task_handler_metadata=task_handler_metadata,
+                    task=db_obj,
+                )
+            else:
+                # set it as failure due to timeout
+                await self.aon_task_timeout(amongo=amongo, task=db_obj)
+
+    async def aget_tasks_by_stage_since(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        stages: Optional[List[TaskStage]] = None,
+        hours: int = 1,
+    ) -> List[TaskMongoObject]:
+        """Get tasks from the last N hours with specified statuses."""
+        cutoff_time = utc_now() - timedelta(hours=hours)
+        return await self.crud_task.aget_multi(
+            amongo,
+            query=QueryBase(
+                filters={
+                    "stage": {
+                        "$in": stages or [TaskStage.PENDING, TaskStage.PROCESSING]
+                    },
+                    "created": {"$gt": cutoff_time},
+                },
+                order_by=[("created", DESCENDING)],
+            ),
+        )
+
+    async def _aget_task_for_event(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        task_id: Optional[ObjectId],
+        task: Optional[TaskMongoObject],
+    ):
+        """Get a task for an event asynchronously."""
+        if task_id is not None:
+            db_obj = await self.aget(amongo=amongo, task_id=task_id, task_type=None)
+            if db_obj is None:
+                raise keble_exceptions.ServerSideInvalidParams(
+                    admin_note={"task_id": task_id},
+                    alert_admin=True,
+                    but_got="report is None",
+                    expected="report exists",
+                    invalid_params=f"task_id={task_id}",
+                )
+            return db_obj
+        elif task is not None:
+            return task
+        else:
+            raise keble_exceptions.ServerSideInvalidParams(
+                admin_note={},
+                alert_admin=True,
+                but_got="no params",
+                expected="at least one of task_id or task",
+                invalid_params="task_id or task",
+            )
+
+    async def _aemit_task_lifecycle_event(
+        self,
+        *,
+        task: TaskMongoObject,
+        event_emitter: AgenticEventEmitter | None,
+    ) -> None:
+        """Emit one canonical task lifecycle event after a persisted mutation.
+
+        Step by step:
+        1. do nothing when the caller did not provide an event emitter;
+        2. build the lifecycle payload from the reloaded task row;
+        3. scope the event to the root task room while identifying the changed task.
+        """
+
+        if event_emitter is None:
+            return
+        event = TaskLifecycleEvent(
+            action_type=TaskEventType.TASK_STAGE_CHANGED.value,
+            status=_task_lifecycle_event_status(stage=task.stage),
+            payload=TaskLifecycleEventPayload(task=task),
+            root_id=str(task.root_task or task.id),
+            object_id=str(task.id),
+            correlation_id=task.progress_key,
+        )
+        await event_emitter.aemit(event)
+        # Drain at the lifecycle boundary: the helpers emitter detaches callbacks, but
+        # lifecycle callbacks carry DOMAIN side effects that the workflow reads right
+        # after (spawn-on-terminal, terminal-subtree settle, chat-memory writes). The
+        # task-row mutation is already persisted before this emit, so joining the
+        # detached chain here makes each aon_task_* atomic w.r.t. its own callbacks and
+        # guarantees those side effects complete before the lifecycle method returns.
+        # Side effect if changes: removing this re-opens the spawn-on-terminal /
+        # subtree-settle race in keble.backend (subagent_workflow lifecycle callback).
+        await event_emitter.adrain()
+
+    async def aon_task_success(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        extended_aredis: ExtendedAsyncRedis,
+        aneo4j: Optional[Neo4jAsyncDriver] = None,
+        task_id: Optional[ObjectId] = None,
+        task: Optional[TaskMongoObject] = None,
+        consuming_token: int,  # how many token should be consume (actually)
+        metadata: Optional[TaskMetadata] = None,
+        event_emitter: AgenticEventEmitter | None = None,
+    ):
+        """Mark a task successful and optionally emit a lifecycle event."""
+        db_obj = await self._aget_task_for_event(
+            amongo=amongo, task=task, task_id=task_id
+        )
+        update_consumed_token = db_obj.consumed_token
+        if consuming_token < db_obj.consumed_token:
+            # do a recovery
+            recover_amount = db_obj.consumed_token - consuming_token
+            recovered = await self._execute_token_consumption_handler(
+                TokenConsumptionPayload(
+                    consumption_type=TokenConsumptionType.RECOVER,
+                    owner=db_obj.owner,
+                    task_id=db_obj.id,
+                    token=recover_amount,
+                    amongo=amongo,
+                    extended_aredis=extended_aredis,
+                    aneo4j=aneo4j,
+                    metadata=metadata,
+                )
+            )
+            if recovered:
+                update_consumed_token = consuming_token
+        elif consuming_token > db_obj.consumed_token:
+            # do a token consumption
+            consume_amount = consuming_token - db_obj.consumed_token
+            consumed = await self._execute_token_consumption_handler(
+                TokenConsumptionPayload(
+                    token=consume_amount,
+                    owner=db_obj.owner,
+                    consumption_type=TokenConsumptionType.CONSUME,
+                    task_id=db_obj.id,
+                    amongo=amongo,
+                    extended_aredis=extended_aredis,
+                    aneo4j=aneo4j,
+                    metadata=metadata,
+                )
+            )
+            if consumed:
+                update_consumed_token = consuming_token
+
+        # Mark as done
+        now = utc_now()
+        await self.crud_task.aupdate(
+            amongo,
+            _id=db_obj.id,
+            obj_in={
+                "stage": TaskStage.SUCCESS,
+                "success_ts": int(now.timestamp()),
+                "consumed_token": update_consumed_token,
+            },
+        )
+
+        result = await self.crud_task.afirst_by_id(amongo, _id=db_obj.id)
+        assert result is not None, f"Task with ID {db_obj.id} not found after update"
+        await self._aemit_task_lifecycle_event(
+            task=result,
+            event_emitter=event_emitter,
+        )
+        return result
+
+    async def aon_task_processing(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        task_id: Optional[ObjectId] = None,
+        task: Optional[TaskMongoObject] = None,
+        event_emitter: AgenticEventEmitter | None = None,
+    ) -> TaskMongoObject:
+        """Mark a task processing and optionally emit a lifecycle event."""
+        db_obj = await self._aget_task_for_event(
+            amongo=amongo, task=task, task_id=task_id
+        )
+        if not db_obj.stage.allow_to_start:
+            logger.warning(
+                f"[Task] Can not switch to processing for task: {db_obj.id}, current stage = {db_obj.stage}"
+            )
+            return db_obj
+        # Mark as processing
+        now = utc_now()
+        await self.crud_task.aupdate(
+            amongo,
+            _id=db_obj.id,
+            obj_in={
+                "stage": TaskStage.PROCESSING,
+                "started_ts": int(now.timestamp()),
+            },
+        )
+        result = await self.crud_task.afirst_by_id(amongo, _id=db_obj.id)
+        assert result is not None, f"Task with ID {db_obj.id} not found after update"
+        await self._aemit_task_lifecycle_event(
+            task=result,
+            event_emitter=event_emitter,
+        )
+        return result
+
+    async def aon_task_failure(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        extended_aredis: ExtendedAsyncRedis,
+        aneo4j: Optional[Neo4jAsyncDriver] = None,
+        task: Optional[TaskMongoObject] = None,
+        task_id: Optional[ObjectId] = None,
+        error: Optional[str] = None,
+        exception_type: TaskExceptionType,
+        metadata: Optional[TaskMetadata] = None,
+        event_emitter: AgenticEventEmitter | None = None,
+    ) -> TaskMongoObject:
+        """Mark a task failed and optionally emit a lifecycle event."""
+        db_obj = await self._aget_task_for_event(
+            amongo=amongo, task=task, task_id=task_id
+        )
+
+        # Handle token recovery if needed
+        update_obj = {
+            "stage": TaskStage.FAILURE,
+            "failure_ts": int(utc_now().timestamp()),
+            "error": error,
+            "exception_type": exception_type,
+        }
+
+        if db_obj.consumed_token > 0:
+            # recover tokens
+            recovered = await self._execute_token_consumption_handler(
+                TokenConsumptionPayload(
+                    consumption_type=TokenConsumptionType.RECOVER,
+                    owner=db_obj.owner,
+                    token=db_obj.consumed_token,
+                    task_id=db_obj.id,
+                    amongo=amongo,
+                    extended_aredis=extended_aredis,
+                    aneo4j=aneo4j,
+                    metadata=metadata,
+                )
+            )
+            if recovered:
+                update_obj["consumed_token"] = 0
+
+        # Mark as failure
+        await self.crud_task.aupdate(
+            amongo,
+            _id=db_obj.id,
+            obj_in=update_obj,
+        )
+
+        result = await self.crud_task.afirst_by_id(amongo, _id=db_obj.id)
+        assert result is not None, f"Task with ID {db_obj.id} not found after update"
+        await self._aemit_task_lifecycle_event(
+            task=result,
+            event_emitter=event_emitter,
+        )
+        return result
+
+    async def aon_task_timeout(
+        self,
+        *,
+        amongo: AsyncIOMotorClient,
+        task: Optional[TaskMongoObject] = None,
+        task_id: Optional[ObjectId] = None,
+        event_emitter: AgenticEventEmitter | None = None,
+    ):
+        """Mark a task timed out and optionally emit a lifecycle event."""
+        db_obj = await self._aget_task_for_event(
+            amongo=amongo, task=task, task_id=task_id
+        )
+        error_message = (
+            f"Task timeout after {db_obj.timeout_mins} min. Created at {db_obj.created}"
+        )
+
+        # Mark as failure
+        now = utc_now()
+        await self.crud_task.aupdate(
+            amongo,
+            _id=db_obj.id,
+            obj_in={
+                "stage": TaskStage.FAILURE,
+                "failure_ts": int(now.timestamp()),
+                "error": error_message,
+                "exception_type": TaskExceptionType.TIMEOUT,
+            },
+        )
+
+        result = await self.crud_task.afirst_by_id(amongo, _id=db_obj.id)
+        assert result is not None, f"Task with ID {db_obj.id} not found after update"
+        await self._aemit_task_lifecycle_event(
+            task=result,
+            event_emitter=event_emitter,
+        )
+        return result