pipane 0.1.5 → 0.1.7

package/dist/server/server/process-pool.js

@@ -1,29 +1,43 @@
 /**
- * CWD-aware process pool for pi RPC processes.
+ * CWD-aware, leased process pool for pi RPC processes.
  *
- * Processes are grouped by cwd. When acquiring a process for a session,
- * the caller provides the session's cwd and gets a process that was
- * spawned in that directory. This ensures bash/read/edit/write tools
- * operate in the correct project directory.
- *
- * Features:
- * - Lazy spawning per-cwd with configurable pre-warming
- * - Readiness check via get_state RPC (replaces setTimeout(500) hack)
- * - Dead process cleanup on exit
- * - Global process cap
+ * The pool is the sole owner of process availability. Callers receive an
+ * idempotent lease instead of passing a separate busy set, so a process is
+ * reserved atomically before control returns to asynchronous application code.
  */
 import { spawn } from "node:child_process";
-import * as readline from "node:readline";
 import { existsSync } from "node:fs";
+import { attachStrictJsonlReader, decodePiRpcLine, encodePiRpcCommand, } from "./pi-rpc-protocol.js";
+/** Exclusive ownership of one pooled process. */
+export class RpcProcessLease {
+    process;
+    pool;
+    _released = false;
+    constructor(pool, process) {
+        this.pool = pool;
+        this.process = process;
+    }
+    get released() { return this._released; }
+    release() {
+        this.pool.releaseLease(this);
+    }
+    /** @internal */
+    markReleased() {
+        this._released = true;
+    }
+}
 export class ProcessPool {
-    /** All live processes, keyed by cwd */
     pools = new Map();
+    leases = new Map();
+    decommissioning = new Set();
+    finalizedProcesses = new WeakSet();
     nextProcId = 0;
     spawnConfig;
     maxProcesses;
     prewarmCount;
     rpcTimeout;
     onProcessExit;
+    eventListeners = new Set();
     constructor(spawnConfig, options) {
         this.spawnConfig = spawnConfig;
         this.maxProcesses = options?.maxProcesses ?? 6;
@@ -31,40 +45,50 @@ export class ProcessPool {
         this.rpcTimeout = options?.rpcTimeout ?? 30000;
         this.onProcessExit = options?.onProcessExit;
     }
-    /** Get the total count of live processes across all cwds */
     get totalProcesses() {
         let count = 0;
         for (const procs of this.pools.values()) {
-            count += procs.filter((p) => p.process.exitCode === null).length;
+            count += procs.filter((proc) => proc.process.exitCode === null).length;
         }
         return count;
     }
-    /** Get all processes (for debug endpoint) */
+    subscribeEvents(listener) {
+        this.eventListeners.add(listener);
+        return () => this.eventListeners.delete(listener);
+    }
+    emitEvent(proc, event) {
+        for (const listener of this.eventListeners) {
+            try {
+                listener(proc, event);
+            }
+            catch (err) {
+                console.error(`[pool] pi#${proc.id} event listener failed:`, err);
+            }
+        }
+    }
     getAllProcesses() {
         const all = [];
-        for (const procs of this.pools.values()) {
+        for (const procs of this.pools.values())
             all.push(...procs);
-        }
         return all;
     }
-    /** Return recent stderr lines for a process (tail). */
+    isLeased(proc) {
+        return this.leases.has(proc);
+    }
+    isDecommissioning(proc) {
+        return this.decommissioning.has(proc);
+    }
     getRecentStderr(proc, maxLines = 12) {
         if (!proc?.recentStderr?.length)
             return [];
         return proc.recentStderr.slice(Math.max(0, proc.recentStderr.length - maxLines));
     }
-    /**
-     * Spawn a new RPC process for the given cwd.
-     * Does not wait for readiness — call waitForReady() separately.
-     */
     spawn(cwd) {
         if (!existsSync(cwd)) {
             throw new Error(`Cannot spawn pi process: directory does not exist: ${cwd}`);
         }
         const procId = ++this.nextProcId;
         console.log(`[pool] Spawning pi process #${procId} (cwd: ${cwd})...`);
-        // Strip NODE_ENV from the child environment so tools spawned by pi
-        // (e.g. bash) don't inherit pipane's "production" setting.
         const { NODE_ENV: _, ...parentEnv } = process.env;
         const baseArgs = typeof this.spawnConfig.baseArgs === "function"
             ? this.spawnConfig.baseArgs()
@@ -90,71 +114,46 @@ export class ProcessPool {
                     recentStderr.splice(0, recentStderr.length - 60);
             }
         });
-        const rl = readline.createInterface({ input: child.stdout, terminal: false });
         const proc = {
             id: procId,
             cwd,
             process: child,
-            rl,
+            stopReadingStdout: () => { },
             pendingRequests: new Map(),
             requestId: 0,
             lastResponseTime: Date.now(),
             recentStderr,
         };
-        // Set up response handler
-        rl.on("line", (line) => {
-            let data;
-            try {
-                data = JSON.parse(line);
-            }
-            catch {
+        proc.stopReadingStdout = attachStrictJsonlReader(child.stdout, (line) => {
+            const decoded = decodePiRpcLine(line);
+            if (!decoded.ok) {
+                console.warn(`[pool] pi#${proc.id} ignored invalid RPC output: ${decoded.error.message}`);
                 return;
             }
-            if (data.type === "response" && data.id && proc.pendingRequests.has(data.id)) {
-                const pending = proc.pendingRequests.get(data.id);
-                proc.pendingRequests.delete(data.id);
+            const message = decoded.value;
+            if (message.type === "response") {
+                const pending = message.id ? proc.pendingRequests.get(message.id) : undefined;
+                if (!pending || !message.id)
+                    return;
+                proc.pendingRequests.delete(message.id);
                 proc.lastResponseTime = Date.now();
-                pending.resolve(data);
+                if (message.command !== pending.command) {
+                    pending.reject(new Error(`Mismatched RPC response for ${pending.command}: received ${message.command}`));
+                    return;
+                }
+                pending.resolve(message);
+                return;
             }
+            this.emitEvent(proc, message);
         });
         child.on("error", (err) => {
             console.error(`[pool] pi#${proc.id} spawn error: ${err.message}`);
-            // Remove from pool
-            const poolForCwd = this.pools.get(cwd);
-            if (poolForCwd) {
-                const idx = poolForCwd.indexOf(proc);
-                if (idx !== -1)
-                    poolForCwd.splice(idx, 1);
-                if (poolForCwd.length === 0)
-                    this.pools.delete(cwd);
-            }
-            // Reject any pending requests
-            for (const [, pending] of proc.pendingRequests) {
-                pending.reject(new Error(`pi process #${proc.id} failed to spawn: ${err.message}`));
-            }
-            proc.pendingRequests.clear();
-            this.onProcessExit?.(proc);
+            this.finalizeProcess(proc, new Error(`pi process #${proc.id} failed to spawn: ${err.message}`));
         });
         child.on("exit", (code) => {
             console.log(`[pool] pi#${proc.id} exited (code ${code})`);
-            // Remove from pool
-            const poolForCwd = this.pools.get(cwd);
-            if (poolForCwd) {
-                const idx = poolForCwd.indexOf(proc);
-                if (idx !== -1)
-                    poolForCwd.splice(idx, 1);
-                if (poolForCwd.length === 0)
-                    this.pools.delete(cwd);
-            }
-            // Reject any pending requests
-            for (const [, pending] of proc.pendingRequests) {
-                pending.reject(new Error(`pi process #${proc.id} exited unexpectedly (code ${code})`));
-            }
-            proc.pendingRequests.clear();
-            // Notify caller
-            this.onProcessExit?.(proc);
+            this.finalizeProcess(proc, new Error(`pi process #${proc.id} exited unexpectedly (code ${code})`));
         });
-        // Add to pool
         let poolForCwd = this.pools.get(cwd);
         if (!poolForCwd) {
             poolForCwd = [];
@@ -164,10 +163,29 @@ export class ProcessPool {
         console.log(`[pool] pi#${procId} spawned (total: ${this.totalProcesses})`);
         return proc;
     }
-    /**
-     * Wait for a process to be ready by sending a get_state RPC.
-     * Falls back to a timeout if the process doesn't respond.
-     */
+    finalizeProcess(proc, error) {
+        if (this.finalizedProcesses.has(proc))
+            return;
+        this.finalizedProcesses.add(proc);
+        const poolForCwd = this.pools.get(proc.cwd);
+        if (poolForCwd) {
+            const index = poolForCwd.indexOf(proc);
+            if (index !== -1)
+                poolForCwd.splice(index, 1);
+            if (poolForCwd.length === 0)
+                this.pools.delete(proc.cwd);
+        }
+        const lease = this.leases.get(proc);
+        if (lease)
+            lease.markReleased();
+        this.leases.delete(proc);
+        this.decommissioning.delete(proc);
+        proc.stopReadingStdout();
+        for (const pending of proc.pendingRequests.values())
+            pending.reject(error);
+        proc.pendingRequests.clear();
+        this.onProcessExit?.(proc);
+    }
     async waitForReady(proc, timeoutMs = 5000) {
         try {
             await this.sendRpc(proc, { type: "get_state" }, timeoutMs);
@@ -177,116 +195,141 @@ export class ProcessPool {
             return false;
         }
     }
-    /**
-     * Get an idle process for the given cwd, or spawn one.
-     * An "idle" process is one that's alive and not in the busySet.
-     */
-    acquire(cwd, busySet) {
+    /** Atomically reserve a matching process, spawning one when capacity allows. */
+    acquire(cwd) {
         const poolForCwd = this.pools.get(cwd);
-        if (poolForCwd) {
-            const idle = poolForCwd.find((p) => p.process.exitCode === null && !busySet.has(p));
-            if (idle)
-                return idle;
-        }
-        // Check global cap
-        if (this.totalProcesses >= this.maxProcesses) {
+        const idle = poolForCwd?.find((proc) => this.isAvailable(proc));
+        if (idle)
+            return this.reserve(idle);
+        if (this.totalProcesses >= this.maxProcesses)
             return null;
-        }
-        return this.spawn(cwd);
+        return this.reserve(this.spawn(cwd));
     }
-    /**
-     * Get any live process (idle preferred). For model queries that don't
-     * need a specific cwd.
-     */
-    getAny(busySet) {
-        // Prefer idle
+    /** Reserve any idle live process, regardless of cwd. */
+    acquireAny() {
         for (const procs of this.pools.values()) {
-            const idle = procs.find((p) => p.process.exitCode === null && !busySet.has(p));
+            const idle = procs.find((proc) => this.isAvailable(proc));
             if (idle)
-                return idle;
-        }
-        // Fall back to any live process
-        for (const procs of this.pools.values()) {
-            const live = procs.find((p) => p.process.exitCode === null);
-            if (live)
-                return live;
+                return this.reserve(idle);
         }
         return null;
     }
-    /**
-     * Evict one idle process from a different cwd to free capacity.
-     * Returns the evicted process, or null if none can be evicted.
-     */
-    evictIdleDifferentCwd(targetCwd, busySet) {
+    isAvailable(proc) {
+        return proc.process.exitCode === null
+            && !this.leases.has(proc)
+            && !this.decommissioning.has(proc);
+    }
+    reserve(proc) {
+        if (!this.isAvailable(proc))
+            throw new Error(`pi process #${proc.id} is not available`);
+        const lease = new RpcProcessLease(this, proc);
+        this.leases.set(proc, lease);
+        return lease;
+    }
+    /** @internal Called only by RpcProcessLease.release(). */
+    releaseLease(lease) {
+        if (lease.released)
+            return;
+        const proc = lease.process;
+        if (this.leases.get(proc) !== lease) {
+            lease.markReleased();
+            return;
+        }
+        this.leases.delete(proc);
+        lease.markReleased();
+        if (this.decommissioning.has(proc) && proc.process.exitCode === null && !proc.process.killed) {
+            // Keep the marker until the exit event removes the process. Otherwise a
+            // concurrent acquisition could lease the dying process before SIGTERM lands.
+            console.log(`[pool] Decommissioning pi#${proc.id} after completed operation`);
+            proc.process.kill("SIGTERM");
+        }
+    }
+    decommission(proc) {
+        if (proc.process.exitCode !== null || this.decommissioning.has(proc))
+            return "ignored";
+        this.decommissioning.add(proc);
+        if (this.leases.has(proc))
+            return "draining";
+        proc.process.kill("SIGTERM");
+        return "killed";
+    }
+    /** Evict an idle process from a different cwd to free capacity. */
+    evictIdleDifferentCwd(targetCwd) {
         for (const [cwd, procs] of this.pools) {
             if (cwd === targetCwd)
                 continue;
-            const victim = procs.find((p) => p.process.exitCode === null && !busySet.has(p));
+            const victim = procs.find((proc) => this.isAvailable(proc));
             if (!victim)
                 continue;
             console.log(`[pool] Evicting idle pi#${victim.id} from cwd ${cwd} to make room for ${targetCwd}`);
-            victim.process.kill("SIGTERM");
+            this.decommission(victim);
             return victim;
         }
         return null;
     }
-    /**
-     * Pre-warm the pool with processes for the given cwd.
-     */
-    /**
-     * Pre-warm the pool with processes for the given cwd.
-     * Spawns are staggered: each process must be ready (via get_state RPC)
-     * before the next one is spawned. This prevents lock contention on
-     * shared resources like auth.json during startup.
-     */
+    /** Immediately kill one process while keeping it unavailable until exit. */
+    forceKill(proc) {
+        if (proc.process.exitCode !== null)
+            return false;
+        this.decommissioning.add(proc);
+        return proc.process.kill("SIGKILL");
+    }
+    /** Kill idle processes now and mark leased processes for kill on release. */
+    decommissionAll() {
+        let killed = 0;
+        let draining = 0;
+        for (const proc of this.getAllProcesses()) {
+            const result = this.decommission(proc);
+            if (result === "killed")
+                killed += 1;
+            if (result === "draining")
+                draining += 1;
+        }
+        return { killed, draining };
+    }
     async prewarm(cwd) {
-        const existing = this.pools.get(cwd)?.filter((p) => p.process.exitCode === null).length ?? 0;
+        const existing = this.pools.get(cwd)?.filter((proc) => proc.process.exitCode === null).length ?? 0;
         const needed = Math.min(this.prewarmCount, this.maxProcesses) - existing;
         for (let i = 0; i < needed; i++) {
             if (this.totalProcesses >= this.maxProcesses)
                 break;
             const proc = this.spawn(cwd);
-            if (i < needed - 1) {
-                // Wait for this process to be ready before spawning the next one
+            if (i < needed - 1)
                 await this.waitForReady(proc);
-            }
         }
     }
-    /**
-     * Send an RPC command to a process and wait for a response.
-     */
     sendRpc(proc, command, timeoutMs) {
         if (!proc.process || proc.process.exitCode !== null) {
             return Promise.reject(new Error("RPC process is dead"));
         }
         const timeout = timeoutMs ?? this.rpcTimeout;
         const id = `req_${++proc.requestId}`;
-        const fullCommand = { ...command, id };
         return new Promise((resolve, reject) => {
             const timer = setTimeout(() => {
                 proc.pendingRequests.delete(id);
+                // A timed-out RPC may still be executing. Never return that process to
+                // the idle pool where a new session command could overlap with it.
+                this.decommission(proc);
                 reject(new Error(`Timeout waiting for RPC response to ${command.type}`));
             }, timeout);
             proc.pendingRequests.set(id, {
-                resolve: (data) => {
+                command: command.type,
+                resolve: (response) => {
                     clearTimeout(timer);
-                    resolve(data);
+                    resolve(response);
                 },
-                reject: (err) => {
+                reject: (error) => {
                     clearTimeout(timer);
-                    reject(err);
+                    reject(error);
                 },
             });
-            proc.process.stdin.write(JSON.stringify(fullCommand) + "\n");
+            proc.process.stdin.write(encodePiRpcCommand(command, id));
         });
     }
-    /**
-     * Send an RPC command and throw if it fails.
-     */
     async sendRpcChecked(proc, command) {
         const response = await this.sendRpc(proc, command);
-        if (!response?.success) {
-            throw new Error(response?.error || `RPC command failed: ${command.type}`);
+        if (!response.success) {
+            throw new Error(response.error || `RPC command failed: ${command.type}`);
         }
         return response;
     }