agentic-mobile-control/apps/device-host-agent/host_agent/app.py

from __future__ import annotations

import asyncio
from contextlib import suppress
from dataclasses import dataclass

from cloud.internal_api.models import AssignmentModel
from device.manager import DeviceManager
from driver.registry import build_driver_factory
from host_agent.assignment import AssignmentExecutor
from host_agent.client import HostAgentClient, HostAgentEnrollmentClient
from host_agent.config import HostAgentConfig, load_host_agent_config
from host_agent.enrollment import resolve_host_identity
from host_agent.execution import create_execution_factories
from host_agent.heartbeat import HeartbeatSynchronizer
from host_agent.identity import HostIdentityStore
from host_agent.lease import ActiveAssignmentRunner
from host_agent.processor import AssignmentProcessingResult, AssignmentProcessor
from storage.device_config import DeviceConfigStore


@dataclass
class HostAgentApplication:
    client: HostAgentClient
    heartbeat: HeartbeatSynchronizer
    processor: AssignmentProcessor

    def run(self) -> None:
        asyncio.run(self.run_async())

    async def run_async(self, stop: asyncio.Event | None = None) -> None:
        stop_requested = stop or asyncio.Event()
        heartbeat_stop = asyncio.Event()
        heartbeat_task = asyncio.create_task(self.heartbeat.run(heartbeat_stop))
        active_processing: asyncio.Task[AssignmentProcessingResult] | None = None
        try:
            while not stop_requested.is_set():
                assignment = await self._claim_until_stopped(stop_requested)
                if assignment is None:
                    continue
                active_processing = asyncio.create_task(
                    self.processor.process(assignment)
                )
                stopped = asyncio.create_task(stop_requested.wait())
                done, _ = await asyncio.wait(
                    {active_processing, stopped},
                    return_when=asyncio.FIRST_COMPLETED,
                )
                if stopped in done:
                    self.processor.request_stop()
                else:
                    stopped.cancel()
                with suppress(asyncio.CancelledError):
                    await stopped
                await asyncio.shield(active_processing)
                active_processing = None
        finally:
            self.processor.request_stop()
            if active_processing is not None:
                with suppress(Exception):
                    await asyncio.shield(active_processing)
            heartbeat_stop.set()
            try:
                await asyncio.gather(heartbeat_task, return_exceptions=True)
                with suppress(Exception):
                    await self.heartbeat.sync_once()
            finally:
                await self.client.aclose()

    async def _claim_until_stopped(
        self,
        stop: asyncio.Event,
    ) -> AssignmentModel | None:
        claim = asyncio.create_task(self.client.claim())
        stopped = asyncio.create_task(stop.wait())
        try:
            done, _ = await asyncio.wait(
                {claim, stopped},
                return_when=asyncio.FIRST_COMPLETED,
            )
        except asyncio.CancelledError:
            claim.cancel()
            stopped.cancel()
            await asyncio.gather(claim, stopped, return_exceptions=True)
            raise
        if stopped in done:
            claim.cancel()
            with suppress(asyncio.CancelledError):
                await claim
            return None
        stopped.cancel()
        with suppress(asyncio.CancelledError):
            await stopped
        return await claim


def create_application(
    *,
    config: HostAgentConfig | None = None,
    manager: DeviceManager | None = None,
    device_config_store: DeviceConfigStore | None = None,
    identity_store: HostIdentityStore | None = None,
    enrollment_client: HostAgentEnrollmentClient | None = None,
) -> HostAgentApplication:
    startup_config = config or load_host_agent_config()
    config_store = device_config_store or DeviceConfigStore()
    owned_enrollment_client = enrollment_client is None
    bootstrap_client = enrollment_client or HostAgentEnrollmentClient(startup_config)
    try:
        resolved_config = resolve_host_identity(
            startup_config,
            identity_store=identity_store
            or HostIdentityStore(startup_config.identity_path),
            client=bootstrap_client,
        )
        bootstrap_client.config = resolved_config
        resolved_manager = manager or _configured_device_manager(
            config_store,
            config=resolved_config,
            enrollment_client=bootstrap_client,
        )
    finally:
        if owned_enrollment_client:
            bootstrap_client.close()
    client = HostAgentClient(resolved_config)
    heartbeat = HeartbeatSynchronizer(resolved_manager, client, resolved_config)
    executor = AssignmentExecutor(create_execution_factories(resolved_manager))
    active_runner = ActiveAssignmentRunner(client, executor)
    return HostAgentApplication(
        client=client,
        heartbeat=heartbeat,
        processor=AssignmentProcessor(client, active_runner),
    )


def _configured_device_manager(
    config_store: DeviceConfigStore,
    *,
    config: HostAgentConfig,
    enrollment_client: HostAgentEnrollmentClient,
) -> DeviceManager:
    manager = DeviceManager()
    for device in config_store.list():
        runtime_device_id = device["device_id"]
        if config.enrollment_managed:
            enrollment = enrollment_client.enroll_device(
                local_device_id=device["device_id"],
                driver_type=device["driver_type"],
                name=device["name"],
                capability_tags=[],
            )
            runtime_device_id = enrollment.device_id
            config_store.set_cloud_device_id(device["device_id"], runtime_device_id)
        manager.register_device(
            runtime_device_id,
            build_driver_factory(
                device["driver_type"],
                device["connection_info"],
            ),
            name=device["name"],
            driver_type=device["driver_type"],
            connection_info=device["connection_info"],
        )
    return manager