agentic-mobile-control/api/mcp.py

import base64
from collections.abc import Callable
from typing import Any

from api.errors import call_with_semantic_errors
from device.manager import DEFAULT_MANAGER, DeviceManager
from tools.describe_screen import describe_screen
from tools.find_icon import find_icon_on_screen
from tools.find_text import find_text_on_screen
from tools.input_text import input_text
from tools.launch_app import launch_app
from tools.screenshot import take_screenshot
from tools.swipe import swipe
from tools.tap import tap
from tools.ui_tree import get_ui_tree


def tool_handlers(
    *,
    manager: DeviceManager | None = None,
) -> dict[str, Callable[..., Any]]:
    device_manager = manager or DEFAULT_MANAGER

    def _screenshot(device_id: str | None = None) -> dict[str, Any]:
        image = take_screenshot(device_id, manager=device_manager)
        return {
            "ok": True,
            "image_base64": base64.b64encode(image).decode("ascii"),
            "mime_type": "image/png",
        }

    return {
        "take_screenshot": lambda device_id=None: call_with_semantic_errors(
            _screenshot,
            device_id,
        ),
        "tap": lambda x, y, device_id=None: call_with_semantic_errors(
            tap,
            x,
            y,
            device_id=device_id,
            manager=device_manager,
        ),
        "swipe": lambda start_x, start_y, end_x, end_y, duration_ms=500, device_id=None: call_with_semantic_errors(
            swipe,
            start_x,
            start_y,
            end_x,
            end_y,
            duration_ms=duration_ms,
            device_id=device_id,
            manager=device_manager,
        ),
        "input_text": lambda text, device_id=None: call_with_semantic_errors(
            input_text,
            text,
            device_id=device_id,
            manager=device_manager,
        ),
        "launch_app": lambda app_id, device_id=None: call_with_semantic_errors(
            launch_app,
            app_id,
            device_id=device_id,
            manager=device_manager,
        ),
        "find_text": lambda query, device_id=None: call_with_semantic_errors(
            find_text_on_screen,
            query,
            device_id=device_id,
            manager=device_manager,
        ),
        "find_icon": lambda name, device_id=None: call_with_semantic_errors(
            find_icon_on_screen,
            name,
            device_id=device_id,
            manager=device_manager,
        ),
        "get_ui_tree": lambda device_id=None: call_with_semantic_errors(
            get_ui_tree,
            device_id,
            manager=device_manager,
        ),
        "describe_screen": lambda device_id=None: call_with_semantic_errors(
            lambda: describe_screen(device_id, manager=device_manager).to_dict()
        ),
        "list_devices": lambda: [
            device.to_dict() for device in device_manager.list_devices()
        ],
        "device_status": lambda device_id: call_with_semantic_errors(
            lambda: {"device_id": device_id, "status": device_manager.status(device_id)}
        ),
    }


def create_mcp_server(*, manager: DeviceManager | None = None) -> Any:
    try:
        from mcp.server.fastmcp import FastMCP
    except ImportError as exc:
        raise RuntimeError("mcp SDK is not installed") from exc

    handlers = tool_handlers(manager=manager)
    server = FastMCP("apex-agent")

    @server.tool(name="take_screenshot")
    def _take_screenshot(device_id: str | None = None) -> dict[str, Any]:
        return handlers["take_screenshot"](device_id=device_id)

    @server.tool(name="tap")
    def _tap(x: float, y: float, device_id: str | None = None) -> dict[str, Any]:
        return handlers["tap"](x=x, y=y, device_id=device_id)

    @server.tool(name="swipe")
    def _swipe(
        start_x: float,
        start_y: float,
        end_x: float,
        end_y: float,
        duration_ms: int = 500,
        device_id: str | None = None,
    ) -> dict[str, Any]:
        return handlers["swipe"](
            start_x=start_x,
            start_y=start_y,
            end_x=end_x,
            end_y=end_y,
            duration_ms=duration_ms,
            device_id=device_id,
        )

    @server.tool(name="input_text")
    def _input_text(text: str, device_id: str | None = None) -> dict[str, Any]:
        return handlers["input_text"](text=text, device_id=device_id)

    @server.tool(name="launch_app")
    def _launch_app(app_id: str, device_id: str | None = None) -> dict[str, Any]:
        return handlers["launch_app"](app_id=app_id, device_id=device_id)

    @server.tool(name="find_text")
    def _find_text(query: str, device_id: str | None = None) -> dict[str, Any]:
        return handlers["find_text"](query=query, device_id=device_id)

    @server.tool(name="find_icon")
    def _find_icon(name: str, device_id: str | None = None) -> dict[str, Any]:
        return handlers["find_icon"](name=name, device_id=device_id)

    @server.tool(name="get_ui_tree")
    def _get_ui_tree(device_id: str | None = None) -> Any:
        return handlers["get_ui_tree"](device_id=device_id)

    @server.tool(name="describe_screen")
    def _describe_screen(device_id: str | None = None) -> dict[str, Any]:
        return handlers["describe_screen"](device_id=device_id)

    @server.tool(name="list_devices")
    def _list_devices() -> list[dict[str, Any]]:
        return handlers["list_devices"]()

    @server.tool(name="device_status")
    def _device_status(device_id: str) -> dict[str, Any]:
        return handlers["device_status"](device_id=device_id)

    return server