@struktoai/mirage-core/dist/rs/diff_pool.js

A Unified Virtual Filesystem For AI Agents

import { Worker } from "../concurrency/worker.js";
import { chunks } from "../io/cooperative.js";
import { abortable, makeAbortError } from "../utils/abort.js";
import { DiffReply } from "./diff_worker.js";
export const CHUNK_SIZE = 64 * 1024;
export const WORKER_LIMIT = 2;
export const IDLE_TTL = 1000;
export const QUEUE_LIMIT = 32;
export const RETIRE_BYTES = 16 * 1024 * 1024;
export class DiffPool {
    entries = new Set();
    queue = [];
    active = new Set();
    cleanupErrors = [];
    stopped = false;
    closing;
    compare(options, old, next, signal) {
        const comparison = this.run(options, old, next, signal);
        this.active.add(comparison);
        void comparison.then(() => this.active.delete(comparison), () => this.active.delete(comparison));
        return comparison;
    }
    stop() {
        this.stopped = true;
        for (const pending of this.queue.splice(0)) {
            pending.detach();
            pending.reject(new Error('diff pool is closed'));
        }
        for (const entry of this.entries)
            if (entry.state.kind === 'leased')
                entry.state.abort.abort();
    }
    close() {
        if (this.closing !== undefined)
            return this.closing;
        this.stop();
        this.closing = this.closeAll();
        return this.closing;
    }
    async closeAll() {
        await Promise.allSettled([...this.entries]
            .filter((entry) => entry.state.kind !== 'stop_failed')
            .map((entry) => this.retire(entry)));
        await Promise.allSettled([...this.entries]
            .filter((entry) => entry.state.kind === 'stop_failed')
            .map((entry) => this.retire(entry, true)));
        await Promise.allSettled([...this.active]);
        if (this.cleanupErrors.length > 0)
            throw new AggregateError(this.cleanupErrors, 'diff worker cleanup failed');
    }
    acquire(signal) {
        if (this.stopped)
            return Promise.reject(new Error('diff pool is closed'));
        if (signal?.aborted === true)
            return Promise.reject(makeAbortError(signal));
        if (this.queue.length >= QUEUE_LIMIT)
            return Promise.reject(new Error('diff worker queue is full'));
        return new Promise((resolve, reject) => {
            const onAbort = () => {
                const index = this.queue.indexOf(pending);
                if (index < 0)
                    return;
                this.queue.splice(index, 1);
                pending.detach();
                reject(makeAbortError(signal));
            };
            const pending = {
                resolve,
                reject,
                detach: () => signal?.removeEventListener('abort', onAbort),
            };
            signal?.addEventListener('abort', onAbort, { once: true });
            this.queue.push(pending);
            this.pump();
        });
    }
    pump() {
        if (this.stopped)
            return;
        while (this.queue.length > 0) {
            let entry = [...this.entries].find((item) => item.state.kind === 'idle');
            if (entry?.state.kind === 'idle')
                clearTimeout(entry.state.timer);
            else {
                if (this.entries.size >= WORKER_LIMIT) {
                    if ([...this.entries].every((item) => item.state.kind === 'stop_failed')) {
                        for (const pending of this.queue.splice(0)) {
                            pending.detach();
                            pending.reject(new Error('diff worker capacity is unavailable after failed cleanup'));
                        }
                    }
                    return;
                }
                entry = {
                    worker: new Worker(),
                    inFlight: new Set(),
                    state: { kind: 'leased', abort: new AbortController() },
                };
                this.entries.add(entry);
            }
            const pending = this.queue.shift();
            if (pending === undefined)
                return;
            entry.state = { kind: 'leased', abort: new AbortController() };
            pending.detach();
            pending.resolve(entry);
        }
    }
    async run(options, old, next, signal) {
        const entry = await this.acquire(signal);
        if (entry.state.kind !== 'leased')
            throw new Error('diff worker lease is invalid');
        const combined = signal === undefined
            ? entry.state.abort.signal
            : AbortSignal.any([signal, entry.state.abort.signal]);
        let success = false;
        let bytes = 0;
        const output = [];
        const stderr = [];
        const call = async (request) => {
            const pending = entry.worker.call('diffDispatch', [request], combined);
            entry.inFlight.add(pending);
            let raw;
            try {
                raw = await pending;
            }
            finally {
                entry.inFlight.delete(pending);
            }
            const reply = DiffReply.parse(raw);
            if (reply.kind === 'unavailable' || reply.kind === 'refusal')
                throw new Error(reply.message);
            if (reply.kind === 'output') {
                if (reply.output.byteLength > CHUNK_SIZE)
                    throw new Error('diff worker output exceeds limit');
                if (reply.output.byteLength > 0)
                    output.push(reply.output);
                bytes += reply.output.byteLength;
            }
            else {
                if (reply.stderr.byteLength > CHUNK_SIZE)
                    throw new Error('diff worker stderr exceeds limit');
                if (reply.stderr.byteLength > 0)
                    stderr.push(reply.stderr);
                bytes += reply.stderr.byteLength;
            }
            return reply;
        };
        try {
            await call({ kind: 'open', options });
            for (const [index, factory] of [
                [0, old],
                [1, next],
            ]) {
                const source = await abortable(Promise.resolve().then(factory), combined);
                for await (const chunk of chunks(source, combined)) {
                    bytes += chunk.byteLength;
                    await call({ kind: 'feed', index, chunk: new Uint8Array(chunk) });
                }
                await call({ kind: 'end', index });
            }
            for (;;) {
                const state = await call({ kind: 'inspect' });
                if (state.kind !== 'state')
                    throw new Error('diff worker returned an invalid state');
                if (!state.busy) {
                    if (state.status !== 2 && state.status !== (state.outcome === 'equal' ? 0 : 1))
                        throw new Error('diff worker outcome disagrees with status');
                    success = true;
                    return { outcome: state.outcome, status: state.status, output, stderr };
                }
                await call({ kind: 'drain' });
            }
        }
        finally {
            await this.release(entry, success, bytes);
        }
    }
    async release(entry, success, bytes) {
        if (entry.state.kind === 'leased') {
            if (!success ||
                bytes > RETIRE_BYTES ||
                this.stopped ||
                (this.queue.length === 0 && [...this.entries].some((item) => item.state.kind === 'idle')))
                await this.retire(entry);
            else {
                const timer = setTimeout(() => {
                    if (entry.state.kind === 'idle' && entry.state.timer === timer)
                        void this.retire(entry).catch(() => undefined);
                }, IDLE_TTL);
                entry.state = { kind: 'idle', timer };
                this.pump();
            }
        }
        else if (entry.state.kind === 'stopping')
            await entry.state.completion;
    }
    retire(entry, retry = false) {
        if (entry.state.kind === 'stopping')
            return entry.state.completion;
        if (entry.state.kind === 'stop_failed' && !retry)
            return Promise.reject(entry.state.error);
        if (entry.state.kind === 'idle')
            clearTimeout(entry.state.timer);
        if (entry.state.kind === 'leased')
            entry.state.abort.abort();
        const completion = Promise.resolve()
            .then(async () => {
            const [stop] = await Promise.allSettled([entry.worker.close(retry), ...entry.inFlight]);
            if (stop.status === 'rejected')
                throw stop.reason;
            this.entries.delete(entry);
        })
            .catch((error) => {
            const failure = error instanceof Error ? error : new Error('diff worker cleanup failed', { cause: error });
            this.cleanupErrors.push(failure);
            entry.state = { kind: 'stop_failed', error: failure };
            throw failure;
        })
            .finally(() => {
            this.pump();
        });
        entry.state = { kind: 'stopping', completion };
        return completion;
    }
}
//# sourceMappingURL=diff_pool.js.map