9router/open-sse/utils/streamErrorPeek.js

import { matchStreamErrorPatterns } from "./streamErrorPatterns.js";
import { HTTP_STATUS } from "../config/runtimeConfig.js";

const DEFAULT_TIMEOUT_MS = (() => {
	const raw = process.env.STREAM_ERROR_PEEK_TIMEOUT_MS;
	const n = raw ? parseInt(raw, 10) : NaN;
	return Number.isFinite(n) && n > 0 ? n : 3000;
})();

const DEFAULT_MAX_BYTES = (() => {
	const raw = process.env.STREAM_ERROR_PEEK_MAX_BYTES;
	const n = raw ? parseInt(raw, 10) : NaN;
	return Number.isFinite(n) && n > 0 ? n : 8192;
})();

function makeAbortError(reason) {
	const error = new Error(reason?.message || reason || "Request aborted");
	error.name = "AbortError";
	return error;
}

/**
 * Read the first bytes of a 200 response and reject it with a 502 when a
 * configured stream-error pattern matches — BEFORE any byte reaches the
 * client, so account/model fallback can still kick in for streaming too.
 * On no-match/timeout/abort the response is re-emitted (buffered bytes +
 * rest of stream) unchanged in spirit. Fail-open: never throws.
 */
export async function maybeRejectEarlyStreamError(
	response,
	patterns,
	{
		signal = null,
		timeoutMs = DEFAULT_TIMEOUT_MS,
		maxBytes = DEFAULT_MAX_BYTES,
	} = {},
) {
	if (!Array.isArray(patterns) || patterns.length === 0) return response;
	const reader = response.body.getReader();
	const decoder = new TextDecoder();
	const abortController = new AbortController();
	const forwardAbort = () => abortController.abort(signal?.reason);
	if (signal?.aborted) abortController.abort(signal?.reason);
	else if (signal)
		signal.addEventListener("abort", forwardAbort, { once: true });

	// Raw bytes for lossless re-emission; decoded text ONLY for pattern matching.
	// Never re-encode decoded text: TextDecoder holds a split multi-byte char
	// internally and flush() would replace it with U+FFFD, corrupting the stream.
	const rawChunks = [];
	let peekedText = "";
	let total = 0;

	const readWithTimeout = (ms) => {
		if (abortController.signal.aborted)
			return Promise.reject(makeAbortError(abortController.signal.reason));
		const timeoutPromise = new Promise((_, reject) => {
			const t = setTimeout(
				() => reject(new Error("stream error peek timeout")),
				ms,
			);
			t.unref?.();
		});
		const abortPromise = new Promise((_, reject) => {
			abortController.signal.addEventListener(
				"abort",
				() => reject(makeAbortError(abortController.signal.reason)),
				{ once: true },
			);
		});
		return Promise.race([reader.read(), timeoutPromise, abortPromise]);
	};

	try {
		const deadline = Date.now() + timeoutMs;
		while (total < maxBytes && Date.now() < deadline) {
			const { done, value } = await readWithTimeout(
				Math.max(deadline - Date.now(), 1),
			);
			if (done) break;
			rawChunks.push(value);
			peekedText += decoder.decode(value, { stream: true });
			total += value.byteLength;
			const matched = matchStreamErrorPatterns(patterns, peekedText);
			if (matched) {
				await reader.cancel("stream error pattern matched").catch(() => {});
				return new Response(
					JSON.stringify({
						error: {
							message: `Stream error pattern matched: ${matched}`,
							type: "upstream_error",
						},
					}),
					{
						status: HTTP_STATUS.BAD_GATEWAY,
						statusText: String(matched).slice(0, 200),
						headers: { "Content-Type": "application/json" },
					},
				);
			}
		}
	} catch {
		// timeout / abort / read failure → commit; downstream stall/abort handling takes over.
	}

	if (signal) signal.removeEventListener("abort", forwardAbort);

	const remaining = new ReadableStream({
		start(controller) {
			(async () => {
				try {
					// Re-emit RAW bytes (never re-encoded decoded text) so split
					// multi-byte UTF-8 sequences survive the peek untouched.
					for (const c of rawChunks) controller.enqueue(c);
					while (true) {
						const { done, value } = await reader.read();
						if (done) break;
						controller.enqueue(value);
					}
					controller.close();
				} catch (err) {
					controller.error(err);
				}
			})();
		},
		cancel() {
			reader.cancel("stream cancelled during peek commit").catch(() => {});
		},
	});

	return new Response(remaining, {
		status: response.status,
		statusText: response.statusText,
		headers: response.headers,
	});
}