agent-runtime-sdk 0.1.0__py3-none-any.whl

agent_runtime/runtime_engine/task_flow.py

@@ -0,0 +1,897 @@
+"""任务生命周期状态机。
+
+这个文件处理一条任务从开始到结束的全过程：
+
+- 启动任务
+- 进入 input_required / auth_required 等待态
+- 用户恢复任务
+- 超时与取消
+- 最终完成或失败
+
+它是 runtime 里最像“调度器”的一层。
+"""
+
+from __future__ import annotations
+
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
+from contextvars import copy_context
+from functools import partial
+import logging
+import os
+from typing import Any
+
+from .a2a_task_io import (
+    emit_wait_state,
+    publish_error,
+    publish_message_completion,
+    publish_result,
+)
+from .context import (
+    TaskContext,
+    TaskPhase,
+    TaskUpdaterProtocol,
+    WAIT_TYPE_AUTH_REQUIRED,
+    WAIT_TYPE_INPUT_REQUIRED,
+    wait_state_payload,
+    wait_state_type,
+)
+from .observability import (
+    A2ATaskObservation,
+    a2a_task_observation,
+    update_observation_output,
+)
+from ..exceptions import (
+    TaskCancelledError,
+    TaskExecutionTimeoutError,
+    TaskWaitTimeoutError,
+    UserCancelledError,
+)
+
+
+logger = logging.getLogger(__name__)
+
+
+class RuntimeTaskFlow:
+    """负责任务生命周期调度与 wait/resume 状态转换。"""
+
+    _AUTH_DECISION_KEYS = (
+        "approve",
+        "approved",
+        "confirm",
+        "confirmed",
+        "accept",
+        "accepted",
+        "allow",
+        "allowed",
+        "authorize",
+        "authorized",
+        "grant",
+        "granted",
+        "decision",
+    )
+    _AUTH_DECISION_KEY_SET = set(_AUTH_DECISION_KEYS)
+
+    @staticmethod
+    def _format_timeout_seconds(value: float) -> str:
+        if float(value).is_integer():
+            return str(int(value))
+        return f"{value:g}"
+
+    def pop_failed_task_id(self, task_id: str) -> bool:
+        if task_id not in self._failed_task_ids:
+            return False
+        self._failed_task_ids.remove(task_id)
+        return True
+
+    @staticmethod
+    def _normalize_auth_decision(value: object) -> bool | None:
+        if isinstance(value, bool):
+            return value
+        if isinstance(value, int) and not isinstance(value, bool):
+            if value == 1:
+                return True
+            if value == 0:
+                return False
+        if value is None:
+            return None
+        text = str(value).strip().lower()
+        if text in {
+            "true",
+            "1",
+            "yes",
+            "y",
+            "ok",
+            "approve",
+            "approved",
+            "confirm",
+            "confirmed",
+            "accept",
+            "accepted",
+            "allow",
+            "allowed",
+            "authorize",
+            "authorized",
+            "grant",
+            "granted",
+            "同意",
+            "确认",
+            "批准",
+            "允许",
+            "可以",
+            "是",
+            "好的",
+            "继续",
+            "通过",
+        }:
+            return True
+        if text in {
+            "false",
+            "0",
+            "no",
+            "n",
+            "deny",
+            "denied",
+            "reject",
+            "rejected",
+            "cancel",
+            "cancelled",
+            "取消",
+            "拒绝",
+            "不同意",
+            "否",
+            "不行",
+            "停止",
+        }:
+            return False
+        return None
+
+    @classmethod
+    def _extract_auth_decision(cls, parsed_dict: dict[str, Any]) -> bool | None:
+        for key in cls._AUTH_DECISION_KEYS:
+            if key in parsed_dict:
+                return cls._normalize_auth_decision(parsed_dict.get(key))
+        return None
+
+    @classmethod
+    def _auth_payload(cls, parsed_dict: dict[str, Any]) -> dict[str, Any]:
+        for key in ("data", "payload"):
+            value = parsed_dict.get(key)
+            if not isinstance(value, dict):
+                continue
+            if cls._AUTH_DECISION_KEY_SET.intersection(value) or any(
+                arg_key in value for arg_key in ("args", "tool_args", "arguments")
+            ):
+                return value
+        return parsed_dict
+
+    @classmethod
+    def _extract_auth_override(
+        cls,
+        parsed_dict: dict[str, Any] | None,
+    ) -> tuple[str | None, Any]:
+        if parsed_dict is None:
+            return None, None
+
+        tool_name = parsed_dict.get("tool_name") or parsed_dict.get("name")
+        if "args" in parsed_dict:
+            return tool_name, parsed_dict.get("args")
+        if "tool_args" in parsed_dict:
+            return tool_name, parsed_dict.get("tool_args")
+        if "arguments" in parsed_dict:
+            return tool_name, parsed_dict.get("arguments")
+        if not cls._AUTH_DECISION_KEY_SET.intersection(parsed_dict):
+            return tool_name, parsed_dict
+        return tool_name, None
+
+    @staticmethod
+    def _cancel_wait_timeout(task_info: TaskContext) -> None:
+        handle = task_info.timeout_handle
+        if handle is not None:
+            handle.cancel()
+        task_info.timeout_handle = None
+
+    @staticmethod
+    def _cancel_task_timeout(task_info: TaskContext) -> None:
+        handle = task_info.control.task_timeout_handle
+        if handle is not None:
+            handle.cancel()
+        task_info.control.task_timeout_handle = None
+
+    def _cancel_all_timeouts(self, task_info: TaskContext) -> None:
+        self._cancel_wait_timeout(task_info)
+        self._cancel_task_timeout(task_info)
+
+    def _clear_wait_state(
+        self, task_info: TaskContext, *, clear_resume_payload: bool = True
+    ) -> None:
+        self._cancel_wait_timeout(task_info)
+        if clear_resume_payload:
+            task_info.clear_wait()
+        else:
+            task_info.clear_wait_item()
+
+    @staticmethod
+    def _shutdown_task_executor(task_info: TaskContext) -> None:
+        executor = task_info.task_executor
+        if executor is None:
+            return
+        task_info.task_executor = None
+        try:
+            executor.shutdown(wait=False, cancel_futures=True)
+        except TypeError:
+            executor.shutdown(wait=False)
+
+    @staticmethod
+    def _release_event(task_info: TaskContext) -> None:
+        event = task_info.event
+        if event:
+            event.set()
+        task_info.event = asyncio.Event()
+
+    @staticmethod
+    def _interrupt_agent(task_info: TaskContext) -> None:
+        agent = task_info.agent
+        if agent is not None:
+            try:
+                agent.interrupt_switch = True
+            except Exception:
+                pass
+
+        agent_task = task_info.agent_task
+        if agent_task is not None and not agent_task.done():
+            try:
+                agent_task.cancel()
+            except Exception:
+                pass
+
+    def _abort_active_tool_call(self, task_info: TaskContext) -> None:
+        active_tool = task_info.active_tool
+        if active_tool is None:
+            return
+        try:
+            active_tool.cancel()
+        except Exception:
+            logger.warning(
+                "Failed to cancel active tool call agent_id=%s tool=%s mcp=%s",
+                self._agent.agent_id,
+                active_tool.tool_name,
+                active_tool.mcp_name,
+                exc_info=True,
+            )
+        finally:
+            task_info.clear_active_tool()
+
+    @staticmethod
+    def _phase_for_error(error: Exception) -> TaskPhase:
+        if isinstance(error, (TaskWaitTimeoutError, TaskExecutionTimeoutError)):
+            return TaskPhase.TIMED_OUT
+        if isinstance(error, (UserCancelledError, TaskCancelledError)):
+            return TaskPhase.CANCELLED
+        return TaskPhase.FAILED
+
+    async def _start_new_task(
+        self,
+        task_id: str,
+        initial_text: str,
+        request_headers: dict[str, str],
+        main_loop: asyncio.AbstractEventLoop,
+        task_updater: TaskUpdaterProtocol | Any,
+        *,
+        context_id: str | None = None,
+        task_store=None,
+    ) -> None:
+        from a2a.types import TaskState
+
+        await task_updater.update_status(TaskState.submitted)
+        task_info = TaskContext(
+            event=asyncio.Event(),
+            loop=main_loop,
+            updater=task_updater,
+            phase=TaskPhase.SUBMITTED,
+        )
+        task_observation = A2ATaskObservation(
+            agent_id=self._agent.agent_id,
+            task_id=task_id,
+            context_id=context_id,
+            request_headers=request_headers,
+            task_input=initial_text,
+        )
+        self.task_pool[task_id] = task_info
+
+        try:
+            agent = self.build_agent(mcp_headers=request_headers)
+        except Exception:
+            self.task_pool.pop(task_id, None)
+            raise
+
+        logger.info(
+            "Starting new task agent_id=%s task_id=%s has_text=%s",
+            self._agent.agent_id,
+            task_id,
+            bool(initial_text),
+        )
+        task_executor = ThreadPoolExecutor(
+            max_workers=1,
+            thread_name_prefix=f"runtime-task-{task_id[:8]}",
+        )
+        run_context = copy_context()
+        agent_task = main_loop.run_in_executor(
+            task_executor,
+            partial(
+                run_context.run,
+                self._run_agent_with_cleanup,
+                agent,
+                initial_text,
+                task_observation,
+            ),
+        )
+        agent_task.add_done_callback(
+            lambda task: None if task.cancelled() else task.exception()
+        )
+
+        task_info.agent = agent
+        task_info.agent_task = agent_task
+        task_info.task_executor = task_executor
+        task_info.set_phase(TaskPhase.RUNNING)
+        self._schedule_task_timeout(
+            task_id,
+            task_info,
+            task_updater,
+            task_store=task_store,
+        )
+
+    def _schedule_task_timeout(
+        self,
+        task_id: str,
+        task_info: TaskContext,
+        task_updater: TaskUpdaterProtocol | Any,
+        *,
+        task_store=None,
+    ) -> None:
+        self._cancel_task_timeout(task_info)
+        task_timeout = float(os.getenv("MCP_AGENT_TASK_TIMEOUT_SECONDS", "0"))
+        if task_timeout <= 0:
+            return
+
+        loop = task_info.loop
+        if not loop or loop.is_closed():
+            return
+
+        def on_timeout() -> None:
+            current_info = self.task_pool.get(task_id)
+            if current_info is None or current_info.finalized:
+                return
+
+            reason = (
+                f"任务执行超过 {self._format_timeout_seconds(task_timeout)} 秒，已自动取消"
+            )
+            logger.warning(
+                "Task execution timeout agent_id=%s task_id=%s reason=%s",
+                self._agent.agent_id,
+                task_id,
+                reason,
+            )
+            current_info.control.task_timeout_handle = None
+            loop.create_task(
+                self._request_task_stop(
+                    task_id,
+                    task_updater,
+                    error=TaskExecutionTimeoutError(reason),
+                    reason=reason,
+                    phase=TaskPhase.TIMED_OUT,
+                    task_store=task_store,
+                )
+            )
+
+        task_info.control.task_timeout_handle = loop.call_later(task_timeout, on_timeout)
+
+    def _schedule_wait_timeout(
+        self,
+        task_id: str,
+        task_info: TaskContext,
+        task_updater: TaskUpdaterProtocol | Any,
+        *,
+        task_store=None,
+    ) -> None:
+        self._cancel_wait_timeout(task_info)
+        wait_timeout = float(os.getenv("MCP_AGENT_WAIT_TIMEOUT_SECONDS", "1800"))
+        if wait_timeout <= 0:
+            return
+
+        loop = task_info.loop
+        if not loop or loop.is_closed():
+            return
+
+        def on_timeout() -> None:
+            current_info = self.task_pool.get(task_id)
+            if (
+                current_info is None
+                or current_info.finalized
+                or current_info.wait_item is None
+                or current_info.timed_out
+            ):
+                return
+            reason = (
+                "任务等待用户输入或授权超过 "
+                f"{self._format_timeout_seconds(wait_timeout)} 秒，已自动失败"
+            )
+            logger.warning(
+                "Task wait timeout agent_id=%s task_id=%s reason=%s",
+                self._agent.agent_id,
+                task_id,
+                reason,
+            )
+            current_info.timeout_handle = None
+            loop.create_task(
+                self._request_task_stop(
+                    task_id,
+                    task_updater,
+                    error=TaskWaitTimeoutError(reason),
+                    reason=reason,
+                    phase=TaskPhase.TIMED_OUT,
+                    task_store=task_store,
+                )
+            )
+
+        task_info.timeout_handle = loop.call_later(wait_timeout, on_timeout)
+
+    async def _request_task_stop(
+        self,
+        task_id: str,
+        task_updater: TaskUpdaterProtocol | Any,
+        *,
+        error: Exception,
+        reason: str,
+        phase: TaskPhase,
+        task_store=None,
+    ) -> None:
+        """统一处理任务停止。
+
+        所有 stop path 都收口到这里：
+
+        - 用户取消
+        - 等待超时
+        - 整体任务超时
+        """
+
+        task_info = self.task_pool.get(task_id)
+        if not task_info or task_info.finalized:
+            return
+        if task_info.control.stop_requested and task_info.control.stop_error is not None:
+            return
+
+        logger.info(
+            "Stopping task agent_id=%s task_id=%s phase=%s reason=%s",
+            self._agent.agent_id,
+            task_id,
+            phase.value,
+            reason,
+        )
+        task_info.request_stop(
+            error=error,
+            reason=reason,
+            timed_out=phase == TaskPhase.TIMED_OUT,
+        )
+        task_info.set_phase(TaskPhase.CANCELLING)
+        self._cancel_all_timeouts(task_info)
+        self._clear_wait_state(task_info)
+        self._abort_active_tool_call(task_info)
+        self._interrupt_agent(task_info)
+        self._release_event(task_info)
+        await self._finalize_task(
+            task_id,
+            task_updater,
+            error=error,
+            task_store=task_store,
+        )
+
+    async def _delete_failed_task_from_store(
+        self,
+        task_id: str,
+        task_updater: TaskUpdaterProtocol | Any,
+        *,
+        task_store=None,
+    ) -> None:
+        if task_store is None:
+            return
+
+        event_queue = getattr(task_updater, "event_queue", None)
+        if event_queue is not None and not event_queue.is_closed():
+            self._failed_task_ids.add(task_id)
+            return
+
+        try:
+            await task_store.delete(task_id)
+        except Exception:
+            logger.warning(
+                "Failed to delete failed task directly from task_store task_id=%s",
+                task_id,
+                exc_info=True,
+            )
+
+    def _run_agent_with_cleanup(
+        self,
+        agent: Any,
+        task_text: str,
+        observation: A2ATaskObservation | None = None,
+    ):
+        try:
+            with a2a_task_observation(observation) as span:
+                result = agent.run(task=task_text)
+                update_observation_output(span, result)
+                return result
+        finally:
+            mcp_clients = list(getattr(agent, "_runtime_mcp_clients", []) or [])
+            for client in mcp_clients:
+                try:
+                    client.disconnect()
+                except Exception:
+                    logger.warning(
+                        "Failed to disconnect MCP client during task cleanup agent_id=%s",
+                        self._agent.agent_id,
+                        exc_info=True,
+                    )
+            try:
+                agent._runtime_mcp_clients = []
+            except Exception:
+                pass
+
+    async def _handle_wait_resume(
+        self,
+        task_id: str,
+        task_info: TaskContext,
+        task_updater: TaskUpdaterProtocol | Any,
+        user_input: str | dict | bool | None,
+        *,
+        task_store=None,
+    ) -> bool:
+        wait_item = task_info.wait_item
+        if wait_item is None or user_input is None:
+            return False
+
+        wait_type = wait_state_type(wait_item)
+        if wait_type == WAIT_TYPE_INPUT_REQUIRED:
+            task_info.user_input = user_input
+            task_info.set_phase(TaskPhase.RUNNING)
+            logger.debug("Resuming task with user input task_id=%s", task_id)
+            self._release_event(task_info)
+            self._clear_wait_state(task_info, clear_resume_payload=False)
+            return True
+
+        if wait_type != WAIT_TYPE_AUTH_REQUIRED:
+            return False
+
+        approved: bool | None = None
+        override_tool_name = None
+        override_args = None
+        parsed_dict = user_input if isinstance(user_input, dict) else None
+        if parsed_dict is None:
+            approved = self._normalize_auth_decision(user_input)
+        else:
+            auth_payload = self._auth_payload(parsed_dict)
+            approved = self._extract_auth_decision(auth_payload)
+            override_tool_name, override_args = self._extract_auth_override(auth_payload)
+
+        if approved is None:
+            return False
+        if approved:
+            task_info.auth_denied = False
+            task_info.tool_name_override = override_tool_name
+            task_info.tool_args_override = override_args
+            task_info.set_phase(TaskPhase.RUNNING)
+            logger.debug(
+                "Resuming task with auth approved task_id=%s tool=%s",
+                task_id,
+                override_tool_name,
+            )
+            self._release_event(task_info)
+            self._clear_wait_state(task_info)
+            return True
+
+        logger.info(
+            "Task auth denied agent_id=%s task_id=%s",
+            self._agent.agent_id,
+            task_id,
+        )
+        task_info.auth_denied = True
+        self._release_event(task_info)
+        self._clear_wait_state(task_info)
+        await self._request_task_stop(
+            task_id,
+            task_updater,
+            error=UserCancelledError("user denied auth request"),
+            reason="user denied auth request",
+            phase=TaskPhase.CANCELLED,
+            task_store=task_store,
+        )
+        return True
+
+    async def _listen_for_events(self, task_id: str):
+        while True:
+            task_info = self.task_pool.get(task_id)
+            if not task_info or task_info.finalized:
+                return
+            agent_task = task_info.agent_task
+            if agent_task and agent_task.done():
+                return
+            if task_info.wait_item is not None:
+                yield task_info.wait_item
+            await asyncio.sleep(0.2)
+
+    async def _emit_wait_state(
+        self,
+        task_id: str,
+        task_info: TaskContext,
+        task_updater: TaskUpdaterProtocol | Any,
+        *,
+        task_store=None,
+    ) -> None:
+        async for pending_item in self._listen_for_events(task_id):
+            payload = wait_state_payload(pending_item)
+            wait_type = payload["type"]
+            wait_data = payload["data"]
+            if wait_type == WAIT_TYPE_INPUT_REQUIRED:
+                task_info.set_phase(TaskPhase.WAITING_INPUT)
+                logger.info(
+                    "Task entered input_required agent_id=%s task_id=%s",
+                    self._agent.agent_id,
+                    task_id,
+                )
+                self._schedule_wait_timeout(
+                    task_id,
+                    task_info,
+                    task_updater,
+                    task_store=task_store,
+                )
+                await emit_wait_state(task_updater, payload)
+                task_info.wait_item_emitted = True
+                return
+            if wait_type == WAIT_TYPE_AUTH_REQUIRED:
+                task_info.set_phase(TaskPhase.WAITING_AUTH)
+                tool_name = wait_data.get("tool_name")
+                logger.info(
+                    "Task entered auth_required agent_id=%s task_id=%s tool=%s",
+                    self._agent.agent_id,
+                    task_id,
+                    tool_name,
+                )
+                self._schedule_wait_timeout(
+                    task_id,
+                    task_info,
+                    task_updater,
+                    task_store=task_store,
+                )
+                await emit_wait_state(task_updater, payload)
+                task_info.wait_item_emitted = True
+                return
+
+    async def _process_task_request(
+        self,
+        task_id: str,
+        task_updater: TaskUpdaterProtocol | Any,
+        user_input: str | dict | bool | None = None,
+        *,
+        task_store=None,
+    ) -> None:
+        task_info = self.task_pool.get(task_id)
+        if not task_info:
+            return
+
+        wait_item = task_info.wait_item
+        if wait_item and user_input is None and task_info.wait_item_emitted:
+            return
+
+        if wait_item and user_input is not None:
+            handled = await self._handle_wait_resume(
+                task_id,
+                task_info,
+                task_updater,
+                user_input,
+                task_store=task_store,
+            )
+            if handled:
+                return
+
+        await self._emit_wait_state(
+            task_id,
+            task_info,
+            task_updater,
+            task_store=task_store,
+        )
+
+    async def _run_task_cycle(
+        self,
+        task_id: str,
+        task_updater: TaskUpdaterProtocol | Any,
+        user_input: str | dict | bool | None = None,
+        *,
+        task_store=None,
+    ) -> None:
+        """驱动一次任务执行周期。
+
+        A2A bridge 只负责把请求转发进来；真正的等待、恢复、完成、失败都在这里收口。
+        """
+
+        await self._process_task_request(
+            task_id,
+            task_updater,
+            user_input,
+            task_store=task_store,
+        )
+
+        task_info = self.task_pool.get(task_id)
+        if not task_info or task_info.finalized:
+            return
+        if task_info.control.stop_requested:
+            return
+
+        agent_task = task_info.agent_task
+        if not agent_task:
+            return
+
+        while True:
+            if task_info.wait_item is not None:
+                await self._process_task_request(
+                    task_id,
+                    task_updater,
+                    None,
+                    task_store=task_store,
+                )
+                return
+            if task_info.control.stop_requested:
+                return
+            if agent_task.done():
+                break
+            await asyncio.sleep(0.1)
+
+        try:
+            result = await agent_task
+        except asyncio.CancelledError:
+            if task_info.finalized:
+                return
+            error = task_info.control.stop_error or TaskCancelledError(
+                "task cancelled by runtime"
+            )
+            await self._finalize_task(
+                task_id,
+                task_updater,
+                error=error,
+                task_store=task_store,
+            )
+        except Exception as exc:
+            if task_info.finalized:
+                return
+            error = task_info.control.stop_error or exc
+            await self._finalize_task(
+                task_id,
+                task_updater,
+                error=error,
+                task_store=task_store,
+            )
+        else:
+            if task_info.finalized:
+                return
+            if (
+                task_info.control.stop_requested
+                and task_info.control.stop_error is not None
+            ):
+                await self._finalize_task(
+                    task_id,
+                    task_updater,
+                    error=task_info.control.stop_error,
+                    task_store=task_store,
+                )
+                return
+            await self._finalize_task(
+                task_id,
+                task_updater,
+                result=result,
+                task_store=task_store,
+            )
+
+    async def _cancel_task(
+        self,
+        task_id: str,
+        task_updater: TaskUpdaterProtocol | Any,
+        *,
+        task_store=None,
+    ) -> None:
+        """终止任务并统一走 finalize 清理路径。"""
+
+        task_info = self.task_pool.get(task_id)
+        if not task_info or task_info.finalized:
+            return
+
+        logger.info(
+            "Cancelling task agent_id=%s task_id=%s",
+            self._agent.agent_id,
+            task_id,
+        )
+        await self._request_task_stop(
+            task_id,
+            task_updater,
+            error=TaskCancelledError("task cancelled by client"),
+            reason="task cancelled by client",
+            phase=TaskPhase.CANCELLED,
+            task_store=task_store,
+        )
+
+    async def _finalize_task(
+        self,
+        task_id: str,
+        task_updater: TaskUpdaterProtocol | Any,
+        result: Any | None = None,
+        error: Exception | None = None,
+        message_text: str | None = None,
+        *,
+        task_store=None,
+    ) -> None:
+        task_info = self.task_pool.get(task_id)
+        if not task_info or task_info.finalized:
+            return
+        task_info.finalized = True
+
+        if message_text is not None:
+            task_info.set_phase(TaskPhase.COMPLETED)
+            logger.debug(
+                "Task finalized with message agent_id=%s task_id=%s",
+                self._agent.agent_id,
+                task_id,
+            )
+            await publish_message_completion(task_updater, message_text)
+        elif error is not None:
+            task_info.set_phase(self._phase_for_error(error))
+            if isinstance(error, TaskWaitTimeoutError):
+                logger.warning(
+                    "Task timed out agent_id=%s task_id=%s error=%s: %s",
+                    self._agent.agent_id,
+                    task_id,
+                    type(error).__name__,
+                    error,
+                )
+            elif isinstance(error, TaskExecutionTimeoutError):
+                logger.warning(
+                    "Task execution timed out agent_id=%s task_id=%s error=%s: %s",
+                    self._agent.agent_id,
+                    task_id,
+                    type(error).__name__,
+                    error,
+                )
+            elif isinstance(error, (UserCancelledError, TaskCancelledError)):
+                logger.info(
+                    "Task cancelled agent_id=%s task_id=%s error=%s: %s",
+                    self._agent.agent_id,
+                    task_id,
+                    type(error).__name__,
+                    error,
+                )
+            else:
+                logger.error(
+                    "Task failed agent_id=%s task_id=%s error=%s: %s",
+                    self._agent.agent_id,
+                    task_id,
+                    type(error).__name__,
+                    error,
+                )
+            await publish_error(task_updater, error)
+            await self._delete_failed_task_from_store(
+                task_id,
+                task_updater,
+                task_store=task_store,
+            )
+        else:
+            task_info.set_phase(TaskPhase.COMPLETED)
+            logger.info(
+                "Task completed agent_id=%s task_id=%s",
+                self._agent.agent_id,
+                task_id,
+            )
+            final_result = self.format_result(result)
+            await publish_result(task_updater, final_result)
+
+        self._cancel_all_timeouts(task_info)
+        self._abort_active_tool_call(task_info)
+        self._shutdown_task_executor(task_info)
+        task_info.set_phase(TaskPhase.FINALIZED)
+        self.task_pool.pop(task_id, None)