mixdog 0.9.85 → 0.9.86

package/src/standalone/agent-shard/shard-pool.mjs

@@ -0,0 +1,442 @@
+// Lead-side agent shard pool (P2/P3 — see docs/agent-shard-design.md).
+//
+// Topology: up to `MIXDOG_AGENT_SHARD_POOL` lazily-forked shard children
+// (default max(2, cpus/4)); worker spawns go to the LEAST-LOADED live shard
+// (a fresh slot opens once every live shard already runs a spawn) and a
+// shard process is REUSED across spawns (runtime import in the child costs
+// ~1-2s and must not be paid per spawn). A codex-style execution window
+// (`MIXDOG_AGENT_SHARD_WINDOW`, default 16, 0=off) caps CONCURRENT shard
+// spawns — overflow queues FIFO in the lead and stays cancellable by
+// tag/sessionId. Shards idle past `MIXDOG_AGENT_SHARD_IDLE_MS` (default 5m)
+// are reaped (~174MB each) and re-fork on demand. cancel/enqueue route to
+// the shard that owns the tag/session; provider 429 cooldowns and credential
+// resets are bridged both ways so one account's quota window parks every
+// process. Flag-gated: MIXDOG_AGENT_SHARD=1 opts in; default keeps the
+// in-process path byte-identical.
+import { fork } from 'node:child_process';
+import { cpus } from 'node:os';
+import { fileURLToPath } from 'node:url';
+import { startChildGuardian } from '../../runtime/shared/child-guardian.mjs';
+import { scrubLoaderVars } from '../../runtime/agent/orchestrator/tools/env-scrub.mjs';
+import { providerAdmissionScheduler } from '../../runtime/agent/orchestrator/providers/admission-scheduler.mjs';
+import { getMcpTools, executeMcpTool } from '../../runtime/agent/orchestrator/mcp/client.mjs';
+import { getInternalTools, executeInternalTool } from '../../runtime/agent/orchestrator/internal-tools.mjs';
+
+const CHILD_ENTRY = fileURLToPath(new URL('./shard-child.mjs', import.meta.url));
+const READY_TIMEOUT_MS = 30_000;
+const MAX_SHARDS = 16;
+const IDLE_REAP_SWEEP_MS = 60_000;
+
+export function agentShardEnabled() {
+    const value = process.env.MIXDOG_AGENT_SHARD;
+    if (value === '1') return true;
+    if (value === '0') return false;
+    // P3 default: ON for lead-like hosts. Worker-mode children (channels
+    // worker, shard child itself) never recursively shard their own spawns
+    // unless explicitly opted in.
+    return process.env.MIXDOG_WORKER_MODE !== '1';
+}
+
+function agentShardPoolSize() {
+    const configured = Math.floor(Number(process.env.MIXDOG_AGENT_SHARD_POOL));
+    if (Number.isFinite(configured) && configured >= 1) return Math.min(configured, MAX_SHARDS);
+    let cores = 4;
+    try { cores = cpus()?.length || 4; } catch { /* keep fallback */ }
+    return Math.max(2, Math.min(MAX_SHARDS, Math.floor(cores / 4)));
+}
+
+// Concurrent shard-spawn window — REACTIVE brake, not a preemptive cap.
+// Concurrent shard-spawn window. UNBOUNDED unless explicitly configured
+// (workstream rule: no local concurrency limits and no auto-brake heuristics
+// — provider 429 cooldown parking in the admission scheduler is the ONLY
+// limiter). `MIXDOG_AGENT_SHARD_WINDOW=N` is a pure operator opt-in to a
+// fixed window with a bounded FIFO overflow queue.
+function _currentWindow() {
+    const configured = Math.floor(Number(process.env.MIXDOG_AGENT_SHARD_WINDOW));
+    if (Number.isFinite(configured) && configured >= 1) return configured;
+    return 0; // unbounded
+}
+
+function agentShardIdleMs() {
+    const configured = Math.floor(Number(process.env.MIXDOG_AGENT_SHARD_IDLE_MS));
+    if (Number.isFinite(configured) && configured >= 0) return configured;
+    return 300_000;
+}
+
+// Overflow-queue bound (provider-admission parity: fail fast with a visible
+// error instead of silently accumulating unbounded deferred spawns).
+function agentShardQueueMax() {
+    const configured = Math.floor(Number(process.env.MIXDOG_AGENT_SHARD_QUEUE_MAX));
+    if (Number.isFinite(configured) && configured >= 1) return configured;
+    return 1_024;
+}
+
+const _shards = new Map(); // slot -> shard record
+let _callSeq = 0;
+// tag / sessionId -> owning shard, so cancel + tag-reuse enqueue reach the
+// process that hosts the live session. Entries are optimistic at spawn
+// dispatch (tag), confirmed by the `spawned` progress event (sessionId), and
+// dropped when the spawn settles (shard sessions are one-shot) or the shard
+// exits.
+const _ownerByKey = new Map();
+let _cooldownUnsub = null;
+// Execution-window state: dispatched-but-unsettled spawn count + FIFO
+// overflow queue of { payload, onProgress, resolve, reject, tag }.
+let _activeSpawns = 0;
+const _spawnQueue = [];
+let _idleReapTimer = null;
+
+function _liveShards() {
+    return [..._shards.values()].filter((shard) => !shard.closed);
+}
+
+function _broadcast(msg, { except = null } = {}) {
+    for (const shard of _liveShards()) {
+        if (shard === except) continue;
+        try { shard.child.send(msg); } catch { /* exit handler reaps */ }
+    }
+}
+
+// Lead-recorded 429/reset → every shard. Registered once, on first fork.
+function _ensureCooldownBridge() {
+    if (_cooldownUnsub || typeof providerAdmissionScheduler.onCooldownEvent !== 'function') return;
+    _cooldownUnsub = providerAdmissionScheduler.onCooldownEvent((event) => {
+        _broadcast({ type: 'cooldown', payload: event });
+    });
+}
+
+// Shard-reported 429/reset → lead scheduler (apply, no re-emit) + sibling
+// shards. applyExternalCooldown never emits, so there is no feedback loop;
+// resets are loop-terminated by the resetCount>0 emit gate.
+function _applyShardCooldownReport(payload, sourceShard) {
+    try {
+        if (payload?.type === 'reset') {
+            providerAdmissionScheduler.resetCooldowns(payload.provider || null);
+        } else if (payload?.key) {
+            providerAdmissionScheduler.applyExternalCooldown?.(payload.key, payload.cooldownUntil);
+            _broadcast({ type: 'cooldown', payload }, { except: sourceShard });
+        }
+    } catch { /* best-effort */ }
+}
+
+function _dropOwnerKeysFor(shard) {
+    for (const [key, owner] of [..._ownerByKey.entries()]) {
+        if (owner === shard) _ownerByKey.delete(key);
+    }
+}
+
+// Lead-tool proxy (P2 MCP proxy, generalized): shard children have neither
+// MCP connections nor the internal-tool executor (both are lead singletons),
+// so worker calls to those surfaces arrive here over IPC and execute in the
+// lead. The child receives the tool DEFINITIONS with each spawn payload
+// (below) so its session schema stays byte-identical to an in-process worker.
+function _leadToolDefs() {
+    let mcp = [];
+    let internal = [];
+    try { mcp = getMcpTools() || []; } catch { /* defs best-effort */ }
+    try { internal = getInternalTools() || []; } catch { /* defs best-effort */ }
+    return { mcp, internal };
+}
+
+async function _handleLeadToolCall(shard, msg) {
+    const { kind, name, args, callerCtx } = msg.payload || {};
+    let reply;
+    try {
+        const value = kind === 'mcp'
+            ? await executeMcpTool(name, args)
+            : await executeInternalTool(name, args, callerCtx || {});
+        reply = { type: 'lead-tool-result', callId: msg.callId, result: value };
+    } catch (error) {
+        reply = { type: 'lead-tool-result', callId: msg.callId, error: error?.message || String(error) };
+    }
+    try { shard.child.send(reply); } catch { /* child gone — its timeout reaps */ }
+}
+
+function _ownerFor(payload) {
+    const bySession = _ownerByKey.get(String(payload?.sessionId || '').trim());
+    if (bySession && !bySession.closed) return bySession;
+    const byTag = _ownerByKey.get(String(payload?.tag || '').trim());
+    return byTag && !byTag.closed ? byTag : null;
+}
+
+function _ensureShardAt(slot) {
+    const existing = _shards.get(slot);
+    if (existing && !existing.closed && existing.child.exitCode === null && existing.child.signalCode === null) {
+        return existing;
+    }
+    const env = { ...process.env };
+    scrubLoaderVars(env);
+    const child = fork(CHILD_ENTRY, [], {
+        stdio: ['ignore', 'ignore', 'pipe', 'ipc'],
+        // Never inherit the parent's execArgv (a driver started with
+        // --input-type/--test would poison the child's module resolution).
+        execArgv: [],
+        windowsHide: true,
+        detached: false,
+        env: {
+            ...env,
+            MIXDOG_STANDALONE: '1',
+            MIXDOG_WORKER_MODE: '1',
+            MIXDOG_CLI_OWNED: '0',
+            MIXDOG_QUIET_SESSION_LOG: process.env.MIXDOG_QUIET_SESSION_LOG ?? '1',
+            // The shard must never recursively shard its own nested spawns.
+            MIXDOG_AGENT_SHARD: '0',
+        },
+    });
+    const shard = {
+        slot,
+        child,
+        closed: false,
+        pending: new Map(),
+        inflight: 0,
+        lastActiveAt: Date.now(),
+        readyResolve: null,
+        ready: null,
+    };
+    shard.ready = new Promise((resolve, reject) => {
+        shard.readyResolve = resolve;
+        const t = setTimeout(() => reject(new Error('agent shard did not become ready within 30s')), READY_TIMEOUT_MS);
+        t.unref?.();
+    });
+    shard.ready.catch(() => {});
+    child.stderr?.on('data', (chunk) => {
+        const text = String(chunk || '').trimEnd();
+        if (text) { try { process.stderr.write(`[agent-shard:${slot}] ${text}\n`); } catch { /* best-effort */ } }
+    });
+    child.on('message', (msg) => {
+        if (!msg || typeof msg !== 'object') return;
+        if (msg.type === 'ready') {
+            shard.readyResolve?.(true);
+            shard.readyResolve = null;
+            return;
+        }
+        if (msg.type === 'rate-limit') {
+            _applyShardCooldownReport(msg.payload, shard);
+            return;
+        }
+        if (msg.type === 'lead-tool' && msg.callId) {
+            void _handleLeadToolCall(shard, msg);
+            return;
+        }
+        if (msg.type === 'progress' && msg.callId) {
+            if (msg.event === 'spawned' && msg.payload) {
+                // Confirm ownership under the sessionId (cancel-by-id) in
+                // addition to the optimistic tag mapping set at dispatch.
+                const sessionId = String(msg.payload.sessionId || '').trim();
+                if (sessionId) _ownerByKey.set(sessionId, shard);
+                const tag = String(msg.payload.tag || '').trim();
+                if (tag) _ownerByKey.set(tag, shard);
+            }
+            const item = shard.pending.get(msg.callId);
+            if (item?.onProgress) {
+                // A lead-side hook error must never kill the IPC pump.
+                try { item.onProgress(msg.event, msg.payload); } catch { /* best-effort */ }
+            }
+            return;
+        }
+        if (msg.type === 'result' && msg.callId) {
+            const item = shard.pending.get(msg.callId);
+            if (!item) return;
+            shard.pending.delete(msg.callId);
+            if (item.cleanupKeys) {
+                for (const key of item.cleanupKeys) {
+                    if (_ownerByKey.get(key) === shard) _ownerByKey.delete(key);
+                }
+                const settledSessionId = String(msg.result?.sessionId || '').trim();
+                if (settledSessionId && _ownerByKey.get(settledSessionId) === shard) {
+                    _ownerByKey.delete(settledSessionId);
+                }
+            }
+            if (msg.error) item.reject(new Error(msg.error));
+            else item.resolve(msg.result);
+        }
+    });
+    child.on('exit', (code, signal) => {
+        shard.closed = true;
+        const error = new Error(`agent shard exited (${signal || (code ?? 'unknown')})`);
+        for (const [, item] of shard.pending) item.reject(error);
+        shard.pending.clear();
+        _dropOwnerKeysFor(shard);
+        if (_shards.get(slot) === shard) _shards.delete(slot);
+    });
+    child.on('error', () => { /* exit handler owns cleanup */ });
+    startChildGuardian({ childPid: child.pid, label: `agent-shard-${slot}`, orphanGraceMs: 8000, forceGraceMs: 3000 });
+    _shards.set(slot, shard);
+    _ensureCooldownBridge();
+    _ensureIdleReapTimer();
+    return shard;
+}
+
+// ── Idle shard reap ─────────────────────────────────────────────────────────
+// A warmed shard holds ~174MB of session runtime. One shared sweep timer
+// kills shards with no inflight work past the idle window; the next spawn
+// re-forks on demand (~1-2s boot, acceptable after minutes of idleness).
+function _ensureIdleReapTimer() {
+    if (_idleReapTimer) return;
+    // Sweep at half the idle window (floor 5s, cap 60s) so short windows
+    // (tests, aggressive configs) actually observe the reap.
+    const idleMsNow = agentShardIdleMs();
+    const sweepMs = idleMsNow > 0
+        ? Math.max(5_000, Math.min(IDLE_REAP_SWEEP_MS, Math.floor(idleMsNow / 2)))
+        : IDLE_REAP_SWEEP_MS;
+    _idleReapTimer = setInterval(() => {
+        const idleMs = agentShardIdleMs();
+        if (idleMs <= 0) return;
+        const now = Date.now();
+        for (const shard of _liveShards()) {
+            if (shard.inflight > 0 || shard.pending.size > 0) continue;
+            if (now - shard.lastActiveAt < idleMs) continue;
+            shard.closed = true;
+            try { shard.child.kill(); } catch { /* exit handler reaps */ }
+        }
+        if (_shards.size === 0 && _idleReapTimer) {
+            clearInterval(_idleReapTimer);
+            _idleReapTimer = null;
+        }
+    }, sweepMs);
+    _idleReapTimer.unref?.();
+}
+
+// Least-loaded assignment: reuse the emptiest live shard; open a fresh slot
+// only when every live shard is already busy and the pool has headroom.
+function _shardForSpawn() {
+    const poolSize = agentShardPoolSize();
+    const live = _liveShards();
+    let best = null;
+    for (const shard of live) {
+        if (!best || shard.inflight < best.inflight) best = shard;
+    }
+    if (best && best.inflight === 0) return best;
+    if (live.length < poolSize) {
+        for (let slot = 0; slot < poolSize; slot += 1) {
+            const existing = _shards.get(slot);
+            if (!existing || existing.closed) return _ensureShardAt(slot);
+        }
+    }
+    return best || _ensureShardAt(0);
+}
+
+async function _requestOn(shard, type, payload, onProgress) {
+    await shard.ready;
+    return new Promise((resolve, reject) => {
+        const callId = `c${++_callSeq}`;
+        // One-shot shard sessions: when the spawn settles, its routing keys
+        // are stale (the child closes the session) and must not capture later
+        // same-tag traffic once the tag is reused on another shard.
+        const cleanupKeys = type === 'spawn'
+            ? [String(payload?.tag || '').trim()].filter(Boolean)
+            : null;
+        shard.pending.set(callId, { resolve, reject, onProgress, cleanupKeys });
+        try {
+            shard.child.send({ type, callId, payload });
+        } catch (error) {
+            shard.pending.delete(callId);
+            reject(error);
+        }
+    });
+}
+
+function _startSpawn(entry) {
+    const shard = _shardForSpawn();
+    if (entry.tag) _ownerByKey.set(entry.tag, shard);
+    // Fresh lead-tool snapshot per spawn: MCP servers reconnect/toggle at
+    // runtime and the child rebuilds its proxy surface from this payload.
+    const payload = { ...entry.payload, leadTools: _leadToolDefs() };
+    _activeSpawns += 1;
+    shard.inflight += 1;
+    shard.lastActiveAt = Date.now();
+    _requestOn(shard, 'spawn', payload, entry.onProgress)
+        .then(entry.resolve, entry.reject)
+        .finally(() => {
+            _activeSpawns = Math.max(0, _activeSpawns - 1);
+            shard.inflight = Math.max(0, shard.inflight - 1);
+            shard.lastActiveAt = Date.now();
+            _drainSpawnQueue();
+        });
+}
+
+function _drainSpawnQueue() {
+    const window = _currentWindow();
+    while (_spawnQueue.length > 0 && (window === 0 || _activeSpawns < window)) {
+        _startSpawn(_spawnQueue.shift());
+    }
+}
+
+// A cancel for a spawn still WAITING in the window queue must not fan out to
+// the shards (nothing was dispatched): drop the entry and fail its spawn
+// promise the same way a closed session fails a live one.
+function _cancelQueuedSpawn(payload) {
+    const tag = String(payload?.tag || '').trim();
+    const sessionId = String(payload?.sessionId || '').trim();
+    if (!tag && !sessionId) return null;
+    const index = _spawnQueue.findIndex((entry) => (tag && entry.tag === tag));
+    if (index < 0) return null;
+    const [entry] = _spawnQueue.splice(index, 1);
+    entry.reject(new Error('Session closed: agent shard spawn cancelled while queued'));
+    return { cancelled: true, queued: true };
+}
+
+/**
+ * One request/response round-trip to the shard pool. `spawn` goes to the
+ * least-loaded shard, bounded by the execution window (overflow queues FIFO,
+ * cancellable); `cancel`/`enqueue` route to the owning shard (fan-out to
+ * every live shard when ownership is unknown); `ping` warms shard 0.
+ * `onProgress` receives interim `progress` messages (e.g. `spawned`).
+ */
+export async function shardRequest(type, payload = {}, { onProgress } = {}) {
+    if (type === 'cancel' || type === 'enqueue') {
+        const fallback = type === 'cancel' ? { cancelled: false } : { queued: false };
+        if (type === 'cancel') {
+            const queuedCancel = _cancelQueuedSpawn(payload);
+            if (queuedCancel) return queuedCancel;
+        }
+        const owner = _ownerFor(payload);
+        if (owner) return _requestOn(owner, type, payload, onProgress);
+        const live = _liveShards();
+        if (live.length === 0) return fallback;
+        const settled = await Promise.allSettled(live.map((shard) => _requestOn(shard, type, payload, onProgress)));
+        const values = settled.filter((r) => r.status === 'fulfilled').map((r) => r.value);
+        return values.find((v) => v?.cancelled === true || v?.queued === true) || values[0] || fallback;
+    }
+    if (type === 'spawn') {
+        return new Promise((resolve, reject) => {
+            const entry = {
+                payload,
+                onProgress,
+                resolve,
+                reject,
+                tag: String(payload?.tag || '').trim() || null,
+            };
+            const window = _currentWindow();
+            if (window > 0 && _activeSpawns >= window) {
+                const queueMax = agentShardQueueMax();
+                if (_spawnQueue.length >= queueMax) {
+                    reject(new Error(`agent shard spawn queue full (${queueMax} deferred spawns; window=${window}) — `
+                        + 'wait for running workers to finish, or raise MIXDOG_AGENT_SHARD_QUEUE_MAX / MIXDOG_AGENT_SHARD_WINDOW'));
+                    return;
+                }
+                _spawnQueue.push(entry);
+            } else {
+                _startSpawn(entry);
+            }
+        });
+    }
+    const shard = _ensureShardAt(0);
+    shard.lastActiveAt = Date.now();
+    return _requestOn(shard, type, payload, onProgress);
+}
+
+/**
+ * Boot prewarm: fork the shard child now and let it import the session
+ * runtime, so the FIRST sharded spawn skips the ~1-2s child boot. No-op when
+ * sharding is disabled. Best-effort; resolves true when the shard answered.
+ */
+export async function prewarmAgentShard() {
+    if (!agentShardEnabled()) return false;
+    try {
+        const pong = await shardRequest('ping', {});
+        return pong?.pong === true;
+    } catch {
+        return false;
+    }
+}