feat: add Android task file evidence and finalize plugin design

2026-10-15 18:27:11 +08:00
1 parent 96fedc816f
commit 72ff342605
@@ -0,0 +1,348 @@
+"""Task-scoped Android files and native-resolution screenshot artifacts."""
+
+import asyncio
+import ctypes
+from datetime import datetime, UTC
+import hashlib
+from io import BytesIO
+import json
+import os
+from pathlib import Path, PurePosixPath
+import re
+import shlex
+from typing import Any, Literal
+import uuid
+
+from PIL import Image
+from pydantic import BaseModel, Field
+
+from artemis.core.tool_failure import ToolFailure
+from artemis.runtime.adb_endpoint import AdbSession, current_adb_endpoint
+from artemis.tools.base import ArtemisTool
+
+
+# Bound memory, model context and subprocess lifetime independently of task timeout.
+MAX_FILE_BYTES = 64 * 1024 * 1024
+MAX_TEXT_BYTES = 64 * 1024
+ADB_TIMEOUT_SECONDS = 60
+
+
+def android_files_available(ctx: Any) -> bool:
+    engine = getattr(ctx, "data_engine", None)
+    return bool(
+        getattr(ctx, "execution_target", "android") != "computer"
+        and getattr(getattr(ctx, "device", None), "device_id", None)
+        and engine is not None
+        and getattr(engine, "current_session_id", None)
+    )
+
+
+def desktop_directory() -> Path:
+    """Resolve the actual Windows desktop, including redirected/OneDrive desktops."""
+    if os.name == "nt":
+        buffer = ctypes.create_unicode_buffer(32768)
+        if ctypes.windll.shell32.SHGetFolderPathW(None, 0x10, None, 0, buffer) != 0:
+            raise OSError("无法定位当前用户的桌面目录")
+        path = Path(buffer.value)
+    else:
+        path = Path.home() / "Desktop"
+    if not path.is_dir():
+        raise OSError("桌面目录不存在；请使用 destination=trace")
+    return path.resolve()
+
+
+def _filename(value: str) -> str:
+    if (
+        not value
+        or len(value) > 180
+        or value in {".", ".."}
+        or value != value.strip(" .")
+        or re.search(r'[\\/:*?"<>|\x00-\x1f]', value)
+        or re.match(r"^(CON|PRN|AUX|NUL|COM[0-9]|LPT[0-9])(?:\.|$)", value, re.I)
+        or value.startswith(".artemis-")
+    ):
+        raise ValueError("filename 必须是普通文件名，不能包含路径或 Windows 保留名称")
+    return value
+
+
+def _device_path(value: str) -> str:
+    path = PurePosixPath(value)
+    if (
+        not value.startswith("/")
+        or value.startswith("//")
+        or path == PurePosixPath("/")
+        or ".." in path.parts
+        or any(ord(char) < 32 for char in value)
+    ):
+        raise ValueError("device_path 必须是 Android 绝对路径，不能包含 .. 或控制字符")
+    return str(path)
+
+
+def _host_path(root: Path, name: str) -> Path:
+    path = root / _filename(name)
+    if path.is_symlink() or path.resolve().parent != root.resolve():
+        raise ValueError("文件不能是符号链接或指向目标目录之外")
+    return path
+
+
+class AndroidFiles:
+    """Use only the task's selected serial and its process-pinned ADB endpoint."""
+
+    def __init__(self, ctx: Any):
+        if not android_files_available(ctx):
+            raise ValueError("安卓文件工具需要当前任务的 Android 设备和 trace 会话")
+        self.serial = str(ctx.device.device_id)
+        self.root = Path(ctx.data_engine.base_dir) / "artifacts" / "android"
+        self.adb = AdbSession(current_adb_endpoint())
+
+    async def _run(self, *args: str, limit: int = MAX_TEXT_BYTES) -> bytes:
+        process = await asyncio.create_subprocess_exec(
+            *self.adb.command(["-s", self.serial, *args]),
+            env=self.adb.environment(),
+            stdin=asyncio.subprocess.DEVNULL,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+            **({"creationflags": 0x08000000} if os.name == "nt" else {}),
+        )
+
+        async def bounded_read(stream, maximum):
+            output = bytearray()
+            while chunk := await stream.read(65536):
+                output.extend(chunk)
+                if len(output) > maximum:
+                    raise ValueError(f"ADB 输出超过 {maximum} 字节限制")
+            return bytes(output)
+
+        readers = [
+            asyncio.create_task(bounded_read(process.stdout, limit)),
+            asyncio.create_task(bounded_read(process.stderr, MAX_TEXT_BYTES)),
+        ]
+        try:
+            async with asyncio.timeout(ADB_TIMEOUT_SECONDS):
+                stdout, stderr = await asyncio.gather(*readers)
+                await process.wait()
+            if process.returncode:
+                detail = (stderr or stdout).decode("utf-8", errors="replace")[:2000]
+                raise OSError(f"ADB 文件操作失败（{process.returncode}）：{detail}")
+            return stdout
+        finally:
+            if process.returncode is None:
+                try:
+                    process.kill()
+                except ProcessLookupError:
+                    pass
+                await process.wait()
+            for reader in readers:
+                if not reader.done():
+                    reader.cancel()
+            await asyncio.gather(*readers, return_exceptions=True)
+
+    async def _shell(self, command: str) -> str:
+        return (await self._run("shell", command)).decode("utf-8", errors="strict").strip()
+
+    async def _hash(self, path: str) -> str:
+        result = await self._shell(f"sha256sum {shlex.quote(path)}")
+        digest = result.split(maxsplit=1)[0] if result else ""
+        if not re.fullmatch(r"[0-9a-fA-F]{64}", digest):
+            raise OSError("设备未返回有效 SHA-256，无法确认文件内容")
+        return digest.lower()
+
+    async def _read(self, path: str, limit: int) -> bytes:
+        quoted = shlex.quote(path)
+        size = int(await self._shell(f"test -f {quoted} && stat -c %s {quoted}"))
+        if not 0 <= size <= limit:
+            raise ValueError(f"文件大小 {size} 超过上限 {limit} 字节，请缩小文件或使用 pull")
+        before = await self._hash(path)
+        data = await self._run("exec-out", f"cat {quoted}", limit=limit)
+        if len(data) != size or hashlib.sha256(data).hexdigest() != before:
+            raise OSError("文件传输不完整或传输期间内容发生变化，未保存")
+        if await self._hash(path) != before:
+            raise OSError("传输期间设备文件发生变化，未保存")
+        return data
+
+    def _save(self, data: bytes, name: str, destination: str, overwrite: bool) -> dict:
+        self.root.mkdir(parents=True, exist_ok=True)
+        trace_path = _host_path(self.root, name)
+        targets = [trace_path]
+        if destination == "desktop":
+            targets.append(_host_path(desktop_directory(), name))
+        for path in targets:
+            if path.exists() and not overwrite:
+                raise FileExistsError(f"文件已存在：{path}；换名或在获准覆盖后使用 overwrite=true")
+        digest = hashlib.sha256(data).hexdigest()
+        for path in targets:
+            temporary = path.parent / f".artemis-{uuid.uuid4().hex}.tmp"
+            try:
+                with temporary.open("xb") as handle:
+                    handle.write(data)
+                    handle.flush()
+                    os.fsync(handle.fileno())
+                if hashlib.sha256(temporary.read_bytes()).hexdigest() != digest:
+                    raise OSError(f"保存后读回校验失败：{path}")
+                if overwrite:
+                    os.replace(temporary, path)
+                elif os.name == "nt":
+                    # Windows rename is exclusive and also works on exFAT/FAT.
+                    os.rename(temporary, path)
+                else:
+                    # Publish only a complete file, without clobbering a racing writer.
+                    os.link(temporary, path)
+            finally:
+                temporary.unlink(missing_ok=True)
+        return {
+            "path": str(targets[-1]),
+            "trace_path": str(trace_path),
+            "bytes": len(data),
+            "sha256": digest,
+            "verified": True,
+        }
+
+    async def _write(self, path: str, data: bytes, overwrite: bool) -> dict:
+        quoted = shlex.quote(path)
+        parent = str(PurePosixPath(path).parent)
+        remote_temp = f"{parent}/.artemis-{uuid.uuid4().hex}.tmp"
+        temp_quoted = shlex.quote(remote_temp)
+        self.root.mkdir(parents=True, exist_ok=True)
+        local_temp = self.root / f".artemis-{uuid.uuid4().hex}.tmp"
+        if not overwrite:
+            exists = await self._shell(
+                f"if test -e {quoted} || test -L {quoted}; then printf exists; fi"
+            )
+            if exists:
+                raise FileExistsError(f"设备文件已存在：{path}；换名或获准后使用 overwrite=true")
+        await self._shell(f"mkdir -p {shlex.quote(parent)}")
+        digest = hashlib.sha256(data).hexdigest()
+        local_temp.write_bytes(data)
+        try:
+            await self._run("push", str(local_temp), remote_temp)
+            if await self._hash(remote_temp) != digest:
+                raise OSError("写入设备临时文件后校验失败，未替换目标文件")
+            # Toybox mv -n returns success when it skipped an existing target;
+            # the temporary file's continued existence makes that a failure.
+            flag = "-f" if overwrite else "-n"
+            await self._shell(
+                f"test ! -d {quoted} && test ! -L {quoted} && "
+                f"mv {flag} {temp_quoted} {quoted} && test ! -e {temp_quoted}"
+            )
+            if await self._hash(path) != digest:
+                raise OSError("设备目标文件读回校验失败")
+            return {"device_path": path, "bytes": len(data), "sha256": digest, "verified": True}
+        finally:
+            local_temp.unlink(missing_ok=True)
+            # Remove only this operation's own temporary file, even on failure.
+            try:
+                await self._shell(f"rm -f {temp_quoted}")
+            except (OSError, ValueError, TimeoutError):
+                pass
+
+    async def execute(self, args: "AndroidFileArgs") -> dict:
+        if args.action == "screenshot":
+            name = _filename(args.filename or f"android-{uuid.uuid4().hex[:12]}.png")
+            if Path(name).suffix.lower() != ".png":
+                raise ValueError("原生截图保存为 PNG，filename 必须以 .png 结尾")
+            data = await self._run("exec-out", "screencap", "-p", limit=MAX_FILE_BYTES)
+            with Image.open(BytesIO(data)) as image:
+                if image.format != "PNG":
+                    raise ValueError("设备没有返回有效 PNG 截图")
+                image.load()
+                dimensions = {"width": image.width, "height": image.height}
+            return {
+                **self._save(data, name, args.destination, args.overwrite),
+                **dimensions,
+                "captured_at": datetime.now(UTC).isoformat(),
+                "device_serial": self.serial,
+            }
+        path = _device_path(args.device_path)
+        if args.action == "list":
+            return {
+                "device_path": path,
+                "entries": await self._shell(f"ls -la {shlex.quote(path)}"),
+            }
+        if args.action == "read_text":
+            data = await self._read(path, MAX_TEXT_BYTES)
+            return {"device_path": path, "content": data.decode("utf-8-sig"), "bytes": len(data)}
+        if args.action == "pull":
+            data = await self._read(path, MAX_FILE_BYTES)
+            return self._save(
+                data, args.filename or PurePosixPath(path).name, args.destination, args.overwrite
+            )
+        if args.action == "write_text":
+            if args.content is None:
+                raise ValueError("write_text 需要 content（UTF-8 文本，可为空字符串）")
+            data = args.content.encode("utf-8")
+            if len(data) > MAX_TEXT_BYTES:
+                raise ValueError("文本超过 64 KiB，请使用 push 传输现有文件")
+        else:  # push: only a named artifact or a file explicitly selected on the desktop
+            root = self.root if args.destination == "trace" else desktop_directory()
+            source = _host_path(root, args.filename)
+            if not source.is_file() or source.stat().st_size > MAX_FILE_BYTES:
+                raise ValueError("本机源文件不存在或超过 64 MiB")
+            with source.open("rb") as handle:
+                data = handle.read(MAX_FILE_BYTES + 1)
+            if len(data) > MAX_FILE_BYTES:
+                raise ValueError("本机源文件超过 64 MiB")
+        return await self._write(path, data, args.overwrite)
+
+
+class AndroidFileArgs(BaseModel):
+    action: Literal["screenshot", "list", "read_text", "write_text", "pull", "push"] = Field(
+        description="screenshot captures Android to host; list/read_text/write_text use Android paths; pull downloads, push uploads."
+    )
+    device_path: str = Field(
+        default="",
+        description="Absolute Android path, e.g. /sdcard/Download/result.txt. Not a PC path. Unused for screenshot.",
+    )
+    filename: str = Field(
+        default="",
+        description="Host file name only, no directories. Screenshot: .png (generated if empty); pull: defaults to device basename; push: required existing file.",
+    )
+    destination: Literal["trace", "desktop"] = Field(
+        default="trace",
+        description="Host location. Default: current task trace/artifacts/android. desktop: current user's actual desktop, only if requested; downloads also keep a trace copy. For push this selects the source location.",
+    )
+    content: str | None = Field(
+        default=None,
+        description="UTF-8 text for write_text only, max 64 KiB; empty string creates an empty file.",
+    )
+    overwrite: bool = Field(
+        default=False,
+        description="Only true when the user authorized replacing an existing file; default never overwrites.",
+    )
+
+
+async def _android_file(ctx: Any, **kwargs: Any) -> str:
+    try:
+        args = AndroidFileArgs.model_validate(kwargs)
+        result = await AndroidFiles(ctx).execute(args)
+        return json.dumps({"ok": True, "action": args.action, **result}, ensure_ascii=False)
+    except TimeoutError:
+        return ToolFailure(f"安卓文件操作超时（{ADB_TIMEOUT_SECONDS} 秒），ADB 子进程已停止")
+    except (OSError, ValueError) as exc:
+        return ToolFailure(f"安卓文件操作失败：{exc or type(exc).__name__}")
+
+
+class AndroidFileTool(ArtemisTool):
+    def is_available(self, ctx: Any = None) -> bool:
+        return android_files_available(ctx)
+
+    async def execute(self, driver=None, ctx=None, **kwargs):
+        return await _android_file(ctx, **kwargs)
+
+
+android_file = AndroidFileTool(
+    name="android_file",
+    description=(
+        "Save Android screenshots and read/write/transfer Android files on the task's selected device. "
+        "Use action=screenshot to save the CURRENT phone screen as native PNG, without UI shortcuts "
+        "or the computer-use plugin. Default host location is this task's trace/artifacts/android; "
+        "use destination=desktop only when the user requests the PC desktop. Returns verified paths, "
+        "size and SHA-256; never claim success without a successful result. Screenshot capture does "
+        "not prove the page is correct. Capture only after intended UI actions have executed and "
+        "the resulting screen has been observed, not alongside pending UI actions. "
+        "list/read_text/write_text operate on absolute ANDROID paths; "
+        "text is UTF-8, max 64 KiB. pull/push transfer a single file up to 64 MiB. No root or app-private "
+        "storage bypass; report Android permission errors. Existing files are protected by default."
+    ),
+    args_schema=AndroidFileArgs,
+    category="system",
+)