@sema-agent/core/dist/core/mcp-connection-pool.js

Stateless, task-oriented AI agent core

import { canonicalize } from "./canonical-json.js";
export class McpConnectionPool {
    entries = new Map();
    dialing = new Map();
    generation = new Map();
    shut = false;
    static keyOf(spec, principal, kind, startupMs) {
        const rest = Object.fromEntries(Object.keys(spec.transport).filter((k) => k !== "kind").map((k) => [k, Reflect.get(spec.transport, k)]));
        const fingerprint = spec.fingerprint;
        const identity = [spec.name, kind ?? null, rest, principal ?? null];
        if (fingerprint !== undefined)
            identity.push(fingerprint);
        if (startupMs !== undefined)
            identity.push({ startupTimeoutMs: startupMs });
        return canonicalize(identity);
    }
    async acquire(key, dial, opts) {
        if (this.shut)
            return undefined;
        if (opts?.fresh === true)
            this.detach(key);
        let entry = this.entries.get(key);
        if (entry === undefined || entry.closed) {
            let pending = this.dialing.get(key);
            if (pending === undefined) {
                const started = this.dialNew(key, this.generation.get(key) ?? 0, dial);
                pending = started;
                this.dialing.set(key, started);
                const settled = () => {
                    if (this.dialing.get(key) === started)
                        this.dialing.delete(key);
                };
                void started.then(settled, settled);
            }
            entry = await pending;
        }
        return this.lease(entry);
    }
    detach(key) {
        this.generation.set(key, (this.generation.get(key) ?? 0) + 1);
        const e = this.entries.get(key);
        if (e !== undefined) {
            e.current = false;
            this.entries.delete(key);
        }
        this.dialing.delete(key);
    }
    close() {
        this.shut = true;
    }
    get live() {
        let n = 0;
        for (const e of this.entries.values())
            if (!e.closed)
                n++;
        return n;
    }
    async dialNew(key, generation, dial) {
        const client = await dial();
        const entry = { client, refs: 0, listeners: new Set(), holders: new Map(), meta: new Map(), nextVersion: 0, current: true, closed: false };
        client.onclose = () => {
            entry.closed = true;
            if (this.entries.get(key) === entry)
                this.entries.delete(key);
            for (const fn of [...entry.listeners]) {
                try {
                    fn();
                }
                catch {
                }
            }
        };
        if ((this.generation.get(key) ?? 0) === generation)
            this.entries.set(key, entry);
        else
            entry.current = false;
        return entry;
    }
    lease(entry) {
        entry.refs++;
        const mine = new Set();
        let released = false;
        const cached = () => {
            const held = new Set([...entry.holders.values()].flat());
            return [...entry.meta].filter(([name]) => held.has(name)).map(([, r]) => r.tool);
        };
        const lease = {
            client: entry.client,
            onClose: (fn) => {
                mine.add(fn);
                entry.listeners.add(fn);
            },
            openListing: () => {
                const version = entry.nextVersion++;
                return (tools, apply, retained) => {
                    for (const tool of tools) {
                        const prior = entry.meta.get(tool.name);
                        if (retained?.has(tool.name) === true) {
                            if (prior === undefined)
                                entry.meta.set(tool.name, { tool, version: -1 });
                            continue;
                        }
                        if (prior !== undefined && prior.version > version)
                            continue;
                        entry.meta.delete(tool.name);
                        entry.meta.set(tool.name, { tool, version });
                    }
                    if (!released)
                        entry.holders.set(lease, tools.map((t) => t.name));
                    entry.apply = apply;
                    apply(cached());
                };
            },
            release: async () => {
                if (released)
                    return;
                released = true;
                for (const fn of mine)
                    entry.listeners.delete(fn);
                entry.refs--;
                if (entry.holders.delete(lease) && entry.refs > 0)
                    entry.apply?.(cached());
                if (entry.refs > 0)
                    return;
                entry.current = false;
                for (const [k, v] of this.entries)
                    if (v === entry)
                        this.entries.delete(k);
                await entry.client.close().catch(() => undefined);
            },
        };
        return lease;
    }
}