auto_agent/backend/app/api/edge/tasks.py

from __future__ import annotations

import json
from typing import Annotated
from uuid import uuid4

from fastapi import APIRouter, Depends, Header, HTTPException, status
from sqlalchemy.orm import Session

from app.core.config import get_settings
from app.core.constants import ERROR_CODE_OK
from app.core.time import format_now
from app.db.session import get_db
from app.schemas.common import ApiResponse
from app.schemas.edge import (
    EdgeEventData,
    EdgeEventRequest,
    EdgeHeartbeatData,
    EdgeHeartbeatRequest,
    EdgePullTasksData,
    EdgePullTasksRequest,
    EdgeTaskItem,
    EdgeTaskReportData,
    EdgeTaskReportRequest,
)
from app.services.edge_service import EdgeService, EdgeTaskConflictError, EdgeTaskNotFoundError

router = APIRouter(prefix="/api/agent/edge", tags=["agent-edge"])


def build_request_id(header_value: str | None) -> str:
    return header_value or f"req-{uuid4().hex[:12]}"


@router.post("/heartbeat", response_model=ApiResponse[EdgeHeartbeatData])
def heartbeat(
    payload: EdgeHeartbeatRequest,
    db: Annotated[Session, Depends(get_db)],
    x_request_id: Annotated[str | None, Header(alias="X-Request-Id")] = None,
) -> ApiResponse[EdgeHeartbeatData]:
    settings = get_settings()
    request_id = build_request_id(x_request_id)
    node = EdgeService(db, settings.default_timezone).heartbeat(
        edge_id=payload.edge_id,
        hostname=payload.hostname,
        os_type=payload.os_type,
        agent_version=payload.agent_version,
        capabilities=payload.capabilities,
    )
    return ApiResponse[EdgeHeartbeatData](
        request_id=request_id,
        success=True,
        code=ERROR_CODE_OK,
        message="success",
        data=EdgeHeartbeatData(
            edge_id=node.edge_id,
            node_status=node.node_status,
            last_heartbeat_at=node.last_heartbeat_at,
        ),
        timestamp=format_now(settings.default_timezone),
    )


@router.post("/tasks/pull", response_model=ApiResponse[EdgePullTasksData])
def pull_tasks(
    payload: EdgePullTasksRequest,
    db: Annotated[Session, Depends(get_db)],
    x_request_id: Annotated[str | None, Header(alias="X-Request-Id")] = None,
) -> ApiResponse[EdgePullTasksData]:
    settings = get_settings()
    request_id = build_request_id(x_request_id)
    tasks = EdgeService(db, settings.default_timezone).pull_tasks(payload.edge_id, payload.max_tasks)
    return ApiResponse[EdgePullTasksData](
        request_id=request_id,
        success=True,
        code=ERROR_CODE_OK,
        message="success",
        data=EdgePullTasksData(
            tasks=[
                EdgeTaskItem(
                    task_id=item.task_id,
                    step_id=item.step_id,
                    tool_name=item.tool_name,
                    params=json.loads(item.params_json),
                    expire_at=item.expire_at,
                )
                for item in tasks
            ]
        ),
        timestamp=format_now(settings.default_timezone),
    )


@router.post("/tasks/report", response_model=ApiResponse[EdgeTaskReportData])
def report_task(
    payload: EdgeTaskReportRequest,
    db: Annotated[Session, Depends(get_db)],
    x_request_id: Annotated[str | None, Header(alias="X-Request-Id")] = None,
) -> ApiResponse[EdgeTaskReportData]:
    settings = get_settings()
    request_id = build_request_id(x_request_id)
    try:
        edge_task, task_status = EdgeService(db, settings.default_timezone).report_task(
            edge_id=payload.edge_id,
            step_id=payload.step_id,
            success=payload.success,
            message=payload.message,
            data=payload.data,
            evidence=payload.evidence,
            started_at=payload.started_at,
            finished_at=payload.finished_at,
        )
    except EdgeTaskNotFoundError as exc:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail={"code": exc.code, "message": "edge step not found"}) from exc
    except EdgeTaskConflictError as exc:
        message = exc.args[0] if exc.args else "edge task conflict"
        raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail={"code": exc.code, "message": message}) from exc

    if payload.task_id != edge_task.task_id or payload.tool_name != edge_task.tool_name:
        raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail={"code": "CONFLICT", "message": "task_id or tool_name mismatch"})

    return ApiResponse[EdgeTaskReportData](
        request_id=request_id,
        success=True,
        code=ERROR_CODE_OK,
        message="success",
        data=EdgeTaskReportData(
            task_id=edge_task.task_id,
            step_id=edge_task.step_id,
            step_status=edge_task.step_status,
            task_status=task_status,
        ),
        timestamp=format_now(settings.default_timezone),
    )


@router.post("/events", response_model=ApiResponse[EdgeEventData])
def report_event(
    payload: EdgeEventRequest,
    db: Annotated[Session, Depends(get_db)],
    x_request_id: Annotated[str | None, Header(alias="X-Request-Id")] = None,
) -> ApiResponse[EdgeEventData]:
    settings = get_settings()
    request_id = build_request_id(x_request_id)
    EdgeService(db, settings.default_timezone).record_event(
        edge_id=payload.edge_id,
        event_type=payload.event_type,
        message=payload.message,
        detail=payload.detail,
    )
    return ApiResponse[EdgeEventData](
        request_id=request_id,
        success=True,
        code=ERROR_CODE_OK,
        message="success",
        data=EdgeEventData(
            edge_id=payload.edge_id,
            event_type=payload.event_type,
            accepted=True,
        ),
        timestamp=format_now(settings.default_timezone),
    )