script-build-and-framework-tmp/agent/agent/orchestration/api_v2.py

"""Optional FastAPI adapter for orchestration V2.

FastAPI is imported only when :func:`create_orchestration_router` is called,
so importing the core orchestration package does not require server extras.
"""

from __future__ import annotations

import asyncio
from typing import Any, Awaitable, Callable, Optional

from pydantic import ValidationError

from .coordinator import OrchestrationError, TaskConflict, TaskCoordinator
from .store import (
    ArtifactConflict,
    RevisionConflict,
    TaskStoreError,
    TaskStoreNotFound,
)
from .wire import (
    ArtifactSnapshotView,
    AttemptList,
    AttemptView,
    DispatchOperationRequest,
    EventPageView,
    EventView,
    MissionSnapshotView,
    OperationRequest,
    OperationView,
    PlannerDecisionView,
    ProblemDetails,
    RootCompletionView,
    TaskList,
    TaskView,
    ValidationList,
    ValidationView,
)


_ERROR_RESPONSES = {
    status: {"model": ProblemDetails}
    for status in (404, 409, 422, 503)
}


class _ResourceNotFound(LookupError):
    pass


def create_orchestration_router(coordinator: TaskCoordinator) -> Any:
    """Build an isolated router backed by the supplied coordinator instance."""

    try:
        from fastapi import APIRouter, Header
        from fastapi.exceptions import RequestValidationError
        from fastapi.responses import JSONResponse
        from fastapi.routing import APIRoute
    except ImportError as exc:  # pragma: no cover - depends on optional extra
        raise RuntimeError(
            "FastAPI is optional; install cyber-agent[server] to create the V2 router"
        ) from exc

    class ProblemRoute(APIRoute):
        def get_route_handler(self) -> Callable[..., Awaitable[Any]]:
            original = super().get_route_handler()

            async def route_handler(request: Any) -> Any:
                try:
                    return await original(request)
                except RequestValidationError as exc:
                    value = ProblemDetails(
                        title="Invalid Request",
                        status=422,
                        detail=str(exc),
                        code="invalid_request",
                        instance=request.url.path,
                    )
                    return JSONResponse(
                        status_code=422, content=value.model_dump(mode="json")
                    )

            return route_handler

    router = APIRouter(
        prefix="/api/v2", tags=["orchestration-v2"], route_class=ProblemRoute
    )

    @router.get(
        "/roots/{root_trace_id}/snapshot",
        response_model=MissionSnapshotView,
        responses=_ERROR_RESPONSES,
    )
    async def get_mission_snapshot(root_trace_id: str) -> Any:
        path = f"/api/v2/roots/{root_trace_id}/snapshot"

        async def action() -> MissionSnapshotView:
            # This is deliberately the only load in the endpoint: every
            # collection must describe one immutable ledger revision.
            ledger = await coordinator.task_store.load(root_trace_id)
            return MissionSnapshotView(
                revision=ledger.revision,
                root_trace_id=ledger.root_trace_id,
                root_task_id=ledger.root_task_id,
                root_objective=ledger.root_objective,
                focused_task_id=ledger.focused_task_id,
                tasks=[
                    TaskView.from_domain(item)
                    for item in sorted(
                        ledger.tasks.values(),
                        key=lambda value: (value.created_at, value.task_id),
                    )
                ],
                attempts=[
                    AttemptView.from_domain(item)
                    for item in sorted(
                        ledger.attempts.values(),
                        key=lambda value: (value.created_at, value.attempt_id),
                    )
                ],
                validations=[
                    ValidationView.from_domain(item)
                    for item in sorted(
                        ledger.validations.values(),
                        key=lambda value: (value.created_at, value.validation_id),
                    )
                ],
                decisions=[
                    PlannerDecisionView.from_domain(item)
                    for item in sorted(
                        ledger.decisions.values(),
                        key=lambda value: (value.created_at, value.decision_id),
                    )
                ],
                operations=[
                    OperationView.from_domain(item)
                    for item in sorted(
                        ledger.operations.values(),
                        key=lambda value: (value.created_at, value.operation_id),
                    )
                ],
            )

        return await call(action, path)

    @router.get(
        "/roots/{root_trace_id}/completion",
        response_model=RootCompletionView,
        responses=_ERROR_RESPONSES,
    )
    async def get_root_completion(root_trace_id: str) -> Any:
        path = f"/api/v2/roots/{root_trace_id}/completion"

        async def action() -> RootCompletionView:
            value = await coordinator.root_completion(root_trace_id)
            return RootCompletionView(
                root_task_id=value["root_task_id"],
                root_objective=value["root_objective"],
                status=value["status"],
                blocked_reason=value.get("blocked_reason"),
                result_summary=value.get("result_summary"),
            )

        return await call(action, path)

    @router.get(
        "/roots/{root_trace_id}/snapshots/{snapshot_id}",
        response_model=ArtifactSnapshotView,
        responses=_ERROR_RESPONSES,
    )
    async def get_artifact_snapshot(root_trace_id: str, snapshot_id: str) -> Any:
        path = f"/api/v2/roots/{root_trace_id}/snapshots/{snapshot_id}"

        async def action() -> ArtifactSnapshotView:
            try:
                value = await coordinator.artifact_store.get(
                    root_trace_id, snapshot_id
                )
            except FileNotFoundError as exc:
                raise _ResourceNotFound(
                    f"Artifact snapshot not found: {snapshot_id}"
                ) from exc
            return ArtifactSnapshotView.from_domain(value)

        return await call(action, path)

    def problem(exc: Exception, instance: str) -> JSONResponse:
        status, title, code, detail = _problem_fields(exc)
        value = ProblemDetails(
            title=title,
            status=status,
            detail=detail,
            code=code,
            instance=instance,
        )
        return JSONResponse(status_code=status, content=value.model_dump(mode="json"))

    async def call(
        action: Callable[[], Awaitable[Any]],
        instance: str,
    ) -> Any:
        try:
            return await action()
        except Exception as exc:  # endpoint boundary owns stable error mapping
            return problem(exc, instance)

    @router.post(
        "/roots/{root_trace_id}/operations",
        status_code=202,
        response_model=OperationView,
        responses=_ERROR_RESPONSES,
    )
    async def start_operation(
        root_trace_id: str,
        body: OperationRequest,
        idempotency_key: Optional[str] = Header(
            None,
            alias="Idempotency-Key",
            min_length=1,
            max_length=200,
            pattern=r"^[\x21-\x7e]+$",
        ),
    ) -> Any:
        path = f"/api/v2/roots/{root_trace_id}/operations"

        async def action() -> OperationView:
            if isinstance(body, DispatchOperationRequest):
                request = body
                operation = await coordinator.start_operation(
                    root_trace_id,
                    request.kind,
                    task_ids=request.task_ids,
                    worker_presets=request.worker_presets,
                    deadline_at=request.deadline_at,
                    idempotency_key=idempotency_key,
                )
            else:
                request = body
                operation = await coordinator.start_operation(
                    root_trace_id,
                    request.kind,
                    task_id=request.task_id,
                    attempt_id=request.attempt_id,
                    deadline_at=request.deadline_at,
                    idempotency_key=idempotency_key,
                )
            return OperationView.from_domain(operation)

        return await call(action, path)

    @router.get(
        "/roots/{root_trace_id}/operations/{operation_id}",
        response_model=OperationView,
        responses=_ERROR_RESPONSES,
    )
    async def get_operation(root_trace_id: str, operation_id: str) -> Any:
        path = f"/api/v2/roots/{root_trace_id}/operations/{operation_id}"

        async def action() -> OperationView:
            operation = await coordinator.get_operation(root_trace_id, operation_id)
            return OperationView.from_domain(operation)

        return await call(action, path)

    @router.post(
        "/roots/{root_trace_id}/operations/{operation_id}/stop",
        response_model=OperationView,
        responses=_ERROR_RESPONSES,
    )
    async def stop_operation(
        root_trace_id: str,
        operation_id: str,
        idempotency_key: Optional[str] = Header(
            None,
            alias="Idempotency-Key",
            min_length=1,
            max_length=200,
            pattern=r"^[\x21-\x7e]+$",
        ),
    ) -> Any:
        path = f"/api/v2/roots/{root_trace_id}/operations/{operation_id}/stop"

        async def action() -> OperationView:
            operation = await coordinator.stop_operation(
                root_trace_id, operation_id, idempotency_key=idempotency_key
            )
            return OperationView.from_domain(operation)

        return await call(action, path)

    @router.post(
        "/roots/{root_trace_id}/operations/{operation_id}/resume",
        response_model=OperationView,
        responses=_ERROR_RESPONSES,
    )
    async def resume_operation(
        root_trace_id: str,
        operation_id: str,
        idempotency_key: Optional[str] = Header(
            None,
            alias="Idempotency-Key",
            min_length=1,
            max_length=200,
            pattern=r"^[\x21-\x7e]+$",
        ),
    ) -> Any:
        path = f"/api/v2/roots/{root_trace_id}/operations/{operation_id}/resume"

        async def action() -> OperationView:
            operation = await coordinator.resume_operation(
                root_trace_id, operation_id, idempotency_key=idempotency_key
            )
            return OperationView.from_domain(operation)

        return await call(action, path)

    @router.get(
        "/roots/{root_trace_id}/tasks",
        response_model=TaskList,
        responses=_ERROR_RESPONSES,
    )
    async def list_tasks(root_trace_id: str) -> Any:
        path = f"/api/v2/roots/{root_trace_id}/tasks"

        async def action() -> TaskList:
            ledger = await coordinator.task_store.load(root_trace_id)
            return TaskList(
                revision=ledger.revision,
                root_task_id=ledger.root_task_id,
                items=[
                    TaskView.from_domain(item)
                    for item in sorted(
                        ledger.tasks.values(),
                        key=lambda value: (value.created_at, value.task_id),
                    )
                ],
            )

        return await call(action, path)

    @router.get(
        "/roots/{root_trace_id}/tasks/{task_id}",
        response_model=TaskView,
        responses=_ERROR_RESPONSES,
    )
    async def get_task(root_trace_id: str, task_id: str) -> Any:
        path = f"/api/v2/roots/{root_trace_id}/tasks/{task_id}"

        async def action() -> TaskView:
            ledger = await coordinator.task_store.load(root_trace_id)
            item = ledger.tasks.get(task_id)
            if item is None:
                raise _ResourceNotFound(f"Task not found: {task_id}")
            return TaskView.from_domain(item)

        return await call(action, path)

    @router.get(
        "/roots/{root_trace_id}/attempts",
        response_model=AttemptList,
        responses=_ERROR_RESPONSES,
    )
    async def list_attempts(root_trace_id: str) -> Any:
        path = f"/api/v2/roots/{root_trace_id}/attempts"

        async def action() -> AttemptList:
            ledger = await coordinator.task_store.load(root_trace_id)
            return AttemptList(
                revision=ledger.revision,
                items=[
                    AttemptView.from_domain(item)
                    for item in sorted(
                        ledger.attempts.values(),
                        key=lambda value: (value.created_at, value.attempt_id),
                    )
                ],
            )

        return await call(action, path)

    @router.get(
        "/roots/{root_trace_id}/attempts/{attempt_id}",
        response_model=AttemptView,
        responses=_ERROR_RESPONSES,
    )
    async def get_attempt(root_trace_id: str, attempt_id: str) -> Any:
        path = f"/api/v2/roots/{root_trace_id}/attempts/{attempt_id}"

        async def action() -> AttemptView:
            ledger = await coordinator.task_store.load(root_trace_id)
            item = ledger.attempts.get(attempt_id)
            if item is None:
                raise _ResourceNotFound(f"Attempt not found: {attempt_id}")
            return AttemptView.from_domain(item)

        return await call(action, path)

    @router.get(
        "/roots/{root_trace_id}/validations",
        response_model=ValidationList,
        responses=_ERROR_RESPONSES,
    )
    async def list_validations(root_trace_id: str) -> Any:
        path = f"/api/v2/roots/{root_trace_id}/validations"

        async def action() -> ValidationList:
            ledger = await coordinator.task_store.load(root_trace_id)
            return ValidationList(
                revision=ledger.revision,
                items=[
                    ValidationView.from_domain(item)
                    for item in sorted(
                        ledger.validations.values(),
                        key=lambda value: (value.created_at, value.validation_id),
                    )
                ],
            )

        return await call(action, path)

    @router.get(
        "/roots/{root_trace_id}/validations/{validation_id}",
        response_model=ValidationView,
        responses=_ERROR_RESPONSES,
    )
    async def get_validation(root_trace_id: str, validation_id: str) -> Any:
        path = f"/api/v2/roots/{root_trace_id}/validations/{validation_id}"

        async def action() -> ValidationView:
            ledger = await coordinator.task_store.load(root_trace_id)
            item = ledger.validations.get(validation_id)
            if item is None:
                raise _ResourceNotFound(f"Validation not found: {validation_id}")
            return ValidationView.from_domain(item)

        return await call(action, path)

    @router.get(
        "/roots/{root_trace_id}/events",
        response_model=EventPageView,
        responses=_ERROR_RESPONSES,
    )
    async def list_events(
        root_trace_id: str,
        cursor: Optional[str] = None,
        limit: str = "100",
        wait_seconds: str = "0",
    ) -> Any:
        path = f"/api/v2/roots/{root_trace_id}/events"

        async def action() -> EventPageView:
            try:
                parsed_limit = int(limit)
                parsed_wait = float(wait_seconds)
            except ValueError as exc:
                raise ValueError("limit and wait_seconds must be numeric") from exc
            if not 1 <= parsed_limit <= 1000:
                raise ValueError("limit must be between 1 and 1000")
            if not 0 <= parsed_wait <= 30:
                raise ValueError("wait_seconds must be between 0 and 30")
            deadline = asyncio.get_running_loop().time() + parsed_wait
            while True:
                page = await coordinator.task_store.list_events(
                    root_trace_id, cursor=cursor, limit=parsed_limit
                )
                if page.events or asyncio.get_running_loop().time() >= deadline:
                    return EventPageView(
                        events=[EventView.from_domain(item) for item in page.events],
                        next_cursor=page.next_cursor,
                        has_more=page.has_more,
                    )
                await asyncio.sleep(min(0.1, max(deadline - asyncio.get_running_loop().time(), 0)))

        return await call(action, path)

    return router


def _problem_fields(exc: Exception) -> tuple[int, str, str, str]:
    if isinstance(exc, (TaskStoreNotFound, _ResourceNotFound)):
        return 404, "Not Found", "not_found", str(exc)
    if isinstance(exc, ValueError) and "not found" in str(exc).lower():
        return 404, "Not Found", "not_found", str(exc)
    if isinstance(exc, (TaskConflict, RevisionConflict, ArtifactConflict)):
        return 409, "Conflict", "conflict", str(exc)
    if isinstance(exc, ValidationError):
        return 422, "Invalid Request", "invalid_request", str(exc)
    if isinstance(exc, ValueError):
        return 422, "Invalid Request", "invalid_request", str(exc)
    if isinstance(exc, RuntimeError) and "AgentExecutor" in str(exc):
        return 503, "Service Unavailable", "service_unavailable", str(exc)
    if isinstance(exc, TaskStoreError):
        return 503, "Service Unavailable", "storage_unavailable", str(exc)
    if isinstance(exc, OrchestrationError):
        return 409, "Conflict", "orchestration_conflict", str(exc)
    return 500, "Internal Server Error", "internal_error", "Internal server error"


__all__ = ["create_orchestration_router"]