agentic-mobile-control/api/rest.py

import os
from pathlib import Path
from typing import Any

from api.console import create_console_router
from api.errors import semantic_error
from core.models import Task
from device.manager import DEFAULT_MANAGER, DeviceManager
from driver.registry import build_driver_factory
from runtime.executor import Executor, default_tool_registry
from runtime.task import TaskRunner, TaskRunnerConfig
from storage.device_config import DEFAULT_MAX_STEPS, DeviceConfigStore
from storage.task_metadata import TaskMetadataStore
from storage.timeline import Timeline
from tools.launch_app import launch_app
from tools.screenshot import take_screenshot
from tools.tap import tap


def create_app(
    *,
    manager: DeviceManager | None = None,
    task_runner: TaskRunner | None = None,
    metadata_store: TaskMetadataStore | None = None,
    device_config_store: DeviceConfigStore | None = None,
    timeline: Timeline | None = None,
) -> Any:
    from fastapi import BackgroundTasks, FastAPI, HTTPException
    from fastapi.middleware.cors import CORSMiddleware
    from fastapi.responses import RedirectResponse
    from fastapi.staticfiles import StaticFiles
    from pydantic import BaseModel
    from starlette.exceptions import HTTPException as StarletteHTTPException
    from starlette.types import Scope

    class SpaStaticFiles(StaticFiles):
        """``StaticFiles`` variant that falls back to ``index.html`` for SPA routes.

        Mirrors ``apps/cloud-api/cloud_api/app.py``'s implementation: an unknown
        path like ``/ui/tasks/abc`` would otherwise 404 instead of letting the
        SPA's client-side router handle it.
        """

        async def get_response(self, path: str, scope: Scope) -> Any:
            try:
                return await super().get_response(path, scope)
            except StarletteHTTPException as exc:
                if exc.status_code == 404 and path != "index.html":
                    return await super().get_response("index.html", scope)
                raise

    device_manager = manager or DEFAULT_MANAGER
    store = metadata_store or TaskMetadataStore()
    config_store = device_config_store or DeviceConfigStore()
    timeline_store = timeline or Timeline()
    max_steps = _load_max_steps(config_store)
    _reload_device_configs(device_manager, config_store)
    runner = task_runner or TaskRunner(
        metadata_store=store,
        executor=Executor(tools=default_tool_registry(manager=device_manager)),
        timeline=timeline_store,
        config=TaskRunnerConfig(max_steps=max_steps),
    )
    _apply_max_steps(runner, max_steps)
    app = FastAPI(title="Apex Agent API")
    app.add_middleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_methods=["*"],
        allow_headers=["*"],
    )

    class TapRequest(BaseModel):
        x: float
        y: float

    class ScreenshotResponse(BaseModel):
        image_base64: str
        mime_type: str = "image/png"

    class LaunchRequest(BaseModel):
        app_id: str

    class AgentTaskRequest(BaseModel):
        goal: str
        device_id: str

    @app.get("/devices")
    def devices() -> list[dict[str, Any]]:
        return [device.to_dict() for device in device_manager.list_devices()]

    @app.post("/devices/{device_id}/tap")
    def tap_device(device_id: str, request: TapRequest) -> dict[str, Any]:
        return _raise_semantic(
            lambda: tap(
                request.x,
                request.y,
                device_id=device_id,
                manager=device_manager,
            )
        )

    @app.post("/devices/{device_id}/screenshot")
    def screenshot_device(device_id: str) -> dict[str, str]:
        import base64

        image = _raise_semantic(
            lambda: take_screenshot(device_id, manager=device_manager)
        )
        response = ScreenshotResponse(
            image_base64=base64.b64encode(image).decode("ascii")
        )
        return response.model_dump()

    @app.post("/devices/{device_id}/launch")
    def launch_device(device_id: str, request: LaunchRequest) -> dict[str, Any]:
        return _raise_semantic(
            lambda: launch_app(
                request.app_id,
                device_id=device_id,
                manager=device_manager,
            )
        )

    @app.post("/agent/task")
    def start_task(
        request: AgentTaskRequest,
        background_tasks: BackgroundTasks,
    ) -> dict[str, str]:
        task = Task(goal=request.goal, device_id=request.device_id)
        store.create_task(task)
        background_tasks.add_task(runner.run, task)
        return {"task_id": task.id, "status": task.status}

    @app.get("/task/{task_id}")
    def get_task(task_id: str) -> dict[str, Any]:
        task = store.get_task(task_id)
        if task is None:
            raise HTTPException(status_code=404, detail="task not found")
        return task

    app.include_router(
        create_console_router(
            device_manager=device_manager,
            metadata_store=store,
            timeline=timeline_store,
            config_store=config_store,
            task_runner=runner,
        )
    )

    console_static_dir = os.environ.get("RUNTIME_CONSOLE_STATIC_DIR")
    if console_static_dir:
        dist_dir = Path(console_static_dir)
        if not dist_dir.is_dir():
            raise ValueError(
                f"RUNTIME_CONSOLE_STATIC_DIR is not a directory: {dist_dir}"
            )

        @app.get("/", include_in_schema=False)
        async def _redirect_to_console() -> RedirectResponse:
            return RedirectResponse(url="/ui/")

        app.mount(
            "/ui",
            SpaStaticFiles(directory=str(dist_dir), html=True),
            name="console",
        )

    return app


def _raise_semantic(func: Any) -> Any:
    from fastapi import HTTPException

    try:
        return func()
    except Exception as exc:
        raise HTTPException(status_code=400, detail=semantic_error(exc)) from exc


def _load_max_steps(config_store: DeviceConfigStore) -> int:
    raw_value = config_store.get_setting("max_steps")
    try:
        max_steps = int(raw_value) if raw_value is not None else DEFAULT_MAX_STEPS
    except ValueError:
        return DEFAULT_MAX_STEPS
    if max_steps <= 0:
        return DEFAULT_MAX_STEPS
    return max_steps


def _apply_max_steps(task_runner: Any, max_steps: int) -> None:
    config = getattr(task_runner, "config", None)
    if config is not None and hasattr(config, "max_steps"):
        config.max_steps = max_steps


def _reload_device_configs(
    device_manager: DeviceManager,
    config_store: DeviceConfigStore,
) -> None:
    for config in config_store.list():
        device_manager.register_device(
            config["device_id"],
            build_driver_factory(config["driver_type"], config["connection_info"]),
            name=config["name"],
            driver_type=config["driver_type"],
            connection_info=config["connection_info"],
        )