assistant-stream/dist/core/utils/stream/merge.js

Streaming utilities for AI assistants

import { Counter } from "../Counter.js";
import { createPathMergeMapper } from "./path-utils.js";
import { promiseWithResolvers } from "../../../utils/promiseWithResolvers.js";
// src/core/utils/stream/merge.ts
const createMergeStream = () => {
	const list = [];
	const pendingPartIndexes = /* @__PURE__ */ new Set();
	const contentCounter = new Counter();
	let sealed = false;
	let cancelled = false;
	let errored = false;
	let controller;
	let rawChunkBatch;
	let rawBarrier;
	let pendingRawBatches = 0;
	let currentPull;
	let cleanupPromise;
	const releaseHolds = (item) => {
		for (const hold of item.holds) hold.reached.resolve();
		item.holds = [];
	};
	const cancelAllReaders = () => {
		rawChunkBatch = void 0;
		for (const partIndex of pendingPartIndexes) {
			partIndex.reference.value = -1;
			partIndex.resolve(-1);
		}
		pendingPartIndexes.clear();
		if (!cleanupPromise) {
			const items = list.splice(0);
			for (const item of items) {
				item.orderedFinish?.delivered.resolve();
				releaseHolds(item);
			}
			cleanupPromise = Promise.all(items.map(async (item) => {
				try {
					await item.reader.cancel().catch(() => void 0);
					await item.pipeTask;
				} finally {
					item.reader.releaseLock();
				}
			})).then(() => void 0);
		}
		return cleanupPromise;
	};
	const handleError = (e) => {
		if (cancelled || errored) return;
		errored = true;
		console.error(e);
		cancelAllReaders();
		controller.error(e);
		currentPull?.reject(e);
		currentPull = void 0;
	};
	const handlePull = (item) => {
		if (!item.promise) item.promise = item.reader.read().then(async ({ done, value }) => {
			if (cancelled || errored) return;
			if (done) {
				item.orderedFinish?.delivered.resolve();
				list.splice(list.indexOf(item), 1);
				item.reader.releaseLock();
				releaseHolds(item);
				if (sealed && list.length === 0 && pendingRawBatches === 0) controller.close();
			} else {
				if (item.gate) {
					await item.gate;
					item.gate = void 0;
					if (cancelled || errored) return;
				}
				while (item.holds[0] && item.delivered >= item.holds[0].writtenBefore) {
					await item.holds[0].flushed;
					item.holds.shift();
					if (cancelled || errored) return;
				}
				if (value.type === "part-finish" && item.orderedFinish) {
					await item.orderedFinish.previous();
					if (cancelled || errored) return;
				}
				controller.enqueue(item.pathMapper?.(value) ?? value);
				item.delivered++;
				for (const hold of item.holds) if (item.delivered >= hold.writtenBefore) hold.reached.resolve();
				if (value.type === "part-finish") item.orderedFinish?.delivered.resolve();
			}
			item.promise = void 0;
			currentPull?.resolve();
			currentPull = void 0;
		}).catch((error) => {
			item.promise = void 0;
			handleError(error);
		});
	};
	const readable = new ReadableStream({
		start(c) {
			controller = c;
		},
		pull() {
			currentPull = promiseWithResolvers();
			list.forEach((item) => {
				handlePull(item);
			});
			return currentPull.promise;
		},
		async cancel() {
			cancelled = true;
			const cleanup = cancelAllReaders();
			currentPull?.resolve();
			currentPull = void 0;
			await cleanup;
		}
	});
	const holdChild = (item, flushed) => {
		const hold = {
			writtenBefore: item.written(),
			reached: promiseWithResolvers(),
			flushed
		};
		if (item.delivered >= hold.writtenBefore) hold.reached.resolve();
		item.holds.push(hold);
		return hold.reached.promise;
	};
	const enqueueRawChunk = (chunk, partIndex) => {
		const waitForPartWrites = chunk.type === "part-start" || chunk.type === "step-finish" || chunk.type === "message-finish" || chunk.type === "error";
		if (list.length > 0 || waitForPartWrites) rawChunkBatch = void 0;
		if (!rawChunkBatch) {
			const batch = [];
			rawChunkBatch = batch;
			pendingRawBatches++;
			const previous = rawBarrier;
			const parts = waitForPartWrites ? list.filter((item) => item.written) : [];
			const flushed = previous || parts.length > 0 ? promiseWithResolvers() : void 0;
			const reached = flushed ? parts.map((item) => holdChild(item, flushed.promise)) : [];
			if (flushed) rawBarrier = flushed.promise;
			Promise.resolve().then(async () => {
				if (previous) await previous;
				if (reached.length > 0) await Promise.all(reached);
				pendingRawBatches--;
				if (rawChunkBatch === batch) rawChunkBatch = void 0;
				if (cancelled || errored) return;
				for (const { chunk: rawChunk, partIndex } of batch) {
					if (rawChunk.type === "part-start" && rawChunk.path.length === 0) {
						const index = contentCounter.up();
						if (partIndex) {
							pendingPartIndexes.delete(partIndex);
							partIndex.reference.value = index;
							partIndex.resolve(index);
						}
					}
					controller.enqueue(rawChunk);
				}
				if (sealed && list.length === 0 && pendingRawBatches === 0) controller.close();
				currentPull?.resolve();
				currentPull = void 0;
			}).catch(handleError).finally(() => {
				if (!flushed) return;
				flushed.resolve();
				if (rawBarrier === flushed.promise) rawBarrier = void 0;
			});
		}
		rawChunkBatch.push({
			chunk,
			partIndex
		});
		if (waitForPartWrites) rawChunkBatch = void 0;
	};
	const addStream = (stream, pipeTask, options) => {
		const handledPipeTask = pipeTask?.catch(() => void 0);
		if (cancelled || errored) {
			stream.cancel().catch(() => void 0);
			return;
		}
		if (sealed) {
			stream.cancel().catch(() => void 0);
			throw new Error("Cannot add streams after the run callback has settled.");
		}
		rawChunkBatch = void 0;
		const item = {
			reader: stream.getReader(),
			pipeTask: handledPipeTask,
			orderedFinish: options?.orderedFinish,
			written: options?.written,
			delivered: 0,
			gate: rawBarrier,
			holds: [],
			pathMapper: options?.pathTransform === "merge" ? createPathMergeMapper(() => contentCounter.up()) : void 0
		};
		list.push(item);
		handlePull(item);
	};
	return {
		readable,
		isSealed() {
			return sealed;
		},
		isCancelled() {
			return cancelled;
		},
		isErrored() {
			return errored;
		},
		seal() {
			if (sealed || cancelled || errored) return;
			sealed = true;
			if (list.length === 0 && pendingRawBatches === 0) controller.close();
		},
		addStream,
		enqueuePartStart(chunk) {
			if (cancelled || errored) return {
				value: -1,
				promise: Promise.resolve(-1)
			};
			if (sealed) throw new Error("Cannot add streams after the run callback has settled.");
			const deferred = promiseWithResolvers();
			const reference = {
				value: void 0,
				promise: deferred.promise
			};
			const reservation = {
				reference,
				resolve: deferred.resolve
			};
			pendingPartIndexes.add(reservation);
			enqueueRawChunk(chunk, reservation);
			return reference;
		},
		enqueue(chunk) {
			if (cancelled || errored) return;
			if (sealed) throw new Error("Cannot add streams after the run callback has settled.");
			enqueueRawChunk(chunk);
		}
	};
};
export { createMergeStream };

//# sourceMappingURL=merge.js.map