Pamudu13/gemma-3-chat

1import os2import json, asyncio, socket3from pathlib import Path4from typing import Dict, Optional5from py.get_setting import EXT_DIR, IS_DOCKER6 7PORT_RANGE = (3100, 13999)8 9# 获取环境变量（由 Docker 或 Electron 注入）10ELECTRON_NODE = os.environ.get("ELECTRON_NODE_EXEC")11ELECTRON_NPM_CLI = os.environ.get("ELECTRON_NPM_CLI")12 13class NodeExtension:14    def __init__(self, ext_id: str):15        self.ext_id   = ext_id16        self.proc: Optional[asyncio.subprocess.Process] = None17        self.port: Optional[int] = None18        self.root     = Path(EXT_DIR) / ext_id19        self.pkg      = json.loads((self.root / "package.json").read_text(encoding="utf-8"))20 21    def _get_exec_cmds(self):22        """智能生成 node 和 npm 的执行命令列表"""23        if IS_DOCKER or not ELECTRON_NODE:24            # Docker 或原生环境：直接使用系统全局的 node 和 npm25            npm_exe = "npm.cmd" if os.name == "nt" else "npm"26            return ["node"], [npm_exe]27        else:28            # Electron 桌面端环境：29            # Node 命令: electron.exe30            # NPM 命令: electron.exe /path/to/npm-cli.js31            return [ELECTRON_NODE], [ELECTRON_NODE, ELECTRON_NPM_CLI]32 33    def _get_env(self):34        """生成带 ELECTRON_RUN_AS_NODE 标记的环境变量"""35        env = os.environ.copy()36        if not IS_DOCKER and ELECTRON_NODE:37            env["ELECTRON_RUN_AS_NODE"] = "1"38        return env39 40    async def start(self) -> int:41        if self.proc and self.proc.returncode is None:42            return self.port43 44        pkg_file = self.root / "package.json"45        nm_folder = self.root / "node_modules"46        47        node_cmd, npm_cmd = self._get_exec_cmds()48        run_env = self._get_env()49        run_env["NODE_EXTENSION_ID"] = self.ext_id 50        # 0. 快速判断：node_modules 存在且比 package.json 新51        if nm_folder.is_dir() and nm_folder.stat().st_mtime >= pkg_file.stat().st_mtime:52            print(f"[{self.ext_id}] node_modules 已存在，跳过 npm install")53        else:54            print(f"[{self.ext_id}] 首次/依赖变更，执行 npm install")55            # 1. 启动 npm install56            # 注意这里使用 *npm_cmd 解包列表57            proc = await asyncio.create_subprocess_exec(58                *npm_cmd, "install", "--production",59                cwd=self.root,60                env=run_env,  # 必须传入修改后的环境变量61                stdout=asyncio.subprocess.PIPE,62                stderr=asyncio.subprocess.STDOUT63            )64            stdout, _ = await proc.communicate()65            if proc.returncode != 0:66                raise RuntimeError(f"npm install 失败:\n{stdout.decode('utf-8', errors='ignore')}")67            # 刷新时间戳68            nm_folder.touch(exist_ok=True)69 70        # 2. 选端口71        want = self.pkg.get("nodePort", 0)72        self.port = want if want else _free_port()73 74        # 3. 起进程75        self.proc = await asyncio.create_subprocess_exec(76            *node_cmd, "index.js", str(self.port),77            cwd=self.root,78            env=run_env, # 必须传入修改后的环境变量79            stdout=asyncio.subprocess.PIPE,80            stderr=asyncio.subprocess.STDOUT81        )82        83        # 4. 等健康84        await _wait_port(self.port)85        return self.port86    87    async def stop(self):88        if self.proc:89            self.proc.terminate()90            await self.proc.wait()91            self.proc = None92 93# ---------- 工具 ----------94def _free_port() -> int:95    with socket.socket() as s:96        s.bind(("", 0))97        return s.getsockname()[1]98 99async def _wait_port(port: int, timeout=10):100    for _ in range(timeout * 10):101        try:102            _, w = await asyncio.wait_for(asyncio.open_connection("127.0.0.1", port), 1)103            w.close()104            return105        except:106            await asyncio.sleep(0.1)107    raise RuntimeError("端口未就绪")108 109# ---------- 全局管理器 ----------110class NodeManager:111    def __init__(self):112        self.exts: Dict[str, NodeExtension] = {}113 114    async def start(self, ext_id: str) -> int:115        if ext_id not in self.exts:116            self.exts[ext_id] = NodeExtension(ext_id)117        return await self.exts[ext_id].start()118 119    async def stop(self, ext_id: str):120        if ext_id in self.exts:121            await self.exts[ext_id].stop()122            del self.exts[ext_id]123 124node_mgr = NodeManager()