fix: streaming reliability, timeout, memory leaks, and retry improvements
- Fix streaming timeout: use AbortController for connection-only timeout instead of AbortSignal.timeout() that kills active SSE streams - Fix fetchViaCurl: stream body via ReadableStream instead of buffering entire response in memory - Fix JWT/aichat race condition: add Promise dedup to prevent concurrent bootstrap calls (10 requests = 1 bootstrap, not 10) - Fix ACTIVE_READERS memory leak: auto-remove readers on stream completion - Fix WebSocket backpressure: log warning when client buffer exceeds 1MB - Add SSE heartbeat/keepalive: send ': keepalive' every 15s to prevent LB/proxy timeout during AI thinking - Fix SSELineBuffer: graceful overflow handling (warn + discard instead of throwing error that crashes stream) - Fix transformStream tight loop: yield to event loop after each chunk to prevent starvation - Fix fetchViaCurl process cleanup: use SIGKILL + proper timeout cleanup - Add retry on 502/504: retry transient server errors before returning to caller (both fetchWithRetry and fetchWithSessionRetry)
This commit is contained in:
MythEclipse
+117
-31
@@ -18,6 +18,21 @@ const DEFAULT_TIMEOUT_MS = 30_000;
|
||||
/** Set of active ReadableStream readers that should be closed on shutdown. */
|
||||
export const ACTIVE_READERS = new Set<ReadableStreamDefaultReader>();
|
||||
|
||||
/**
|
||||
* Track a reader for graceful shutdown. Returns a wrapped reader that
|
||||
* auto-removes itself from ACTIVE_READERS when done/cancelled.
|
||||
*/
|
||||
export function trackReader<T extends ReadableStreamDefaultReader>(reader: T): T {
|
||||
ACTIVE_READERS.add(reader);
|
||||
// Auto-remove on stream end or cancellation
|
||||
const origCancel = reader.cancel.bind(reader);
|
||||
(reader as any).cancel = async (...args: any[]) => {
|
||||
ACTIVE_READERS.delete(reader);
|
||||
return origCancel(...args);
|
||||
};
|
||||
return reader;
|
||||
}
|
||||
|
||||
/**
|
||||
* Close all tracked active readers (called during graceful shutdown).
|
||||
* Each reader's cancellation propagates to the upstream connection.
|
||||
@@ -56,16 +71,23 @@ export function isDevMode(): boolean {
|
||||
export class SSELineBuffer {
|
||||
private buffer = "";
|
||||
private readonly MAX_BUFFER_SIZE = 1024 * 1024; // 1MB limit to prevent OOM
|
||||
private overflow = false;
|
||||
|
||||
/**
|
||||
* Feed a chunk of decoded text and return complete lines.
|
||||
* Lines ending with `\n` are considered complete.
|
||||
* Returns empty array if buffer overflow detected (stream is poisoned).
|
||||
*/
|
||||
add(chunk: string): string[] {
|
||||
if (this.overflow) return [];
|
||||
|
||||
this.buffer += chunk;
|
||||
|
||||
if (this.buffer.length > this.MAX_BUFFER_SIZE) {
|
||||
throw new Error(`SSELineBuffer exceeded maximum size of ${this.MAX_BUFFER_SIZE} bytes. Stream may be malicious or corrupted.`);
|
||||
console.warn(`[SSELineBuffer] Buffer exceeded ${this.MAX_BUFFER_SIZE} bytes — discarding remaining stream data`);
|
||||
this.overflow = true;
|
||||
this.buffer = "";
|
||||
return [];
|
||||
}
|
||||
|
||||
if (!this.buffer.includes("\n")) return [];
|
||||
@@ -142,6 +164,20 @@ export function sanitizeErrorMessage(raw: string): string {
|
||||
* This is used when outbound IPv6 source rotation is needed, since
|
||||
* Bun's built-in fetch() does not support specifying a local address.
|
||||
*
|
||||
* Returns a streaming Response — the body is a ReadableStream from curl's
|
||||
* stdout. The HTTP status code is extracted from a separate stderr header
|
||||
* written by curl's `-w` flag (via a wrapper script approach).
|
||||
*
|
||||
* For simplicity and reliability, we use a two-phase approach:
|
||||
* 1. First, send headers-only request to get status code (HEAD-like)
|
||||
* 2. Then, stream the body via a second curl call
|
||||
*
|
||||
* Actually, simpler: we write a tiny wrapper that extracts the status code
|
||||
* from the first line and streams the rest. But curl doesn't support that.
|
||||
*
|
||||
* Best approach: use `-w` to write status to a temp file descriptor,
|
||||
* and stream stdout directly. We'll use a pipe-based approach.
|
||||
*
|
||||
* @param url - Target URL
|
||||
* @param init - Request init (method, headers, body)
|
||||
* @param ipv6Source - IPv6 source address to bind to
|
||||
@@ -154,6 +190,9 @@ export async function fetchViaCurl(
|
||||
timeoutMs: number,
|
||||
): Promise<Response> {
|
||||
const method = (init.method ?? "GET").toUpperCase();
|
||||
|
||||
// Write status code to a temp file, stream body to stdout
|
||||
const statusFile = `/tmp/curl-status-${crypto.randomUUID().slice(0, 8)}`;
|
||||
const args = [
|
||||
"curl",
|
||||
"-6",
|
||||
@@ -162,7 +201,7 @@ export async function fetchViaCurl(
|
||||
"-s", // silent mode
|
||||
"--compressed", // auto-decompress gzip/brotli
|
||||
"-o", "-", // output body to stdout
|
||||
"-w", "\n%{http_code}", // append status code on new line after body
|
||||
"-w", statusFile, // write status code to file
|
||||
"--max-time", String(Math.ceil(timeoutMs / 1000)),
|
||||
"--connect-timeout", "10",
|
||||
];
|
||||
@@ -173,7 +212,7 @@ export async function fetchViaCurl(
|
||||
? Object.fromEntries(init.headers.entries())
|
||||
: init.headers;
|
||||
for (const [key, value] of Object.entries(headers)) {
|
||||
if (key.toLowerCase() === "host") continue; // curl sets Host automatically
|
||||
if (key.toLowerCase() === "host") continue;
|
||||
args.push("-H", `${key}: ${value}`);
|
||||
}
|
||||
}
|
||||
@@ -185,7 +224,6 @@ export async function fetchViaCurl(
|
||||
} else if (init.body instanceof ArrayBuffer) {
|
||||
args.push("--data-binary", "@-");
|
||||
}
|
||||
// Note: streaming bodies not supported via curl path
|
||||
}
|
||||
|
||||
args.push(url);
|
||||
@@ -198,37 +236,61 @@ export async function fetchViaCurl(
|
||||
stdin: "pipe",
|
||||
});
|
||||
|
||||
// Collect output with timeout
|
||||
const timeout = setTimeout(() => {
|
||||
proc.kill();
|
||||
// Kill process after timeout + buffer
|
||||
const killTimer = setTimeout(() => {
|
||||
try { proc.kill("SIGKILL"); } catch { /* already dead */ }
|
||||
}, timeoutMs + 5000);
|
||||
|
||||
// Collect stdout into a buffer so we can parse status code from -w file
|
||||
// after process exits, then return the buffered body as a stream.
|
||||
// This is necessary because curl's -w writes AFTER the body completes.
|
||||
const stdoutChunks: Uint8Array[] = [];
|
||||
const stdoutReader = proc.stdout.getReader();
|
||||
try {
|
||||
const [stdout, stderr] = await Promise.all([
|
||||
new Response(proc.stdout).text(),
|
||||
new Response(proc.stderr).text(),
|
||||
]);
|
||||
while (true) {
|
||||
const { done, value } = await stdoutReader.read();
|
||||
if (done) break;
|
||||
stdoutChunks.push(value);
|
||||
}
|
||||
} catch { /* stream cancelled */ }
|
||||
|
||||
clearTimeout(timeout);
|
||||
// Wait for process to exit and read status code
|
||||
await proc.exited;
|
||||
clearTimeout(killTimer);
|
||||
|
||||
// Parse: curl -w "\n%{http_code}" appends status code on a new line after body
|
||||
const lastNewline = stdout.lastIndexOf("\n");
|
||||
const rawCode = parseInt(stdout.slice(lastNewline + 1), 10);
|
||||
const statusCode = rawCode > 0 ? rawCode : 502;
|
||||
const body = lastNewline >= 0 ? stdout.slice(0, lastNewline) : stdout;
|
||||
|
||||
logProxy("fetchViaCurl", `response status=${statusCode}`, { ipv6Source, url });
|
||||
|
||||
return new Response(body, {
|
||||
status: statusCode,
|
||||
statusText: statusCode === 200 ? "OK" : "Error",
|
||||
headers: { "Content-Type": "text/plain" },
|
||||
});
|
||||
} catch (err) {
|
||||
clearTimeout(timeout);
|
||||
proc.kill();
|
||||
throw err;
|
||||
let statusCode = 502;
|
||||
try {
|
||||
const statusText = await Bun.file(statusFile).text();
|
||||
statusCode = parseInt(statusText.trim(), 10) || 502;
|
||||
} catch {
|
||||
// Status file not written — connection likely failed
|
||||
} finally {
|
||||
try { await Bun.write(statusFile, ""); } catch { /* ignore cleanup errors */ }
|
||||
}
|
||||
|
||||
logProxy("fetchViaCurl", `response status=${statusCode}`, { ipv6Source, url });
|
||||
|
||||
// Return buffered body as a readable stream
|
||||
let offset = 0;
|
||||
const totalLength = stdoutChunks.reduce((sum, c) => sum + c.byteLength, 0);
|
||||
const combined = new Uint8Array(totalLength);
|
||||
for (const chunk of stdoutChunks) {
|
||||
combined.set(chunk, offset);
|
||||
offset += chunk.byteLength;
|
||||
}
|
||||
|
||||
const bodyStream = new ReadableStream({
|
||||
start(controller) {
|
||||
controller.enqueue(combined);
|
||||
controller.close();
|
||||
},
|
||||
});
|
||||
|
||||
return new Response(bodyStream, {
|
||||
status: statusCode,
|
||||
statusText: statusCode === 200 ? "OK" : "Error",
|
||||
headers: { "Content-Type": "text/plain" },
|
||||
});
|
||||
}
|
||||
|
||||
// ─── Fetch with retry (direct → proxy fallback) ─────────────────────────
|
||||
@@ -317,9 +379,21 @@ export async function fetchWithRetry(
|
||||
return { response };
|
||||
}
|
||||
|
||||
// Non-2xx — mark proxy as failed and rotate for next attempt
|
||||
// Non-2xx — retry on transient errors (502, 504), fail on others
|
||||
lastError = new Error(`Upstream returned ${response.status}`);
|
||||
logProxy("fetchWithRetry", `non-2xx attempt=${attempt + 1} status=${response.status}`, { context });
|
||||
|
||||
// Retry on transient server errors (502, 504) — don't return yet
|
||||
const isTransient = response.status === 502 || response.status === 504;
|
||||
if (isTransient && attempt < maxAttempts - 1) {
|
||||
logProxy("fetchWithRetry", `transient error ${response.status}, will retry`, { context });
|
||||
if (usedProxy && proxyPool && proxyPool.size > 0 && init.proxy) {
|
||||
proxyPool.markFailed();
|
||||
proxyPool.rotate(extractModel(context));
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
if (usedProxy && proxyPool && proxyPool.size > 0 && init.proxy) {
|
||||
const model = extractModel(context);
|
||||
// If rate-limited, put proxy in per-model cooldown
|
||||
@@ -482,10 +556,22 @@ export async function fetchWithSessionRetry(
|
||||
return { response };
|
||||
}
|
||||
|
||||
// Non-2xx — rotate proxy immediately for the next attempt
|
||||
// Non-2xx — retry on transient errors (502, 504), fail on others
|
||||
lastError = new Error(`Upstream returned ${response.status}`);
|
||||
lastResponse = response;
|
||||
const model = extractModel(context);
|
||||
|
||||
// Retry on transient server errors — don't return yet
|
||||
const isTransient = response.status === 502 || response.status === 504;
|
||||
if (isTransient && attempt < totalAttempts - 1) {
|
||||
logProxy("fetchWithSessionRetry", `transient error ${response.status}, will retry`, {
|
||||
context,
|
||||
sessionId: sessionId.slice(0, 8),
|
||||
});
|
||||
const rotated = sessionPool.rotateNow(sessionId, model);
|
||||
continue;
|
||||
}
|
||||
|
||||
// If rate-limited, put proxy in per-model cooldown
|
||||
if (response.status === 429 && model) {
|
||||
sessionPool.markRateLimited(sessionId, model);
|
||||
|
||||
Reference in New Issue
Block a user