@sema-agent/core/dist/core/mcp-probe.js

Stateless, task-oriented AI agent core

import { ErrorCode, McpError } from "@modelcontextprotocol/sdk/types.js";
import { snapshotMcpEntries } from "./locked-config.js";
import { classifyMcpFailure, livenessOf } from "./mcp-failure.js";
import { MCP_ENV_MS_MAX, materializeMcpTools, mcpStartupTimeoutMs } from "./mcp.js";
import { findNamespacePrefixCollision } from "./protocol-naming.js";
import { MCP_NAMESPACE } from "./protocol-table.js";
import { mcpManifestEntries } from "./runner/prepare-wiring-manifest.js";
const DECLINE_EVERY_ELICITATION = async () => ({ action: "decline" });
function wholeInRange(name, value, min, max) {
    if (value === undefined)
        return undefined;
    if (typeof value === "number" && Number.isInteger(value) && value >= min && value <= max)
        return value;
    const shown = typeof value === "number" ? String(value) : `a ${typeof value}`;
    const e = new Error(`probeMcpServers: \`${name}\` must be a whole number in [${min}, ${max}] — received ${shown}`);
    e.code = "config.mcp_probe_option";
    throw e;
}
function abortCause(signal) {
    const reason = signal.reason;
    if (reason instanceof Error)
        return reason;
    const e = new Error("the MCP probe was aborted by its caller's signal");
    e.name = "AbortError";
    return e;
}
function whenAborted(signal, until) {
    if (signal.aborted)
        return Promise.resolve();
    return new Promise((resolve) => {
        signal.addEventListener("abort", () => resolve(), { once: true, signal: until });
    });
}
function afterMs(ms, until) {
    return new Promise((resolve) => {
        const timer = setTimeout(resolve, ms);
        timer.unref();
        until.addEventListener("abort", () => clearTimeout(timer), { once: true });
    });
}
async function raceDial(dial, ms, signal) {
    if (ms === undefined && signal === undefined)
        return { kind: "settled", handle: await dial };
    const until = new AbortController();
    const legs = [dial.then((handle) => ({ kind: "settled", handle }))];
    if (ms !== undefined)
        legs.push(afterMs(ms, until.signal).then(() => ({ kind: "timeout", afterMs: ms })));
    if (signal !== undefined)
        legs.push(whenAborted(signal, until.signal).then(() => ({ kind: "aborted", signal })));
    try {
        return await Promise.race(legs);
    }
    finally {
        until.abort();
    }
}
function stoppedWaitingStatus(entry, ms) {
    const condition = new McpError(ErrorCode.RequestTimeout, `the probe stopped waiting for this server's connect and initial tool listing after ${ms} ms`);
    const { kind, delivered } = classifyMcpFailure(condition, { phase: "connect", transport: entry.transport.kind });
    const liveness = livenessOf(kind, Date.now());
    return { name: entry.name, status: "failed", errorCode: kind, delivered, error: condition.message, ...(liveness !== undefined ? { liveness } : {}) };
}
export async function probeMcpServers(specs, opts = {}) {
    const declared = snapshotMcpEntries(specs);
    const timeoutMs = wholeInRange("timeoutMs", opts.timeoutMs, 1, MCP_ENV_MS_MAX) ?? mcpStartupTimeoutMs();
    const concurrency = wholeInRange("concurrency", opts.concurrency, 1, Number.MAX_SAFE_INTEGER) ?? 1;
    const onElicit = opts.onElicit ?? DECLINE_EVERY_ELICITATION;
    const signal = opts.signal;
    if (findNamespacePrefixCollision(MCP_NAMESPACE, declared.map((e) => e.name)) !== undefined) {
        await materializeMcpTools(declared, opts.principal, onElicit);
    }
    const statuses = [];
    const servers = new Array(declared.length);
    const teardowns = [];
    let fatal;
    const probeOne = async (i) => {
        const entry = declared[i];
        const started = performance.now();
        const elapsed = () => Math.round(performance.now() - started);
        const dial = materializeMcpTools([entry], opts.principal, onElicit);
        const verdict = await raceDial(dial, timeoutMs, signal);
        if (verdict.kind === "settled") {
            const found = verdict.handle.statuses.find((s) => s.name === entry.name);
            const status = found === undefined ? undefined : { ...found };
            teardowns.push(verdict.handle.dispose());
            if (status !== undefined)
                statuses.push(status);
            servers[i] = { name: entry.name, connectMs: elapsed(), tools: status?.toolNames ?? [], ...(status?.errorCode !== undefined ? { errorCode: status.errorCode } : {}) };
            return;
        }
        teardowns.push(dial.then((h) => h.dispose(), () => undefined));
        if (verdict.kind === "aborted") {
            fatal ??= { cause: abortCause(verdict.signal) };
            return;
        }
        const status = stoppedWaitingStatus(entry, verdict.afterMs);
        statuses.push(status);
        servers[i] = { name: entry.name, connectMs: elapsed(), tools: [], errorCode: status.errorCode };
    };
    if (signal?.aborted === true)
        fatal = { cause: abortCause(signal) };
    let next = 0;
    const lane = async () => {
        for (;;) {
            if (signal?.aborted === true)
                fatal ??= { cause: abortCause(signal) };
            if (fatal !== undefined)
                return;
            const i = next++;
            if (i >= declared.length)
                return;
            try {
                await probeOne(i);
            }
            catch (e) {
                fatal ??= { cause: e };
                return;
            }
        }
    };
    try {
        await Promise.all(Array.from({ length: Math.min(concurrency, declared.length) }, lane));
    }
    finally {
        await Promise.allSettled(teardowns);
    }
    if (fatal === undefined && signal?.aborted === true)
        fatal = { cause: abortCause(signal) };
    if (fatal !== undefined)
        throw fatal.cause;
    return { entries: mcpManifestEntries(declared, statuses), servers };
}