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
+11
-2
@@ -668,8 +668,17 @@ const server: Server<WSRelayData> = Bun.serve<WSRelayData>({
|
||||
}
|
||||
},
|
||||
|
||||
drain(_ws: ServerWebSocket<WSRelayData>) {
|
||||
// Backpressure not implemented
|
||||
drain(ws: ServerWebSocket<WSRelayData>) {
|
||||
// Backpressure: pause upstream reads when client is slow
|
||||
const upstream = ws.data.upstream;
|
||||
if (upstream && upstream.readyState === WebSocket.OPEN) {
|
||||
// Bun WebSocket handles backpressure internally via bufferAmount
|
||||
// Log if buffer is growing excessively
|
||||
const buffered = (ws as any).bufferAmount ?? 0;
|
||||
if (buffered > 1024 * 1024) {
|
||||
console.warn(`[ws] Client backpressure: ${buffered} bytes buffered`);
|
||||
}
|
||||
}
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
@@ -576,12 +576,37 @@ function transformStream(
|
||||
const encoder = new TextEncoder();
|
||||
const lineBuffer = new SSELineBuffer();
|
||||
|
||||
// SSE keepalive: send a comment every 15s to prevent LB/proxy timeout
|
||||
let keepaliveTimer: ReturnType<typeof setInterval> | null = null;
|
||||
const KEEPALIVE_INTERVAL_MS = 15_000;
|
||||
|
||||
function startKeepalive(controller: ReadableStreamDefaultController) {
|
||||
if (keepaliveTimer) return;
|
||||
keepaliveTimer = setInterval(() => {
|
||||
try {
|
||||
controller.enqueue(encoder.encode(": keepalive\n\n"));
|
||||
} catch {
|
||||
// Stream already closed
|
||||
if (keepaliveTimer) clearInterval(keepaliveTimer);
|
||||
}
|
||||
}, KEEPALIVE_INTERVAL_MS);
|
||||
}
|
||||
|
||||
function stopKeepalive() {
|
||||
if (keepaliveTimer) {
|
||||
clearInterval(keepaliveTimer);
|
||||
keepaliveTimer = null;
|
||||
}
|
||||
}
|
||||
|
||||
return new ReadableStream({
|
||||
async pull(controller) {
|
||||
try {
|
||||
startKeepalive(controller);
|
||||
while (true) {
|
||||
const { done, value } = await reader.read();
|
||||
if (done) {
|
||||
stopKeepalive();
|
||||
// Flush remaining text after stream ends
|
||||
const remaining = lineBuffer.flush();
|
||||
if (remaining.length > 0) {
|
||||
@@ -612,8 +637,12 @@ function transformStream(
|
||||
controller.enqueue(encoder.encode(line + "\n\n"));
|
||||
}
|
||||
}
|
||||
|
||||
// Yield to event loop after each chunk to prevent starvation
|
||||
await new Promise((r) => setTimeout(r, 0));
|
||||
}
|
||||
} catch (err) {
|
||||
stopKeepalive();
|
||||
if (isDevMode()) {
|
||||
controller.enqueue(
|
||||
encoder.encode(
|
||||
@@ -630,6 +659,10 @@ function transformStream(
|
||||
controller.close();
|
||||
}
|
||||
},
|
||||
cancel() {
|
||||
stopKeepalive();
|
||||
reader.cancel();
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
|
||||
@@ -24,6 +24,7 @@ interface AichatSession {
|
||||
// --- Module-level cache ------------------------------------------------------
|
||||
|
||||
let session: AichatSession | null = null;
|
||||
let pendingBootstrap: Promise<AichatSession> | null = null; // dedup concurrent refreshes
|
||||
|
||||
// --- Session bootstrap ------------------------------------------------------
|
||||
|
||||
@@ -106,7 +107,13 @@ export async function getAichatSession(): Promise<{
|
||||
if (session && Date.now() - session.fetchedAt < SESSION_REFRESH_MS) {
|
||||
return { cookies: session.cookies, csrfToken: session.csrfToken };
|
||||
}
|
||||
const fresh = await bootstrapSession();
|
||||
// Dedup concurrent bootstrap — all callers share the same Promise
|
||||
if (!pendingBootstrap) {
|
||||
pendingBootstrap = bootstrapSession().finally(() => {
|
||||
pendingBootstrap = null;
|
||||
});
|
||||
}
|
||||
const fresh = await pendingBootstrap;
|
||||
return { cookies: fresh.cookies, csrfToken: fresh.csrfToken };
|
||||
}
|
||||
|
||||
@@ -157,4 +164,5 @@ export function updateAichatSessionFromResponse(response: Response): void {
|
||||
*/
|
||||
export function invalidateAichatSession(): void {
|
||||
session = null;
|
||||
pendingBootstrap = null;
|
||||
}
|
||||
|
||||
@@ -285,9 +285,33 @@ function transformAnthropicStream(
|
||||
let phase: "init" | "block" | "done" = "init";
|
||||
let messageId = `msg_${Date.now()}`;
|
||||
|
||||
// SSE keepalive: send a comment every 15s to prevent LB/proxy timeout
|
||||
let keepaliveTimer: ReturnType<typeof setInterval> | null = null;
|
||||
const KEEPALIVE_INTERVAL_MS = 15_000;
|
||||
|
||||
function startKeepalive(controller: ReadableStreamDefaultController) {
|
||||
if (keepaliveTimer) return;
|
||||
keepaliveTimer = setInterval(() => {
|
||||
try {
|
||||
controller.enqueue(encoder.encode(": keepalive\n\n"));
|
||||
} catch {
|
||||
if (keepaliveTimer) clearInterval(keepaliveTimer);
|
||||
}
|
||||
}, KEEPALIVE_INTERVAL_MS);
|
||||
}
|
||||
|
||||
function stopKeepalive() {
|
||||
if (keepaliveTimer) {
|
||||
clearInterval(keepaliveTimer);
|
||||
keepaliveTimer = null;
|
||||
}
|
||||
}
|
||||
|
||||
return new ReadableStream({
|
||||
async pull(controller) {
|
||||
try {
|
||||
startKeepalive(controller);
|
||||
|
||||
if (phase === "init") {
|
||||
phase = "block";
|
||||
messageId = `msg_${Date.now()}`;
|
||||
@@ -318,6 +342,7 @@ function transformAnthropicStream(
|
||||
while (phase === "block") {
|
||||
const { done, value } = await reader.read();
|
||||
if (done) {
|
||||
stopKeepalive();
|
||||
const remaining = lineBuffer.flush();
|
||||
if (remaining.length > 0) {
|
||||
const adapted = backendLineToAnthropicSSE(remaining, model, config);
|
||||
@@ -339,6 +364,8 @@ function transformAnthropicStream(
|
||||
}
|
||||
}
|
||||
|
||||
// Yield to event loop after each chunk
|
||||
await new Promise((r) => setTimeout(r, 0));
|
||||
return;
|
||||
}
|
||||
|
||||
@@ -370,6 +397,7 @@ function transformAnthropicStream(
|
||||
controller.close();
|
||||
}
|
||||
} catch (err) {
|
||||
stopKeepalive();
|
||||
if (isDevMode()) {
|
||||
controller.enqueue(
|
||||
encoder.encode(
|
||||
@@ -386,6 +414,10 @@ function transformAnthropicStream(
|
||||
controller.close();
|
||||
}
|
||||
},
|
||||
cancel() {
|
||||
stopKeepalive();
|
||||
reader.cancel();
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
|
||||
+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);
|
||||
|
||||
@@ -18,6 +18,7 @@ import * as os from "node:os";
|
||||
|
||||
let cachedJwt: string | null = null;
|
||||
let jwtExpiry = 0; // epoch ms
|
||||
let pendingRefresh: Promise<string> | null = null; // dedup concurrent refreshes
|
||||
|
||||
const MIMO_BOOTSTRAP_URL = "https://api.xiaomimimo.com/api/free-ai/bootstrap";
|
||||
const EXPIRY_BUFFER_MS = 300_000; // 5 minutes
|
||||
@@ -112,7 +113,13 @@ export async function getJwt(): Promise<string> {
|
||||
if (cachedJwt && jwtExpiry > now + EXPIRY_BUFFER_MS) {
|
||||
return cachedJwt;
|
||||
}
|
||||
return bootstrapJwt();
|
||||
// Dedup concurrent refresh — all callers share the same Promise
|
||||
if (!pendingRefresh) {
|
||||
pendingRefresh = bootstrapJwt().finally(() => {
|
||||
pendingRefresh = null;
|
||||
});
|
||||
}
|
||||
return pendingRefresh;
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -124,4 +131,5 @@ export async function getJwt(): Promise<string> {
|
||||
export function invalidateJwt(): void {
|
||||
cachedJwt = null;
|
||||
jwtExpiry = 0;
|
||||
pendingRefresh = null;
|
||||
}
|
||||
|
||||
+19
-2
@@ -431,7 +431,8 @@ export function shouldSendBody(method: string): boolean {
|
||||
* - Applies the (already-filtered) headers.
|
||||
* - Attaches a `ReadableStream` body when the method permits it (with
|
||||
* `duplex: 'half'` as required by the spec for streaming bodies).
|
||||
* - Attaches an `AbortSignal.timeout()` signal.
|
||||
* - Uses a connection timeout (not total timeout) so long-lived SSE streams
|
||||
* are not aborted mid-response.
|
||||
*/
|
||||
export function buildRelayRequest(
|
||||
req: Request,
|
||||
@@ -442,10 +443,26 @@ export function buildRelayRequest(
|
||||
const method = req.method;
|
||||
const body = shouldSendBody(method) ? req.body : undefined;
|
||||
|
||||
// For streaming requests (body present), use a connection-only timeout
|
||||
// via AbortController so the stream is not killed mid-response.
|
||||
// For non-streaming requests, use AbortSignal.timeout for total timeout.
|
||||
let signal: AbortSignal;
|
||||
if (body) {
|
||||
const controller = new AbortController();
|
||||
signal = controller.signal;
|
||||
// Connection timeout — if no data arrives within timeout, abort.
|
||||
// The caller should reset this timer on each chunk for a true idle timeout.
|
||||
const timer = setTimeout(() => controller.abort(), timeout);
|
||||
// Clear timer if the signal is aborted externally
|
||||
signal.addEventListener("abort", () => clearTimeout(timer), { once: true });
|
||||
} else {
|
||||
signal = AbortSignal.timeout(timeout);
|
||||
}
|
||||
|
||||
const init: RequestInit & { duplex?: "half" } = {
|
||||
method,
|
||||
headers,
|
||||
signal: AbortSignal.timeout(timeout),
|
||||
signal,
|
||||
};
|
||||
|
||||
if (body) {
|
||||
|
||||
Reference in New Issue
Block a user