@oh-my-pi/pi-utils 11.2.2 → 11.2.3

package/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "@oh-my-pi/pi-utils",
-	"version": "11.2.2",
+	"version": "11.2.3",
 	"description": "Shared utilities for pi packages",
 	"type": "module",
 	"main": "./src/index.ts",

package/src/stream.ts

@@ -1,5 +1,3 @@
-import { ArrayBufferSink } from "bun";
-
 /**
  * Sanitize binary output for display/storage.
  * Removes characters that crash string-width or cause display issues:
@@ -42,58 +40,70 @@ export function sanitizeText(text: string): string {
 	return sanitizeBinaryOutput(Bun.stripANSI(text)).replace(/\r/g, "");
 }
 
-/**
- * Create a transform stream that splits lines.
- */
-export function createSplitterStream<T>(options: {
-	newLine?: boolean;
-	mapFn: (chunk: Uint8Array) => T;
-}): TransformStream<Uint8Array, T> {
-	const { newLine = false, mapFn } = options;
-	const LF = 0x0a;
-	const sink = new Bun.ArrayBufferSink();
-	sink.start({ asUint8Array: true, stream: true, highWaterMark: 4096 });
-	let pending = false; // whether the sink has unflushed data
-
-	return new TransformStream<Uint8Array, T>({
-		transform(chunk, ctrl) {
-			let pos = 0;
-
-			while (pos < chunk.length) {
-				const nl = chunk.indexOf(LF, pos);
-				if (nl === -1) {
-					sink.write(chunk.subarray(pos));
-					pending = true;
-					break;
-				}
-
-				const slice = chunk.subarray(pos, newLine ? nl + 1 : nl);
+const LF = 0x0a;
 
-				if (pending) {
-					if (slice.length > 0) sink.write(slice);
-					ctrl.enqueue(mapFn(sink.flush() as Uint8Array));
-					pending = false;
-				} else {
-					ctrl.enqueue(mapFn(slice));
-				}
-				pos = nl + 1;
+export async function* readLines(stream: ReadableStream<Uint8Array>, signal?: AbortSignal): AsyncGenerator<Uint8Array> {
+	const buffer = new ConcatSink();
+	const source = signal ? stream.pipeThrough(new TransformStream(), { signal }) : stream;
+	try {
+		for await (const chunk of source) {
+			for (const line of buffer.appendAndFlushLines(chunk)) {
+				yield line;
 			}
-		},
-		flush(ctrl) {
-			if (pending) {
-				const tail = sink.end() as Uint8Array;
-				if (tail.length > 0) ctrl.enqueue(mapFn(tail));
+		}
+		if (!buffer.isEmpty) {
+			const tail = buffer.flush();
+			if (tail) {
+				buffer.clear();
+				yield tail;
 			}
-		},
-	});
+		}
+	} catch (err) {
+		// Abort errors are expected — just stop the generator.
+		if (signal?.aborted) return;
+		throw err;
+	}
 }
 
-export function createTextLineSplitter(sanitize = false): TransformStream<Uint8Array, string> {
-	const dec = new TextDecoder("utf-8", { ignoreBOM: true, fatal: true });
-	if (sanitize) {
-		return createSplitterStream({ mapFn: chunk => sanitizeText(dec.decode(chunk)) });
+export async function* readJsonl<T>(stream: ReadableStream<Uint8Array>, signal?: AbortSignal): AsyncGenerator<T> {
+	const buffer = new ConcatSink();
+	const source = signal ? stream.pipeThrough(new TransformStream(), { signal }) : stream;
+	try {
+		const yieldBuffer: T[] = [];
+		for await (const chunk of source) {
+			buffer.appendAndConsume(chunk, 0, chunk.length, (payload, beg, end) => {
+				const { values, error, read, done } = Bun.JSONL.parseChunk(payload, beg, end);
+				if (values.length > 0) {
+					yieldBuffer.push(...(values as T[]));
+				}
+				if (error) throw error;
+				if (done) return 0;
+				return end - read;
+			});
+			if (yieldBuffer.length > 0) {
+				yield* yieldBuffer;
+				yieldBuffer.length = 0;
+			}
+		}
+		if (!buffer.isEmpty) {
+			const tail = buffer.flush();
+			if (tail) {
+				buffer.clear();
+				const { values, error, done } = Bun.JSONL.parseChunk(tail, 0, tail.length);
+				if (values.length > 0) {
+					yield* values as T[];
+				}
+				if (error) throw error;
+				if (!done) {
+					throw new Error("JSONL stream ended unexpectedly");
+				}
+			}
+		}
+	} catch (err) {
+		// Abort errors are expected — just stop the generator.
+		if (signal?.aborted) return;
+		throw err;
 	}
-	return createSplitterStream({ mapFn: dec.decode.bind(dec) });
 }
 
 /**
@@ -118,28 +128,169 @@ export function createTextDecoderStream(): TransformStream<Uint8Array, string> {
 // SSE (Server-Sent Events)
 // =============================================================================
 
-const LF = 0x0a;
-const CR = 0x0d;
-const SPACE = 0x20;
+class Bitmap {
+	private bits: Uint32Array;
+	constructor(n: number) {
+		this.bits = new Uint32Array((n + 31) >>> 5);
+	}
+
+	set(i: number, value: boolean) {
+		const index = i >>> 5;
+		const mask = 1 << (i & 31);
+		if (value) {
+			this.bits[index] |= mask;
+		} else {
+			this.bits[index] &= ~mask;
+		}
+	}
+	get(i: number) {
+		const index = i >>> 5;
+		const mask = 1 << (i & 31);
+		const word = this.bits[index];
+		return word !== undefined && (word & mask) !== 0;
+	}
+}
+
+const WHITESPACE = new Bitmap(256);
+for (let i = 0; i <= 0x7f; i++) {
+	const c = String.fromCharCode(i);
+	switch (c) {
+		case " ":
+		case "\t":
+		case "\n":
+		case "\r":
+			WHITESPACE.set(i, true);
+			break;
+		default:
+			WHITESPACE.set(i, !c.trim());
+			break;
+	}
+}
+
+const createPattern = (prefix: string) => {
+	const pre = Buffer.from(prefix, "utf-8");
+	return {
+		strip(buf: Uint8Array): number | null {
+			const n = pre.length;
+			if (buf.length < n) return null;
+			if (pre.equals(buf.subarray(0, n))) {
+				return n;
+			}
+			return null;
+		},
+	};
+};
+
+const PAT_DATA = createPattern("data:");
+
+const PAT_DONE = createPattern("[DONE]");
+
+class ConcatSink {
+	#space?: Buffer;
+	#length = 0;
 
-// "data:" = [0x64, 0x61, 0x74, 0x61, 0x3a]
-const DATA_0 = 0x64; // d
-const DATA_1 = 0x61; // a
-const DATA_2 = 0x74; // t
-const DATA_3 = 0x61; // a
-const DATA_4 = 0x3a; // :
+	#ensureCapacity(size: number): Buffer {
+		const space = this.#space;
+		if (space && space.length >= size) return space;
+		const nextSize = space ? Math.max(size, space.length * 2) : size;
+		const next = Buffer.allocUnsafe(nextSize);
+		if (space && this.#length > 0) {
+			space.copy(next, 0, 0, this.#length);
+		}
+		this.#space = next;
+		return next;
+	}
+
+	append(chunk: Uint8Array) {
+		const n = chunk.length;
+		if (!n) return;
+		const offset = this.#length;
+		const space = this.#ensureCapacity(offset + n);
+		space.set(chunk, offset);
+		this.#length += n;
+	}
+
+	reset(chunk: Uint8Array) {
+		const n = chunk.length;
+		if (!n) {
+			this.#length = 0;
+			return;
+		}
+		const space = this.#ensureCapacity(n);
+		space.set(chunk, 0);
+		this.#length = n;
+	}
+
+	get isEmpty(): boolean {
+		return this.#length === 0;
+	}
 
-// "[DONE]" = [0x5b, 0x44, 0x4f, 0x4e, 0x45, 0x5d]
-const DONE = Uint8Array.from([0x5b, 0x44, 0x4f, 0x4e, 0x45, 0x5d]);
+	flush(): Uint8Array | undefined {
+		if (!this.#length) return undefined;
+		return this.#space!.subarray(0, this.#length);
+	}
 
-function isDone(buf: Uint8Array, start: number, end: number): boolean {
-	if (end - start !== 6) return false;
-	for (let i = 0; i < 6; i++) {
-		if (buf[start + i] !== DONE[i]) return false;
+	clear() {
+		this.#length = 0;
+	}
+
+	*appendAndFlushLines(chunk: Uint8Array) {
+		let pos = 0;
+		while (pos < chunk.length) {
+			const nl = chunk.indexOf(LF, pos);
+			if (nl === -1) {
+				this.append(chunk.subarray(pos));
+				return;
+			}
+			const suffix = chunk.subarray(pos, nl);
+			pos = nl + 1;
+			if (this.isEmpty) {
+				yield suffix;
+			} else {
+				this.append(suffix);
+				const payload = this.flush();
+				if (payload) {
+					yield payload;
+					this.clear();
+				}
+			}
+		}
+	}
+
+	appendAndConsume(
+		chunk: Uint8Array,
+		beg: number,
+		end: number,
+		// (slice) => [remaining length]
+		consumer: (payload: Uint8Array, beg: number, end: number) => number,
+	) {
+		if (this.isEmpty) {
+			const rem = consumer(chunk, beg, end);
+			if (!rem) return;
+			this.reset(chunk.subarray(end - rem, end));
+			return;
+		}
+
+		const offset = this.#length;
+		const n = end - beg;
+		const total = offset + n;
+		const space = this.#ensureCapacity(total);
+		space.set(chunk.subarray(beg, end), offset);
+		this.#length = total;
+		const rem = consumer(space.subarray(0, total), 0, total);
+		if (!rem) {
+			this.#length = 0;
+			return;
+		}
+		if (rem < total) {
+			space.copyWithin(0, total - rem, total);
+		}
+		this.#length = rem;
 	}
-	return true;
 }
 
+const kDoneError = new Error("SSE stream done");
+
 /**
  * Stream parsed JSON objects from SSE `data:` lines.
  *
@@ -150,96 +301,113 @@ function isDone(buf: Uint8Array, start: number, end: number): boolean {
  * }
  * ```
  */
-export async function* readSseJson<T>(
-	stream: ReadableStream<Uint8Array>,
-	abortSignal?: AbortSignal,
-): AsyncGenerator<T> {
-	const sink = new ArrayBufferSink();
-	sink.start({ asUint8Array: true, stream: true, highWaterMark: 4096 });
-	let pending = false;
+export async function* readSseJson<T>(stream: ReadableStream<Uint8Array>, signal?: AbortSignal): AsyncGenerator<T> {
+	const lineBuffer = new ConcatSink();
+	const jsonBuffer = new ConcatSink();
 
 	// pipeThrough with { signal } makes the stream abort-aware: the pipe
 	// cancels the source and errors the output when the signal fires,
 	// so for-await-of exits cleanly without manual reader/listener management.
-	const source = abortSignal ? stream.pipeThrough(new TransformStream(), { signal: abortSignal }) : stream;
-
+	stream = signal ? stream.pipeThrough(new TransformStream(), { signal }) : stream;
 	try {
-		for await (const chunk of source) {
-			let pos = 0;
-			while (pos < chunk.length) {
-				const nl = chunk.indexOf(LF, pos);
-				if (nl === -1) {
-					sink.write(chunk.subarray(pos));
-					pending = true;
-					break;
-				}
+		const yieldBuffer: T[] = [];
+		const processLine = (line: Uint8Array) => {
+			// Strip trailing spaces including \r.
+			let end = line.length;
+			while (end && WHITESPACE.get(line[end - 1])) {
+				--end;
+			}
+			if (!end) return; // blank line
+
+			const trimmed = end === line.length ? line : line.subarray(0, end);
+
+			// Check "data:" prefix and optional space afterwards.
+			let beg = PAT_DATA.strip(trimmed);
+			if (beg === null) return;
+			while (beg < end && WHITESPACE.get(trimmed[beg])) {
+				++beg;
+			}
+			if (beg >= end) return;
 
-				let line: Uint8Array;
-				if (pending) {
-					if (nl > pos) sink.write(chunk.subarray(pos, nl));
-					line = sink.flush() as Uint8Array;
-					pending = false;
-				} else {
-					line = chunk.subarray(pos, nl);
+			jsonBuffer.appendAndConsume(trimmed, beg, end, (payload, beg, end) => {
+				const { values, error, read, done } = Bun.JSONL.parseChunk(payload, beg, end);
+				if (values.length > 0) {
+					yieldBuffer.push(...(values as T[]));
+				}
+				if (error) {
+					if (PAT_DONE.strip(payload.subarray(beg, end))) {
+						throw kDoneError;
+					}
+					throw error;
+				}
+				if (done) return 0;
+				return end - read;
+			});
+		};
+		for await (const chunk of stream) {
+			for (const line of lineBuffer.appendAndFlushLines(chunk)) {
+				processLine(line);
+				if (yieldBuffer.length > 0) {
+					yield* yieldBuffer;
+					yieldBuffer.length = 0;
+				}
+			}
+		}
+		if (!lineBuffer.isEmpty) {
+			const tail = lineBuffer.flush();
+			if (tail) {
+				lineBuffer.clear();
+				processLine(tail);
+				if (yieldBuffer.length > 0) {
+					yield* yieldBuffer;
+					yieldBuffer.length = 0;
 				}
-				pos = nl + 1;
-
-				// Strip trailing CR, skip blank/short lines.
-				const len = line.length > 0 && line[line.length - 1] === CR ? line.length - 1 : line.length;
-				if (len < 6) continue; // "data:" + at least 1 byte
-
-				// Check "data:" prefix.
-				if (
-					line[0] !== DATA_0 ||
-					line[1] !== DATA_1 ||
-					line[2] !== DATA_2 ||
-					line[3] !== DATA_3 ||
-					line[4] !== DATA_4
-				)
-					continue;
-
-				// Payload start — skip optional space after colon.
-				const pStart = line[5] === SPACE ? 6 : 5;
-				if (pStart >= len) continue;
-				if (isDone(line, pStart, len)) return;
-
-				// Build payload + \n for JSONL.parse.
-				const pLen = len - pStart;
-				const buf = new Uint8Array(pLen + 1);
-				buf.set(line.subarray(pStart, len));
-				buf[pLen] = LF;
-
-				const [parsed] = Bun.JSONL.parse(buf);
-				if (parsed !== undefined) yield parsed as T;
 			}
 		}
 	} catch (err) {
+		if (err === kDoneError) return;
 		// Abort errors are expected — just stop the generator.
-		if (abortSignal?.aborted) return;
+		if (signal?.aborted) return;
 		throw err;
 	}
+	if (!jsonBuffer.isEmpty) {
+		throw new Error("SSE stream ended unexpectedly");
+	}
+}
 
-	// Trailing line without final newline.
-	if (pending) {
-		const tail = sink.end() as Uint8Array;
-		const len = tail.length > 0 && tail[tail.length - 1] === CR ? tail.length - 1 : tail.length;
-		if (
-			len >= 6 &&
-			tail[0] === DATA_0 &&
-			tail[1] === DATA_1 &&
-			tail[2] === DATA_2 &&
-			tail[3] === DATA_3 &&
-			tail[4] === DATA_4
-		) {
-			const pStart = tail[5] === SPACE ? 6 : 5;
-			if (pStart < len && !isDone(tail, pStart, len)) {
-				const pLen = len - pStart;
-				const buf = new Uint8Array(pLen + 1);
-				buf.set(tail.subarray(pStart, len));
-				buf[pLen] = LF;
-				const [parsed] = Bun.JSONL.parse(buf);
-				if (parsed !== undefined) yield parsed as T;
+/**
+ * Parse a complete JSONL string, skipping malformed lines instead of throwing.
+ *
+ * Uses `Bun.JSONL.parseChunk` internally. On parse errors, the malformed
+ * region is skipped up to the next newline and parsing continues.
+ *
+ * @example
+ * ```ts
+ * const entries = parseJsonlLenient<MyType>(fileContents);
+ * ```
+ */
+export function parseJsonlLenient<T>(buffer: string): T[] {
+	let entries: T[] | undefined;
+
+	while (buffer.length > 0) {
+		const { values, error, read, done } = Bun.JSONL.parseChunk(buffer);
+		if (values.length > 0) {
+			const ext = values as T[];
+			if (!entries) {
+				entries = ext;
+			} else {
+				entries.push(...ext);
 			}
 		}
+		if (error) {
+			const nextNewline = buffer.indexOf("\n", read);
+			if (nextNewline === -1) break;
+			buffer = buffer.substring(nextNewline + 1);
+			continue;
+		}
+		if (read === 0) break;
+		buffer = buffer.substring(read);
+		if (done) break;
 	}
+	return entries ?? [];
 }