Artemis-Eapil/artemis/runtime/diagnostics.py

"""Small, task-owned runtime snapshots; no polling or readiness decisions."""

import asyncio
import ctypes
from ctypes import wintypes
from functools import lru_cache
import hashlib
import json
import logging
import os
from pathlib import Path
import platform
import subprocess
import sys
import time

import psutil

from artemis._version import get_version
from artemis.config.paths import ROOT_DIR
from artemis.runtime.adb_endpoint import current_adb_endpoint
from artemis.toolchain import toolchain

logger = logging.getLogger(__name__)
FILENAME = "runtime_diagnostics.json"


def _command(arguments: list[str]) -> str | None:
    try:
        result = subprocess.run(
            arguments,
            capture_output=True,
            text=True,
            encoding="utf-8",
            errors="replace",
            timeout=3,
            creationflags=subprocess.CREATE_NO_WINDOW if sys.platform == "win32" else 0,
            check=False,
        )
        return result.stdout.strip()[:4096] if result.returncode == 0 else None
    except (OSError, subprocess.SubprocessError):
        return None


def _tool_identity(name: str) -> dict:
    path = toolchain.resolve(name)
    result = {"path": path, "version": None, "sha256": None}
    if path:
        result["version"] = _command([path, "version" if name == "adb" else "--version"])
        try:
            with Path(path).open("rb") as stream:
                result["sha256"] = hashlib.file_digest(stream, "sha256").hexdigest()
        except OSError:
            logger.debug("Could not hash %s", name)
    return result


@lru_cache(maxsize=2)
def runtime_identity(android: bool) -> dict:
    """Read identity once per worker, before later source updates can change it."""
    git = ["git", "-C", str(ROOT_DIR)]
    checkout = (ROOT_DIR / ".git").exists()
    commit = _command([*git, "rev-parse", "--verify", "HEAD"]) if checkout else None
    changes = (
        _command([*git, "status", "--porcelain", "--untracked-files=no"]) if checkout else None
    )
    return {
        "artemis_version": get_version(),
        "source_commit": commit,
        "source_modified": bool(changes) if changes is not None else None,
        "installation_archive_sha256": None,
        "python_version": platform.python_version(),
        "platform": platform.system(),
        "tools": {name: _tool_identity(name) for name in ("adb", "scrcpy")} if android else {},
    }


class _PerformanceInformation(ctypes.Structure):
    # https://learn.microsoft.com/windows/win32/api/psapi/ns-psapi-performance_information
    _fields_ = (
        [("cb", wintypes.DWORD)]
        + [
            (name, ctypes.c_size_t)
            for name in (
                "CommitTotal",
                "CommitLimit",
                "CommitPeak",
                "PhysicalTotal",
                "PhysicalAvailable",
                "SystemCache",
                "KernelTotal",
                "KernelPaged",
                "KernelNonpaged",
                "PageSize",
            )
        ]
        + [(name, wintypes.DWORD) for name in ("HandleCount", "ProcessCount", "ThreadCount")]
    )


def _windows_commit() -> dict | None:
    if sys.platform != "win32":
        return None
    values = _PerformanceInformation()
    values.cb = ctypes.sizeof(values)
    query = ctypes.WinDLL("psapi", use_last_error=True).GetPerformanceInfo
    query.argtypes = [ctypes.POINTER(_PerformanceInformation), wintypes.DWORD]
    query.restype = wintypes.BOOL
    if not query(ctypes.byref(values), values.cb):
        raise ctypes.WinError(ctypes.get_last_error())
    committed = values.CommitTotal * values.PageSize
    limit = values.CommitLimit * values.PageSize
    return {
        "committed_bytes": committed,
        "limit_bytes": limit,
        "available_bytes": max(0, limit - committed),
        "note": "系统提交额度，不能当作页面文件大小。",
    }


def _memory_snapshot() -> dict:
    result = {"physical": None, "system_commit": None, "worker": None}
    try:
        memory = psutil.virtual_memory()
        result["physical"] = {"total_bytes": memory.total, "available_bytes": memory.available}
        result["system_commit"] = _windows_commit()
    except (OSError, psutil.Error):
        logger.debug("Host memory snapshot unavailable")
    try:
        memory = psutil.Process().memory_info()
        result["worker"] = {
            "pid": os.getpid(),
            "rss_bytes": memory.rss,
            "private_bytes": getattr(memory, "private", None),
        }
    except (OSError, psutil.Error):
        logger.debug("Worker memory snapshot unavailable")
    return result


def _adb_server() -> dict:
    endpoint = current_adb_endpoint()
    result = {"endpoint": endpoint.to_dict(), "local_listener": None}
    if endpoint.host.lower() not in {"localhost", "127.0.0.1", "::1"}:
        return result  # A remote server's process cannot be observed on this host.
    try:
        for connection in psutil.net_connections(kind="tcp"):
            if (
                connection.status == psutil.CONN_LISTEN
                and connection.laddr.port == endpoint.port
                and connection.laddr.ip
                in {
                    endpoint.host,
                    "0.0.0.0",
                    "::",
                    *(("127.0.0.1", "::1") if endpoint.host.lower() == "localhost" else ()),
                }
                and connection.pid
            ):
                process = psutil.Process(connection.pid)
                result["local_listener"] = {
                    "pid": process.pid,
                    "name": process.name(),
                    "path": process.exe(),
                }
                break
    except (OSError, psutil.Error):
        logger.debug("ADB listener identity unavailable")
    return result


def write_runtime_diagnostics(directory: Path, phase: str, *, android: bool = True) -> None:
    """Keep a start snapshot and the latest failure snapshot without changing task results."""
    if phase not in {"start", "failure"}:
        raise ValueError("Unknown runtime snapshot phase")
    path = directory / FILENAME
    try:
        payload = {"schema_version": 1}
        if path.is_file() and path.stat().st_size <= 256 * 1024:
            try:
                saved = json.loads(path.read_text(encoding="utf-8"))
                if isinstance(saved, dict):
                    payload = saved
            except (OSError, ValueError):
                logger.debug("Previous runtime snapshot unreadable")
        if phase == "start" and phase in payload:
            return
        payload[phase] = {
            "captured_at": time.time(),
            "runtime": runtime_identity(android),
            "memory": _memory_snapshot(),
            "adb_server": _adb_server() if android else None,
        }
        directory.mkdir(parents=True, exist_ok=True)
        temporary = path.with_suffix(".tmp")
        temporary.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
        temporary.replace(path)
    except (OSError, psutil.Error, ValueError, TypeError):
        logger.warning("Could not save runtime snapshot (%s)", phase, exc_info=True)


async def capture_runtime_diagnostics(directory: Path, phase: str, *, android: bool = True) -> None:
    """Keep optional evidence failures outside the task's execution result."""
    try:
        await asyncio.to_thread(write_runtime_diagnostics, directory, phase, android=android)
    except (OSError, RuntimeError):
        logger.warning("Could not start runtime snapshot (%s)", phase, exc_info=True)