agentic-mobile-control/packages/cloud-platform/cloud/repository.py

from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime
from typing import TYPE_CHECKING, Any, Literal, Protocol

if TYPE_CHECKING:
    from cloud.plugins import PluginManifest
    from cloud.pool import HostRegistration, PooledDevice
    from cloud.scheduler import ScheduledTask, ScheduledTaskStatus


AttemptStatus = Literal["assigned", "dispatched", "done", "failed", "expired"]
TerminalTaskStatus = Literal["done", "failed"]
ResultRecordStatus = Literal["recorded", "already_recorded", "conflict"]
LeaseRenewalStatus = Literal["renewed", "not_found", "conflict", "expired"]


class HostEnrollmentConflictError(RuntimeError):
    """Raised when a Host enrollment cannot preserve its existing binding."""


class EnrollmentTokenConflictError(HostEnrollmentConflictError):
    """Raised when a one-time enrollment token is already bound elsewhere."""


class DeviceEnrollmentConflictError(RuntimeError):
    """Raised when a local device enrollment conflicts with stored identity."""


@dataclass(frozen=True)
class HostEnrollment:
    host_id: str
    agent_instance_id: str
    display_name: str | None
    enrolled_at: datetime
    revoked_at: datetime | None = None


@dataclass(frozen=True)
class DeviceEnrollment:
    device_id: str
    host_id: str
    local_device_id: str
    driver_type: str
    name: str | None
    capability_tags: list[str]
    enrolled_at: datetime
    revoked_at: datetime | None = None


@dataclass(frozen=True)
class TaskAttemptRecord:
    task_id: str
    attempt: int
    lease_id: str
    host_id: str
    device_id: str
    status: AttemptStatus
    lease_expires_at: datetime
    created_at: datetime
    completed_at: datetime | None = None
    failure_reason: str | None = None
    terminal_result: dict[str, Any] | None = None


@dataclass(frozen=True)
class LeasedAssignment:
    task_id: str
    attempt: int
    lease_id: str
    lease_expires_at: datetime
    host_id: str
    device_id: str
    goal: str | None
    workflow_definition_id: str | None


class CloudRepository(Protocol):
    """Persistence port for cloud state and atomic scheduling operations."""

    def enroll_host(
        self,
        *,
        host_id: str,
        agent_instance_id: str,
        credential_digest: str,
        enrollment_token_digest: str,
        display_name: str | None,
        enrolled_at: datetime,
    ) -> HostEnrollment: ...

    def authenticate_enrolled_host(
        self,
        credential_digest: str,
    ) -> str | None: ...

    def revoke_enrolled_host(
        self,
        host_id: str,
        *,
        revoked_at: datetime,
    ) -> bool: ...

    def is_enrollment_managed_host(self, host_id: str) -> bool: ...

    def enroll_device(
        self,
        *,
        device_id: str,
        host_id: str,
        local_device_id: str,
        driver_type: str,
        name: str | None,
        capability_tags: list[str],
        enrolled_at: datetime,
    ) -> DeviceEnrollment: ...

    def get_device_enrollment(
        self,
        device_id: str,
    ) -> DeviceEnrollment | None: ...

    def list_device_enrollments(
        self,
        host_id: str,
    ) -> list[DeviceEnrollment]: ...

    def upsert_host(
        self,
        host_id: str,
        *,
        address: str | None,
        last_seen_at: datetime,
    ) -> None: ...

    def replace_host_devices(
        self,
        host_id: str,
        devices: list[PooledDevice],
        *,
        allow_device_takeover: bool = False,
    ) -> None: ...

    def list_hosts(self) -> list[HostRegistration]: ...

    def get_host(self, host_id: str) -> HostRegistration | None: ...

    def list_devices(self) -> list[PooledDevice]: ...

    def get_device(self, device_id: str) -> PooledDevice | None: ...

    def enqueue_task(self, task: ScheduledTask) -> None: ...

    def list_queued_tasks(self) -> list[ScheduledTask]: ...

    def list_tasks(
        self,
        *,
        status: ScheduledTaskStatus | None = None,
        limit: int = 50,
        offset: int = 0,
    ) -> list[ScheduledTask]: ...

    def count_tasks(self, status: ScheduledTaskStatus | None = None) -> int: ...

    def get_task(self, task_id: str) -> ScheduledTask | None: ...

    def update_task(
        self,
        task_id: str,
        *,
        status: str | None = None,
        assigned_device_id: str | None = None,
        assigned_host_id: str | None = None,
    ) -> None: ...

    def count_queued_tasks(self) -> int: ...

    def save_plugin(self, manifest: PluginManifest, *, wired: bool) -> None: ...

    def list_plugins(self) -> list[tuple[PluginManifest, bool]]: ...

    def get_plugin(self, name: str) -> tuple[PluginManifest, bool] | None: ...

    def list_reserved_device_ids(self, *, now: datetime) -> set[str]: ...

    def assign_task(
        self,
        *,
        task_id: str,
        host_id: str,
        device_id: str,
        lease_id: str,
        lease_expires_at: datetime,
        now: datetime,
    ) -> LeasedAssignment | None: ...

    def claim_assignment(
        self,
        *,
        host_id: str,
        now: datetime,
    ) -> LeasedAssignment | None: ...

    def renew_lease(
        self,
        *,
        task_id: str,
        attempt: int,
        lease_id: str,
        host_id: str,
        lease_expires_at: datetime,
        now: datetime,
    ) -> LeaseRenewalStatus: ...

    def record_task_result(
        self,
        *,
        task_id: str,
        attempt: int,
        lease_id: str,
        host_id: str,
        status: TerminalTaskStatus,
        failure_reason: str | None,
        terminal_result: dict[str, Any] | None,
        completed_at: datetime,
    ) -> ResultRecordStatus: ...

    def reap_expired_leases(
        self,
        *,
        now: datetime,
        max_attempts: int,
    ) -> list[str]: ...

    def list_task_attempts(self, task_id: str) -> list[TaskAttemptRecord]: ...

    def health_check(self) -> None: ...

    def close(self) -> None: ...