Agent/gateway/core/executor/task_manager.py

from __future__ import annotations

import asyncio
import logging
import uuid
from collections import defaultdict
from collections.abc import Callable
from typing import Any, Literal

import httpx

from utils.env_parse import env_float, env_str

from gateway.core.executor.errors import ExecutorError, TaskNotFoundError
from gateway.core.executor.models import TERMINAL_STATUSES, RunMode, TaskRecord
from gateway.core.lifecycle import TraceManager, WorkspaceManager
from gateway.core.lifecycle.errors import LifecycleError

logger = logging.getLogger(__name__)


# --- 执行上下文（路径 + 内存日志） ---


class ExecutionContext:
    """Workspace 路径解析与执行日志（内存）。"""

    def __init__(self, workspace_manager: WorkspaceManager, trace_manager: TraceManager) -> None:
        self._wm = workspace_manager
        self._tm = trace_manager
        self._logs: dict[str, list[dict[str, Any]]] = defaultdict(list)

    async def create_context(self, task_id: str, trace_id: str) -> dict[str, Any]:
        try:
            workspace_id = self._tm.get_workspace_id(trace_id)
        except LifecycleError as e:
            logger.warning("ExecutionContext.create_context: %s", e)
            return {
                "task_id": task_id,
                "trace_id": trace_id,
                "workspace_id": None,
                "workspace_path": None,
                "note": str(e),
            }
        try:
            path = await self._wm.get_workspace_path(workspace_id)
        except Exception as e:
            logger.warning("ExecutionContext.get_workspace_path failed: %s", e)
            path = None
        return {
            "task_id": task_id,
            "trace_id": trace_id,
            "workspace_id": workspace_id,
            "workspace_path": path,
        }

    async def get_workspace_path(self, trace_id: str) -> str | None:
        try:
            workspace_id = self._tm.get_workspace_id(trace_id)
            return await self._wm.get_workspace_path(workspace_id)
        except Exception as e:
            logger.warning("ExecutionContext.get_workspace_path trace_id=%s err=%s", trace_id, e)
            return None

    def log_execution(self, task_id: str, log_entry: dict[str, Any]) -> None:
        self._logs[task_id].append(dict(log_entry))

    def get_logs(self, task_id: str) -> list[dict[str, Any]]:
        return list(self._logs.get(task_id, []))


# --- 内存任务表、按 trace 互斥 ---


class _InMemoryTaskStore:
    def __init__(self) -> None:
        self._by_id: dict[str, TaskRecord] = {}

    def put(self, record: TaskRecord) -> None:
        self._by_id[record.task_id] = record

    def get(self, task_id: str) -> TaskRecord | None:
        return self._by_id.get(task_id)

    def require(self, task_id: str) -> TaskRecord:
        rec = self._by_id.get(task_id)
        if rec is None:
            raise TaskNotFoundError(task_id)
        return rec

    def list_as_dicts(
        self,
        *,
        trace_id: str | None = None,
        status: str | None = None,
    ) -> list[dict[str, Any]]:
        out: list[dict[str, Any]] = []
        for rec in self._by_id.values():
            if trace_id is not None and rec.trace_id != trace_id:
                continue
            if status is not None and rec.status != status:
                continue
            out.append(rec.to_dict())
        out.sort(key=lambda d: d.get("created_at") or "", reverse=True)
        return out

    def contains(self, task_id: str) -> bool:
        return task_id in self._by_id


class _TraceSerialLocks:
    def __init__(self) -> None:
        self._locks: dict[str, asyncio.Lock] = {}

    def lock_for(self, trace_id: str) -> asyncio.Lock:
        if trace_id not in self._locks:
            self._locks[trace_id] = asyncio.Lock()
        return self._locks[trace_id]


# --- Agent HTTP 与 gateway_exec ---


class AgentTraceHttpClient:
    """对 Agent ``/api/traces`` 的最小封装。"""

    def __init__(self, *, base_url: str, timeout: float) -> None:
        self._base = base_url.rstrip("/")
        self._timeout = timeout

    async def post_run(
        self,
        trace_id: str,
        *,
        messages: list[dict[str, Any]],
        gateway_exec: dict[str, Any] | None = None,
    ) -> tuple[int, Any]:
        body: dict[str, Any] = {"messages": messages}
        if gateway_exec:
            body["gateway_exec"] = gateway_exec
        async with httpx.AsyncClient(timeout=self._timeout) as client:
            r = await client.post(f"{self._base}/api/traces/{trace_id}/run", json=body)
        try:
            payload: Any = r.json()
        except Exception:
            payload = r.text
        return r.status_code, payload

    async def get_trace_status(self, trace_id: str) -> str | None:
        async with httpx.AsyncClient(timeout=self._timeout) as client:
            r = await client.get(f"{self._base}/api/traces/{trace_id}")
        if r.status_code != 200:
            return None
        try:
            data = r.json()
        except Exception:
            return None
        trace = data.get("trace")
        if isinstance(trace, dict):
            st = trace.get("status")
            return str(st) if st is not None else None
        return None

    async def post_stop(self, trace_id: str) -> bool:
        async with httpx.AsyncClient(timeout=self._timeout) as client:
            r = await client.post(f"{self._base}/api/traces/{trace_id}/stop")
        return r.status_code < 400


class _GatewayExecResolver:
    def __init__(self, workspace_manager: WorkspaceManager, trace_manager: TraceManager) -> None:
        self._wm = workspace_manager
        self._trace_mgr = trace_manager

    def resolve(self, trace_id: str) -> dict[str, Any] | None:
        try:
            wid = self._trace_mgr.get_workspace_id(trace_id)
        except LifecycleError:
            return None
        cid = self._wm.get_workspace_container_id(wid)
        if not cid:
            return None
        return {
            "docker_container": cid,
            "container_user": "agent",
            "container_workdir": "/home/agent/workspace",
        }


# --- 单任务执行管道 ---


class _TaskExecutionPipeline:
    def __init__(
        self,
        *,
        store: _InMemoryTaskStore,
        trace_client: AgentTraceHttpClient,
        gateway_exec: _GatewayExecResolver,
        execution_context: ExecutionContext,
        poll_interval: float,
        poll_max_seconds: float,
        is_cancelled: Callable[[str], bool],
    ) -> None:
        self._store = store
        self._http = trace_client
        self._gateway_exec = gateway_exec
        self._ctx = execution_context
        self._poll_interval = poll_interval
        self._poll_max_seconds = poll_max_seconds
        self._is_cancelled = is_cancelled

    async def run_after_lock_acquired(self, task_id: str) -> None:
        rec = self._store.get(task_id)
        if rec is None:
            return
        if self._is_cancelled(task_id):
            rec.touch(status="cancelled", error_message="cancelled_before_run")
            return

        rec.touch(status="running")
        self._ctx.log_execution(task_id, {"event": "run_start"})

        gateway_exec = self._gateway_exec.resolve(rec.trace_id)
        messages = [{"role": "user", "content": rec.task_description}]

        try:
            status_code, payload = await self._http.post_run(
                rec.trace_id,
                messages=messages,
                gateway_exec=gateway_exec,
            )
        except Exception as e:
            logger.exception("executor post_run failed task_id=%s", task_id)
            rec.touch(status="failed", error_message=str(e))
            self._ctx.log_execution(task_id, {"event": "http_error", "error": str(e)})
            return

        if status_code == 409:
            rec.touch(status="failed", error_message="trace_already_running")
            return

        if status_code >= 400:
            detail = payload if isinstance(payload, str) else str(payload)
            rec.touch(status="failed", error_message=f"agent_http_{status_code}: {detail[:500]}")
            return

        self._ctx.log_execution(task_id, {"event": "run_accepted", "status_code": status_code})

        poll_kind, terminal, poll_err = await self._poll_until_terminal(task_id, rec)
        if poll_kind == "cancelled":
            return

        if poll_err == "poll_timeout":
            rec.touch(status="failed", error_message="poll_timeout", trace_terminal_status=terminal)
        elif terminal == "failed":
            rec.touch(status="failed", error_message="trace_failed", trace_terminal_status=terminal)
        elif terminal == "stopped":
            rec.touch(status="cancelled", error_message="trace_stopped", trace_terminal_status=terminal)
        elif terminal == "completed":
            rec.touch(status="completed", trace_terminal_status=terminal)
        elif terminal:
            rec.touch(status="completed", trace_terminal_status=terminal)
        else:
            rec.touch(status="failed", error_message="status_unknown")

        self._ctx.log_execution(
            task_id,
            {"event": "finished", "terminal": terminal, "err": poll_err},
        )

    async def _poll_until_terminal(
        self, task_id: str, rec: TaskRecord
    ) -> tuple[Literal["ok", "cancelled"], str | None, str | None]:
        elapsed = 0.0
        terminal: str | None = None
        poll_err: str | None = None
        while elapsed <= self._poll_max_seconds:
            if self._is_cancelled(task_id):
                await self._http.post_stop(rec.trace_id)
                rec.touch(
                    status="cancelled",
                    error_message="cancelled_during_run",
                    trace_terminal_status="stopped",
                )
                self._ctx.log_execution(task_id, {"event": "cancelled_mid_poll"})
                return "cancelled", None, None

            st = await self._http.get_trace_status(rec.trace_id)
            if st and st in TERMINAL_STATUSES:
                terminal = st
                break
            await asyncio.sleep(self._poll_interval)
            elapsed += self._poll_interval
        else:
            poll_err = "poll_timeout"
        return "ok", terminal, poll_err


# --- 编排入口 ---


class TaskManager:
    """
    校验 Trace、内存任务表、按 trace 串行执行、同步等待与取消。
    内部类：存储、锁、HTTP、管道均在本模块实现。
    """

    def __init__(
        self,
        *,
        workspace_manager: WorkspaceManager,
        trace_manager: TraceManager,
        agent_api_base_url: str,
        http_timeout: float,
        poll_interval: float = 2.0,
        poll_max_seconds: float = 86400.0,
    ) -> None:
        self._trace_mgr = trace_manager
        self._store = _InMemoryTaskStore()
        self._trace_locks = _TraceSerialLocks()
        self._http = AgentTraceHttpClient(base_url=agent_api_base_url, timeout=http_timeout)
        self._gateway_exec = _GatewayExecResolver(workspace_manager, trace_manager)
        self._ctx = ExecutionContext(workspace_manager, trace_manager)
        self._pipeline = _TaskExecutionPipeline(
            store=self._store,
            trace_client=self._http,
            gateway_exec=self._gateway_exec,
            execution_context=self._ctx,
            poll_interval=poll_interval,
            poll_max_seconds=poll_max_seconds,
            is_cancelled=self._is_cancelled,
        )
        self._cancelled: set[str] = set()
        self._done_events: dict[str, asyncio.Event] = {}

    def _is_cancelled(self, task_id: str) -> bool:
        return task_id in self._cancelled

    @classmethod
    def from_env(
        cls,
        workspace_manager: WorkspaceManager,
        trace_manager: TraceManager,
    ) -> TaskManager:
        return cls(
            workspace_manager=workspace_manager,
            trace_manager=trace_manager,
            agent_api_base_url=env_str("GATEWAY_AGENT_API_BASE_URL", "http://127.0.0.1:8000"),
            http_timeout=env_float("GATEWAY_AGENT_API_TIMEOUT", 60.0),
            poll_interval=env_float("GATEWAY_EXECUTOR_POLL_INTERVAL", 2.0),
            poll_max_seconds=env_float("GATEWAY_EXECUTOR_POLL_MAX_SECONDS", 86400.0),
        )

    async def submit_task(
        self,
        trace_id: str,
        task_description: str,
        mode: RunMode = "async",
        metadata: dict[str, Any] | None = None,
    ) -> str:
        await self._trace_mgr.get_trace(trace_id)

        task_id = f"gtask-{uuid.uuid4()}"
        rec = TaskRecord(
            task_id=task_id,
            trace_id=trace_id,
            task_description=task_description,
            mode=mode,
            metadata=dict(metadata or {}),
        )
        self._store.put(rec)
        done_ev = asyncio.Event()
        self._done_events[task_id] = done_ev

        self._ctx.log_execution(
            task_id,
            {"event": "submitted", "trace_id": trace_id, "mode": mode},
        )

        asyncio.create_task(self._run_task_pipeline(task_id), name=f"executor:{task_id}")

        if mode == "sync":
            await done_ev.wait()
        return task_id

    async def _run_task_pipeline(self, task_id: str) -> None:
        try:
            rec = self._store.get(task_id)
            if not rec:
                return
            if task_id in self._cancelled:
                rec.touch(status="cancelled", error_message="cancelled_before_start")
                return

            async with self._trace_locks.lock_for(rec.trace_id):
                if task_id in self._cancelled:
                    rec.touch(status="cancelled", error_message="cancelled_before_run")
                    return

                await self._pipeline.run_after_lock_acquired(task_id)
        finally:
            self._finish(task_id)

    def _finish(self, task_id: str) -> None:
        ev = self._done_events.pop(task_id, None)
        if ev is not None:
            ev.set()

    def get_task(self, task_id: str) -> dict[str, Any]:
        return self._store.require(task_id).to_dict()

    def list_tasks(
        self,
        trace_id: str | None = None,
        status: str | None = None,
    ) -> list[dict[str, Any]]:
        return self._store.list_as_dicts(trace_id=trace_id, status=status)

    async def cancel_task(self, task_id: str) -> None:
        rec = self._store.get(task_id)
        if not rec:
            raise TaskNotFoundError(task_id)

        if rec.status in ("completed", "failed", "cancelled"):
            raise ExecutorError(f"任务已结束，无法取消: {rec.status}")

        self._cancelled.add(task_id)

        if rec.status == "running":
            try:
                await self._http.post_stop(rec.trace_id)
            except Exception as e:
                logger.warning("TaskManager cancel stop trace failed: %s", e)

        self._ctx.log_execution(task_id, {"event": "cancel_requested"})

    def get_execution_logs(self, task_id: str) -> list[dict[str, Any]]:
        if not self._store.contains(task_id):
            raise TaskNotFoundError(task_id)
        return self._ctx.get_logs(task_id)