agentic-mobile-control/api/rest.py

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 pydantic import BaseModel

    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,
        )
    )

    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"],
        )