Agent/gateway/core/executor/api.py

from __future__ import annotations

import logging
from typing import Any

from fastapi import APIRouter, HTTPException, Query
from pydantic import BaseModel, Field

from gateway.core.executor.errors import ExecutorError, TaskNotFoundError
from gateway.core.executor.models import RunMode
from gateway.core.executor.task_manager import TaskManager
from gateway.core.lifecycle.errors import LifecycleError

logger = logging.getLogger(__name__)


class SubmitTaskRequest(BaseModel):
    trace_id: str = Field(..., description="Agent Trace ID")
    task_description: str = Field(..., description="用户任务描述，将作为一条 user 消息续跑")
    mode: RunMode = Field("async", description="async：立即返回 task_id；sync：阻塞至 Trace 终态")
    metadata: dict[str, Any] = Field(default_factory=dict)


class SubmitTaskResponse(BaseModel):
    task_id: str


def build_executor_router(task_manager: TaskManager) -> APIRouter:
    router = APIRouter(prefix="/gateway/executor", tags=["gateway-executor"])

    @router.post("/tasks", response_model=SubmitTaskResponse)
    async def submit_task(req: SubmitTaskRequest) -> SubmitTaskResponse:
        try:
            task_id = await task_manager.submit_task(
                req.trace_id,
                req.task_description,
                mode=req.mode,
                metadata=req.metadata,
            )
        except TaskNotFoundError as e:
            raise HTTPException(status_code=404, detail=str(e)) from e
        except LifecycleError as e:
            raise HTTPException(status_code=404, detail=str(e)) from e
        except ExecutorError as e:
            raise HTTPException(status_code=400, detail=str(e)) from e
        except Exception as e:
            logger.exception("executor submit_task")
            raise HTTPException(status_code=502, detail=str(e)) from e
        return SubmitTaskResponse(task_id=task_id)

    @router.get("/tasks/{task_id}")
    async def get_task(task_id: str) -> dict[str, Any]:
        try:
            return task_manager.get_task(task_id)
        except TaskNotFoundError as e:
            raise HTTPException(status_code=404, detail=str(e)) from e

    @router.get("/tasks")
    async def list_tasks(
        trace_id: str | None = Query(None),
        status: str | None = Query(None),
    ) -> dict[str, Any]:
        items = task_manager.list_tasks(trace_id=trace_id, status=status)
        return {"tasks": items, "count": len(items)}

    @router.get("/tasks/{task_id}/logs")
    async def task_logs(task_id: str) -> dict[str, Any]:
        try:
            logs = task_manager.get_execution_logs(task_id)
        except TaskNotFoundError as e:
            raise HTTPException(status_code=404, detail=str(e)) from e
        return {"task_id": task_id, "logs": logs}

    @router.post("/tasks/{task_id}/cancel")
    async def cancel_task(task_id: str) -> dict[str, str]:
        try:
            await task_manager.cancel_task(task_id)
        except TaskNotFoundError as e:
            raise HTTPException(status_code=404, detail=str(e)) from e
        except ExecutorError as e:
            raise HTTPException(status_code=400, detail=str(e)) from e
        return {"task_id": task_id, "status": "cancel_requested"}

    return router