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 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,
) -> dict[str, Callable[..., Any]]:

    def _screenshot(device_id: str | None = None) -> dict[str, Any]:
        image = take_screenshot(device_id, manager=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=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=manager,
            )
        ),
        "input_text": lambda text, device_id=None: call_with_semantic_errors(
            input_text,
            text,
            device_id=device_id,
            manager=manager,
        ),
        "launch_app": lambda app_id, device_id=None: call_with_semantic_errors(
            launch_app,
            app_id,
            device_id=device_id,
            manager=manager,
        ),
        "find_text": lambda query, device_id=None: call_with_semantic_errors(
            find_text_on_screen,
            query,
            device_id=device_id,
            manager=manager,
        ),
        "find_icon": lambda name, device_id=None: call_with_semantic_errors(
            find_icon_on_screen,
            name,
            device_id=device_id,
            manager=manager,
        ),
        "get_ui_tree": lambda device_id=None, include_app_info=False: (
            call_with_semantic_errors(
                get_ui_tree,
                device_id,
                manager=manager,
                include_app_info=include_app_info,
            )
        ),
        "describe_screen": lambda device_id=None: call_with_semantic_errors(
            lambda: describe_screen(device_id, manager=manager).to_dict()
        ),
        "list_devices": lambda: [device.to_dict() for device in manager.list_devices()],
        "device_status": lambda device_id: call_with_semantic_errors(
            lambda: {"device_id": device_id, "status": manager.status(device_id)}
        ),
    }


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

    if manager is None:
        raise ValueError("create_mcp_server requires a non-None manager")

    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,
        include_app_info: bool = False,
    ) -> Any:
        return handlers["get_ui_tree"](
            device_id=device_id,
            include_app_info=include_app_info,
        )

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

    if skill_catalog_store is not None:
        from api.skill_catalog_mcp import register_skill_catalog_tools
        from storage.local_skills import LocalSkillStore

        local_store = skill_local_store or LocalSkillStore()
        register_skill_catalog_tools(
            server,
            store=skill_catalog_store,
            local_store=local_store,
            get_active_subscriptions=lambda: set(skill_active_subscriptions or set()),
            get_registered_tools=lambda: set(handlers.keys()),
        )

    return server