@linxiraos/pi-utils 1.1.14 → 1.1.16

package/src/sqlite.ts

@@ -1,13 +1,248 @@
+/** Shared SQLite opening, error attribution, and result-code classification for persistent stores. */
+import { Database } from "bun:sqlite";
+import * as fs from "node:fs";
+import { getDbBusyTimeoutMs } from "./env";
+import { withFileLockSync } from "./file-lock";
+import { isEnoent } from "./fs-error";
+import * as logger from "./logger";
+
+const BUSY_MAX_ATTEMPTS = 4;
+const BUSY_BASE_DELAY_MS = 100;
+const SQLITE_STORE_SUFFIXES = ["-wal", "-shm", "-journal", ""];
+
+type SqliteFileIdentity = string | null | undefined;
+
+class SqliteAttemptFailure extends Error {
+	readonly original: unknown;
+	readonly identity: SqliteFileIdentity;
+	readonly canRecover: boolean;
+	readonly db?: Database;
+
+	constructor(original: unknown, identity: SqliteFileIdentity, options: { canRecover?: boolean; db?: Database } = {}) {
+		super(original instanceof Error ? original.message : String(original));
+		this.original = original;
+		this.identity = identity;
+		this.canRecover = options.canRecover ?? true;
+		this.db = options.db;
+	}
+}
+
+function sqliteFileIdentity(dbPath: string): SqliteFileIdentity {
+	try {
+		const stat = fs.statSync(dbPath);
+		return `${stat.dev}:${stat.ino}:${stat.birthtimeMs}`;
+	} catch (error) {
+		return isEnoent(error) ? null : undefined;
+	}
+}
+
+function closeFailedDatabase(db: Database | undefined, error: unknown, identity: SqliteFileIdentity): void {
+	try {
+		db?.close();
+	} catch (closeError) {
+		const original = error instanceof Error ? error : new Error(String(error));
+		const detail = closeError instanceof Error ? closeError.message : String(closeError);
+		original.message += `; failed to close the SQLite handle: ${detail}`;
+		throw new SqliteAttemptFailure(original, identity, { canRecover: false });
+	}
+}
+
+/** Controls opt-in replacement of an unrecoverably corrupt SQLite store. */
+export interface SqliteOpenOptions {
+	/**
+	 * Preserve a corrupt store and its sidecars, recreate it, and run the
+	 * initializer once more. Disabled by default.
+	 */
+	recoverCorruption?: boolean;
+	/** Runs after preservation and before the replacement is initialized. */
+	onCorruptionPreserved?: (backupPath: string, error: unknown) => void;
+}
+
+async function openWithBusyRetries<T>(
+	dbPath: string,
+	initialize: (db: Database) => T | Promise<T>,
+	options: SqliteOpenOptions,
+): Promise<T> {
+	for (let attempt = 0; ; attempt++) {
+		let db: Database | undefined;
+		const identity = sqliteFileIdentity(dbPath);
+		try {
+			db = new Database(dbPath);
+			// WAL recovery can bypass the busy handler; both it and retries are needed (#2421).
+			db.run(`PRAGMA busy_timeout = ${getDbBusyTimeoutMs()}`);
+			return await initialize(db);
+		} catch (error) {
+			if (options.recoverCorruption && isSqliteCorruptionError(error)) {
+				throw new SqliteAttemptFailure(error, identity, { db });
+			}
+			closeFailedDatabase(db, error, identity);
+			if (!isSqliteBusyError(error) || attempt + 1 >= BUSY_MAX_ATTEMPTS) {
+				throw new SqliteAttemptFailure(error, identity);
+			}
+			await Bun.sleep(BUSY_BASE_DELAY_MS * 2 ** attempt);
+		}
+	}
+}
+
+function openOnce<T>(dbPath: string, initialize: (db: Database) => T, options: SqliteOpenOptions): T {
+	let db: Database | undefined;
+	const identity = sqliteFileIdentity(dbPath);
+	try {
+		db = new Database(dbPath);
+		db.run(`PRAGMA busy_timeout = ${getDbBusyTimeoutMs()}`);
+		return initialize(db);
+	} catch (error) {
+		if (options.recoverCorruption && isSqliteCorruptionError(error)) {
+			throw new SqliteAttemptFailure(error, identity, { db });
+		}
+		closeFailedDatabase(db, error, identity);
+		throw new SqliteAttemptFailure(error, identity);
+	}
+}
+
+function quarantineCorruptSqliteStore(dbPath: string, db: Database | undefined): string {
+	const backupPath = `${dbPath}.corrupt-${Date.now()}-${crypto.randomUUID()}`;
+	const preserved: string[] = [];
+	// Closing a failed WAL connection can truncate its WAL. Copy evidence
+	// before closing, and remove originals only after every copy succeeds.
+	for (const suffix of SQLITE_STORE_SUFFIXES) {
+		try {
+			fs.chmodSync(`${dbPath}${suffix}`, 0o600);
+			fs.copyFileSync(`${dbPath}${suffix}`, `${backupPath}${suffix}`, fs.constants.COPYFILE_EXCL);
+			preserved.push(suffix);
+		} catch (error) {
+			if (isEnoent(error) && suffix !== "") continue;
+			throw error;
+		}
+	}
+	db?.close();
+
+	const removed: string[] = [];
+	try {
+		// Remove the main file last so a failed sidecar removal cannot leave
+		// a path at which another startup creates an empty database.
+		for (const suffix of preserved) {
+			try {
+				fs.unlinkSync(`${dbPath}${suffix}`);
+				removed.push(suffix);
+			} catch (error) {
+				if (!isEnoent(error)) throw error;
+			}
+		}
+	} catch (error) {
+		for (const suffix of removed) {
+			try {
+				fs.copyFileSync(`${backupPath}${suffix}`, `${dbPath}${suffix}`, fs.constants.COPYFILE_EXCL);
+			} catch (rollbackError) {
+				logger.error("SQLite quarantine rollback failed; original preserved at backup path", {
+					path: `${dbPath}${suffix}`,
+					backupPath: `${backupPath}${suffix}`,
+					error: String(rollbackError),
+				});
+			}
+		}
+		throw error;
+	}
+	return backupPath;
+}
+
+function corruptionPreservationError(corruption: unknown, dbPath: string, preservationError: unknown): Error {
+	const annotated = annotateSqliteError(corruption, dbPath);
+	const detail = preservationError instanceof Error ? preservationError.message : String(preservationError);
+	annotated.message += `; failed to preserve the corrupt database: ${detail}`;
+	return annotated;
+}
+
+function recoverCorruptDatabase(dbPath: string, error: unknown, options: SqliteOpenOptions): void {
+	if (!(error instanceof SqliteAttemptFailure)) throw annotateSqliteError(error, dbPath);
+	const failure = error;
+	if (!options.recoverCorruption || !failure.canRecover || !isSqliteCorruptionError(failure.original)) {
+		throw annotateSqliteError(failure.original, dbPath);
+	}
+
+	let backupPath: string | null;
+	try {
+		try {
+			backupPath = withFileLockSync(`${dbPath}.recovery`, () => {
+				const currentIdentity = sqliteFileIdentity(dbPath);
+				if (failure.identity === undefined || currentIdentity === undefined) {
+					throw new Error("could not verify the corrupt database file identity");
+				}
+				if (currentIdentity !== failure.identity) return null;
+				return quarantineCorruptSqliteStore(dbPath, failure.db);
+			});
+		} finally {
+			closeFailedDatabase(failure.db, failure.original, failure.identity);
+		}
+	} catch (preservationError) {
+		throw corruptionPreservationError(failure.original, dbPath, preservationError);
+	}
+
+	if (backupPath === null) return;
+	logger.warn("SQLite database corrupt; preserved damaged store before recreating it", {
+		path: dbPath,
+		backupPath,
+		warning: "Stored credentials from this database may require re-login.",
+	});
+	options.onCorruptionPreserved?.(backupPath, failure.original);
+}
+
 /**
- * Shared classifiers for `bun:sqlite` error result codes.
+ * Opens and initializes a store, retrying BUSY failures up to four total attempts.
+ * Installs the busy handler before initialization and closes failed connections.
+ * The initializer may run again on a fresh connection; on success it owns the handle.
  *
- * Every omp SQLite store (`agent.db` credential/usage store, `models.db` model
- * cache, `history.db`) needs the same two distinctions: a transient BUSY that
- * clears by retrying, and an unrecoverable corruption that never does. Keeping
- * one implementation here prevents the classifiers from drifting between the
- * credential store and the model cache.
+ * With corruption recovery enabled, recovery is serialized across processes.
+ * The identity observed by the failed handle is checked under that lock, so a
+ * waiter adopts a replacement made by a peer instead of quarantining it.
+ * Final failures retain their SQLite codes and include the database path.
  */
-import type { Database } from "bun:sqlite";
+export async function openSqliteDatabase<T>(
+	dbPath: string,
+	initialize: (db: Database) => T | Promise<T>,
+	options: SqliteOpenOptions = {},
+): Promise<T> {
+	try {
+		return await openWithBusyRetries(dbPath, initialize, options);
+	} catch (error) {
+		recoverCorruptDatabase(dbPath, error, options);
+	}
+
+	try {
+		return await openWithBusyRetries(dbPath, initialize, {});
+	} catch (error) {
+		throw annotateSqliteError(error instanceof SqliteAttemptFailure ? error.original : error, dbPath);
+	}
+}
+
+/**
+ * Synchronous counterpart to {@link openSqliteDatabase}. It performs no BUSY
+ * retry loop; corruption recovery, when enabled, is bounded to one replacement.
+ */
+export function openSqliteDatabaseSync<T>(
+	dbPath: string,
+	initialize: (db: Database) => T,
+	options: SqliteOpenOptions = {},
+): T {
+	try {
+		return openOnce(dbPath, initialize, options);
+	} catch (error) {
+		recoverCorruptDatabase(dbPath, error, options);
+	}
+
+	try {
+		return openOnce(dbPath, initialize, {});
+	} catch (error) {
+		throw annotateSqliteError(error instanceof SqliteAttemptFailure ? error.original : error, dbPath);
+	}
+}
+
+/** Adds the failing store's path to an error without losing SQLite result codes or its original stack. */
+export function annotateSqliteError(error: unknown, dbPath: string): Error {
+	const annotated = error instanceof Error ? error : new Error(String(error));
+	annotated.message = `Database ${JSON.stringify(dbPath)}: ${annotated.message}`;
+	return annotated;
+}
 
 /** Checkpoints committed WAL frames without waiting for concurrent readers. */
 export function checkpointWal(db: Database): void {

package/src/stream.ts

@@ -4,7 +4,14 @@ import { abortableSource } from "./abortable";
 import { parseStreamingJson } from "./json-parse";
 
 const LF = 0x0a;
+const CR = 0x0d;
 
+/**
+ * Split a byte stream on LF boundaries.
+ *
+ * Every yielded line owns its bytes and remains unchanged after the generator
+ * advances or drains. Line terminators are excluded.
+ */
 export async function* readLines(stream: ReadableStream<Uint8Array>, signal?: AbortSignal): AsyncGenerator<Uint8Array> {
 	const buffer = new ConcatSink();
 	const source = abortableSource(stream, signal);
@@ -56,13 +63,19 @@ export async function* readJsonl<T>(stream: ReadableStream<Uint8Array>, signal?:
 	}
 }
 
-// =============================================================================
-// SSE (Server-Sent Events)
-// =============================================================================
-
-class ConcatSink {
+/**
+ * Amortized byte accumulator for chunked stream readers.
+ *
+ * Holds the unconsumed tail of a stream in a single growing `Buffer` so that
+ * appending N chunks costs O(total bytes) instead of re-copying the whole
+ * prefix per chunk. Backs {@link readLines}, {@link readJsonl} and
+ * {@link readSseEvents}; also usable directly when a reader needs its own
+ * framing loop (see `consume` and `flush`).
+ */
+export class ConcatSink {
 	#space?: Buffer;
 	#length = 0;
+	#skipLeadingLf = false;
 
 	#ensureCapacity(size: number): Buffer {
 		const space = this.#space;
@@ -100,16 +113,37 @@ class ConcatSink {
 		return this.#length === 0;
 	}
 
+	/**
+	 * The buffered bytes as a live view — invalidated by the next `append`,
+	 * `reset` or `consume`.
+	 */
 	flush(): Uint8Array | undefined {
 		if (!this.#length) return undefined;
 		return this.#space!.subarray(0, this.#length);
 	}
 
+	/** Drop the first `count` buffered bytes, keeping the remainder. */
+	consume(count: number) {
+		if (count <= 0) return;
+		if (count >= this.#length) {
+			this.#length = 0;
+			return;
+		}
+		this.#space!.copyWithin(0, count, this.#length);
+		this.#length -= count;
+	}
+
 	clear() {
 		this.#length = 0;
 	}
 
-	*appendAndFlushLines(chunk: Uint8Array) {
+	/**
+	 * Append a chunk and yield each complete LF-delimited line.
+	 *
+	 * Yielded lines are owned snapshots. Unlike {@link flush}, they remain
+	 * valid after this sink or the input chunk is mutated.
+	 */
+	*appendAndFlushLines(chunk: Uint8Array): Generator<Uint8Array> {
 		let pos = 0;
 		while (pos < chunk.length) {
 			const nl = chunk.indexOf(LF, pos);
@@ -120,12 +154,12 @@ class ConcatSink {
 			const suffix = chunk.subarray(pos, nl);
 			pos = nl + 1;
 			if (this.isEmpty) {
-				yield suffix;
+				yield new Uint8Array(suffix);
 			} else {
 				this.append(suffix);
 				const payload = this.flush();
 				if (payload) {
-					yield payload;
+					yield new Uint8Array(payload);
 					this.clear();
 				}
 			}
@@ -133,19 +167,27 @@ class ConcatSink {
 	}
 
 	appendAndFlushText(chunk: Uint8Array, decoder: TextDecoder): string | undefined {
-		const lastNewline = chunk.lastIndexOf(LF);
-		if (lastNewline === -1) {
-			this.append(chunk);
+		let start = 0;
+		if (this.#skipLeadingLf) {
+			if (chunk.length === 0) return undefined;
+			this.#skipLeadingLf = false;
+			if (chunk[0] === LF) start = 1;
+		}
+
+		const lastLineEnd = Math.max(chunk.lastIndexOf(LF), chunk.lastIndexOf(CR));
+		if (lastLineEnd < start) {
+			if (start < chunk.length) this.append(chunk.subarray(start));
 			return undefined;
 		}
 
-		const completeEnd = lastNewline + 1;
+		const completeEnd = lastLineEnd + 1;
+		this.#skipLeadingLf = chunk[lastLineEnd] === CR && completeEnd === chunk.length;
 		let text: string;
 		if (this.isEmpty) {
-			const complete = completeEnd === chunk.length ? chunk : chunk.subarray(0, completeEnd);
+			const complete = start === 0 && completeEnd === chunk.length ? chunk : chunk.subarray(start, completeEnd);
 			text = decoder.decode(complete);
 		} else {
-			this.append(completeEnd === chunk.length ? chunk : chunk.subarray(0, completeEnd));
+			this.append(chunk.subarray(start, completeEnd));
 			text = decoder.decode(this.flush());
 			this.clear();
 		}
@@ -197,6 +239,10 @@ class ConcatSink {
 	}
 }
 
+// =============================================================================
+// SSE (Server-Sent Events)
+// =============================================================================
+
 /**
  * Stream parsed JSON objects from SSE `data:` lines.
  *
@@ -236,12 +282,27 @@ function isRecoverableTrailingJson(data: string): boolean {
 	return typeof recovered === "object" && recovered !== null;
 }
 
-export async function* readSseJson<T>(
+/**
+ * One dispatched `data:` frame from {@link readSseFrames}: either the parsed JSON
+ * value, or the text of a frame `JSON.parse` rejected together with the
+ * `SyntaxError` it raised (so the strict reader can rethrow it unchanged).
+ */
+type SseFrame<T> = { ok: true; value: T } | { ok: false; raw: string; error: SyntaxError };
+
+/**
+ * Shared `data:`-line framing for {@link readSseJson} and
+ * {@link readSseJsonOrText}: skips empty events, stops at the OpenAI `[DONE]`
+ * sentinel, notifies the diagnostic observer, and treats a container-shaped
+ * stream tail as a clean end of iteration.
+ */
+async function* readSseFrames<T>(
 	stream: ReadableStream<Uint8Array>,
 	signal?: AbortSignal,
 	onEvent?: SseEventObserver,
-): AsyncGenerator<T> {
-	for await (const sse of readSseEvents(stream, signal)) {
+): AsyncGenerator<SseFrame<T>> {
+	// The diagnostic observer is the only reader of `raw`; capture it exactly
+	// when one is attached so the hot path stays allocation-free.
+	for await (const sse of readSseEvents(stream, signal, onEvent ? { captureRaw: true } : undefined)) {
 		const isTrailing = trailingEvents.has(sse);
 		notifySseEventObserver(onEvent, sse);
 		const data = sse.data;
@@ -250,16 +311,60 @@ export async function* readSseJson<T>(
 			continue;
 		}
 		try {
-			yield JSON.parse(data) as T;
+			yield { ok: true, value: JSON.parse(data) as T };
 		} catch (err) {
 			if (err instanceof SyntaxError && isTrailing && isRecoverableTrailingJson(data)) {
 				return;
 			}
+			if (err instanceof SyntaxError) {
+				yield { ok: false, raw: data, error: err };
+				continue;
+			}
 			throw err;
 		}
 	}
 }
 
+export async function* readSseJson<T>(
+	stream: ReadableStream<Uint8Array>,
+	signal?: AbortSignal,
+	onEvent?: SseEventObserver,
+): AsyncGenerator<T> {
+	for await (const frame of readSseFrames<T>(stream, signal, onEvent)) {
+		if (!frame.ok) throw frame.error;
+		yield frame.value;
+	}
+}
+
+/**
+ * Like {@link readSseJson}, but a `data:` frame that is not valid JSON is yielded
+ * as its raw text instead of raising a `SyntaxError`. Cut-off container-shaped
+ * stream tails stay recoverable, exactly as they are in {@link readSseJson}.
+ *
+ * Consumers that only understand objects must treat a `string` yield as a
+ * transport-level failure (for example a `429 Too Many Requests` or an HTML
+ * throttle page from a reverse proxy that already committed to the stream). This
+ * exists because `readSseJson`'s baseline consumers span unrelated transports
+ * whose error handling a text yield would subtly change; new call sites opt in.
+ *
+ * Note that the text lane is only the frames `JSON.parse` *rejected*: a frame
+ * carrying a JSON-encoded string (`data: "429 Too Many Requests"`) parses, so it
+ * is yielded as that string and is indistinguishable from a rejected frame by
+ * type alone. Consumers branching on `typeof === "string"` therefore see both,
+ * which is the safe direction — each is classified as text rather than trusted as
+ * an event object.
+ */
+export async function* readSseJsonOrText<T>(
+	stream: ReadableStream<Uint8Array>,
+	signal?: AbortSignal,
+	onEvent?: SseEventObserver,
+): AsyncGenerator<T | string> {
+	for await (const frame of readSseFrames<T>(stream, signal, onEvent)) {
+		if (!frame.ok) yield frame.raw;
+		else yield frame.value;
+	}
+}
+
 /**
  * A single Server-Sent Event dispatched on a blank-line boundary.
  *
@@ -276,6 +381,15 @@ export async function* readSseJson<T>(
 export interface ServerSentEvent {
 	event: string | null;
 	data: string;
+	/**
+	 * Decoded wire lines for this event (`event:`/`data:`/etc.), for the
+	 * diagnostic pipeline. Populated only when the reader opts in via
+	 * {@link ReadSseEventsOptions.captureRaw} (or attaches an `onSseEvent`
+	 * observer to the JSON readers, which opt in automatically); otherwise
+	 * `[]`. Direct `readSseEvents` callers that need wire text must pass
+	 * `{ captureRaw: true }` — the field is allocation-free by default so
+	 * the token path pays no per-frame array/slice cost.
+	 */
 	raw: string[];
 	id?: string;
 	retry?: number;
@@ -288,50 +402,48 @@ interface SseEventState {
 	// of buffering an array and joining at flush. `null` means "no data: field
 	// seen yet" (distinct from a `data:` field with an empty value).
 	data: string | null;
-	raw: string[];
+	// Diagnostic wire lines, captured only when a reader asked for them (see
+	// `readSseEventsOptions.captureRaw`): per-frame array+slice allocation on
+	// the token path otherwise. `null` means capture is off.
+	raw: string[] | null;
 	id?: string;
 	retry?: number;
 }
 
 // Complete lines are decoded in one batch per source chunk. Each batch ends on
-// LF, which cannot split a multi-byte UTF-8 sequence.
+// an ASCII line-ending byte, which cannot split a multi-byte UTF-8 sequence.
 const SSE_DECODER = new TextDecoder("utf-8");
 
 function flushSseEvent(state: SseEventState): ServerSentEvent | null {
 	if (state.event === null && state.data === null && state.id === undefined && state.retry === undefined) {
-		state.raw = [];
+		if (state.raw !== null) state.raw = [];
 		return null;
 	}
 	const event: ServerSentEvent = {
 		event: state.event,
 		data: state.data ?? "",
-		raw: state.raw,
+		raw: state.raw ?? [],
 	};
 	if (state.id !== undefined) event.id = state.id;
 	if (state.retry !== undefined) event.retry = state.retry;
 	state.event = null;
 	state.data = null;
-	state.raw = [];
+	if (state.raw !== null) state.raw = [];
 	state.id = undefined;
 	state.retry = undefined;
 	return event;
 }
 
 function pushSseLine(line: string, state: SseEventState): ServerSentEvent | null {
-	// Complete-line batches split on LF only; strip a trailing CR so CRLF sources
-	// don't leak `\r` into field values.
-	if (line.charCodeAt(line.length - 1) === 0x0d /* '\r' */) {
-		line = line.slice(0, -1);
-	}
 	if (line.length === 0) return flushSseEvent(state);
 
 	// Comment line: keep in `raw` for diagnostic context, skip parsing.
 	if (line.charCodeAt(0) === 0x3a /* ':' */) {
-		state.raw.push(line);
+		state.raw?.push(line);
 		return null;
 	}
 
-	state.raw.push(line);
+	state.raw?.push(line);
 
 	const colon = line.indexOf(":");
 	const fieldName = colon === -1 ? line : line.slice(0, colon);
@@ -385,12 +497,24 @@ function pushSseLine(line: string, state: SseEventState): ServerSentEvent | null
  * }
  * ```
  */
+export interface ReadSseEventsOptions {
+	/**
+	 * Capture per-line wire text into `event.raw` for the diagnostic
+	 * pipeline (`onSseEvent` observers, raw-SSE viewer). Off by default:
+	 * every frame otherwise pays an array allocation plus one string slice
+	 * per line on the token path.
+	 */
+	captureRaw?: boolean;
+}
+
 export async function* readSseEvents(
 	stream: ReadableStream<Uint8Array>,
 	signal?: AbortSignal,
+	options?: ReadSseEventsOptions,
 ): AsyncGenerator<ServerSentEvent> {
 	const lineBuffer = new ConcatSink();
-	const state: SseEventState = { event: null, data: null, raw: [] };
+	const captureRaw = options?.captureRaw === true;
+	const state: SseEventState = { event: null, data: null, raw: captureRaw ? [] : null };
 	const source = abortableSource(stream, signal);
 	try {
 		for await (const chunk of source) {
@@ -398,13 +522,21 @@ export async function* readSseEvents(
 			if (text === undefined) continue;
 			let start = 0;
 			while (start < text.length) {
-				const newline = text.indexOf("\n", start);
-				const event = pushSseLine(text.slice(start, newline), state);
+				let lineEnd = start;
+				while (lineEnd < text.length) {
+					const code = text.charCodeAt(lineEnd);
+					if (code === LF || code === CR) break;
+					lineEnd++;
+				}
+				const event = pushSseLine(text.slice(start, lineEnd), state);
 				if (event) yield event;
-				start = newline + 1;
+				if (text.charCodeAt(lineEnd) === CR && text.charCodeAt(lineEnd + 1) === LF) {
+					lineEnd++;
+				}
+				start = lineEnd + 1;
 			}
 		}
-		// Treat any trailing partial line (no terminating LF) as a complete line.
+		// Treat any trailing partial line (no terminating line ending) as complete.
 		if (!lineBuffer.isEmpty) {
 			const tail = lineBuffer.flush();
 			if (tail) {