mixdog 0.9.87 → 0.9.89

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

@@ -1,443 +0,0 @@
-// 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;
-    // P2 gate: sharding stays opt-in (MIXDOG_AGENT_SHARD=1). The P3
-    // default-ON flip (lead-like hosts shard unless MIXDOG_WORKER_MODE=1)
-    // is deferred until the remaining P2 items land — a default-ON pool
-    // routes every spawn through shard children, which breaks hosts that
-    // inject a custom session manager (fake/test mgr, embedded runtimes).
-    return false;
-}
-
-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;
-    }
-}