From 22f2990f56ee1f9a751a12f03f949d56dfee2668 Mon Sep 17 00:00:00 2001 From: Owen Date: Thu, 6 Aug 2026 17:15:11 -0400 Subject: [PATCH] support streaming and closes properly on site targets --- server/lib/aiCapabilities.ts | 7 ++-- server/lib/aiGatewayUpstreamFetch.ts | 14 +++++++ server/routers/aiGateway/pipeline.ts | 39 +++++++++++++++--- server/routers/aiGateway/targetRouting.ts | 50 +++++++++++++++++++---- 4 files changed, 92 insertions(+), 18 deletions(-) diff --git a/server/lib/aiCapabilities.ts b/server/lib/aiCapabilities.ts index b6d2f7984..e899f83c3 100644 --- a/server/lib/aiCapabilities.ts +++ b/server/lib/aiCapabilities.ts @@ -62,9 +62,10 @@ export function joinUpstreamUrl(baseUrl: string, path: string): string { } function pathFromRequest(req: Request): string { - // Prefer originalUrl path (includes mounted path) over req.path when available. - const raw = - req.originalUrl?.split("?")[0] || req.url?.split("?")[0] || req.path; + // Prefer originalUrl (includes mounted path) over req.url when available. + // Query string is preserved - some providers use it to select the + // streaming response format (e.g. Gemini's `?alt=sse`). + const raw = req.originalUrl || req.url || req.path; return raw.startsWith("/") ? raw : `/${raw}`; } diff --git a/server/lib/aiGatewayUpstreamFetch.ts b/server/lib/aiGatewayUpstreamFetch.ts index e41ec23e4..7691f1f95 100644 --- a/server/lib/aiGatewayUpstreamFetch.ts +++ b/server/lib/aiGatewayUpstreamFetch.ts @@ -7,6 +7,7 @@ type UpstreamFetchInit = { headers: Record; body?: string; skipTlsVerification?: boolean; + signal?: AbortSignal; }; const insecureHttpsAgent = new https.Agent({ @@ -25,6 +26,11 @@ export function aiGatewayUpstreamFetch( isHttps && init.skipTlsVerification ? insecureHttpsAgent : undefined; return new Promise((resolve, reject) => { + if (init.signal?.aborted) { + reject(init.signal.reason ?? new Error("Request aborted")); + return; + } + const req = lib.request( url, { @@ -60,6 +66,14 @@ export function aiGatewayUpstreamFetch( req.on("error", reject); + if (init.signal) { + const onAbort = () => req.destroy(init.signal!.reason); + init.signal.addEventListener("abort", onAbort, { once: true }); + req.on("close", () => + init.signal!.removeEventListener("abort", onAbort) + ); + } + if (init.body !== undefined) { req.write(init.body); } diff --git a/server/routers/aiGateway/pipeline.ts b/server/routers/aiGateway/pipeline.ts index 32690b1e1..ff76a5317 100644 --- a/server/routers/aiGateway/pipeline.ts +++ b/server/routers/aiGateway/pipeline.ts @@ -592,15 +592,33 @@ export async function handleAiGatewayProxy( skipTlsVerification: provider.skipTlsVerification }); + // Cancel the upstream request (and, transitively, anything it fans + // out to) if the client goes away before we're done - otherwise a + // client-cancelled streaming chat completion keeps running upstream + // to completion, wasting the connection and any per-token billing. + const abortController = new AbortController(); + const onClientClose = () => { + if (!res.writableEnded) { + abortController.abort(); + } + }; + res.on("close", onClientClose); + let upstreamRes: globalThis.Response; try { upstreamRes = await aiGatewayUpstreamFetch(targetUrl, { method: "POST", headers, body, - skipTlsVerification: provider.skipTlsVerification + skipTlsVerification: provider.skipTlsVerification, + signal: abortController.signal }); } catch (fetchError) { + res.off("close", onClientClose); + if (abortController.signal.aborted) { + // Client already disconnected; nothing left to respond to. + return; + } logger.error({ message: "AI gateway upstream fetch failed", url: targetUrl, @@ -628,14 +646,23 @@ export async function handleAiGatewayProxy( if (isStream && upstreamRes.body) { res.flushHeaders(); const reader = upstreamRes.body.getReader(); - while (true) { - const { done, value } = await reader.read(); - if (done) break; - res.write(value); + try { + while (!abortController.signal.aborted) { + const { done, value } = await reader.read(); + if (done) break; + res.write(value); + } + } finally { + await reader.cancel().catch(() => {}); + res.off("close", onClientClose); } - return res.end(); + if (!res.writableEnded) { + res.end(); + } + return; } + res.off("close", onClientClose); const text = await upstreamRes.text(); return res.send(text); } catch (error) { diff --git a/server/routers/aiGateway/targetRouting.ts b/server/routers/aiGateway/targetRouting.ts index 2e415b819..b0b2456d4 100644 --- a/server/routers/aiGateway/targetRouting.ts +++ b/server/routers/aiGateway/targetRouting.ts @@ -128,8 +128,10 @@ function pickTarget( } function pathFromRequest(req: Request): string { - const raw = - req.originalUrl?.split("?")[0] || req.url?.split("?")[0] || req.path; + // Query string is preserved - some providers use it to select the + // streaming response format (e.g. Gemini's `?alt=sse`), and gerbil's + // /router/* forwards it through untouched. + const raw = req.originalUrl || req.url || req.path; return raw.startsWith("/") ? raw : `/${raw}`; } @@ -205,14 +207,32 @@ export async function proxyAiGatewayToSiteTarget( body: req.body }); + // Cancel the request to gerbil (which cascades to gerbil cancelling its + // proxied request to the actual site target, since gerbil's reverse + // proxy derives the outbound request's context from the inbound one) if + // the client goes away before we're done. + const abortController = new AbortController(); + const onClientClose = () => { + if (!res.writableEnded) { + abortController.abort(); + } + }; + res.on("close", onClientClose); + let upstreamRes: globalThis.Response; try { upstreamRes = await fetch(gerbilUrl, { method: "POST", headers, - body + body, + signal: abortController.signal }); } catch (fetchError) { + res.off("close", onClientClose); + if (abortController.signal.aborted) { + // Client already disconnected; nothing left to respond to. + return; + } logger.error({ message: "AI gateway target proxy request failed", url: gerbilUrl, @@ -232,7 +252,11 @@ export async function proxyAiGatewayToSiteTarget( const contentType = upstreamRes.headers.get("content-type") || ""; const isStream = req.body?.stream === true || - contentType.includes("text/event-stream"); + contentType.includes("text/event-stream") || + pathFromRequest(req).includes("streamGenerateContent") || + pathFromRequest(req).includes("streamRawPredict") || + pathFromRequest(req).includes("converse-stream") || + pathFromRequest(req).includes("invoke-with-response-stream"); res.status(upstreamRes.status); res.setHeader("Content-Type", contentType || "application/json"); @@ -240,15 +264,23 @@ export async function proxyAiGatewayToSiteTarget( if (isStream && upstreamRes.body) { res.flushHeaders(); const reader = upstreamRes.body.getReader(); - while (true) { - const { done, value } = await reader.read(); - if (done) break; - res.write(value); + try { + while (!abortController.signal.aborted) { + const { done, value } = await reader.read(); + if (done) break; + res.write(value); + } + } finally { + await reader.cancel().catch(() => {}); + res.off("close", onClientClose); + } + if (!res.writableEnded) { + res.end(); } - res.end(); return; } + res.off("close", onClientClose); const text = await upstreamRes.text(); res.send(text); }