@struktoai/mirage-core/dist/io/spool.js

A Unified Virtual Filesystem For AI Agents

import { OutputStream, closeQuietly } from "./stream.js";
import { CHUNK_SIZE, chunks } from "./cooperative.js";
import { concat } from "../utils/bytes.js";
export const MEMORY_LIMIT = 1024 * 1024;
export const IO_BUFFER_SIZE = 256 * 1024;
let factory = null;
/** Install the execution host's private scratch storage. */
export function setSpoolFactory(value) {
    factory = value;
}
/**
 * Retain replayable bytes, spilling through the host's scratch storage. An
 * append that fails leaves no telling which of its bytes were kept, so every
 * call after it fails too, rather than read bytes from where the append never
 * put them. Bytes count once the spool holds them, before they flush, as in
 * Python, where a cancelled append keeps the spool usable. Bytes waiting for
 * the store go at the offset they end up at, so a flush that failed writes
 * them there again rather than after what it got through. A reader that `hold`s
 * the spool keeps it open: a `close` while it is held takes effect at the last
 * `unhold`.
 */
export class ByteSpool {
    memoryLimit;
    size = 0;
    closed = false;
    memory = [];
    store = null;
    lost = undefined;
    pendingSize = 0;
    stored = 0;
    holders = 0;
    closing = false;
    cache = new Uint8Array();
    cacheOffset = 0;
    turn = Promise.resolve();
    constructor(memoryLimit = MEMORY_LIMIT) {
        this.memoryLimit = memoryLimit;
        if (!Number.isSafeInteger(memoryLimit) || memoryLimit < 0) {
            throw new RangeError('spool memory limit must be nonnegative');
        }
    }
    /** Run `fn` once every call before it has settled, as Python's lock holds them. */
    locked(fn) {
        const run = this.turn.then(fn);
        this.turn = run.then(() => undefined, () => undefined);
        return run;
    }
    append(data) {
        return this.locked(async () => {
            this.usable();
            try {
                await this.appendNow(data);
            }
            catch (error) {
                this.lost = error;
                throw error;
            }
        });
    }
    /** Overwrite an existing window without changing the spool length. */
    write(offset, data) {
        return this.locked(async () => {
            this.usable();
            if (!Number.isSafeInteger(offset) || offset < 0 || offset + data.length > this.size)
                throw new RangeError('spool write must stay within its size');
            this.cache = new Uint8Array();
            try {
                if (this.store !== null) {
                    await this.flush();
                    await this.store.write(offset, data);
                }
                else {
                    let written = 0;
                    while (written < data.length) {
                        const position = offset + written;
                        const part = this.memory[Math.floor(position / CHUNK_SIZE)];
                        if (part === undefined)
                            throw new Error('spool lost a block');
                        const start = position % CHUNK_SIZE;
                        const count = Math.min(part.length - start, data.length - written);
                        part.set(data.subarray(written, written + count), start);
                        written += count;
                    }
                }
            }
            catch (error) {
                this.lost = error;
                throw error;
            }
        });
    }
    usable() {
        if (this.closed)
            throw new Error('spool is closed');
        if (this.lost !== undefined) {
            throw Object.assign(new Error('spool lost a write', { cause: this.lost }), { code: 'EIO' });
        }
    }
    async appendNow(data) {
        this.cache = new Uint8Array();
        if (this.store === null && this.size + data.length > this.memoryLimit && factory !== null) {
            this.store = await factory();
            for (const part of this.memory) {
                await this.store.write(this.stored, part);
                this.stored += part.length;
            }
            this.memory = [];
            this.pendingSize = 0;
        }
        let offset = 0;
        while (offset < data.length) {
            const tail = this.memory.at(-1);
            if (tail !== undefined && tail.length < CHUNK_SIZE) {
                const count = Math.min(CHUNK_SIZE - tail.length, data.length - offset);
                const extended = new Uint8Array(tail.buffer, tail.byteOffset, tail.length + count);
                extended.set(data.subarray(offset, offset + count), tail.length);
                this.memory[this.memory.length - 1] = extended;
                offset += count;
            }
            else {
                const count = Math.min(CHUNK_SIZE, data.length - offset);
                const block = new Uint8Array(CHUNK_SIZE);
                block.set(data.subarray(offset, offset + count));
                this.memory.push(block.subarray(0, count));
                offset += count;
            }
        }
        this.pendingSize += data.length;
        this.size += data.length;
        if (this.store !== null && this.pendingSize >= IO_BUFFER_SIZE)
            await this.flush();
    }
    async flush() {
        if (this.store !== null && this.pendingSize > 0) {
            await this.store.write(this.stored, concat(this.memory));
            this.stored += this.pendingSize;
            this.memory = [];
            this.pendingSize = 0;
        }
    }
    async fill(source) {
        for await (const chunk of chunks(source))
            await this.append(chunk);
    }
    async read(offset, size) {
        if (!Number.isSafeInteger(offset) || !Number.isSafeInteger(size) || offset < 0 || size < 0)
            throw new RangeError('spool window must be nonnegative');
        return this.locked(() => this.readNow(offset, size));
    }
    async readNow(offset, size) {
        this.usable();
        if (this.store !== null) {
            await this.flush();
            if (size > CHUNK_SIZE)
                return this.store.read(offset, size);
            if (!(this.cacheOffset <= offset && offset + size <= this.cacheOffset + this.cache.length)) {
                this.cacheOffset = Math.floor(offset / IO_BUFFER_SIZE) * IO_BUFFER_SIZE;
                this.cache = await this.store.read(this.cacheOffset, Math.max(IO_BUFFER_SIZE, offset + size - this.cacheOffset));
            }
            const start = offset - this.cacheOffset;
            return this.cache.slice(start, start + size);
        }
        const result = new Uint8Array(Math.max(0, Math.min(size, this.size - offset)));
        let index = Math.floor(offset / CHUNK_SIZE);
        let start = offset - index * CHUNK_SIZE;
        for (let written = 0; written < result.length; index++, start = 0) {
            const part = this.memory[index];
            if (part === undefined)
                break;
            const count = Math.min(part.length - start, result.length - written);
            result.set(part.subarray(start, start + count), written);
            written += count;
        }
        return result;
    }
    /**
     * The spool's bytes from `offset`, `piece` bytes at a time: a bounded reader
     * keeps `CHUNK_SIZE`, a writer of the whole spool takes `IO_BUFFER_SIZE`,
     * one store read per step.
     */
    async *stream(offset = 0, size, piece = CHUNK_SIZE) {
        const end = size === undefined ? this.size : Math.min(this.size, offset + size);
        while (offset < end) {
            const data = await this.read(offset, Math.min(piece, end - offset));
            if (data.length === 0)
                return;
            offset += data.length;
            yield data;
        }
    }
    async output() {
        if (this.size <= this.memoryLimit) {
            const data = await this.read(0, this.size);
            await this.close();
            return data;
        }
        return new OutputStream(this.stream(0, undefined, IO_BUFFER_SIZE), () => this.close());
    }
    body(kind) {
        return this.locked(async () => {
            this.usable();
            await this.flush();
            if (this.store?.body !== undefined)
                return this.store.body(kind);
            if (this.store !== null)
                throw new Error('scratch storage cannot provide an upload body');
            return this.readNow(0, this.size);
        });
    }
    /** Keep the spool open until `unhold`, for a reader that took windows of it. */
    hold() {
        this.holders++;
    }
    /** Let go of a `hold`; the last one runs a `close` it held off. */
    async unhold() {
        this.holders--;
        if (this.closing && this.holders === 0)
            await this.close();
    }
    close() {
        if (this.holders > 0) {
            this.closing = true;
            return Promise.resolve();
        }
        return this.locked(() => this.closeNow());
    }
    async closeNow() {
        this.closed = true;
        this.memory = [];
        this.cache = new Uint8Array();
        const store = this.store;
        this.store = null;
        if (store !== null)
            await store.close();
    }
}
/**
 * Hold `source` in a spool: its bytes when they fit in memory, a stream that
 * owns the spool past that, null when it is empty. Mirrors Python's spooled.
 */
export async function spooled(source) {
    const store = new ByteSpool();
    try {
        await store.fill(source ?? new Uint8Array());
        if (store.size === 0) {
            await store.close();
            return null;
        }
        return await store.output();
    }
    catch (error) {
        await store.close();
        throw error;
    }
}
/**
 * `spooled` over both outputs of a run; when one fails, what the other holds
 * is closed. Mirrors Python's spooled_pair.
 */
export async function spooledPair(stdout, stderr) {
    let out;
    try {
        out = await spooled(stdout);
    }
    catch (error) {
        await closeQuietly(stderr);
        throw error;
    }
    try {
        return [out, await spooled(stderr)];
    }
    catch (error) {
        await closeQuietly(out);
        throw error;
    }
}
/** A window of a spool, sized and sliced as the bytes it stands for. */
export class ByteSlice {
    source;
    offset;
    size;
    constructor(source, offset, size) {
        this.source = source;
        this.offset = offset;
        this.size = size;
    }
    get byteLength() {
        return this.size;
    }
    /** As `Uint8Array.subarray`, for the nonnegative bounds a buffer slices by. */
    subarray(start = 0, end = this.size) {
        const from = Math.min(start, this.size);
        return new ByteSlice(this.source, this.offset + from, Math.max(0, Math.min(end, this.size) - from));
    }
    stream() {
        return this.source.stream(this.offset, this.size);
    }
}
//# sourceMappingURL=spool.js.map