uiautomator2-mcp-server 0.1.2__py3-none-any.whl → 0.2.0__py3-none-any.whl

u2mcp/tools/device.py

@@ -1,293 +1,307 @@
-from __future__ import annotations
-
-import asyncio
-import sys
-from base64 import b64encode
-from collections.abc import AsyncGenerator
-from contextlib import asynccontextmanager
-from io import BytesIO
-from typing import Any, Literal
-
-import uiautomator2 as u2
-from adbutils import adb
-from fastmcp.server.dependencies import get_context
-from fastmcp.utilities.logging import get_logger
-from PIL.Image import Image
-
-from ..mcp import mcp
-
-__all__ = (
-    "device_list",
-    "init",
-    "connect",
-    "disconnect",
-    "disconnect_all",
-    "window_size",
-    "screenshot",
-    "dump_hierarchy",
-    "info",
-)
-
-
-StdoutType = Literal["stdout", "stderr"]
-
-_devices: dict[str, tuple[asyncio.Semaphore, u2.Device]] = {}
-_device_connect_lock = asyncio.Lock()
-
-
-@asynccontextmanager
-async def get_device(serial: str) -> AsyncGenerator[u2.Device]:
-    async with _device_connect_lock:
-        try:
-            semaphore, device = _devices[serial]
-        except KeyError:
-
-            def _connect():
-                _d = u2.connect(serial)
-                _d.info
-                return _d
-
-            device = await asyncio.to_thread(_connect)
-            semaphore = asyncio.Semaphore()
-            _devices[serial] = semaphore, device
-    async with semaphore:
-        yield device
-
-
-@mcp.tool("device_list")
-async def device_list() -> list[dict[str, Any]]:
-    device_list = await asyncio.to_thread(adb.device_list)
-    return [d.info for d in device_list]
-
-
-@mcp.tool("init")
-async def init(serial: str = ""):
-    """Install essential resources to device.
-
-    Important:
-        This tool must be run on the Android device before running operation actions.
-
-    Args:
-        serial (str): Android device serialno to initialize. If empty string, all devices will be initialized.
-
-    Returns:
-        None upon successful completion (exit code 0).
-        Raises an exception if the subprocess returns a non-zero exit code.
-    """
-    logger = get_logger(f"{__name__}.init")
-    args = ["-m", "uiautomator2", "init"]
-    if serial := serial.strip():
-        args.extend(["--serial", serial])
-
-    logger.info("Running uiautomator2 init command: %s %s", sys.executable, args)
-    process = await asyncio.create_subprocess_exec(
-        sys.executable, *args, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
-    )
-    if process.stdout is None:
-        raise RuntimeError("stdout is None")
-    if process.stderr is None:
-        raise RuntimeError("stderr is None")
-
-    output_queue: asyncio.Queue[tuple[StdoutType, str]] = asyncio.Queue()
-
-    async def stream_subprocess(stream: asyncio.streams.StreamReader, tag: StdoutType):
-        while True:
-            line_bytes = await stream.readline()
-            await output_queue.put((tag, line_bytes.decode()))
-            if not line_bytes:  # Reached EOF, the empty string is in the queue
-                break
-
-    # Start the stream reading tasks
-    tasks = [
-        asyncio.create_task(coro)
-        for coro in (
-            stream_subprocess(process.stdout, "stdout"),
-            stream_subprocess(process.stderr, "stderr"),
-        )
-    ]
-
-    completed_streams = 0
-
-    logger.info("read uiautomator2 init command stdio")
-
-    ctx = get_context()
-
-    while True:
-        tag, line = await output_queue.get()
-
-        if not line:  # This was the EOF sentinel (empty string)
-            completed_streams += 1
-            if completed_streams == len(tasks):
-                output_queue.task_done()
-                break  # Both streams are done, exit the main consumer loop
-
-        # Process the actual line data
-        if line := line.strip():
-            logger.info("%s: %s", tag, line)
-            if tag == "stdout":
-                await ctx.info(line)
-            else:
-                await ctx.warning(line)
-
-        output_queue.task_done()
-
-    # Wait for the tasks to formally complete and the process to exit
-    logger.info("waiting for uiautomator2 init command to complete")
-    await asyncio.gather(*tasks)
-    exit_code = await process.wait()
-    logger.info("uiautomator2 init command exited with code: %s", exit_code)
-    if exit_code != 0:
-        raise RuntimeError(f"uiautomator2 init command exited with non-zero code: {exit_code}")
-
-
-@mcp.tool("connect")
-async def connect(serial: str = ""):
-    """Connect to an Android device
-
-    Args:
-        serial (str): Android device serial number. If empty string, connects to the unique device if only one device is connected.
-
-
-    Returns:
-        dict[str,Any]: Device information
-    """
-    global _devices
-    device: u2.Device | None = None
-
-    logger = get_logger(f"{__name__}.connect")
-
-    if serial := serial.strip():
-        try:
-            async with get_device(serial) as device_1:
-                # Found, then check if it's still connected
-                try:
-                    return await asyncio.to_thread(lambda: device_1.device_info | device_1.info)
-                except u2.ConnectError as e:
-                    # Found, but not connected, delete it
-                    logger.warning("Device %s is no longer connected, delete it!", serial)
-                    del _devices[serial]
-                    raise e from None
-        except KeyError:
-            # Not found, need a new connection!
-            logger.info("Cannot find device with serial %s, connecting...")
-
-    # make new connection here!
-    async with _device_connect_lock:
-        device = await asyncio.to_thread(u2.connect, serial)
-        if device is None:
-            raise RuntimeError("Cannot connect to device")
-        logger.info("Connected to device %s", device.serial)
-        result = await asyncio.to_thread(lambda: device.device_info | device.info)
-        _devices[device.serial] = asyncio.Semaphore(), device
-        return result
-
-
-@mcp.tool("disconnect")
-async def disconnect(serial: str):
-    """Disconnect from an Android device
-
-    Args:
-        serial (str): Android device serialno
-
-    Returns:
-        None
-    """
-    if not (serial := serial.strip()):
-        raise ValueError("serial cannot be empty")
-    async with _device_connect_lock:
-        del _devices[serial]
-
-
-@mcp.tool("disconnect_all")
-async def disconnect_all():
-    """Disconnect from all Android devices"""
-    async with _device_connect_lock:
-        _devices.clear()
-
-
-@mcp.tool("window_size")
-async def window_size(serial: str) -> dict[str, int]:
-    """Get window size of an Android device
-
-    Args:
-        serial (str): Android device serialno
-
-    Returns:
-        dict[str,int]: Window size object:
-            - "width" (int): Window width
-            - "height" (int): Window height
-    """
-    async with get_device(serial) as device:
-        width, height = await asyncio.to_thread(device.window_size)
-        return {"width": width, "height": height}
-
-
-@mcp.tool("screenshot")
-async def screenshot(serial: str, display_id: int = -1) -> dict[str, Any]:
-    """
-    Take screenshot of device
-
-    Args:
-        serial (str): Android device serialno
-        display_id (int): use specific display if device has multiple screen. Defaults to -1.
-
-    Returns:
-        dict[str,Any]: Screenshot image JPEG data with the following keys:
-            - "image" (str): Base64 encoded image data in data URL format (data:image/jpeg;base64,...)
-            - "size" (tuple[int,int]): Image dimensions as (width, height)
-    """
-    display_id = int(display_id)
-    async with get_device(serial) as device:
-        im = await asyncio.to_thread(
-            device.screenshot,
-            display_id=display_id if display_id >= 0 else None,
-        )  # type: ignore[arg-type]
-
-    if not isinstance(im, Image):
-        raise RuntimeError("Invalid image")
-
-    with BytesIO() as fp:
-        im.save(fp, "jpeg")
-        im_data = fp.getvalue()
-
-    return {
-        "width": im.width,
-        "height": im.height,
-        "image": "data:image/jpeg;base64," + b64encode(im_data).decode(),
-    }
-
-
-@mcp.tool("dump_hierarchy")
-async def dump_hierarchy(serial: str, compressed: bool = False, pretty: bool = False, max_depth: int = -1) -> str:
-    """
-    Dump window hierarchy
-
-    Args:
-        serial (str): Android device serialno
-        compressed (bool): return compressed xml
-        pretty (bool): pretty print xml
-        max_depth (int): max depth of hierarchy
-
-    Returns:
-        str: xml string of the hierarchy tree
-    """
-    async with get_device(serial) as device:
-        return await asyncio.to_thread(
-            device.dump_hierarchy, compressed=compressed, pretty=pretty, max_depth=max_depth if max_depth > 0 else None
-        )
-
-
-@mcp.tool("info")
-async def info(serial: str) -> dict[str, Any]:
-    """
-    Get device info
-
-    Args:
-        serial (str): Android device serialno
-
-    Returns:
-        dict[str,Any]: Device info
-    """
-
-    async with get_device(serial) as device:
-        return await asyncio.to_thread(lambda: device.info)
+from __future__ import annotations
+
+import argparse
+from base64 import b64encode
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager, closing
+from io import BytesIO
+from pathlib import Path
+from typing import Any
+
+import uiautomator2 as u2
+from adbutils import adb
+from anyio import Lock, to_thread
+from fastmcp.utilities.logging import get_logger
+from PIL.Image import Image
+
+from ..mcp import mcp
+
+__all__ = (
+    "device_list",
+    "shell_command",
+    "init",
+    "connect",
+    "disconnect",
+    "disconnect_all",
+    "window_size",
+    "screenshot",
+    "save_screenshot",
+    "dump_hierarchy",
+    "info",
+)
+
+
+_devices: dict[str, tuple[Lock, u2.Device]] = {}
+_global_device_connection_lock = Lock()
+
+
+@asynccontextmanager
+async def get_device(serial: str) -> AsyncGenerator[u2.Device]:
+    async with _global_device_connection_lock:
+        try:
+            lock, device = _devices[serial]
+        except KeyError:
+
+            def _connect():
+                _d = u2.connect(serial)
+                _d.info
+                return _d
+
+            device = await to_thread.run_sync(_connect)
+            lock = Lock()
+            _devices[serial] = lock, device
+
+    async with lock:
+        yield device
+
+
+@mcp.tool("init", tags={"device:manage"})
+async def init(serial: str = ""):
+    """Install essential resources (minicap, minitouch, uiautomator ...) to device.
+
+    Important:
+        This tool must be run on the Android device before running operation actions.
+
+    Args:
+        serial (str): Android device serialno to initialize. If empty string, all devices will be initialized.
+
+    Returns:
+        None upon successful completion
+        Raises an exception if the subprocess returns a non-zero exit code.
+    """
+    from uiautomator2.__main__ import cmd_init
+
+    args = argparse.Namespace(serial=serial, serial_optional=None)
+    return await to_thread.run_sync(cmd_init, args)
+
+
+@mcp.tool("purge", tags={"device:manage"})
+async def purge(serial: str = ""):
+    """Purge all resources (minicap, minitouch, uiautomator ...) from device.
+
+    Important:
+        This tool must be run on the Android device before running operation actions.
+
+    Args:
+        serial (str): Android device serialno to purge. If empty string, all devices will be purged.
+
+    Returns:
+        None upon successful completion
+        Raises an exception if the subprocess returns a non-zero exit code.
+    """
+    from uiautomator2.__main__ import cmd_purge
+
+    args = argparse.Namespace(serial=serial)
+    return await to_thread.run_sync(cmd_purge, args)
+
+
+@mcp.tool("shell_command", tags={"device:shell"})
+async def shell_command(serial: str, command: str, timeout: float = 60) -> tuple[int, str]:
+    """Run a shell command on an Android device
+
+    Args:
+        serial (str): Android device serialno
+        command (str): Shell command to run
+        timeout (float): Seconds to wait for command to complete.
+
+    Returns:
+        tuple[int,str]: Return code and output of the command
+    """
+    async with get_device(serial) as device:
+        return_value = await to_thread.run_sync(device.adb_device.shell2, command, timeout)
+        return return_value.returncode, return_value.output
+
+
+@mcp.tool("device_list", tags={"device:info"})
+async def device_list() -> list[dict[str, Any]]:
+    """List of Adb Device with state:device
+
+    Returns:
+        list[dict[str,Any]]: List Adb Device information
+    """
+    device_list = await to_thread.run_sync(adb.device_list)
+    return [d.info for d in device_list]
+
+
+@mcp.tool("connect", tags={"device:manage"})
+async def connect(serial: str = "") -> dict[str, Any]:
+    """Connect to an Android device
+
+    Args:
+        serial (str): Android device serial number. If empty string, connects to the unique device if only one device is connected.
+
+    Returns:
+        dict[str,Any]: Device information
+    """
+    global _devices
+    device: u2.Device | None = None
+
+    logger = get_logger(f"{__name__}.connect")
+
+    if serial := serial.strip():
+        try:
+            async with get_device(serial) as device_1:
+                # Found, then check if it's still connected
+                try:
+                    return await to_thread.run_sync(lambda: device_1.device_info | device_1.info)
+                except u2.ConnectError as e:
+                    # Found, but not connected, delete it
+                    logger.warning("Device %s is no longer connected, delete it!", serial)
+                    del _devices[serial]
+                    raise e from None
+        except KeyError:
+            # Not found, need a new connection!
+            logger.info("Cannot find device with serial %s, connecting...")
+
+    # make new connection here!
+    async with _global_device_connection_lock:
+        device = await to_thread.run_sync(u2.connect, serial)
+        if device is None:
+            raise RuntimeError("Cannot connect to device")
+        logger.info("Connected to device %s", device.serial)
+        result = await to_thread.run_sync(lambda: device.device_info | device.info)
+        _devices[device.serial] = Lock(), device
+        return result
+
+
+@mcp.tool("disconnect", tags={"device:manage"})
+async def disconnect(serial: str):
+    """Disconnect from an Android device
+
+    Args:
+        serial (str): Android device serialno
+
+    Returns:
+        None
+    """
+    if not (serial := serial.strip()):
+        raise ValueError("serial cannot be empty")
+    async with _global_device_connection_lock:
+        _devices.pop(serial, None)
+
+
+@mcp.tool("disconnect_all", tags={"device:manage"})
+async def disconnect_all():
+    """Disconnect from all Android devices"""
+    async with _global_device_connection_lock:
+        _devices.clear()
+
+
+@mcp.tool("window_size", tags={"device:info"})
+async def window_size(serial: str) -> dict[str, int]:
+    """Get window size of an Android device
+
+    Args:
+        serial (str): Android device serialno
+
+    Returns:
+        dict[str,int]: Window size object:
+            - "width" (int): Window width
+            - "height" (int): Window height
+    """
+    async with get_device(serial) as device:
+        width, height = await to_thread.run_sync(device.window_size)
+        return {"width": width, "height": height}
+
+
+@mcp.tool("screenshot", tags={"device:capture", "screen:capture"})
+async def screenshot(serial: str, format: str = "jpeg", display_id: int = -1) -> dict[str, Any]:
+    """
+    Take screenshot of device
+
+    Args:
+        serial (str): Android device serialno.
+        format (str): Image format. Defaults to "jpeg".
+        display_id (int): use specific display if device has multiple screen. Defaults to -1.
+
+    Returns:
+        dict[str,Any]: Screenshot image JPEG data with the following keys:
+            - image (str): Base64 encoded image data in data URL format (data:image/jpeg;base64,...)
+            - height (int): Image height
+            - width (int): Image width
+    """
+    display_id = int(display_id)
+    async with get_device(serial) as device:
+        im = await to_thread.run_sync(lambda: device.screenshot(display_id=display_id if display_id >= 0 else None))
+
+    if not isinstance(im, Image):
+        raise RuntimeError("Invalid image")
+
+    with closing(im):
+        with BytesIO() as fp:
+            im.save(fp, format)
+            im_data = fp.getvalue()
+
+        return {
+            "image": "data:image/jpeg;base64," + b64encode(im_data).decode(),
+            "height": im.height,
+            "width": im.width,
+        }
+
+
+@mcp.tool("save_screenshot", tags={"device:capture", "screen:capture"})
+async def save_screenshot(serial: str, file: str, display_id: int = -1) -> str:
+    """
+    Save screenshot of device to file
+
+    Args:
+        serial(str): Android device serial number.
+        file(str): File path to save the screenshot. Supports both absolute and relative paths.
+        display_id(int): Use specific display if device has multiple screens. Defaults to -1 (default display).
+
+    Returns:
+        str: Screenshot save file path
+    """
+    display_id = int(display_id)
+
+    async with get_device(serial) as device:
+        im = await to_thread.run_sync(lambda: device.screenshot(display_id=display_id if display_id >= 0 else None))
+
+    if not isinstance(im, Image):
+        raise RuntimeError("Invalid image")
+
+    with closing(im):
+        # Convert path to Path object and resolve
+        file_path = Path(file)
+        # Create parent directory if it doesn't exist
+        file_path.parent.mkdir(parents=True, exist_ok=True)
+        # Save the image
+        im.save(file_path)
+
+    return file_path.resolve().as_posix()
+
+
+@mcp.tool("dump_hierarchy", tags={"device:capture"})
+async def dump_hierarchy(serial: str, compressed: bool = False, pretty: bool = False, max_depth: int = -1) -> str:
+    """
+    Dump window hierarchy
+
+    Args:
+        serial (str): Android device serialno
+        compressed (bool): return compressed xml
+        pretty (bool): pretty print xml
+        max_depth (int): max depth of hierarchy
+
+    Returns:
+        str: xml string of the hierarchy tree
+    """
+    async with get_device(serial) as device:
+        return await to_thread.run_sync(
+            lambda: device.dump_hierarchy(compressed=compressed, pretty=pretty, max_depth=max_depth if max_depth > 0 else None)
+        )
+
+
+@mcp.tool("info", tags={"device:info"})
+async def info(serial: str) -> dict[str, Any]:
+    """
+    Get device info
+
+    Args:
+        serial (str): Android device serialno
+
+    Returns:
+        dict[str,Any]: Device info
+    """
+
+    async with get_device(serial) as device:
+        return await to_thread.run_sync(lambda: device.info)