opencode-mem-guard 0.1.0__py3-none-any.whl

opencode_mem_guard/__init__.py

@@ -0,0 +1,3 @@
+"""OpenCode Memory Guard - 内存泄漏监控与僵尸进程回收"""
+
+__version__ = "0.1.0"

opencode_mem_guard/__main__.py

@@ -0,0 +1,5 @@
+"""允许 python -m opencode_mem_guard 启动"""
+from opencode_mem_guard.app import main
+
+if __name__ == "__main__":
+    main()

opencode_mem_guard/app.py

@@ -0,0 +1,351 @@
+"""
+OpenCode MemGuard 主应用编排器
+
+职责：
+- 启动监控线程（collector + leak_detector + reclaimer 循环）
+- 启动 360 式悬浮球（主线程 Win32 Layered Window）
+- 启动系统托盘（后台线程 pystray）
+- 线程安全的共享状态
+"""
+import argparse
+import sys
+import time
+import threading
+import logging
+from pathlib import Path
+from typing import Optional
+
+from .collector import collect_snapshot, MemSnapshot
+from .leak_detector import LeakDetector, LeakAlert
+from .reclaimer import ProcessReclaimer, ReclaimResult
+from .datastore import DataStore
+from .autostart import is_autostart_enabled, enable_autostart, disable_autostart
+
+
+# ─── 共享状态 ───────────────────────────────────────────────
+
+class SharedState:
+    """线程安全的共享监控状态"""
+
+    def __init__(self):
+        self._lock = threading.Lock()
+        self._data = {
+            "ram_pct": 0.0,
+            "ram_used_gb": 0.0,
+            "ram_total_gb": 0.0,
+            "node_count": 0,
+            "node_ram_mb": 0.0,
+            "node_handles": 0,
+            "total_killed": 0,
+            "total_freed_mb": 0.0,
+            "auto_reclaim": True,
+            "status": "normal",
+            "uptime_min": 0.0,
+            "alerts": [],
+        }
+        self._start_time = time.time()
+
+    def update_from_snapshot(self, snap: MemSnapshot):
+        with self._lock:
+            if snap.system:
+                self._data["ram_pct"] = snap.system.usage_pct
+                self._data["ram_used_gb"] = snap.system.used_ram_gb
+                self._data["ram_total_gb"] = snap.system.total_ram_gb
+            self._data["node_count"] = snap.node_process_count
+            self._data["node_ram_mb"] = snap.node_total_ram_mb
+            self._data["node_handles"] = snap.node_total_handles
+            self._data["uptime_min"] = round((time.time() - self._start_time) / 60, 1)
+
+    def update_alerts(self, alerts: list):
+        with self._lock:
+            # 确定状态等级
+            status = "normal"
+            for a in alerts:
+                if a.level in ("critical", "emergency"):
+                    status = "critical"
+                    break
+                elif a.level == "warning":
+                    status = "warning"
+            self._data["status"] = status
+            self._data["alerts"] = [
+                {"level": a.level, "metric": a.metric, "msg": a.message}
+                for a in alerts[-5:]  # 只保留最近5条
+            ]
+
+    def update_reclaim(self, reclaimer: ProcessReclaimer):
+        with self._lock:
+            self._data["total_killed"] = reclaimer.total_killed
+            self._data["total_freed_mb"] = round(reclaimer.total_freed_mb, 1)
+
+    def set_auto_reclaim(self, enabled: bool):
+        with self._lock:
+            self._data["auto_reclaim"] = enabled
+
+    def get_auto_reclaim(self) -> bool:
+        with self._lock:
+            return self._data["auto_reclaim"]
+
+    def get_status_dict(self) -> dict:
+        with self._lock:
+            return dict(self._data)
+
+    def get_tray_status(self) -> str:
+        with self._lock:
+            return self._data["status"]
+
+    def get_tray_tooltip(self) -> str:
+        with self._lock:
+            d = self._data
+            return (
+                f"MemGuard - RAM: {d['ram_pct']}% | "
+                f"Node: {d['node_count']}进程 | "
+                f"{d['node_ram_mb']:.0f}MB"
+            )
+
+
+# ─── 监控线程 ───────────────────────────────────────────────
+
+class MonitorThread(threading.Thread):
+    """后台监控线程"""
+
+    def __init__(
+        self,
+        state: SharedState,
+        interval: float = 5.0,
+        no_reclaim: bool = False,
+        dry_run: bool = False,
+        data_dir: str = "",
+    ):
+        super().__init__(daemon=True, name="monitor-loop")
+        self.state = state
+        self.interval = interval
+        self.no_reclaim = no_reclaim
+        self._stop_event = threading.Event()
+
+        # 数据目录
+        if not data_dir:
+            data_dir = str(Path.home() / ".opencode-mem-guard" / "data")
+        self.store = DataStore(data_dir=data_dir)
+        self.detector = LeakDetector()
+        self.reclaimer = ProcessReclaimer(
+            dry_run=dry_run,
+            logger=self.store.logger,
+        )
+        self.store.logger.info("监控线程初始化完成")
+
+    def run(self):
+        self.store.logger.info(
+            f"监控启动 interval={self.interval}s "
+            f"reclaim={'off' if self.no_reclaim else 'on'}"
+        )
+        while not self._stop_event.is_set():
+            try:
+                self._tick()
+            except Exception as e:
+                self.store.logger.error(f"监控循环异常: {e}")
+            self._stop_event.wait(self.interval)
+
+        self.store.close()
+        self.store.logger.info("监控线程已停止")
+
+    def _tick(self):
+        # 1. 采集
+        snap = collect_snapshot()
+        self.state.update_from_snapshot(snap)
+        self.store.save_snapshot(snap)
+
+        # 2. 泄漏检测
+        alerts = self.detector.feed(snap)
+        self.state.update_alerts(alerts)
+        if alerts:
+            self.store.save_alerts(alerts)
+
+        # 3. 回收
+        if not self.no_reclaim and self.state.get_auto_reclaim():
+            if self.reclaimer.should_run():
+                result = self.reclaimer.reclaim()
+                if result and result.actions:
+                    self.store.save_reclaim(result)
+                    self.state.update_reclaim(self.reclaimer)
+        else:
+            # 即使不回收也更新统计
+            self.state.update_reclaim(self.reclaimer)
+
+    def force_reclaim(self) -> Optional[ReclaimResult]:
+        """手动触发一次回收（忽略冷却期）"""
+        self.reclaimer._last_reclaim_time = 0  # 重置冷却
+        result = self.reclaimer.reclaim()
+        if result and result.actions:
+            self.store.save_reclaim(result)
+            self.state.update_reclaim(self.reclaimer)
+        return result
+
+    def stop(self):
+        self._stop_event.set()
+
+
+# ─── 主入口 ─────────────────────────────────────────────────
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        prog="opencode-mem-guard",
+        description="OpenCode 内存泄漏监控与僵尸进程回收",
+    )
+    parser.add_argument(
+        "--no-reclaim", action="store_true",
+        help="禁用自动回收",
+    )
+    parser.add_argument(
+        "--dry-run", action="store_true",
+        help="试运行模式（只检测不杀进程）",
+    )
+    parser.add_argument(
+        "--no-tray", action="store_true",
+        help="不显示托盘图标（headless 模式）",
+    )
+    parser.add_argument(
+        "--no-ball", action="store_true",
+        help="不显示悬浮球（仅托盘模式）",
+    )
+    parser.add_argument(
+        "-i", "--interval", type=float, default=5.0,
+        help="采集间隔（秒），默认 5",
+    )
+    parser.add_argument(
+        "--data-dir", type=str, default="",
+        help="数据存储目录，默认 ~/.opencode-mem-guard/data/",
+    )
+    args = parser.parse_args()
+
+    # 共享状态
+    state = SharedState()
+    state.set_auto_reclaim(not args.no_reclaim)
+
+    # 启动监控线程
+    monitor = MonitorThread(
+        state=state,
+        interval=args.interval,
+        no_reclaim=args.no_reclaim,
+        dry_run=args.dry_run,
+        data_dir=args.data_dir,
+    )
+    monitor.start()
+
+    if args.no_tray and args.no_ball:
+        # headless 模式：直接等待
+        print("OpenCode MemGuard 已启动（headless 模式）")
+        print(f"数据目录: {monitor.store.data_dir}")
+        print("按 Ctrl+C 退出")
+        try:
+            while True:
+                time.sleep(1)
+        except KeyboardInterrupt:
+            print("\n正在停止...")
+            monitor.stop()
+            monitor.join(timeout=5)
+        return
+
+    # ── 回调函数 ──
+
+    ball = None  # 浮球引用，后面赋值
+
+    def on_quit():
+        monitor.stop()
+        if ball:
+            ball.stop()
+
+    def on_toggle_reclaim():
+        state.set_auto_reclaim(not state.get_auto_reclaim())
+
+    def on_force_reclaim():
+        monitor.force_reclaim()
+
+    def on_toggle_autostart():
+        if is_autostart_enabled():
+            disable_autostart()
+        else:
+            enable_autostart()
+
+    # ── 启动托盘（后台线程） ──
+
+    tray = None
+
+    if not args.no_tray:
+        from .tray import TrayManager
+
+        def on_toggle_ball():
+            if ball:
+                ball.toggle_visibility()
+
+        tray = TrayManager(
+            on_toggle_ball=on_toggle_ball,
+            on_toggle_reclaim=on_toggle_reclaim,
+            on_force_reclaim=on_force_reclaim,
+            on_toggle_autostart=on_toggle_autostart,
+            on_quit=on_quit,
+            get_reclaim_state=lambda: state.get_auto_reclaim(),
+            get_autostart_state=is_autostart_enabled,
+            get_ball_visible=lambda: ball._visible if ball else True,
+            data_dir=str(monitor.store.data_dir),
+        )
+
+        # 托盘状态更新线程
+        def _update_tray():
+            while monitor.is_alive():
+                try:
+                    tray.update_status(
+                        state.get_tray_status(),
+                        state.get_tray_tooltip(),
+                    )
+                except Exception:
+                    pass
+                time.sleep(3)
+
+        tray_thread = threading.Thread(
+            target=tray.run, daemon=True, name="pystray"
+        )
+        tray_thread.start()
+
+        updater = threading.Thread(
+            target=_update_tray, daemon=True, name="tray-updater"
+        )
+        updater.start()
+
+    # ── 启动悬浮球（主线程 Win32） ──
+
+    if not args.no_ball:
+        from .floating_ball import FloatingBall
+
+        ball = FloatingBall(
+            get_data=lambda: state.get_status_dict(),
+            on_force_reclaim=on_force_reclaim,
+            on_toggle_reclaim=on_toggle_reclaim,
+            on_quit=on_quit,
+            get_reclaim_state=lambda: state.get_auto_reclaim(),
+        )
+
+        try:
+            ball.run()  # 阻塞：tkinter mainloop
+        except KeyboardInterrupt:
+            pass
+        finally:
+            monitor.stop()
+            monitor.join(timeout=5)
+            if tray is not None:
+                tray.stop()
+    else:
+        # 仅托盘模式，主线程等待
+        print("OpenCode MemGuard 已启动（仅托盘模式）")
+        try:
+            while True:
+                time.sleep(1)
+        except KeyboardInterrupt:
+            pass
+        finally:
+            monitor.stop()
+            monitor.join(timeout=5)
+
+
+if __name__ == "__main__":
+    main()

opencode_mem_guard/autostart.py

@@ -0,0 +1,129 @@
+"""
+Windows 开机自启动管理模块
+使用注册表实现开机自启动
+"""
+import sys
+import winreg
+from pathlib import Path
+from typing import Optional
+
+
+REG_PATH = r"Software\Microsoft\Windows\CurrentVersion\Run"
+APP_NAME = "OpenCodeMemGuard"
+
+
+def get_executable_path() -> str:
+    """获取当前可执行文件路径"""
+    if getattr(sys, 'frozen', False):
+        # 打包后的 exe
+        return sys.executable
+    else:
+        # 开发模式：使用 pythonw -m opencode_mem_guard
+        python_exe = sys.executable.replace("python.exe", "pythonw.exe")
+        if not Path(python_exe).exists():
+            python_exe = sys.executable
+        return f'"{python_exe}" -m opencode_mem_guard'
+
+
+def enable_autostart() -> bool:
+    """
+    启用开机自启动
+    
+    Returns:
+        bool: 成功返回 True，失败返回 False
+    """
+    try:
+        key = winreg.OpenKey(
+            winreg.HKEY_CURRENT_USER,
+            REG_PATH,
+            0,
+            winreg.KEY_SET_VALUE
+        )
+        
+        exe_path = get_executable_path()
+        winreg.SetValueEx(key, APP_NAME, 0, winreg.REG_SZ, exe_path)
+        winreg.CloseKey(key)
+        return True
+    except Exception as e:
+        print(f"启用自启动失败: {e}")
+        return False
+
+
+def disable_autostart() -> bool:
+    """
+    禁用开机自启动
+    
+    Returns:
+        bool: 成功返回 True，失败返回 False
+    """
+    try:
+        key = winreg.OpenKey(
+            winreg.HKEY_CURRENT_USER,
+            REG_PATH,
+            0,
+            winreg.KEY_SET_VALUE
+        )
+        
+        try:
+            winreg.DeleteValue(key, APP_NAME)
+        except FileNotFoundError:
+            # 键不存在，视为成功
+            pass
+        
+        winreg.CloseKey(key)
+        return True
+    except Exception as e:
+        print(f"禁用自启动失败: {e}")
+        return False
+
+
+def is_autostart_enabled() -> bool:
+    """
+    检查是否已启用开机自启动
+    
+    Returns:
+        bool: 已启用返回 True，未启用返回 False
+    """
+    try:
+        key = winreg.OpenKey(
+            winreg.HKEY_CURRENT_USER,
+            REG_PATH,
+            0,
+            winreg.KEY_READ
+        )
+        
+        try:
+            value, _ = winreg.QueryValueEx(key, APP_NAME)
+            winreg.CloseKey(key)
+            return bool(value)
+        except FileNotFoundError:
+            winreg.CloseKey(key)
+            return False
+    except Exception:
+        return False
+
+
+def get_autostart_path() -> Optional[str]:
+    """
+    获取注册表中的自启动路径
+    
+    Returns:
+        Optional[str]: 自启动路径，如果未设置则返回 None
+    """
+    try:
+        key = winreg.OpenKey(
+            winreg.HKEY_CURRENT_USER,
+            REG_PATH,
+            0,
+            winreg.KEY_READ
+        )
+        
+        try:
+            value, _ = winreg.QueryValueEx(key, APP_NAME)
+            winreg.CloseKey(key)
+            return value
+        except FileNotFoundError:
+            winreg.CloseKey(key)
+            return None
+    except Exception:
+        return None

opencode_mem_guard/collector.py

@@ -0,0 +1,117 @@
+"""
+内存数据采集模块
+使用 psutil 采集 node.exe 进程内存、系统内存、句柄数等关键指标
+比 PowerShell 快 100x+，不会在系统高负载时超时
+"""
+import time
+import psutil
+from dataclasses import dataclass, field, asdict
+from typing import List, Optional
+
+
+@dataclass
+class ProcessMemInfo:
+    """单个进程的内存信息"""
+    pid: int
+    name: str
+    ram_mb: float  # WorkingSet (物理内存)
+    virtual_mb: float  # 虚拟内存
+    paged_mb: float  # 分页内存
+    peak_ram_mb: float  # 峰值物理内存
+    handle_count: int
+    thread_count: int
+    cpu_seconds: float = 0.0
+    cmdline: str = ""
+
+
+@dataclass
+class SystemMemInfo:
+    """系统级内存信息"""
+    total_ram_gb: float
+    free_ram_gb: float
+    used_ram_gb: float
+    usage_pct: float
+    commit_total_gb: float  # 提交总量(虚拟内存)
+    commit_used_gb: float
+    commit_pct: float
+
+
+@dataclass
+class MemSnapshot:
+    """一次完整的内存快照"""
+    timestamp: float
+    system: Optional[SystemMemInfo] = None
+    node_processes: List[ProcessMemInfo] = field(default_factory=list)
+    node_total_ram_mb: float = 0.0
+    node_total_virtual_gb: float = 0.0
+    node_process_count: int = 0
+    node_total_handles: int = 0
+    node_total_threads: int = 0
+    error: str = ""
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+
+def collect_snapshot() -> MemSnapshot:
+    """使用 psutil 执行一次完整的内存快照采集，毫秒级完成"""
+    snap = MemSnapshot(timestamp=time.time())
+
+    try:
+        # ── 系统内存 ──
+        vm = psutil.virtual_memory()
+        swap = psutil.swap_memory()
+
+        total_gb = vm.total / (1024 ** 3)
+        free_gb = vm.available / (1024 ** 3)
+        used_gb = (vm.total - vm.available) / (1024 ** 3)
+
+        # commit = 物理内存 + swap 已用
+        commit_total_gb = (vm.total + swap.total) / (1024 ** 3)
+        commit_used_gb = (vm.total - vm.available + swap.used) / (1024 ** 3)
+        commit_pct = (commit_used_gb / commit_total_gb * 100) if commit_total_gb > 0 else 0
+
+        snap.system = SystemMemInfo(
+            total_ram_gb=round(total_gb, 2),
+            free_ram_gb=round(free_gb, 2),
+            used_ram_gb=round(used_gb, 2),
+            usage_pct=round(vm.percent, 1),
+            commit_total_gb=round(commit_total_gb, 2),
+            commit_used_gb=round(commit_used_gb, 2),
+            commit_pct=round(commit_pct, 1),
+        )
+    except Exception as e:
+        snap.error = f"系统内存采集失败: {e}"
+
+    # ── Node 进程采集 ──
+    try:
+        for proc in psutil.process_iter(['pid', 'name']):
+            try:
+                if proc.info['name'] and proc.info['name'].lower() == 'node.exe':
+                    mem = proc.memory_info()
+                    cpu_times = proc.cpu_times()
+                    snap.node_processes.append(ProcessMemInfo(
+                        pid=proc.info['pid'],
+                        name="node.exe",
+                        ram_mb=round(mem.rss / 1048576, 1),
+                        virtual_mb=round(mem.vms / 1048576, 1),
+                        paged_mb=round(getattr(mem, 'paged_pool', 0) / 1048576, 1),
+                        peak_ram_mb=round(getattr(mem, 'peak_wset', mem.rss) / 1048576, 1),
+                        handle_count=proc.num_handles() if hasattr(proc, 'num_handles') else 0,
+                        thread_count=proc.num_threads(),
+                        cpu_seconds=round(cpu_times.user + cpu_times.system, 2),
+                    ))
+            except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
+                continue
+    except Exception as e:
+        if not snap.error:
+            snap.error = f"进程采集失败: {e}"
+
+    # 汇总
+    snap.node_process_count = len(snap.node_processes)
+    snap.node_total_ram_mb = round(sum(p.ram_mb for p in snap.node_processes), 1)
+    snap.node_total_virtual_gb = round(sum(p.virtual_mb for p in snap.node_processes) / 1024, 2)
+    snap.node_total_handles = sum(p.handle_count for p in snap.node_processes)
+    snap.node_total_threads = sum(p.thread_count for p in snap.node_processes)
+
+    return snap