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

A Unified Virtual Filesystem For AI Agents

// ========= Copyright 2026 @ Strukto.AI All Rights Reserved. =========
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ========= Copyright 2026 @ Strukto.AI All Rights Reserved. =========
import MIRAGE_RS_WASM from "../generated/mirage_rs.js";
import { chunks } from "../io/cooperative.js";
import { SEALED } from "../io/pipe.js";
import { YieldBudget } from "../io/yield_budget.js";
import { concat } from "../utils/bytes.js";
import { decodeBase64 } from "../utils/base64.js";
import { StreamRefusal } from "./errors.js";
// Linear memory only grows, so an instance a large stream (a big sort) grew
// is dropped once no stream is open and the next stream gets a fresh one.
// Movable streams (the matchers a cache keeps) do not hold it: they close
// with it, and their owners open them again on the next instance. One in the
// middle of a search (busy) does hold it, so the search keeps its progress.
const RECYCLE_BYTES = 32 * 1024 * 1024;
const encoder = new TextEncoder();
const decoder = new TextDecoder();
let compiled;
let compiledSync;
let current;
let synchronous;
const movables = new WeakMap();
// The next stream gets a fresh instance; the movable streams still on this one
// close, and `live` says so to their owners.
function retire(exports, trapped) {
    current = undefined;
    if (synchronous === exports)
        synchronous = undefined;
    for (const stream of movables.get(exports) ?? [])
        stream.drop(!trapped);
    movables.delete(exports);
}
/** The compiled module, for the page to hand the workers it starts without a copy. */
export function compiledModule() {
    compiled ??= WebAssembly.compile(decodeBase64(MIRAGE_RS_WASM));
    return compiled;
}
/** The compiled module, compiled now if no one has yet, for a caller that cannot wait. */
export function moduleSync() {
    compiledSync ??= new WebAssembly.Module(decodeBase64(MIRAGE_RS_WASM));
    compiled ??= Promise.resolve(compiledSync);
    return compiledSync;
}
/** Run on the module the page compiled: a browser worker's bundle leaves the bytes out. */
export function adoptModule(module) {
    compiledSync ??= module;
    compiled ??= Promise.resolve(module);
}
async function instance() {
    if (synchronous !== undefined)
        return synchronous;
    current ??= compiledModule().then((module) => {
        compiledSync = module;
        synchronous ??= new WebAssembly.Instance(module, {}).exports;
        return synchronous;
    });
    return current;
}
function bytesAt(exports, ptr, len) {
    const bytes = new Uint8Array(exports.memory.buffer, ptr, len).slice();
    SEALED.add(bytes.buffer);
    return bytes;
}
/** One mirage-rs stream: chunks go in, the bytes they produce come out. Mirrors Python's OpenStream. */
export class OpenStream {
    movable;
    exports;
    handle;
    capacity;
    constructor(exports, handle, movable = false) {
        this.movable = movable;
        this.exports = exports;
        this.handle = handle;
        this.capacity = exports.mirage_rs_input_capacity();
        if (movable) {
            const held = movables.get(exports) ?? new Set();
            held.add(this);
            movables.set(exports, held);
        }
    }
    /** Whether the stream is still open; a movable one closes when its instance retires. */
    get live() {
        return this.handle !== 0;
    }
    /** Push one chunk and return what it produced. */
    feed(chunk) {
        return this.guard(() => this.push(chunk, (len) => this.exports.mirage_rs_feed(this.handle, len)));
    }
    /** End the input and return what the stream held back. */
    finish() {
        return this.guard(() => this.collect(this.exports.mirage_rs_finish(this.handle), []));
    }
    /** Whether `drain` has work or output left. */
    busy() {
        return this.guard(() => this.exports.mirage_rs_busy(this.handle) === 1);
    }
    /** Stop processing and report whether bounded retirement is available. */
    cancel() {
        return this.guard(() => this.exports.mirage_rs_cancel(this.handle) === 1);
    }
    /** Bounded steps of the work `busy` reports, up to a piece of output, and that output. */
    drain() {
        return this.guard(() => this.collect(this.exports.mirage_rs_drain(this.handle), []));
    }
    /** The input a multi-input stream needs next, null for none. */
    wants() {
        const index = this.guard(() => this.exports.mirage_rs_wants(this.handle));
        return index < 0 ? null : index;
    }
    /** Push one chunk of input `index` to a multi-input stream. */
    feedInput(index, chunk) {
        return this.guard(() => this.push(chunk, (len) => this.exports.mirage_rs_feed_input(this.handle, index, len)));
    }
    /** End input `index` of a multi-input stream. */
    endInput(index) {
        return this.guard(() => this.collect(this.exports.mirage_rs_end_input(this.handle, index), []));
    }
    /** Start the next input of a stream that reads one after another. */
    begin(config) {
        const bytes = encoder.encode(JSON.stringify(config));
        if (bytes.byteLength > this.capacity)
            throw new StreamRefusal('stream input configuration too large', new Uint8Array());
        this.guard(() => {
            const window = this.exports.mirage_rs_input(this.handle);
            new Uint8Array(this.exports.memory.buffer, window, bytes.byteLength).set(bytes);
            this.collect(this.exports.mirage_rs_begin(this.handle, bytes.byteLength), []);
        });
    }
    /** Whether the input fed since the last answer, then `chunk`, passes the stream's test, in one call. */
    test(chunk) {
        return this.guard(() => {
            const cut = Math.max(0, chunk.byteLength - this.capacity);
            if (cut > 0)
                this.push(chunk.subarray(0, cut), (len) => this.exports.mirage_rs_feed(this.handle, len));
            const tail = chunk.subarray(cut);
            const window = this.exports.mirage_rs_input(this.handle);
            new Uint8Array(this.exports.memory.buffer, window, tail.byteLength).set(tail);
            const answer = this.exports.mirage_rs_test(this.handle, tail.byteLength);
            if (answer < 0)
                this.collect(1, []);
            return answer === 1;
        });
    }
    /** Whether the stream wants no more of the current input. */
    done() {
        return this.guard(() => this.exports.mirage_rs_done(this.handle) === 1);
    }
    /** The command's exit status for the input so far. */
    status() {
        return this.guard(() => this.exports.mirage_rs_status(this.handle));
    }
    /** The stream's current progress as its typed JSON record. */
    metadata() {
        return this.guard(() => this.collect(this.exports.mirage_rs_metadata(this.handle), []));
    }
    /** What the stream wrote to stderr since the last call. */
    takeStderr() {
        return this.guard(() => this.collect(this.exports.mirage_rs_take_stderr(this.handle), []));
    }
    /** The host operation the stream waits on and the bytes it carries, null while it waits on none. */
    request() {
        return this.guard(() => {
            const len = this.exports.mirage_rs_request(this.handle);
            if (len < 0)
                return null;
            const taken = this.collect(0, []);
            const header = JSON.parse(decoder.decode(taken.subarray(0, len)));
            return [header, taken.subarray(len)];
        });
    }
    /** Answer the request taken last with the host's reply and the bytes it carries. */
    respond(reply, data) {
        const header = encoder.encode(JSON.stringify(reply));
        return this.guard(() => {
            for (const part of [header, ...(data instanceof Uint8Array ? [data] : data)]) {
                for (let offset = 0; offset < part.byteLength; offset += this.capacity) {
                    const piece = part.subarray(offset, offset + this.capacity);
                    const window = this.exports.mirage_rs_input(this.handle);
                    new Uint8Array(this.exports.memory.buffer, window, piece.byteLength).set(piece);
                    this.exports.mirage_rs_stage(this.handle, piece.byteLength);
                }
            }
            return this.collect(this.exports.mirage_rs_respond(this.handle, header.byteLength), []);
        });
    }
    close() {
        if (this.handle === 0)
            return;
        this.drop(true);
        const held = movables.get(this.exports);
        if (this.exports.mirage_rs_open_count() === (held?.size ?? 0) &&
            this.exports.memory.buffer.byteLength > RECYCLE_BYTES &&
            ![...(held ?? [])].some((stream) => stream.busy()))
            retire(this.exports, false);
    }
    /** Close the handle, on the instance itself unless it trapped. */
    drop(release) {
        if (this.handle === 0)
            return;
        if (release)
            this.exports.mirage_rs_close(this.handle);
        this.handle = 0;
        movables.get(this.exports)?.delete(this);
    }
    // A trap (a Rust panic) stops the instance mid-call with its state still
    // borrowed, so every later call on it would trap too: this stream is
    // done, and the next stream gets a fresh instance.
    guard(call) {
        try {
            return call();
        }
        catch (err) {
            if (err instanceof WebAssembly.RuntimeError) {
                this.drop(false);
                retire(this.exports, true);
            }
            throw err;
        }
    }
    push(chunk, feed) {
        // An empty chunk still reaches the stream: a sort starts a buffered run.
        if (chunk.byteLength === 0)
            return this.collect(feed(0), []);
        const parts = [];
        for (let offset = 0; offset < chunk.byteLength; offset += this.capacity) {
            const piece = chunk.subarray(offset, offset + this.capacity);
            const window = this.exports.mirage_rs_input(this.handle);
            new Uint8Array(this.exports.memory.buffer, window, piece.byteLength).set(piece);
            parts.push(this.collect(feed(piece.byteLength), parts));
        }
        const [only] = parts;
        return parts.length === 1 && only !== undefined ? only : concat(parts);
    }
    collect(status, before) {
        const { exports, handle } = this;
        const output = bytesAt(exports, exports.mirage_rs_output(handle), exports.mirage_rs_output_len(handle));
        if (status === 0)
            return output;
        const message = decoder.decode(bytesAt(exports, exports.mirage_rs_error(handle), exports.mirage_rs_error_len(handle)));
        throw new StreamRefusal(message, concat([...before, output]), exports.mirage_rs_error_status(handle));
    }
}
/**
 * Run a stream until it waits on neither work nor its host, a bounded call
 * at a time, serving each request it makes. Mirrors Python's serve_requests.
 */
export async function* serveRequests(stream, host, signal) {
    const budget = new YieldBudget(signal);
    for (;;) {
        while (stream.busy()) {
            await budget.run();
            const output = stream.drain();
            await host.errors.append(stream.takeStderr());
            if (output.byteLength > 0)
                yield output;
        }
        const taken = stream.request();
        if (taken === null)
            break;
        await budget.run();
        const served = await host.serve(...taken);
        if (served.output !== undefined && served.output !== null)
            yield* chunks(served.output, signal);
        const output = stream.respond(served.reply, served.data ?? new Uint8Array());
        // Each diagnostic is collected before the output that follows it, so a
        // reader that closes on that output still leaves it reported.
        await host.errors.append(stream.takeStderr());
        if (output.byteLength > 0)
            yield output;
    }
    const tail = stream.finish();
    await host.errors.append(stream.takeStderr());
    if (tail.byteLength > 0)
        yield tail;
}
/** Open the mirage-rs stream `kind`. Mirrors Python's open_stream. */
export async function openStream(kind, config) {
    const exports = await instance();
    return open(exports, kind, config);
}
/**
 * Open a stream for synchronous pure utilities, initializing Wasm on first use.
 * A `movable` stream does not keep a grown instance alive: it closes when the
 * instance retires, and its owner opens it again.
 */
export function openStreamSync(kind, config, movable = false) {
    synchronous ??= new WebAssembly.Instance(moduleSync(), {}).exports;
    return open(synchronous, kind, config, movable);
}
const movableFinalizer = new FinalizationRegistry((stream) => {
    stream.close();
});
/**
 * A movable synchronous stream for an object a cache keeps, such as a compiled
 * matcher: it opens again on the live instance once the one it was on retired,
 * and closes when its owner is collected. A retired stream is unregistered, so
 * it never holds its instance's memory.
 */
export class MovableStream {
    kind;
    config;
    stream;
    constructor(kind, config) {
        this.kind = kind;
        this.config = config;
        this.stream = this.open();
    }
    /** Close the stream now rather than when its owner is collected. */
    close() {
        movableFinalizer.unregister(this);
        this.stream.close();
    }
    /** The stream, opened again first when its instance retired. */
    get live() {
        if (!this.stream.live) {
            movableFinalizer.unregister(this);
            this.stream = this.open();
        }
        return this.stream;
    }
    open() {
        const stream = openStreamSync(this.kind, this.config, true);
        movableFinalizer.register(this, stream, this);
        return stream;
    }
}
function open(exports, kind, config, movable = false) {
    const kindBytes = encoder.encode(kind);
    const configBytes = encoder.encode(JSON.stringify(config));
    const args = exports.mirage_rs_args(kindBytes.byteLength + configBytes.byteLength);
    const view = new Uint8Array(exports.memory.buffer, args, kindBytes.byteLength + configBytes.byteLength);
    view.set(kindBytes);
    view.set(configBytes, kindBytes.byteLength);
    const handle = exports.mirage_rs_open(kindBytes.byteLength, configBytes.byteLength);
    if (handle === 0) {
        const message = decoder.decode(bytesAt(exports, exports.mirage_rs_error(0), exports.mirage_rs_error_len(0)));
        throw new StreamRefusal(message, new Uint8Array(), exports.mirage_rs_error_status(0));
    }
    return new OpenStream(exports, handle, movable);
}
/**
 * Stream `source` through `kind`, one bounded chunk at a time. A refusal
 * throws StreamRefusal holding the output written before it, so the caller
 * records the failure before it passes that output on: a reader that stops
 * at the output would otherwise never see it. Mirrors Python's run_stream.
 */
export async function* runStream(kind, config, source, signal) {
    const stream = await openStream(kind, config);
    yield* streamOutputChunks(stream, source, signal);
}
/**
 * Feed a multi-input stream the input its `wants` names until it wants none,
 * handing on its output a bounded call at a time. Mirrors Python's
 * merge_output.
 */
export async function* mergeOutput(stream, read, signal) {
    const budget = new YieldBudget(signal);
    for (;;) {
        while (stream.busy()) {
            await budget.run();
            const output = stream.drain();
            if (output.byteLength > 0)
                yield output;
        }
        const index = stream.wants();
        if (index === null)
            break;
        await budget.run();
        const chunk = await read(index);
        const output = chunk === null ? stream.endInput(index) : stream.feedInput(index, chunk);
        if (output.byteLength > 0)
            yield output;
    }
    const tail = stream.finish();
    if (tail.byteLength > 0)
        yield tail;
    while (stream.busy()) {
        await budget.run();
        const output = stream.drain();
        if (output.byteLength > 0)
            yield output;
    }
}
/**
 * Feed a stream that reads its inputs one after another each source in turn,
 * through mergeOutput, and close the sources. Mirrors Python's chain_output.
 */
export async function* chainOutput(stream, sources, signal) {
    const inputs = sources.map((source) => chunks(source, signal));
    try {
        yield* mergeOutput(stream, async (index) => {
            const next = await inputs[index]?.next();
            return next === undefined || next.done === true ? null : next.value;
        }, signal);
    }
    finally {
        for (const [index, input] of inputs.entries()) {
            await input.return?.(undefined);
            await sources[index].return?.();
        }
    }
}
export async function* streamOutputChunks(stream, source, signal, ownsStream = true) {
    const budget = new YieldBudget(signal);
    try {
        for await (const chunk of chunks(source, signal)) {
            const output = stream.feed(chunk);
            if (output.byteLength > 0)
                yield output;
            while (stream.busy()) {
                await budget.run();
                const output = stream.drain();
                if (output.byteLength > 0)
                    yield output;
            }
            if (stream.done())
                break;
        }
        const tail = stream.finish();
        if (tail.byteLength > 0)
            yield tail;
        while (stream.busy()) {
            await budget.run();
            const output = stream.drain();
            if (output.byteLength > 0)
                yield output;
        }
    }
    finally {
        if (ownsStream)
            await closeStream(stream);
    }
}
/** Retire cancelled records through bounded calls before releasing the handle. */
export async function closeStream(stream) {
    if (!stream.live)
        return;
    try {
        if (stream.cancel()) {
            const budget = new YieldBudget();
            while (stream.busy()) {
                await budget.run();
                stream.drain();
            }
        }
    }
    finally {
        stream.close();
    }
}
//# sourceMappingURL=stream.js.map