Artemis-Eapil/artemis/service_registry.py

"""Lightweight, local discovery of user-started services. Never starts processes.

This module deliberately avoids config/runtime imports: it is also used by the
stdio MCP bridge, which must not load agents, device controllers or model SDKs.
"""

from __future__ import annotations

import hashlib
import http.client
import json
import os
from pathlib import Path
import tempfile
from typing import Any
import urllib.error
import urllib.request

from filelock import FileLock, Timeout
import psutil

PROTOCOL_VERSION = 1
ROOT = Path(__file__).resolve().parent.parent
MCP_PATH = "/api/mcp/"
TOKEN_HEADER = "X-Artemis-MCP-Token"


class ServiceUnavailableError(RuntimeError):
    """The user must start/select a compatible Artemis service."""


def canonical_path(path: str | Path) -> str:
    return os.path.normcase(str(Path(path).expanduser().resolve()))


def installation_id(root: Path = ROOT) -> str:
    return hashlib.sha256(canonical_path(root).encode()).hexdigest()[:24]


def registry_directory(root: Path = ROOT) -> Path:
    user_id = hashlib.sha256(canonical_path(Path.home()).encode()).hexdigest()[:16]
    return Path(tempfile.gettempdir()) / "artemis-services" / user_id / installation_id(root)


def _alive(pid: Any) -> bool:
    return isinstance(pid, int) and pid > 0 and psutil.pid_exists(pid)


def _selected(info: dict, host: str | None, port: int | None) -> bool:
    if host is not None and host not in {"localhost", "127.0.0.1", "::1"}:
        return False
    if port is not None and info.get("port") != port:
        return False
    app_dir = os.getenv("ARTEMIS_APP_DIR") or os.getenv("ANTIGRAVITY_APP_DIR")
    traces_dir = os.getenv("ARTEMIS_TRACES_DIR")
    return not (
        (app_dir and canonical_path(app_dir) != info.get("app_dir"))
        or (traces_dir and canonical_path(traces_dir) != info.get("traces_dir"))
    )


def probe_service(info: dict, timeout: float = 1.5) -> bool:
    port = info.get("port")
    if not isinstance(port, int) or not 0 < port < 65536 or not _alive(info.get("pid")):
        return False
    token = info.get("token")
    if not isinstance(token, str) or not token:
        return False
    request = urllib.request.Request(
        f"http://127.0.0.1:{port}/api/system/shared-service",
        headers={TOKEN_HEADER: token},
    )
    try:
        # A local service connection must not go through the user's HTTP proxy.
        opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))
        with opener.open(request, timeout=timeout) as response:
            result = json.loads(response.read())
    except (OSError, ValueError, urllib.error.URLError, http.client.HTTPException):
        return False
    return (
        isinstance(result, dict)
        and result.get("protocol") == PROTOCOL_VERSION
        and result.get("pid") == info.get("pid")
        and result.get("installation") == info.get("installation")
        and result.get("instance") == info.get("instance")
    )


def discover_service(root: Path = ROOT) -> dict:
    """Choose one live service in this installation, respecting explicit scope.

    An unspecified data directory follows the sole running instance (including a
    custom web port). Ambiguous instances require an explicit daemon port; we never
    guess or cross installation boundaries.
    """
    host = os.getenv("ARTEMIS_DAEMON_HOST")
    try:
        port = int(os.environ["ARTEMIS_DAEMON_PORT"]) if os.getenv("ARTEMIS_DAEMON_PORT") else None
    except ValueError as exc:
        raise ServiceUnavailableError("ARTEMIS_DAEMON_PORT 必须是有效端口号。") from exc
    candidates = []
    for path in registry_directory(root).glob("*.json"):
        try:
            info = json.loads(path.read_text(encoding="utf-8"))
        except (OSError, ValueError):
            continue
        if (
            isinstance(info, dict)
            and info.get("installation") == installation_id(root)
            and _selected(info, host, port)
            and probe_service(info)
        ):
            candidates.append(info)
    if len(candidates) == 1:
        return candidates[0]
    if len(candidates) > 1:
        raise ServiceUnavailableError(
            "检测到多个 ARTEMIS 数据空间，请设置 ARTEMIS_DAEMON_PORT 选择要连接的网页服务。"
        )
    raise ServiceUnavailableError(
        "未找到可连接的 ARTEMIS 主程序。请先运行 start.bat / start-cn.bat 或 artemis ui；"
        "若网页已启动，请确认使用同一安装目录的新版本，并重新启动网页服务。"
        "MCP 仅连接主程序，不会自动启动后台或独立执行任务。"
    )


class ServiceRegistration:
    """One server owner per installation and trace directory, released by the OS."""

    def __init__(
        self, *, port: int, token: str, app_dir: Path, traces_dir: Path, root: Path = ROOT
    ):
        import uuid

        directory = registry_directory(root)
        directory.mkdir(parents=True, exist_ok=True, mode=0o700)
        scope = hashlib.sha256(canonical_path(traces_dir).encode()).hexdigest()[:24]
        self.lock = FileLock(directory / f"{scope}.lock", thread_local=False)
        self.path = directory / f"{scope}.json"
        self.info = {
            "protocol": PROTOCOL_VERSION,
            "installation": installation_id(root),
            "instance": uuid.uuid4().hex,
            "pid": os.getpid(),
            "port": port,
            "token": token,
            "app_dir": canonical_path(app_dir),
            "traces_dir": canonical_path(traces_dir),
        }

    def acquire(self) -> None:
        try:
            self.lock.acquire(timeout=0)
        except Timeout as exc:
            raise ServiceUnavailableError(
                "此 ARTEMIS 数据目录已有主程序运行，请使用现有网页服务；不会启动第二份后台。"
            ) from exc
        try:
            self.path.write_text(json.dumps(self.info), encoding="utf-8")
            self.path.chmod(0o600)
        except OSError:
            self.lock.release()
            raise

    def close(self) -> None:
        if self.lock.is_locked:
            try:
                self.path.unlink(missing_ok=True)
            finally:
                self.lock.release()