@agentunion/kite 1.0.7 → 1.2.0

package/core/launcher/process_manager.py

@@ -8,8 +8,10 @@ stdio lifecycle:
 - Structured messages: JSON lines with "kite" field, dispatched via callback
 """
 
+import glob
 import json
 import os
+import re
 import subprocess
 import sys
 import threading
@@ -17,6 +19,9 @@ import time
 from dataclasses import dataclass, asdict
 from typing import Callable
 
+if sys.platform == "win32":
+    import ctypes
+    from ctypes import wintypes
 
 IS_WINDOWS = sys.platform == "win32"
 IS_MACOS = sys.platform == "darwin"
@@ -45,57 +50,311 @@ class ProcessManager:
         self.instance_id = instance_id or str(os.getpid())
         self._on_kite_message = on_kite_message
         # Use KITE_MODULE_DATA (set by Launcher for itself)
-        self.data_dir = os.environ.get("KITE_MODULE_DATA") or os.path.join(os.environ["KITE_INSTANCE_DIR"], "launcher")
-        self.records_path = os.path.join(self.data_dir, "processes.json")
+        launcher_dir = os.environ.get("KITE_MODULE_DATA") or os.path.join(os.environ["KITE_INSTANCE_DIR"], "launcher")
+        self.data_dir = os.path.join(launcher_dir, "state")
+        os.makedirs(self.data_dir, exist_ok=True)
+        self._instance_num = self._allocate_instance_num()
+        suffix = "" if self._instance_num == 1 else f"~{self._instance_num}"
+        self._instance_suffix = suffix
+        self.records_path = os.path.join(self.data_dir, f"processes{suffix}.json")
         self._processes: dict[str, subprocess.Popen] = {}   # name -> Popen
         self._records: dict[str, ProcessRecord] = {}         # name -> record
-        os.makedirs(self.data_dir, exist_ok=True)
+
+    @property
+    def instance_suffix(self) -> str:
+        """Return the instance suffix (empty string for instance 1, '~N' for N>1)."""
+        return self._instance_suffix
+
+    @property
+    def instance_num(self) -> int:
+        """Return the instance number (1-based)."""
+        return self._instance_num
+
+    def get_alive_instances(self) -> list[dict]:
+        """Scan processes*.json and return info about all alive instances (including self).
+
+        Returns list of dicts: {"num": int, "launcher_pid": int, "module_count": int, "is_self": bool, "debug": bool, "cwd": str}.
+        """
+        instances = []
+        pattern = os.path.join(self.data_dir, "processes*.json")
+        for filepath in glob.glob(pattern):
+            basename = os.path.basename(filepath)
+            num = self._parse_instance_num(basename)
+            if num is None:
+                continue
+            pid = self._read_launcher_pid(filepath)
+            is_self = (pid == os.getpid())
+            if not is_self and (not pid or not self._is_pid_alive(pid)):
+                continue
+            # Read module count, debug flag, and cwd
+            module_count = 0
+            debug = False
+            cwd = ""
+            try:
+                with open(filepath, "r", encoding="utf-8") as f:
+                    data = json.load(f)
+                if isinstance(data, dict):
+                    module_count = len(data.get("records", []))
+                    debug = data.get("debug", False)
+                    cwd = data.get("cwd", "")
+                elif isinstance(data, list):
+                    module_count = len(data)
+            except Exception:
+                pass
+            instances.append({
+                "num": num,
+                "launcher_pid": pid or 0,
+                "module_count": module_count,
+                "is_self": is_self,
+                "debug": debug,
+                "cwd": cwd,
+            })
+        instances.sort(key=lambda x: x["num"])
+        return instances
+
+    def get_global_instances(self) -> list[dict]:
+        """Scan all instance dirs under KITE_WORKSPACE for alive launcher instances.
+
+        Returns list of dicts:
+            {"instance_dir": str, "cwd": str, "num": int, "launcher_pid": int,
+             "module_count": int, "is_self": bool, "debug": bool}
+        """
+        workspace = os.environ.get("KITE_WORKSPACE", "")
+        if not workspace or not os.path.isdir(workspace):
+            return []
+
+        my_pid = os.getpid()
+        my_instance_dir = os.environ.get("KITE_INSTANCE_DIR", "")
+        results = []
+
+        try:
+            entries = os.listdir(workspace)
+        except OSError:
+            return []
+
+        for entry in entries:
+            inst_path = os.path.join(workspace, entry)
+            if not os.path.isdir(inst_path):
+                continue
+            state_dir = os.path.join(inst_path, "launcher", "state")
+            if not os.path.isdir(state_dir):
+                continue
+
+            # Read .cwd for display (fallback if processes.json lacks cwd)
+            cwd_from_file = ""
+            cwd_file = os.path.join(inst_path, ".cwd")
+            try:
+                with open(cwd_file, "r", encoding="utf-8") as f:
+                    cwd_from_file = f.read().strip()
+            except Exception:
+                pass
+
+            # Scan processes*.json in this instance's state dir
+            pattern = os.path.join(state_dir, "processes*.json")
+            for filepath in glob.glob(pattern):
+                basename = os.path.basename(filepath)
+                num = self._parse_instance_num(basename)
+                if num is None:
+                    continue
+                pid = self._read_launcher_pid(filepath)
+                is_self_dir = (os.path.normcase(os.path.normpath(inst_path))
+                               == os.path.normcase(os.path.normpath(my_instance_dir)))
+                is_self = (pid == my_pid and is_self_dir)
+                if not is_self and (not pid or not self._is_pid_alive(pid)):
+                    continue
+
+                # Read metadata
+                module_count = 0
+                debug = False
+                cwd = ""
+                try:
+                    with open(filepath, "r", encoding="utf-8") as f:
+                        data = json.load(f)
+                    if isinstance(data, dict):
+                        module_count = len(data.get("records", []))
+                        debug = data.get("debug", False)
+                        cwd = data.get("cwd", "")
+                    elif isinstance(data, list):
+                        module_count = len(data)
+                except Exception:
+                    pass
+
+                results.append({
+                    "instance_dir": entry,  # basename of instance dir
+                    "cwd": cwd or cwd_from_file,
+                    "num": num,
+                    "launcher_pid": pid or 0,
+                    "module_count": module_count,
+                    "is_self": is_self,
+                    "debug": debug,
+                })
+
+        results.sort(key=lambda x: (x["instance_dir"], x["num"]))
+        return results
+
+    def _allocate_instance_num(self) -> int:
+        """Scan processes*.json in data_dir and allocate the smallest available instance number.
+
+        An instance number is "available" if:
+        - No corresponding processes file exists, OR
+        - The file exists but its launcher_pid is dead (stale file from a crashed instance).
+        """
+        # Collect occupied instance numbers (whose launcher is still alive)
+        occupied: set[int] = set()
+        pattern = os.path.join(self.data_dir, "processes*.json")
+        for filepath in glob.glob(pattern):
+            basename = os.path.basename(filepath)
+            num = self._parse_instance_num(basename)
+            if num is None:
+                continue
+            pid = self._read_launcher_pid(filepath)
+            if pid and pid != os.getpid() and self._is_pid_alive(pid):
+                occupied.add(num)
+
+        # Allocate smallest available number starting from 1
+        num = 1
+        while num in occupied:
+            num += 1
+        return num
+
+    @staticmethod
+    def _parse_instance_num(basename: str) -> int | None:
+        """Extract instance number from a processes filename.
+
+        'processes.json' -> 1, 'processes~3.json' -> 3, other -> None.
+        """
+        if basename == "processes.json":
+            return 1
+        m = re.match(r"^processes~(\d+)\.json$", basename)
+        return int(m.group(1)) if m else None
+
+    @staticmethod
+    def _read_launcher_pid(filepath: str) -> int | None:
+        """Read launcher_pid from a processes JSON file. Returns None on error or old format."""
+        try:
+            with open(filepath, "r", encoding="utf-8") as f:
+                data = json.load(f)
+            if isinstance(data, dict):
+                return data.get("launcher_pid")
+        except Exception:
+            pass
+        return None
 
     # ── Leftover cleanup ──
 
-    def cleanup_leftovers(self):
-        """Read processes.json in current instance dir, kill leftovers from previous runs."""
-        if not os.path.isfile(self.records_path):
-            return
+    def cleanup_leftovers(self) -> dict[str, int]:
+        """Scan all processes*.json in data_dir (current instance dir only), clean up dead instances' child processes.
+
+        Returns dict: {instance_dir_basename: killed_count}.
+        """
+        my_pid = os.getpid()
+        inst_name = os.path.basename(os.environ.get("KITE_INSTANCE_DIR", "")) or "unknown"
+        killed = 0
+        pattern = os.path.join(self.data_dir, "processes*.json")
+        for filepath in glob.glob(pattern):
+            killed += self._cleanup_one_leftover_file(filepath, my_pid)
+        return {inst_name: killed} if killed else {}
+
+    def cleanup_global_leftovers(self) -> dict[str, int]:
+        """Scan all instance dirs under KITE_WORKSPACE for dead launcher files.
+
+        Same logic as cleanup_leftovers() but across all workspace instances.
+        Designed to run via run_in_executor (synchronous, blocking I/O).
+        Returns dict: {instance_dir_basename: killed_count}.
+        """
+        workspace = os.environ.get("KITE_WORKSPACE", "")
+        if not workspace or not os.path.isdir(workspace):
+            return {}
+
+        my_pid = os.getpid()
+        result: dict[str, int] = {}
 
         try:
-            with open(self.records_path, "r", encoding="utf-8") as f:
+            entries = os.listdir(workspace)
+        except OSError:
+            return {}
+
+        for entry in entries:
+            state_dir = os.path.join(workspace, entry, "launcher", "state")
+            if not os.path.isdir(state_dir):
+                continue
+            pattern = os.path.join(state_dir, "processes*.json")
+            for filepath in glob.glob(pattern):
+                killed = self._cleanup_one_leftover_file(filepath, my_pid)
+                if killed:
+                    result[entry] = result.get(entry, 0) + killed
+        return result
+
+    def _cleanup_one_leftover_file(self, filepath: str, my_pid: int) -> int:
+        """Process a single processes*.json file for leftover cleanup.
+
+        - Skip if launcher_pid matches my_pid (our own records).
+        - Skip if launcher_pid is still alive (parallel instance).
+        - Otherwise clean up child processes and delete the file.
+        Returns number of processes killed.
+        """
+        try:
+            with open(filepath, "r", encoding="utf-8") as f:
                 saved = json.load(f)
         except Exception:
             try:
-                os.remove(self.records_path)
+                os.remove(filepath)
             except OSError:
                 pass
-            return
+            return 0
+
+        # New format: {"launcher_pid": N, "records": [...]}
+        # Old format: plain array [...]
+        if isinstance(saved, dict):
+            launcher_pid = saved.get("launcher_pid")
+            records = saved.get("records", [])
+        elif isinstance(saved, list):
+            launcher_pid = None
+            records = saved
+        else:
+            try:
+                os.remove(filepath)
+            except OSError:
+                pass
+            return 0
+
+        # Skip our own records (written by current process)
+        if launcher_pid == my_pid:
+            return 0
 
-        for entry in saved:
+        # If the owning launcher is still alive, skip this file (parallel instance)
+        if launcher_pid and self._is_pid_alive(launcher_pid):
+            return 0
+
+        # Dead launcher (or old format) — clean up its child processes
+        killed = 0
+        for entry in records:
             pid = entry.get("pid", 0)
             cmd = entry.get("cmd", [])
             name = entry.get("name", "?")
             if not pid:
                 continue
             if self._is_pid_alive(pid) and self._cmd_matches(pid, cmd):
-                print(f"[launcher] 正在清理残留进程: {name} (PID {pid})")
+                print(f"[launcher] 正在清理残留进程: {name} (PID {pid}) [{os.path.basename(filepath)}]")
                 self._force_kill(pid)
+                killed += 1
             else:
                 print(f"[launcher] 残留进程 {name} (PID {pid}) 已不存在")
 
         try:
-            os.remove(self.records_path)
+            os.remove(filepath)
         except OSError:
             pass
+        return killed
 
     def _is_pid_alive(self, pid: int) -> bool:
-        """Check if a process with given PID exists."""
+        """Check if a process with given PID exists.
+
+        Windows: Uses OpenProcess API via ctypes (fast, ~0.1ms per check).
+        Linux/macOS: Uses os.kill(pid, 0) signal check.
+        """
         if IS_WINDOWS:
-            try:
-                result = subprocess.run(
-                    ["tasklist", "/FI", f"PID eq {pid}", "/NH"],
-                    capture_output=True, text=True, timeout=5,
-                )
-                return str(pid) in result.stdout
-            except Exception:
-                return False
+            return self._win_is_pid_alive_ctypes(pid)
         else:
             try:
                 os.kill(pid, 0)
@@ -103,6 +362,95 @@ class ProcessManager:
             except OSError:
                 return False
 
+    def _win_get_all_alive_pids(self) -> set:
+        """Windows: Get all alive PIDs in one batch call (fallback optimization).
+
+        Returns set of PIDs. Returns empty set on failure.
+        """
+        if not IS_WINDOWS:
+            return set()
+
+        try:
+            result = subprocess.run(
+                ["tasklist", "/FO", "CSV", "/NH"],
+                capture_output=True, timeout=10,
+            )
+            # Use errors='replace' to handle any encoding (GBK, UTF-8, etc.)
+            stdout = result.stdout.decode(errors='replace')
+            pids = set()
+            for line in stdout.strip().split('\n'):
+                if not line:
+                    continue
+                # CSV format: "image_name","pid","session_name","session#","mem_usage"
+                parts = line.split(',')
+                if len(parts) >= 2:
+                    try:
+                        pid_str = parts[1].strip('" \r')
+                        pids.add(int(pid_str))
+                    except (ValueError, IndexError):
+                        continue
+            return pids
+        except Exception:
+            return set()
+
+    def _win_is_pid_alive_ctypes(self, pid: int) -> bool:
+        """Windows: Check if PID exists and is still running using OpenProcess API.
+
+        Returns True only if process exists AND has not exited.
+        Handles zombie processes (exited but handle still open) correctly
+        by checking GetExitCodeProcess after OpenProcess succeeds.
+
+        Falls back to batch tasklist on ctypes failure (rare).
+        """
+        if not IS_WINDOWS:
+            return False
+
+        # Constants from Windows API
+        PROCESS_QUERY_LIMITED_INFORMATION = 0x1000
+        ERROR_INVALID_PARAMETER = 87
+        ERROR_ACCESS_DENIED = 5
+        STILL_ACTIVE = 259
+
+        try:
+            # Try to open the process with minimal permissions
+            handle = ctypes.windll.kernel32.OpenProcess(
+                PROCESS_QUERY_LIMITED_INFORMATION,
+                False,  # bInheritHandle
+                pid
+            )
+
+            if handle:
+                try:
+                    # OpenProcess succeeds even for zombie processes (exited but handle
+                    # not yet closed by parent). Check actual exit status.
+                    exit_code = wintypes.DWORD()
+                    if ctypes.windll.kernel32.GetExitCodeProcess(handle, ctypes.byref(exit_code)):
+                        return exit_code.value == STILL_ACTIVE
+                    # GetExitCodeProcess failed — assume alive to be safe
+                    return True
+                finally:
+                    ctypes.windll.kernel32.CloseHandle(handle)
+
+            # OpenProcess failed, check why
+            error_code = ctypes.windll.kernel32.GetLastError()
+
+            if error_code == ERROR_ACCESS_DENIED:
+                # Process exists but we don't have permission (e.g., SYSTEM process)
+                # For our use case, this means it's NOT our child process
+                return False
+            elif error_code == ERROR_INVALID_PARAMETER:
+                # Invalid PID (process doesn't exist)
+                return False
+            else:
+                # Other errors (e.g., process exited between check and open)
+                return False
+
+        except Exception:
+            # Fallback: batch tasklist (fetch once, cache for subsequent calls)
+            if not hasattr(self, '_alive_pids_cache'):
+                self._alive_pids_cache = self._win_get_all_alive_pids()
+            return pid in self._alive_pids_cache
+
     def _cmd_matches(self, pid: int, expected_cmd: list) -> bool:
         """Verify the process command line matches what we expect (prevent killing recycled PIDs).
         Uses multi-strategy fallback per platform for robustness.
@@ -150,9 +498,9 @@ class ProcessManager:
             r = subprocess.run(
                 ["powershell", "-NoProfile", "-Command",
                  f"(Get-CimInstance Win32_Process -Filter 'ProcessId={pid}').CommandLine"],
-                capture_output=True, text=True, timeout=8,
+                capture_output=True, timeout=8,
             )
-            return r.stdout.strip()
+            return r.stdout.decode(errors='replace').strip()
         except Exception:
             return ""
 
@@ -162,22 +510,31 @@ class ProcessManager:
             r = subprocess.run(
                 ["wmic", "process", "where", f"ProcessId={pid}",
                  "get", "CommandLine", "/value"],
-                capture_output=True, text=True, timeout=5,
+                capture_output=True, timeout=5,
             )
-            return r.stdout.strip()
+            return r.stdout.decode(errors='replace').strip()
         except Exception:
             return ""
 
     def _exe_name_matches(self, pid: int, expected_cmd: list) -> bool:
-        """Last resort: check if process image name matches expected executable."""
+        """Last resort: check if process image name matches expected executable.
+
+        Windows: Uses OpenProcess + QueryFullProcessImageName for fast check.
+        Falls back to tasklist if ctypes fails.
+        """
         exe = os.path.basename(expected_cmd[0]).lower()
         if exe.endswith(".exe"):
             exe = exe[:-4]
         if IS_WINDOWS:
+            # Try fast ctypes method first
+            image_name = self._win_get_image_name_ctypes(pid)
+            if image_name:
+                return exe in image_name.lower()
+            # Fallback to tasklist
             try:
                 r = subprocess.run(["tasklist", "/FI", f"PID eq {pid}", "/FO", "CSV", "/NH"],
-                                   capture_output=True, text=True, timeout=5)
-                return exe in r.stdout.lower()
+                                   capture_output=True, timeout=5)
+                return exe in r.stdout.decode(errors='replace').lower()
             except Exception:
                 return False
         else:
@@ -188,6 +545,50 @@ class ProcessManager:
             except Exception:
                 return False
 
+    def _win_get_image_name_ctypes(self, pid: int) -> str:
+        """Windows: Get process image name using OpenProcess + QueryFullProcessImageName.
+
+        Returns empty string on failure.
+        """
+        if not IS_WINDOWS:
+            return ""
+
+        PROCESS_QUERY_LIMITED_INFORMATION = 0x1000
+
+        try:
+            handle = ctypes.windll.kernel32.OpenProcess(
+                PROCESS_QUERY_LIMITED_INFORMATION,
+                False,
+                pid
+            )
+
+            if not handle:
+                return ""
+
+            try:
+                # Buffer for image name (MAX_PATH = 260)
+                buffer_size = wintypes.DWORD(260)
+                buffer = ctypes.create_unicode_buffer(buffer_size.value)
+
+                # QueryFullProcessImageNameW
+                success = ctypes.windll.kernel32.QueryFullProcessImageNameW(
+                    handle,
+                    0,  # dwFlags (0 = Win32 path format)
+                    buffer,
+                    ctypes.byref(buffer_size)
+                )
+
+                if success:
+                    return os.path.basename(buffer.value)
+                else:
+                    return ""
+
+            finally:
+                ctypes.windll.kernel32.CloseHandle(handle)
+
+        except Exception:
+            return ""
+
     def _force_kill(self, pid: int):
         """Force kill a process by PID."""
         try:
@@ -296,8 +697,10 @@ class ProcessManager:
             return False
 
     def close_stdio(self, name: str):
-        """Close stdin and stdout for a module. Called by Launcher after module.ready → module.started.
-        stdout reader thread will exit naturally when stdout is closed.
+        """Close stdin for a module. Called by Launcher after module.ready → module.started.
+        Only closes stdin (no more messages to send). stdout is left open so the reader
+        thread can continue relaying logs — it exits naturally when the process dies.
+        On Windows, closing stdout while a reader thread is blocked on readline() can deadlock.
         """
         proc = self._processes.get(name)
         if not proc:
@@ -307,11 +710,6 @@ class ProcessManager:
                 proc.stdin.close()
             except Exception:
                 pass
-        if proc.stdout and not proc.stdout.closed:
-            try:
-                proc.stdout.close()
-            except Exception:
-                pass
 
     def _build_cmd(self, info) -> list:
         """Build the command list based on module runtime.
@@ -398,11 +796,18 @@ class ProcessManager:
         if not self._processes:
             return
 
-        names = list(self._processes.keys())
-        print(f"[launcher] 正在停止所有模块: {', '.join(names)}")
+        # Filter to only still-running processes
+        alive = [n for n, p in self._processes.items() if p.poll() is None]
+        if not alive:
+            # All already exited — just clean up bookkeeping
+            self._processes.clear()
+            self._records.clear()
+            return
+
+        print(f"[launcher] 正在停止所有模块: {', '.join(alive)}")
 
         # Phase 1: close stdio and send terminate to all
-        for name in names:
+        for name in alive:
             self.close_stdio(name)
             proc = self._processes.get(name)
             if proc and proc.poll() is None:
@@ -410,7 +815,7 @@ class ProcessManager:
 
         # Phase 2: wait for all to exit
         deadline = time.time() + timeout
-        for name in names:
+        for name in alive:
             proc = self._processes.get(name)
             if not proc:
                 continue
@@ -421,7 +826,7 @@ class ProcessManager:
                 pass
 
         # Phase 3: force kill survivors
-        for name in names:
+        for name in alive:
             proc = self._processes.get(name)
             if proc and proc.poll() is None:
                 print(f"[launcher] 强制终止 {name} (PID {proc.pid})")
@@ -438,8 +843,13 @@ class ProcessManager:
     # ── Persistence & monitoring ──
 
     def persist_records(self):
-        """Write current process records to processes.json."""
-        data = [asdict(r) for r in self._records.values()]
+        """Write current process records to processes.json (with launcher_pid for multi-instance)."""
+        data = {
+            "launcher_pid": os.getpid(),
+            "debug": os.environ.get("KITE_DEBUG") == "1",
+            "cwd": os.environ.get("KITE_CWD", ""),
+            "records": [asdict(r) for r in self._records.values()],
+        }
         self._write_records_file(data)
 
     def check_exited(self) -> list[tuple[str, int]]:
@@ -461,10 +871,10 @@ class ProcessManager:
         proc = self._processes.get(name)
         return proc is not None and proc.poll() is None
 
-    def _write_records_file(self, data: list):
-        """Write data list to processes.json."""
+    def _write_records_file(self, data):
+        """Write data to the instance's processes JSON file."""
         try:
             with open(self.records_path, "w", encoding="utf-8") as f:
                 json.dump(data, f, indent=2)
         except Exception as e:
-            print(f"[launcher] 警告: 写入 processes.json 失败: {e}")
+            print(f"[launcher] 警告: 写入 {os.path.basename(self.records_path)} 失败: {e}")