agentic-mobile-control/packages/cloud-platform/cloud/internal_api/api.py

from __future__ import annotations

import asyncio
import base64
import json
import logging
from collections.abc import Awaitable, Callable
from datetime import timedelta
from time import monotonic
from typing import TYPE_CHECKING
from uuid import uuid4

from fastapi import APIRouter, HTTPException, Request, Response, status
from fastapi.responses import JSONResponse

from cloud.auth import (
    AuthProvider,
    HostAuthorizationError,
    digest_token,
)
from cloud.internal_api.models import (
    AssignmentModel,
    ClaimRequest,
    ClaimResponse,
    DeviceEnrollmentRequest,
    DeviceEnrollmentResponse,
    HeartbeatRequest,
    HeartbeatResponse,
    HostGovernancePolicyModel,
    HostEnrollmentRequest,
    HostEnrollmentResponse,
    HostTaskCancellationResponse,
    HostTaskSubmissionRequest,
    HostTaskSubmissionResponse,
    LeaseRenewalRequest,
    LeaseRenewalResponse,
    PlannerDecisionError,
    PlannerDecisionRequest,
    PlannerDecisionResponse,
    StaleLeaseConflict,
    TerminalResultRequest,
    TerminalResultResponse,
)
from cloud.llm_providers import LlmProviderResolutionError, LlmProviderService
from cloud.planner_config import build_cloud_planner_client
from cloud.provider_secrets import ProviderSecretConfigurationError
from cloud.repository import (
    AssignmentProgressSnapshot,
    DeviceEnrollmentConflictError,
    HostEnrollmentConflictError,
)
from cloud.governance import TokenBudgetExceededError
from core.models import Device, utc_now
from runtime.tool_calling_client import ToolCallingClient, ToolCallUnavailable
from runtime.tool_specs import ToolSpec

if TYPE_CHECKING:
    from cloud.pool import DevicePool
    from cloud.scheduler import TaskScheduler

logger = logging.getLogger(__name__)


def create_internal_router(
    *,
    pool: DevicePool,
    auth_provider: AuthProvider,
    version_prefix: str = "/internal/v1",
    claim_poll_interval_seconds: float = 0.1,
    lease_duration_seconds: float = 60.0,
    sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
    planner_client_factory: Callable[[], ToolCallingClient] | None = None,
    planner_provider_service: LlmProviderService | None = None,
    scheduler: TaskScheduler | None = None,
    planner_token_reservation_ceiling: int = 4096,
    planner_token_reservation_ttl_seconds: float = 300.0,
) -> APIRouter:
    if claim_poll_interval_seconds <= 0:
        raise ValueError("claim_poll_interval_seconds must be greater than zero")
    if lease_duration_seconds <= 0:
        raise ValueError("lease_duration_seconds must be greater than zero")
    if planner_token_reservation_ceiling <= 0:
        raise ValueError("planner_token_reservation_ceiling must be greater than zero")
    if planner_token_reservation_ttl_seconds <= 0:
        raise ValueError(
            "planner_token_reservation_ttl_seconds must be greater than zero"
        )
    router = APIRouter(prefix=version_prefix, tags=["host-agent"])

    def authorize_host(request: Request, host_id: str) -> None:
        principal = auth_provider.authenticate(request)
        if principal is None:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail="unauthorized",
                headers={"WWW-Authenticate": "Bearer"},
            )
        try:
            principal.require_host(host_id)
        except HostAuthorizationError as exc:
            raise HTTPException(
                status_code=status.HTTP_403_FORBIDDEN,
                detail=str(exc),
            ) from exc

    @router.post(
        "/enrollments",
        response_model=HostEnrollmentResponse,
        status_code=status.HTTP_201_CREATED,
    )
    def enroll_host(
        payload: HostEnrollmentRequest,
    ) -> HostEnrollmentResponse:
        try:
            enrollment = pool.store.enroll_host(
                host_id=f"host-{uuid4().hex}",
                agent_instance_id=payload.agent_instance_id,
                credential_digest=digest_token(payload.host_token),
                enrollment_token_digest=None,
                display_name=payload.display_name,
                enrolled_at=utc_now(),
            )
        except HostEnrollmentConflictError as exc:
            raise HTTPException(
                status_code=status.HTTP_409_CONFLICT,
                detail=str(exc),
            ) from exc
        return HostEnrollmentResponse(host_id=enrollment.host_id)

    @router.post(
        "/hosts/{host_id}/devices/enroll",
        response_model=DeviceEnrollmentResponse,
        status_code=status.HTTP_201_CREATED,
    )
    def enroll_device(
        host_id: str,
        payload: DeviceEnrollmentRequest,
        request: Request,
    ) -> DeviceEnrollmentResponse:
        authorize_host(request, host_id)
        try:
            enrollment = pool.store.enroll_device(
                device_id=f"device-{uuid4().hex}",
                host_id=host_id,
                local_device_id=payload.local_device_id,
                driver_type=payload.driver_type,
                name=payload.name,
                capability_tags=list(payload.capability_tags),
                enrolled_at=utc_now(),
            )
        except DeviceEnrollmentConflictError as exc:
            raise HTTPException(
                status_code=status.HTTP_409_CONFLICT,
                detail=str(exc),
            ) from exc
        return DeviceEnrollmentResponse(device_id=enrollment.device_id)

    @router.put(
        "/hosts/{host_id}/heartbeat",
        response_model=HeartbeatResponse,
    )
    def heartbeat(
        host_id: str,
        payload: HeartbeatRequest,
        request: Request,
    ) -> HeartbeatResponse:
        authorize_host(request, host_id)
        allow_device_takeover = _validate_snapshot(
            pool,
            host_id=host_id,
            payload=payload,
        )
        devices = [
            Device(
                id=device.device_id,
                driver_type=device.driver_type,
                status=device.status,
                capability_tags=list(device.capability_tags),
            )
            for device in payload.devices
        ]
        pool.sync_host_devices(
            host_id,
            devices,
            address=payload.address,
            allow_device_takeover=allow_device_takeover,
            planner_transport=payload.planner_transport,
        )
        policy = pool.store.get_host_governance_policy(host_id)
        policy_revision = policy.revision if policy is not None else 0
        return HeartbeatResponse(
            host_id=host_id,
            accepted_devices=len(devices),
            received_at=utc_now(),
            policy_revision=policy_revision,
            policy=(
                HostGovernancePolicyModel(
                    revision=policy.revision,
                    self_submission_enabled=policy.self_submission_enabled,
                    max_active_tasks=policy.max_active_tasks,
                    daily_token_budget=policy.daily_token_budget,
                )
                if policy is not None and payload.policy_revision != policy.revision
                else None
            ),
        )

    @router.post(
        "/hosts/{host_id}/tasks",
        response_model=HostTaskSubmissionResponse,
        status_code=status.HTTP_201_CREATED,
    )
    def submit_host_task(
        host_id: str,
        payload: HostTaskSubmissionRequest,
        request: Request,
    ) -> HostTaskSubmissionResponse:
        authorize_host(request, host_id)
        if payload.host_id != host_id:
            raise HTTPException(
                status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
                detail="task host_id must match the request path",
            )
        if scheduler is None:
            raise HTTPException(
                status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
                detail="task submission is unavailable",
            )
        policy = pool.store.get_host_governance_policy(host_id)
        if policy is not None and not policy.self_submission_enabled:
            raise HTTPException(
                status_code=status.HTTP_403_FORBIDDEN,
                detail="Host self-submission is disabled",
            )
        if payload.device_id is not None and not any(
            device.host_id == host_id and device.device_id == payload.device_id
            for device in pool.list_devices()
        ):
            raise HTTPException(
                status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
                detail="target device is not owned by this Host",
            )
        from cloud.scheduler import TaskConstraints

        task_id = scheduler.submit(
            goal=payload.goal,
            constraints=TaskConstraints(
                target_host_id=host_id,
                target_device_id=payload.device_id,
            ),
        )
        return HostTaskSubmissionResponse(task_id=task_id)

    @router.post(
        "/hosts/{host_id}/tasks/{task_id}/cancel",
        response_model=HostTaskCancellationResponse,
        responses={
            status.HTTP_202_ACCEPTED: {"model": HostTaskCancellationResponse},
        },
    )
    def cancel_host_task(
        host_id: str,
        task_id: str,
        request: Request,
        response: Response,
    ) -> HostTaskCancellationResponse:
        authorize_host(request, host_id)
        if scheduler is None:
            raise HTTPException(
                status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
                detail="task cancellation is unavailable",
            )
        task = scheduler.store.get_task(task_id)
        if task is None or task.constraints.target_host_id != host_id:
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail=f"task {task_id!r} not found",
            )
        result = scheduler.store.request_task_cancellation(
            task_id, requested_at=utc_now()
        )
        if result == "not_found":
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail=f"task {task_id!r} not found",
            )
        if result == "already_terminal":
            raise HTTPException(
                status_code=status.HTTP_409_CONFLICT,
                detail=f"task {task_id!r} has already reached a terminal state",
            )
        task = scheduler.store.get_task(task_id)
        assert task is not None
        if result == "requested" and task.status != "cancelled":
            response.status_code = status.HTTP_202_ACCEPTED
        return HostTaskCancellationResponse(task_id=task_id, status=task.status)

    @router.post(
        "/hosts/{host_id}/assignments/claim",
        response_model=ClaimResponse,
    )
    async def claim_assignment(
        host_id: str,
        payload: ClaimRequest,
        request: Request,
    ) -> ClaimResponse:
        authorize_host(request, host_id)
        if payload.host_id != host_id:
            raise HTTPException(
                status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
                detail="claim host_id must match the request path",
            )

        deadline = monotonic() + payload.timeout_seconds
        while True:
            assignment = pool.store.claim_assignment(host_id=host_id, now=utc_now())
            if assignment is not None:
                return ClaimResponse(
                    assignment=AssignmentModel(
                        task_id=assignment.task_id,
                        attempt=assignment.attempt,
                        lease_id=assignment.lease_id,
                        lease_expires_at=assignment.lease_expires_at,
                        host_id=assignment.host_id,
                        device_id=assignment.device_id,
                        goal=assignment.goal,
                        workflow_definition_id=assignment.workflow_definition_id,
                    )
                )

            remaining = deadline - monotonic()
            if remaining <= 0:
                return ClaimResponse(timed_out=True)
            await sleep(min(claim_poll_interval_seconds, remaining))

    @router.post(
        "/hosts/{host_id}/assignments/{task_id}/renew",
        response_model=LeaseRenewalResponse,
        responses={status.HTTP_409_CONFLICT: {"model": StaleLeaseConflict}},
    )
    def renew_assignment(
        host_id: str,
        task_id: str,
        payload: LeaseRenewalRequest,
        request: Request,
    ):
        authorize_host(request, host_id)
        _validate_assignment_identity(
            host_id=host_id,
            task_id=task_id,
            payload_host_id=payload.host_id,
            payload_task_id=payload.task_id,
        )
        now = utc_now()
        lease_expires_at = now + timedelta(seconds=lease_duration_seconds)
        progress_snapshot: AssignmentProgressSnapshot | None = None
        if payload.progress is not None:
            progress_snapshot = AssignmentProgressSnapshot(
                step_index=payload.progress.step_index,
                step_status=payload.progress.step_status,
                summary=payload.progress.summary[:500],
                updated_at=now,
            )
        renewal = pool.store.renew_lease(
            task_id=task_id,
            attempt=payload.attempt,
            lease_id=payload.lease_id,
            host_id=host_id,
            lease_expires_at=lease_expires_at,
            now=now,
            progress=progress_snapshot,
        )
        if renewal.status == "not_found":
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail="assignment not found",
            )
        if renewal.status != "renewed":
            return _stale_lease_conflict("assignment lease is stale or expired")
        return LeaseRenewalResponse(
            status="renewed",
            lease_expires_at=lease_expires_at,
            cancel_requested=renewal.cancel_requested,
        )

    @router.post(
        "/hosts/{host_id}/assignments/{task_id}/result",
        response_model=TerminalResultResponse,
        responses={status.HTTP_409_CONFLICT: {"model": StaleLeaseConflict}},
    )
    def report_result(
        host_id: str,
        task_id: str,
        payload: TerminalResultRequest,
        request: Request,
    ):
        authorize_host(request, host_id)
        _validate_assignment_identity(
            host_id=host_id,
            task_id=task_id,
            payload_host_id=payload.host_id,
            payload_task_id=payload.task_id,
        )
        result_status = pool.store.record_task_result(
            task_id=task_id,
            attempt=payload.attempt,
            lease_id=payload.lease_id,
            host_id=host_id,
            status=payload.status,
            failure_reason=payload.failure_reason,
            terminal_result=payload.result,
            completed_at=utc_now(),
        )
        if result_status == "conflict":
            return _stale_lease_conflict("assignment lease is stale or superseded")
        return TerminalResultResponse(status=result_status)

    @router.post(
        "/hosts/{host_id}/planner/decide",
        response_model=PlannerDecisionResponse,
        response_model_exclude_none=True,
        responses={
            status.HTTP_502_BAD_GATEWAY: {"model": PlannerDecisionError},
            status.HTTP_429_TOO_MANY_REQUESTS: {"model": PlannerDecisionError},
        },
    )
    def decide_planner_call(
        host_id: str,
        payload: PlannerDecisionRequest,
        request: Request,
    ):
        authorize_host(request, host_id)
        if payload.host_id != host_id:
            raise HTTPException(
                status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
                detail="planner-decision host_id must match the request path",
            )
        _validate_planner_context(pool, host_id=host_id, payload=payload)

        screenshot: bytes | None = None
        if payload.screenshot_base64 is not None:
            try:
                screenshot = base64.b64decode(payload.screenshot_base64)
            except (ValueError, TypeError):
                raise HTTPException(
                    status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
                    detail="screenshot_base64 is not valid base64",
                ) from None

        tools = [
            ToolSpec(
                name=tool.name,
                description=tool.description,
                parameters=tool.parameters,
            )
            for tool in payload.tools
        ]

        try:
            if planner_client_factory is not None:
                client = planner_client_factory()
                resolved_provider = "test"
                resolved_model = "test"
                planner_timeout = payload.timeout_seconds
            elif planner_provider_service is not None:
                resolved = planner_provider_service.resolve_active_profile()
                client = build_cloud_planner_client(resolved)
                resolved_provider = resolved.profile.provider_type
                resolved_model = resolved.profile.model
                planner_timeout = resolved.profile.timeout_seconds
            else:
                raise LlmProviderResolutionError(
                    "no database Provider resolver configured"
                )
        except (LlmProviderResolutionError, ProviderSecretConfigurationError) as exc:
            logger.info(
                "planner-decision request failed",
                extra={"host_id": host_id, "error_class": type(exc).__name__},
            )
            return JSONResponse(
                status_code=status.HTTP_502_BAD_GATEWAY,
                content=PlannerDecisionError(detail=str(exc)).model_dump(),
            )

        now = utc_now()
        reservation = None
        try:
            reservation = pool.store.reserve_host_token_budget(
                reservation_id=uuid4().hex,
                host_id=host_id,
                usage_day=now.date().isoformat(),
                reserved_tokens=planner_token_reservation_ceiling,
                task_id=payload.task_id,
                attempt=payload.attempt,
                created_at=now,
                expires_at=now
                + timedelta(seconds=planner_token_reservation_ttl_seconds),
            )
        except TokenBudgetExceededError as exc:
            return JSONResponse(
                status_code=status.HTTP_429_TOO_MANY_REQUESTS,
                content=PlannerDecisionError(detail=str(exc)).model_dump(),
            )

        started_at = monotonic()
        try:
            decision = client.decide(
                system_prompt=payload.system_prompt,
                user_prompt=payload.user_prompt,
                screenshot=screenshot,
                tools=tools,
                timeout=planner_timeout,
            )
        except ToolCallUnavailable as exc:
            logger.info(
                "planner-decision request failed",
                extra={
                    "host_id": host_id,
                    "latency_seconds": monotonic() - started_at,
                    "error_class": type(exc).__name__,
                },
            )
            return JSONResponse(
                status_code=status.HTTP_502_BAD_GATEWAY,
                content=PlannerDecisionError(detail=str(exc)).model_dump(),
            )
        usage = decision.usage
        if (
            reservation is not None
            and usage is not None
            and usage.total_tokens is not None
        ):
            pool.store.settle_host_token_reservation(
                reservation_id=reservation.id,
                event_id=uuid4().hex,
                provider=resolved_provider,
                model=resolved_model,
                input_tokens=usage.input_tokens,
                output_tokens=usage.output_tokens,
                total_tokens=usage.total_tokens,
                occurred_at=utc_now(),
            )
        if payload.task_id and payload.attempt:
            pool.store.record_planner_decision(
                host_id=host_id,
                task_id=payload.task_id,
                attempt=payload.attempt,
                system_prompt=payload.system_prompt,
                user_prompt=payload.user_prompt,
                tool_name=decision.tool_name,
                arguments_json=json.dumps(decision.arguments),
                now=utc_now(),
                rationale=getattr(decision, "text_output", None),
                thinking=getattr(decision, "thinking", None),
            )
        logger.info(
            "planner-decision request resolved",
            extra={
                "host_id": host_id,
                "tool_name": decision.tool_name,
                "latency_seconds": monotonic() - started_at,
            },
        )
        return PlannerDecisionResponse(
            tool_name=decision.tool_name,
            arguments=dict(decision.arguments),
            input_tokens=(decision.usage.input_tokens if decision.usage else None),
            output_tokens=(decision.usage.output_tokens if decision.usage else None),
            total_tokens=(decision.usage.total_tokens if decision.usage else None),
        )

    return router


def _validate_assignment_identity(
    *,
    host_id: str,
    task_id: str,
    payload_host_id: str,
    payload_task_id: str,
) -> None:
    if payload_host_id != host_id or payload_task_id != task_id:
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
            detail="assignment identity must match the request path",
        )


def _validate_planner_context(
    pool, *, host_id: str, payload: PlannerDecisionRequest
) -> None:
    context_values = (payload.task_id, payload.attempt, payload.lease_id)
    if not any(value is not None for value in context_values):
        return
    if any(value is None for value in context_values):
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
            detail="planner context requires task_id, attempt, and lease_id together",
        )
    attempts = pool.store.list_task_attempts(payload.task_id or "")
    if not any(
        attempt.attempt == payload.attempt
        and attempt.host_id == host_id
        and attempt.lease_id == payload.lease_id
        for attempt in attempts
    ):
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
            detail="planner context does not match a Host assignment",
        )


def _stale_lease_conflict(detail: str) -> JSONResponse:
    return JSONResponse(
        status_code=status.HTTP_409_CONFLICT,
        content=StaleLeaseConflict(detail=detail).model_dump(),
    )


def _validate_snapshot(
    pool: DevicePool,
    *,
    host_id: str,
    payload: HeartbeatRequest,
) -> bool:
    if payload.host_id != host_id:
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
            detail="heartbeat host_id must match the request path",
        )
    device_ids = [device.device_id for device in payload.devices]
    if len(device_ids) != len(set(device_ids)):
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
            detail="heartbeat snapshot contains duplicate device ids",
        )

    if pool.store.is_enrollment_managed_host(host_id):
        enrollments = {
            enrollment.device_id: enrollment
            for enrollment in pool.store.list_device_enrollments(host_id)
            if enrollment.revoked_at is None
        }
        invalid = [
            device.device_id
            for device in payload.devices
            if device.device_id not in enrollments
            or enrollments[device.device_id].driver_type != device.driver_type
        ]
        if invalid:
            raise HTTPException(
                status_code=status.HTTP_409_CONFLICT,
                detail=f"device enrollment conflict: {sorted(set(invalid))}",
            )
        return False

    now = utc_now()
    hosts = {host.host_id: host for host in pool.store.list_hosts()}
    conflicts: list[str] = []
    stale_owner_found = False
    requested_ids = set(device_ids)
    for device in pool.store.list_devices():
        if device.device_id not in requested_ids or device.host_id == host_id:
            continue
        owner = hosts.get(device.host_id)
        if owner is None:
            continue
        age_seconds = (now - owner.last_seen_at).total_seconds()
        if age_seconds <= pool.config.stale_after_seconds:
            conflicts.append(device.device_id)
        else:
            stale_owner_found = True
    if conflicts:
        raise HTTPException(
            status_code=status.HTTP_409_CONFLICT,
            detail=f"device ownership conflict: {sorted(set(conflicts))}",
        )
    return stale_owner_found