From c62de0899a48daf2ce04901a219fc7b5a1e3acf6 Mon Sep 17 00:00:00 2001 From: MythEclipse Date: Sat, 20 Jun 2026 17:13:32 +0700 Subject: [PATCH] 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) --- .mcp.json | 12 --- docker-compose.yml | 2 +- src/index.ts | 13 +++- src/lib/ai-proxy.ts | 33 +++++++++ src/lib/aichat-auth.ts | 10 ++- src/lib/anthropic-proxy.ts | 32 ++++++++ src/lib/fetch-utils.ts | 148 +++++++++++++++++++++++++++++-------- src/lib/mimo-auth.ts | 10 ++- src/lib/relay-utils.ts | 21 +++++- 9 files changed, 231 insertions(+), 50 deletions(-) delete mode 100644 .mcp.json diff --git a/.mcp.json b/.mcp.json deleted file mode 100644 index c942808..0000000 --- a/.mcp.json +++ /dev/null @@ -1,12 +0,0 @@ -{ - "mcpServers": { - "code-review-graph": { - "command": "uvx", - "args": [ - "code-review-graph", - "serve" - ], - "type": "stdio" - } - } -} diff --git a/docker-compose.yml b/docker-compose.yml index 9b146aa..fac8e40 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -6,7 +6,7 @@ services: network_mode: host environment: HOST: "::" - PORT: "3000" + PORT: "4060" IPV6_SOURCES: "2001:df4:c140:1f::8,2001:df4:c140:1f::395,2001:df4:c140:1f::d9,2001:df4:c140:1f::db,2001:df4:c140:1f::dd,2001:df4:c140:1f::de,2001:df4:c140:1f::df,2001:df4:c140:1f:ffff:ffff:ffff:1,2001:df4:c140:1f:ffff:ffff:ffff:2,2001:df4:c140:1f:ffff:ffff:ffff:ff,2001:df4:c140:1f:ffff:ffff:ffff:100,2001:df4:c140:1f:ffff:ffff:ffff:f000,2001:df4:c140:1f:ffff:ffff:ffff:f0e6,2001:df4:c140:1f:ffff:ffff:ffff:f229,2001:df4:c140:1f:ffff:ffff:ffff:f340,2001:df4:c140:1f:ffff:ffff:ffff:f34b,2001:df4:c140:1f:ffff:ffff:ffff:f3ef,2001:df4:c140:1f:ffff:ffff:ffff:f4ba" RELAY_TIMEOUT_MS: "30000" RATE_LIMIT_MAX: "1000" diff --git a/src/index.ts b/src/index.ts index 61757e2..468aeb2 100644 --- a/src/index.ts +++ b/src/index.ts @@ -668,8 +668,17 @@ const server: Server = Bun.serve({ } }, - drain(_ws: ServerWebSocket) { - // Backpressure not implemented + drain(ws: ServerWebSocket) { + // 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`); + } + } }, }, }); diff --git a/src/lib/ai-proxy.ts b/src/lib/ai-proxy.ts index 6559f59..648ea44 100644 --- a/src/lib/ai-proxy.ts +++ b/src/lib/ai-proxy.ts @@ -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 | 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(); + }, }); } diff --git a/src/lib/aichat-auth.ts b/src/lib/aichat-auth.ts index 40e36bc..898a2fa 100644 --- a/src/lib/aichat-auth.ts +++ b/src/lib/aichat-auth.ts @@ -24,6 +24,7 @@ interface AichatSession { // --- Module-level cache ------------------------------------------------------ let session: AichatSession | null = null; +let pendingBootstrap: Promise | 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; } diff --git a/src/lib/anthropic-proxy.ts b/src/lib/anthropic-proxy.ts index 0e435e8..96b0347 100644 --- a/src/lib/anthropic-proxy.ts +++ b/src/lib/anthropic-proxy.ts @@ -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 | 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(); + }, }); } diff --git a/src/lib/fetch-utils.ts b/src/lib/fetch-utils.ts index dd316bb..d837e5e 100644 --- a/src/lib/fetch-utils.ts +++ b/src/lib/fetch-utils.ts @@ -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(); +/** + * Track a reader for graceful shutdown. Returns a wrapped reader that + * auto-removes itself from ACTIVE_READERS when done/cancelled. + */ +export function trackReader(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 { 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); diff --git a/src/lib/mimo-auth.ts b/src/lib/mimo-auth.ts index 3f884bf..69741fc 100644 --- a/src/lib/mimo-auth.ts +++ b/src/lib/mimo-auth.ts @@ -18,6 +18,7 @@ import * as os from "node:os"; let cachedJwt: string | null = null; let jwtExpiry = 0; // epoch ms +let pendingRefresh: Promise | 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 { 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 { export function invalidateJwt(): void { cachedJwt = null; jwtExpiry = 0; + pendingRefresh = null; } diff --git a/src/lib/relay-utils.ts b/src/lib/relay-utils.ts index c4ecf81..2754a7f 100644 --- a/src/lib/relay-utils.ts +++ b/src/lib/relay-utils.ts @@ -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) {