front/xhs-agent-desktop/engine/xhs_bridge/skills_process_manager.py

from __future__ import annotations

import json
import subprocess
import sys
import time
from pathlib import Path

import websockets.sync.client as ws_client

from xhs_bridge.config import AppConfig
from xhs_bridge.exceptions import ProcessError
from xhs_bridge.subprocess_win import popen_kwargs_hide_console


class SkillsProcessManager:
    """管理 skills bridge_server 子进程与扩展连接状态。"""

    def __init__(self, config: AppConfig) -> None:
        self.config = config

    @property
    def _bridge_url(self) -> str:
        return self.config.skills.bridge_ws_url

    def _read_pid(self) -> int | None:
        pid_file = self.config.bridge_pid_file()
        if not pid_file.is_file():
            return None
        try:
            pid = int(pid_file.read_text(encoding="utf-8").strip())
        except ValueError:
            return None
        return pid if self._is_process_running(pid) else None

    @staticmethod
    def _is_process_running(pid: int) -> bool:
        if pid <= 0:
            return False
        try:
            import ctypes

            handle = ctypes.windll.kernel32.OpenProcess(0x1000, False, pid)
            if handle:
                ctypes.windll.kernel32.CloseHandle(handle)
                return True
            return False
        except Exception:
            return False

    def _write_pid(self, pid: int) -> None:
        pid_file = self.config.bridge_pid_file()
        pid_file.parent.mkdir(parents=True, exist_ok=True)
        pid_file.write_text(str(pid), encoding="utf-8")

    def _clear_pid(self) -> None:
        pid_file = self.config.bridge_pid_file()
        if pid_file.is_file():
            pid_file.unlink(missing_ok=True)

    def ping_bridge(self) -> dict[str, bool]:
        """查询 bridge server 与扩展连接状态。"""
        try:
            with ws_client.connect(self._bridge_url, max_size=32 * 1024 * 1024, open_timeout=3) as ws:
                ws.send(json.dumps({"role": "cli", "method": "ping_server"}))
                raw = ws.recv(timeout=5)
            data = json.loads(raw)
            result = data.get("result") or {}
            return {
                "server_running": True,
                "extension_connected": bool(result.get("extension_connected")),
            }
        except Exception:
            return {"server_running": False, "extension_connected": False}

    def is_engine_healthy(self) -> bool:
        ping = self.ping_bridge()
        return ping["server_running"] and ping["extension_connected"]

    def is_bridge_running(self) -> bool:
        return self.ping_bridge()["server_running"]

    def wait_engine_ready(self, timeout_sec: int | None = None) -> bool:
        timeout = timeout_sec or self.config.skills.startup_timeout_sec
        deadline = time.time() + timeout
        while time.time() < deadline:
            if self.is_engine_healthy():
                return True
            time.sleep(0.5)
        return False

    def start_engine(self, force: bool = False) -> int:
        if not force and self.is_engine_healthy():
            return self._read_pid() or 0

        if force:
            self.stop_engine()

        port = self.config.skills.bridge_port
        cmd = self.config.skills.bridge_launcher(port=port)
        if len(cmd) >= 2 and cmd[0].endswith(".py") and not Path(cmd[0]).is_file():
            raise ProcessError(f"Bridge 脚本不存在: {cmd[0]}")
        if cmd[0].endswith(".exe") and not Path(cmd[0]).is_file():
            raise ProcessError(f"Bridge 引擎不存在: {cmd[0]}")

        proc = subprocess.Popen(
            cmd,
            cwd=str(self.config.skills.project_dir),
            **popen_kwargs_hide_console(new_process_group=True),
        )
        self._write_pid(proc.pid)

        if not self.wait_engine_ready():
            raise ProcessError(
                f"Skills 引擎启动超时（{self.config.skills.startup_timeout_sec}s）。\n"
                "请确认：\n"
                "1. Chrome 已安装 XHS Bridge 扩展（extension/）\n"
                "2. 扩展已启用并连接到 ws://localhost:{port}"
            )
        return proc.pid

    def stop_engine(self) -> None:
        pid = self._read_pid()
        if pid:
            subprocess.run(
                ["taskkill", "/PID", str(pid), "/T", "/F"],
                check=False,
                capture_output=True,
            )
        self._clear_pid()

    def open_chrome(self) -> None:
        """尝试启动 Chrome（扩展未连接时）。"""
        import os

        candidates = [
            Path(r"C:\Program Files\Google\Chrome\Application\chrome.exe"),
            Path(r"C:\Program Files (x86)\Google\Chrome\Application\chrome.exe"),
            Path(os.path.expandvars(r"%LOCALAPPDATA%\Google\Chrome\Application\chrome.exe")),
        ]
        for path in candidates:
            if path.is_file():
                subprocess.Popen([str(path)])
                return