@yemi33/minions 0.1.2372 → 0.1.2373

package/engine/agent-worker-pool.js

@@ -0,0 +1,367 @@
+/**
+ * engine/agent-worker-pool.js — Fixed-size ACP worker pool for fleet agent
+ * dispatches (P-4c6e8a72, W-mre7fhxz001ka5fa / P-9b3d5f61 plan).
+ *
+ * Modeled on engine/cc-worker-pool.js but with a DIFFERENT ownership model:
+ * cc-worker-pool.js gives each CC/doc-chat tab a permanently-owned worker
+ * (tabId → Worker, kept warm for the tab's lifetime); this pool instead
+ * leases a worker to exactly one active fleet dispatch at a time and
+ * returns it to a shared free set on release — 1 worker : 1 ACTIVE dispatch,
+ * not 1 worker : 1 dispatch for life. Both pools share the same underlying
+ * spawn + JSON-RPC transport (engine/acp-transport.js) — no duplicated
+ * framing code.
+ *
+ * Public API:
+ *
+ *   pool.acquireWorker({ dispatchId, cwd, model, effort, mcpServers, hermeticDirs })
+ *     → Promise<WorkerHandle>. Blocks (FIFO queue) when no free/matching
+ *       worker exists and the pool is already at its configured size.
+ *       WorkerHandle: {
+ *         dispatchId,
+ *         stream(promptText, opts) → Promise,  // ACP session/prompt turn
+ *         cancel(),                            // ACP session/cancel
+ *         release(),                           // return the worker to the
+ *                                               // free set (or evict it, if
+ *                                               // it died during the lease)
+ *       }
+ *   pool.setPoolSize(n) / getPoolSize()
+ *   pool.shutdown()                  // close every worker, reject the queue
+ *
+ * Assignment strategy (acquireWorker), per the P-7a1c2e4f spike's findings:
+ *
+ *   1. A FREE worker whose mcpServersHash already matches the request wins
+ *      first — reused via a fresh `session/new` (worker.newSession), no
+ *      respawn, no wasted initialize handshake.
+ *   2. Otherwise any OTHER free worker is evicted (closed) and a replacement
+ *      is spawned fresh with the requested mcpServers/cwd — MCP servers are
+ *      resolved at `copilot --acp` process boot (via $COPILOT_HOME), so a
+ *      changed local/stdio server set requires a full respawn, not just a
+ *      new session (spike (c)).
+ *   3. Otherwise, if the pool has room (fewer than `poolSize` workers
+ *      allocated across free+busy+in-flight-spawn), spawn a brand new
+ *      worker.
+ *   4. Otherwise the request queues in FIFO order until a `release()` or a
+ *      worker crash frees a slot.
+ *
+ * A worker that exits its OS process after the ACP handshake (WORKER_DIED)
+ * is evicted immediately — from wherever it currently lives (free set or an
+ * active lease) — and is NEVER returned to the free set; the freed slot
+ * immediately unblocks the next queued acquireWorker (pump re-runs on every
+ * eviction and every release).
+ *
+ * `hermeticDirs` (workspace-manifest `allowed_repos` beyond the primary
+ * `cwd`) has no `session/new` parameter (spike (b) confirmed this) — it is
+ * applied by sending the ACP `/add-dir <path>` runtime slash command as a
+ * `session/prompt` turn for each dir, immediately after the session is
+ * established and before the handle is returned to the caller.
+ *
+ * No new npm deps — Node built-ins only.
+ */
+
+const transport = require('./acp-transport');
+const { Worker, hashMcpServers } = transport;
+
+// Test seam — mirrors engine/cc-worker-pool.js's `_internals` pattern, but is
+// its OWN object (constructor-injected into Worker per acquisition) so this
+// pool's test stubs never collide with cc-worker-pool.js's.
+const _internals = {
+  spawnAcp: transport.spawnAcp,
+  killImmediate: transport.killImmediate,
+  now() { return Date.now(); },
+};
+
+// P-9b3d5f61 — falls back to resolveAgentAcpPoolSize(engine)'s own default
+// chain (engine.agentAcpPoolSize > engine.maxConcurrent >
+// ENGINE_DEFAULTS.maxConcurrent) once the spawn-routing item (P-1d8f0b93)
+// wires `setPoolSize` in on config load. ENGINE_DEFAULTS.maxConcurrent is 5
+// (engine/shared.js) — mirrored here as the pre-wiring default so a caller
+// exercising this module directly (e.g. tests) gets a sane cap.
+const DEFAULT_POOL_SIZE = 5;
+let _poolSize = DEFAULT_POOL_SIZE;
+
+function setPoolSize(n) {
+  const parsed = Number(n);
+  if (!Number.isFinite(parsed) || parsed <= 0) return _poolSize;
+  _poolSize = Math.floor(parsed);
+  _pump();
+  return _poolSize;
+}
+
+function getPoolSize() {
+  return _poolSize;
+}
+
+// AGENT_POOL_TRACE-gated structured trace logger — mirrors cc-worker-pool.js's
+// CC_POOL_TRACE convention but under its own env var so the two pools can be
+// traced independently. NO PII — dispatchId (caller-supplied), sessionIds
+// (opaque ACP ids), and protocol flags only.
+function _trace(...parts) {
+  if (!process.env.AGENT_POOL_TRACE) return;
+  try { process.stderr.write('[agent-pool] ' + parts.join(' ') + '\n'); }
+  catch { /* swallow telemetry errors */ }
+}
+
+let _nextWorkerId = 1;
+
+const _free = [];              // Worker[] — idle, unleased
+const _busy = new Map();       // dispatchId → Worker — currently leased
+const _queue = [];             // pending acquire requests, FIFO
+// Slots reserved for a worker that's mid-reuse/respawn/fresh-spawn — pulled
+// out of `_free` (or not yet spawned) but not yet placed in `_busy`. Counted
+// toward `_poolSize` so a burst of concurrent acquireWorker calls can't
+// overshoot the configured pool size while spawns are still in flight.
+let _reserved = 0;
+
+function _occupied() {
+  return _free.length + _busy.size + _reserved;
+}
+
+// ── Queue pump ──────────────────────────────────────────────────────────
+//
+// Fully synchronous per pass: each iteration either satisfies the head of
+// the queue by reserving capacity synchronously (splice out of `_free`, or
+// bump `_reserved`) and kicking off the (async, fire-and-forget from here)
+// reuse/respawn/spawn work, or — if the queue's head can't be satisfied yet
+// — stops. The async work calls `_pump()` again itself on completion
+// (success moves the reservation into `_busy` and the freed capacity is a
+// no-op re-pump; failure frees the reservation and DOES let another queued
+// request through).
+function _pump() {
+  while (_queue.length > 0) {
+    const req = _queue[0];
+
+    const matchIdx = _free.findIndex((w) => w.mcpServersHash === req.mcpServersHash);
+    if (matchIdx >= 0) {
+      _queue.shift();
+      const worker = _free.splice(matchIdx, 1)[0];
+      _reserved++;
+      _reuseWorker(worker, req);
+      continue;
+    }
+
+    if (_free.length > 0) {
+      _queue.shift();
+      const stale = _free.shift();
+      _reserved++;
+      _respawnWorker(stale, req);
+      continue;
+    }
+
+    if (_occupied() < _poolSize) {
+      _queue.shift();
+      _reserved++;
+      _spawnFreshWorker(req);
+      continue;
+    }
+
+    // Pool is full and no free worker exists — leave req at the head of the
+    // queue; a release() or a worker-death eviction will re-pump.
+    break;
+  }
+}
+
+async function _reuseWorker(worker, req) {
+  try {
+    await worker.newSession({ mcpServers: req.mcpServers, model: req.model, effort: req.effort, cwd: req.cwd });
+    _trace(`dispatch=${req.dispatchId} reuse: worker=${worker.id} sessionId=${worker.sessionId}`);
+    _reserved--;
+    req.resolve(worker);
+  } catch (err) {
+    _reserved--;
+    // The worker died mid-newSession (or the daemon rejected the call) —
+    // it's already unusable; don't return it to `_free`.
+    try { worker.close(); } catch { /* already torn down */ }
+    req.reject(err);
+    _pump();
+  }
+}
+
+async function _respawnWorker(stale, req) {
+  try { stale.close(); } catch { /* already torn down */ }
+  try {
+    const worker = await _bootWorker(req);
+    _reserved--;
+    req.resolve(worker);
+  } catch (err) {
+    _reserved--;
+    req.reject(err);
+    _pump();
+  }
+}
+
+async function _spawnFreshWorker(req) {
+  try {
+    const worker = await _bootWorker(req);
+    _reserved--;
+    req.resolve(worker);
+  } catch (err) {
+    _reserved--;
+    req.reject(err);
+    _pump();
+  }
+}
+
+// Spawn + initialize + session/new a brand-new Worker for `req`. Attaches
+// the pool's own proc-exit eviction listener once the handshake succeeds —
+// Worker's own internal exit handler (engine/acp-transport.js) already marks
+// `.killed` and fails any pending/inflight calls; this listener runs after
+// it (registered later on the same EventEmitter) and removes the worker from
+// whichever pool collection it's currently sitting in, then re-pumps so a
+// queued request can spawn its replacement.
+async function _bootWorker(req) {
+  const worker = new Worker({
+    id: `agent-pool-${_nextWorkerId++}`,
+    model: req.model,
+    effort: req.effort,
+    mcpServers: req.mcpServers,
+    mcpServersHash: req.mcpServersHash,
+    cwd: req.cwd,
+    internals: _internals,
+    trace: _trace,
+  });
+  const initPromise = worker._spawnAndInit();
+  worker.initPromise = initPromise;
+  try {
+    await initPromise;
+  } catch (err) {
+    try { worker.close(); } catch { /* already torn down */ }
+    throw err;
+  } finally {
+    worker.initPromise = null;
+  }
+  if (worker.proc) {
+    worker.proc.once('exit', () => _onWorkerExit(worker));
+  }
+  _trace(`dispatch=${req.dispatchId} spawn: worker=${worker.id} sessionId=${worker.sessionId}`);
+  return worker;
+}
+
+function _onWorkerExit(worker) {
+  let removed = false;
+  const freeIdx = _free.indexOf(worker);
+  if (freeIdx >= 0) {
+    _free.splice(freeIdx, 1);
+    removed = true;
+  }
+  for (const [dispatchId, w] of _busy) {
+    if (w === worker) {
+      _busy.delete(dispatchId);
+      removed = true;
+      break;
+    }
+  }
+  if (removed) _trace(`worker=${worker.id} exited — evicted from pool`);
+  // A slot may now be free (whether or not this worker was tracked in
+  // `_free`/`_busy` — e.g. it may have died mid-reuse/respawn, in which case
+  // the reserving call's own catch already decremented `_reserved` and
+  // re-pumped). Re-pumping here is always safe/idempotent.
+  _pump();
+}
+
+// Apply workspace-manifest `allowed_repos` beyond `cwd` (P-4c6e8a72 /
+// P-7a1c2e4f spike (b)) — there is no `session/new` param for this; the only
+// mechanism is the ACP `/add-dir <path>` runtime slash command sent as an
+// ordinary `session/prompt` turn. Best-effort per dir: a failure surfaces via
+// `_trace` but does not fail the whole acquisition — the dispatch still gets
+// a working session scoped to `cwd`, just without the extra directory.
+async function _applyHermeticDirs(worker, hermeticDirs) {
+  if (!Array.isArray(hermeticDirs) || hermeticDirs.length === 0) return;
+  for (const dir of hermeticDirs) {
+    if (!dir) continue;
+    await new Promise((resolve) => {
+      worker.stream(`/add-dir ${dir}`, {
+        onDone: () => { _trace(`worker=${worker.id} add-dir ok: ${dir}`); resolve(); },
+        onError: (err) => { _trace(`worker=${worker.id} add-dir failed: ${dir} (${err && err.message})`); resolve(); },
+      }).catch(() => resolve());
+    });
+  }
+}
+
+// ── Public API ────────────────────────────────────────────────────────────
+
+async function acquireWorker({ dispatchId, cwd, model, effort, mcpServers, hermeticDirs } = {}) {
+  if (!dispatchId) throw new Error('agent-worker-pool.acquireWorker: dispatchId is required');
+  if (_busy.has(dispatchId)) {
+    throw new Error(
+      `agent-worker-pool.acquireWorker: dispatchId ${dispatchId} already has an active worker ` +
+      '(1 worker : 1 active dispatch invariant — call release() before acquiring again)'
+    );
+  }
+
+  const mcpServersHash = hashMcpServers(mcpServers);
+  const req = { dispatchId, cwd, model, effort, mcpServers, mcpServersHash };
+
+  const worker = await new Promise((resolve, reject) => {
+    req.resolve = resolve;
+    req.reject = reject;
+    _queue.push(req);
+    _pump();
+  });
+
+  await _applyHermeticDirs(worker, hermeticDirs);
+  _busy.set(dispatchId, worker);
+
+  return {
+    dispatchId,
+    // P-1d8f0b93 — exposed so engine.js's PooledAgentProcess facade (the
+    // spawnAgent integration) can write the same PID-file/log identity a
+    // cold-spawned child_process would, and stamp the ACP sessionId into the
+    // synthesized `result` event for parseOutput parity. Both are read-time
+    // snapshots of the underlying worker's current values (a getter, not a
+    // captured copy), since sessionId can rotate under a reused worker.
+    get pid() { return (worker.proc && worker.proc.pid) || null; },
+    get sessionId() { return worker.sessionId || null; },
+    stream: (promptText, opts) => worker.stream(promptText, opts),
+    cancel: () => worker.cancel(),
+    release: () => release(dispatchId),
+  };
+}
+
+// Return a leased worker to the free set — unless it died during the lease
+// (WORKER_DIED), in which case it's simply dropped (already evicted from
+// `_busy` by `_onWorkerExit` if the crash happened before release(), or
+// dropped here if release() is the first to notice `.killed`). Returns
+// false if `dispatchId` has no active lease (already released/evicted or
+// never acquired) — safe to call defensively.
+function release(dispatchId) {
+  const worker = _busy.get(dispatchId);
+  if (!worker) return false;
+  _busy.delete(dispatchId);
+  if (worker.killed) {
+    _trace(`dispatch=${dispatchId} release: worker=${worker.id} was dead — evicted, not returned to free set`);
+    _pump();
+    return true;
+  }
+  _free.push(worker);
+  _trace(`dispatch=${dispatchId} release: worker=${worker.id} returned to free set`);
+  _pump();
+  return true;
+}
+
+function shutdown() {
+  for (const req of _queue.splice(0)) {
+    try { req.reject(new Error('agent-worker-pool: shutdown')); } catch { /* swallow */ }
+  }
+  for (const worker of _free.splice(0)) {
+    try { worker.close(); } catch { /* swallow */ }
+  }
+  for (const [, worker] of _busy) {
+    try { worker.close(); } catch { /* swallow */ }
+  }
+  _busy.clear();
+  _reserved = 0;
+}
+
+module.exports = {
+  acquireWorker,
+  release,
+  setPoolSize,
+  getPoolSize,
+  shutdown,
+  DEFAULT_POOL_SIZE,
+  // Exposed for unit tests; engine code MUST go through the public API.
+  _internals,
+  _free,
+  _busy,
+  _queue,
+};