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

from __future__ import annotations

import asyncio
import base64
import json
from pathlib import Path
from typing import Any

import jinja2
from fastapi import Depends, FastAPI, HTTPException, Request
from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response

from device.manager import DeviceManager
from host_agent.assignment import AssignmentExecutor
from host_agent.client import HostAgentEnrollmentClient
from host_agent.config import HostAgentConfig
from host_agent.devices import register_local_device, unregister_local_device
from host_agent.history import ConsoleHistoryStore
from host_agent.identity import HostIdentityStore
from host_agent.local_account import LocalAccountStore
from host_agent.status import AgentStatusTracker
from host_agent.web.auth import (
    SessionManager,
    SessionState,
    attempt_login,
    change_password,
)
from storage.device_config import DeviceConfigStore
from storage.task_metadata import TaskMetadataStore
from storage.timeline import Timeline

SESSION_COOKIE_NAME = "host_console_session"
CSRF_HEADER_NAME = "X-CSRF-Token"
CSRF_FORM_FIELD = "csrf_token"
_LOOPBACK_BIND_HOSTS = frozenset({"127.0.0.1", "localhost", "::1"})

_ENV = jinja2.Environment(
    loader=jinja2.FileSystemLoader(Path(__file__).parent / "templates"),
    autoescape=jinja2.select_autoescape(["html", "xml"]),
)


def _render(
    template_name: str,
    *,
    status_code: int = 200,
    **context: Any,
) -> HTMLResponse:
    html = _ENV.get_template(template_name).render(**context)
    return HTMLResponse(html, status_code=status_code)


def _device_display_status(device: Any, *, busy_device_id: str | None) -> str:
    """Connected-but-idle devices report "busy" at the DeviceManager layer
    (an active Appium/WDA session is required to reuse it), which is not the
    same as a task currently running on that device. Only the device actually
    bound to the current assignment should read as "busy" here.
    """
    if device.status == "busy" and device.id != busy_device_id:
        return "connected"
    return device.status


def _screenshot_data_uri(record: dict[str, Any]) -> str | None:
    """Return a ``data:`` URI for the step's screenshot, or ``None``."""
    screenshot_path = record.get("screenshot_path")
    if not screenshot_path:
        return None
    path = Path(str(screenshot_path))
    if not path.exists():
        return None
    encoded = base64.b64encode(path.read_bytes()).decode("ascii")
    return f"data:image/png;base64,{encoded}"


def _dashboard_texts(*, snapshot: dict[str, Any]) -> dict[str, str]:
    """Pre-compute human-readable text strings for the dashboard template."""
    heartbeat = snapshot.get("last_heartbeat")
    assignment = snapshot.get("current_assignment")
    policy = snapshot.get("host_policy")
    progress = snapshot.get("progress")
    return {
        "heartbeat_text": (
            f"{'ok' if heartbeat['ok'] else 'failed'} at {heartbeat['at']} "
            f"({heartbeat['device_count']} devices)"
            if heartbeat
            else "never"
        ),
        "assignment_text": (
            f"{assignment['task_id']} on {assignment['device_id']} "
            f"(started {assignment['started_at']})"
            if assignment
            else "none"
        ),
        "progress_text": (
            f"step {progress['step_index']} \u2014 {progress['step_status']}: "
            f"{progress['summary']}"
            if progress
            else ""
        ),
        "policy_text": (
            f"revision {policy['revision']}; self-submission "
            f"{'enabled' if policy['self_submission_enabled'] else 'disabled'}; "
            f"max active tasks {policy['max_active_tasks'] or 'unlimited'}; "
            f"daily token budget {policy['daily_token_budget'] or 'unmetered'}"
            if policy
            else "no Cloud policy cached"
        ),
    }


def create_console_app(
    *,
    config: HostAgentConfig,
    manager: DeviceManager,
    config_store: DeviceConfigStore,
    local_account_store: LocalAccountStore,
    identity_store: HostIdentityStore,
    history_store: ConsoleHistoryStore,
    status_tracker: AgentStatusTracker,
    session_manager: SessionManager,
    enrollment_client: HostAgentEnrollmentClient | None,
    metadata_store: TaskMetadataStore | None = None,
    timeline: Timeline | None = None,
    executor: AssignmentExecutor | None = None,
) -> FastAPI:
    app = FastAPI(title="Host Agent Console")
    cookie_secure = config.console_bind_host not in _LOOPBACK_BIND_HOSTS

    def _session_token(request: Request) -> str | None:
        return request.cookies.get(SESSION_COOKIE_NAME)

    def require_session(request: Request) -> SessionState:
        token = _session_token(request)
        state = session_manager.validate(token) if token else None
        if state is None:
            raise HTTPException(status_code=303, headers={"Location": "/login"})
        return state

    async def require_csrf(
        request: Request,
        session: SessionState = Depends(require_session),
    ) -> SessionState:
        provided = request.headers.get(CSRF_HEADER_NAME)
        if not provided:
            form = await request.form()
            raw = form.get(CSRF_FORM_FIELD)
            provided = str(raw) if raw is not None else None
        token = _session_token(request)
        if (
            not token
            or not provided
            or not session_manager.validate_csrf(token, provided)
        ):
            raise HTTPException(status_code=403, detail="invalid CSRF token")
        return session

    @app.get("/login", response_class=HTMLResponse)
    async def login_page() -> HTMLResponse:
        account = await asyncio.to_thread(local_account_store.load)
        return _render("login.html", title="Login", session=None, account=account)

    @app.post("/login")
    async def login_submit(request: Request) -> Response:
        account = await asyncio.to_thread(local_account_store.load)
        if account is None:
            return _render("login.html", title="Login", session=None, account=None)
        form = await request.form()
        username = str(form.get("username", ""))
        password = str(form.get("password", ""))
        ok = await asyncio.to_thread(
            attempt_login, local_account_store, username=username, password=password
        )
        if not ok:
            return _render(
                "login.html",
                title="Login",
                session=None,
                account=account,
                error="Invalid username or password.",
            )
        session_token, _ = session_manager.create_session(username)
        response = RedirectResponse(url="/", status_code=303)
        response.set_cookie(
            key=SESSION_COOKIE_NAME,
            value=session_token,
            httponly=True,
            samesite="strict",
            secure=cookie_secure,
            path="/",
        )
        return response

    @app.post("/logout")
    async def logout(
        request: Request,
        session: SessionState = Depends(require_csrf),
    ) -> Response:
        token = _session_token(request)
        if token:
            session_manager.invalidate(token)
        response = RedirectResponse(url="/login", status_code=303)
        response.delete_cookie(key=SESSION_COOKIE_NAME, path="/")
        return response

    @app.get("/", response_class=HTMLResponse)
    async def dashboard(
        session: SessionState = Depends(require_session),
    ) -> HTMLResponse:
        identity = await asyncio.to_thread(identity_store.load)
        snapshot = status_tracker.snapshot()
        busy_device_id = (
            snapshot["current_assignment"]["device_id"]
            if snapshot.get("current_assignment")
            else None
        )
        devices = [
            {
                "id": d.id,
                "name": d.name,
                "driver_type": d.driver_type,
                "display_status": _device_display_status(
                    d, busy_device_id=busy_device_id
                ),
            }
            for d in manager.list_devices()
        ]
        texts = _dashboard_texts(snapshot=snapshot)
        return _render(
            "dashboard.html",
            title="Status",
            session=session,
            identity=identity,
            devices=devices,
            config=config,
            **texts,
        )

    @app.get("/api/status")
    async def api_status(
        session: SessionState = Depends(require_session),
    ) -> JSONResponse:
        snapshot = status_tracker.snapshot()
        if executor is not None and snapshot.get("progress") is None:
            live = executor.latest_progress()
            if live is not None:
                snapshot["progress"] = {
                    "step_index": live.step_index,
                    "step_status": live.step_status,
                    "summary": live.summary,
                    "updated_at": live.updated_at.isoformat(),
                }
        current_assignment = snapshot.get("current_assignment")
        busy_device_id = current_assignment["device_id"] if current_assignment else None
        devices = [
            {
                "id": device.id,
                "name": device.name,
                "driver_type": device.driver_type,
                "status": _device_display_status(device, busy_device_id=busy_device_id),
            }
            for device in manager.list_devices()
        ]
        return JSONResponse({"status": snapshot, "devices": devices})

    @app.get("/devices", response_class=HTMLResponse)
    async def devices_page(
        request: Request,
        session: SessionState = Depends(require_session),
    ) -> HTMLResponse:
        devices = await asyncio.to_thread(config_store.list)
        edit_id = request.query_params.get("edit")
        edit_record = (
            await asyncio.to_thread(config_store.get, edit_id) if edit_id else None
        )
        connection_info_json = (
            json.dumps(edit_record["connection_info"]) if edit_record else "{}"
        )
        return _render(
            "devices.html",
            title="Devices",
            session=session,
            devices=devices,
            csrf_token=session.csrf_token,
            edit_record=edit_record,
            connection_info_json=connection_info_json,
            error=None,
        )

    @app.post("/devices/save")
    async def devices_save(
        request: Request,
        session: SessionState = Depends(require_csrf),
    ) -> Response:
        form = await request.form()
        device_id = str(form.get("device_id", "")).strip()
        driver_type = str(form.get("driver_type", "")).strip()
        name = str(form.get("name", "")).strip() or None
        connection_info_raw = str(form.get("connection_info", "") or "{}")

        error: str | None = None
        connection_info: dict[str, Any] = {}
        if not device_id or not driver_type:
            error = "Device ID and driver type are required."
        else:
            try:
                parsed = json.loads(connection_info_raw)
            except ValueError:
                error = "Connection info must be valid JSON."
            else:
                if not isinstance(parsed, dict):
                    error = "Connection info must be a JSON object."
                else:
                    connection_info = parsed

        if error is None:
            try:
                await asyncio.to_thread(
                    register_local_device,
                    config_store,
                    manager,
                    device_id=device_id,
                    driver_type=driver_type,
                    connection_info=connection_info,
                    name=name,
                    config=config,
                    enrollment_client=enrollment_client,
                )
            except ValueError as exc:
                error = str(exc)

        if error is not None:
            devices = await asyncio.to_thread(config_store.list)
            return _render(
                "devices.html",
                title="Devices",
                session=session,
                devices=devices,
                csrf_token=session.csrf_token,
                edit_record=None,
                connection_info_json="{}",
                error=error,
                status_code=400,
            )
        return RedirectResponse(url="/devices", status_code=303)

    @app.post("/devices/remove")
    async def devices_remove(
        request: Request,
        session: SessionState = Depends(require_csrf),
    ) -> Response:
        form = await request.form()
        device_id = str(form.get("device_id", "")).strip()
        if device_id:
            await asyncio.to_thread(
                unregister_local_device, config_store, manager, device_id=device_id
            )
        return RedirectResponse(url="/devices", status_code=303)

    @app.get("/account", response_class=HTMLResponse)
    async def account_page(
        session: SessionState = Depends(require_session),
    ) -> HTMLResponse:
        return _render(
            "account.html",
            title="Account",
            session=session,
            csrf_token=session.csrf_token,
            message=None,
            error=None,
        )

    @app.post("/account", response_class=HTMLResponse)
    async def account_submit(
        request: Request,
        session: SessionState = Depends(require_csrf),
    ) -> HTMLResponse:
        form = await request.form()
        current_password = str(form.get("current_password", ""))
        new_password = str(form.get("new_password", ""))
        confirm_password = str(form.get("confirm_password", ""))
        if not new_password or new_password != confirm_password:
            return _render(
                "account.html",
                title="Account",
                session=session,
                csrf_token=session.csrf_token,
                message=None,
                error="New password and confirmation must match.",
                status_code=400,
            )
        ok = await asyncio.to_thread(
            change_password,
            local_account_store,
            current_password=current_password,
            new_password=new_password,
        )
        if not ok:
            return _render(
                "account.html",
                title="Account",
                session=session,
                csrf_token=session.csrf_token,
                message=None,
                error="Current password is incorrect.",
                status_code=400,
            )
        return _render(
            "account.html",
            title="Account",
            session=session,
            csrf_token=session.csrf_token,
            message="Password updated.",
            error=None,
        )

    @app.get("/history", response_class=HTMLResponse)
    async def history_page(
        session: SessionState = Depends(require_session),
    ) -> HTMLResponse:
        entries = await asyncio.to_thread(history_store.list_recent)
        return _render(
            "history.html",
            title="History",
            session=session,
            entries=entries,
        )

    @app.get("/tasks", response_class=HTMLResponse)
    async def tasks_page(
        session: SessionState = Depends(require_session),
    ) -> HTMLResponse:
        if metadata_store is None:
            raise HTTPException(
                status_code=503, detail="task metadata store not configured"
            )
        tasks_list = await asyncio.to_thread(metadata_store.list_tasks)
        return _render(
            "tasks_list.html",
            title="Tasks",
            session=session,
            tasks=tasks_list,
        )

    @app.get("/tasks/{task_id}", response_class=HTMLResponse)
    async def task_detail_page(
        task_id: str,
        session: SessionState = Depends(require_session),
    ) -> HTMLResponse:
        if metadata_store is None:
            raise HTTPException(
                status_code=503, detail="task metadata store not configured"
            )
        task = await asyncio.to_thread(metadata_store.get_task, task_id)
        if task is None:
            raise HTTPException(status_code=404, detail="task not found")
        timeline_records: list[dict[str, Any]] = []
        if timeline is not None:
            timeline_records = await asyncio.to_thread(timeline.read, task_id)
        task_rows = [
            (key, task[key])
            for key in (
                "id",
                "goal",
                "device_id",
                "status",
                "created_at",
                "updated_at",
            )
            if task.get(key) is not None
        ]
        timeline_steps = [
            {
                "index": record.get("index", ""),
                "timestamp": record.get("timestamp", ""),
                "prompt": record.get("prompt") or "",
                "tool_call_text": (
                    json.dumps(record.get("tool_call"), ensure_ascii=False)
                    if record.get("tool_call")
                    else ""
                ),
                "result_text": (
                    json.dumps(record.get("result"), ensure_ascii=False)
                    if record.get("result")
                    else ""
                ),
                "screenshot_src": _screenshot_data_uri(record),
            }
            for record in timeline_records
        ]
        return _render(
            "task_detail.html",
            title=f"Task {task_id}",
            session=session,
            task=task,
            task_rows=task_rows,
            timeline_steps=timeline_steps,
        )

    return app