Skip to content

fix(api): finalize bifrost streaming usage on close #4616

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
KooshaPari wants to merge 1 commit into from
Closed

fix(api): finalize bifrost streaming usage on close #4616

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
98 changes: 76 additions & 22 deletions src/app/api/v1/relay/chat/completions/bifrost/route.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,11 +30,7 @@

import { CORS_HEADERS, handleCorsOptions } from "@/shared/utils/cors";
import { createInjectionGuard } from "@/middleware/promptInjectionGuard";
import {
getRelayTokenByHash,
checkRateLimit,
recordRelayUsage,
} from "@/lib/db/relayProxies";
import { getRelayTokenByHash, checkRateLimit, recordRelayUsage } from "@/lib/db/relayProxies";
import { buildErrorBody } from "@omniroute/open-sse/utils/error";
import { createHash } from "node:crypto";
import { z } from "zod";
Expand Down Expand Up @@ -84,8 +80,48 @@ function hashToken(token: string): string {
return createHash("sha256").update(token).digest("hex");
}

function wrapStreamWithFinalizer(
source: ReadableStream<Uint8Array>,
finalize: (status?: "success" | "error") => void
): ReadableStream<Uint8Array> {
const reader = source.getReader();
let finalized = false;

const finalizeOnce = (status?: "success" | "error") => {
if (finalized) return;
finalized = true;
finalize(status);
};

return new ReadableStream<Uint8Array>({
async pull(controller) {
Comment on lines +83 to +97

@gemini-code-assist gemini-code-assist Bot Jun 22, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Calling source.getReader() immediately when wrapStreamWithFinalizer is called locks the source stream eagerly. If the returned stream is never started or read, the underlying stream remains locked. It is safer and more idiomatic to obtain the reader inside the start method of the returned ReadableStream.

function wrapStreamWithFinalizer(
  source: ReadableStream<Uint8Array>,
  finalize: (status?: \"success\" | \"error\") => void
): ReadableStream<Uint8Array> {
  let reader: ReadableStreamDefaultReader<Uint8Array>;
  let finalized = false;

  const finalizeOnce = (status?: \"success\" | \"error\") => {
    if (finalized) return;
    finalized = true;
    finalize(status);
  };

  return new ReadableStream<Uint8Array>({
    start() {
      reader = source.getReader();
    },
    async pull(controller) {

try {
const chunk = await reader.read();
if (chunk.done) {
finalizeOnce();
controller.close();
return;
}

controller.enqueue(chunk.value);
} catch (err) {
finalizeOnce("error");
controller.error(err);
}
},
async cancel(reason) {
try {
await reader.cancel(reason);
} finally {
finalizeOnce("error");
}
},
Comment on lines +112 to +118

@gemini-code-assist gemini-code-assist Bot Jun 22, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Logging client-side stream cancellations as "error" might skew your error metrics, making normal user actions (like stopping generation or closing the tab) look like system or upstream failures. Consider whether these should be logged as "success" (or a specific "cancelled" status if supported) instead of "error".

Suggested change
async cancel(reason) {
try {
await reader.cancel(reason);
} finally {
finalizeOnce("error");
}
},
async cancel(reason) {
try {
await reader.cancel(reason);
} finally {
finalizeOnce();
}
},

});
}

export async function POST(request: Request) {
const startTime = Date.now();
const requestId = request.headers.get("x-request-id") || undefined;
const clientIp = sanitizeForensicHeader(
request.headers.get("x-forwarded-for")?.split(",")[0]?.trim() ||
request.headers.get("x-real-ip") ||
Expand Down Expand Up @@ -148,7 +184,7 @@ export async function POST(request: Request) {
const token = getRelayTokenByHash(tokenHash);
if (!token) {
recordRelayUsage("unknown", {
requestId: request.headers.get("x-request-id") || undefined,
requestId,
status: "auth_failed",
statusCode: 401,
latencyMs: Date.now() - startTime,
Expand All @@ -171,7 +207,7 @@ export async function POST(request: Request) {
const rateCheck = checkRateLimit(token.id);
if (!rateCheck.allowed) {
recordRelayUsage(token.id, {
requestId: request.headers.get("x-request-id") || undefined,
requestId,
status: "rate_limited",
statusCode: 429,
latencyMs: Date.now() - startTime,
Expand Down Expand Up @@ -201,7 +237,9 @@ export async function POST(request: Request) {
const parsed = BifrostRequestSchema.safeParse(rawBody);
if (!parsed.success) {
return new Response(
JSON.stringify(buildErrorBody(400, parsed.error.issues[0]?.message || "Invalid request body")),
JSON.stringify(
buildErrorBody(400, parsed.error.issues[0]?.message || "Invalid request body")
),
{ status: 400, headers: JSON_CORS_HEADERS }
);
}
Expand All @@ -210,7 +248,7 @@ export async function POST(request: Request) {
const guard = injectionGuard(body);
if (guard.blocked) {
recordRelayUsage(token.id, {
requestId: request.headers.get("x-request-id") || undefined,
requestId,
status: "error",
statusCode: 400,
latencyMs: Date.now() - startTime,
Expand Down Expand Up @@ -264,20 +302,25 @@ export async function POST(request: Request) {
body: JSON.stringify(body),
signal: ac.signal,
});
} finally {
} catch (err) {
clearTimeout(tid);
throw err;
}
Comment on lines 304 to 308

@gemini-code-assist gemini-code-assist Bot Jun 22, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Keeping the BIFROST_TIMEOUT_MS active during the entire streaming response means that any stream taking longer than 30 seconds (the default) will be prematurely aborted by the AbortController. In LLM applications, streaming responses can easily exceed 30 seconds for longer generations.

The timeout should only cover the initial connection and header retrieval (i.e., cleared as soon as fetch resolves), rather than the entire stream transfer. This matches the original behavior where clearTimeout(tid) was called in the finally block of the fetch call.

Suggested change
});
} finally {
} catch (err) {
clearTimeout(tid);
throw err;
}
});
clearTimeout(tid);
} catch (err) {
clearTimeout(tid);
throw err;
}


// 5. Forward response. For streaming we pass the body stream through
// unmodified — Node's fetch streams chunked correctly.
recordRelayUsage(token.id, {
requestId: request.headers.get("x-request-id") || undefined,
status: upstream.status < 500 ? "success" : "error",
statusCode: upstream.status,
latencyMs: Date.now() - startTime,
clientIp,
userAgent,
});
let finalized = false;
const finalizeRelayUsage = (statusOverride?: "success" | "error") => {
if (finalized) return;
finalized = true;
clearTimeout(tid);
recordRelayUsage(token.id, {
requestId,
status: statusOverride ?? (upstream.status < 500 ? "success" : "error"),
statusCode: upstream.status,
latencyMs: Date.now() - startTime,
clientIp,
userAgent,
});
};

const newHeaders = new Headers(upstream.headers);
newHeaders.set("X-Routed-By", "bifrost");
Expand All @@ -286,7 +329,18 @@ export async function POST(request: Request) {
newHeaders.set("Content-Type", upstream.headers.get("Content-Type") ?? "application/json");
}

return new Response(upstream.body, {
// 5. Forward response. Unary responses keep the existing fast path, but
// streaming responses finalize usage only after the body closes/errors/cancels.
const responseBody =
wantsStream && upstream.body
? wrapStreamWithFinalizer(upstream.body, finalizeRelayUsage)
: upstream.body;

if (!wantsStream || !upstream.body) {
finalizeRelayUsage();
}

return new Response(responseBody, {
status: upstream.status,
headers: newHeaders,
});
Expand All @@ -311,4 +365,4 @@ export async function POST(request: Request) {
}
);
}
}
}
109 changes: 106 additions & 3 deletions tests/unit/api/v1/bifrost-sidecar.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ const ORIGINAL_BIFROST_API_KEY = process.env.BIFROST_API_KEY;
const ORIGINAL_BIFROST_OMNI_KEY = process.env.OMNIROUTE_BIFROST_KEY;
const ORIGINAL_BIFROST_TIMEOUT = process.env.BIFROST_TIMEOUT_MS;
const ORIGINAL_BIFROST_STREAMING = process.env.BIFROST_STREAMING_ENABLED;
const encoder = new TextEncoder();

function restoreEnv() {
if (ORIGINAL_BIFROST_BASE_URL === undefined) delete process.env.BIFROST_BASE_URL;
Expand All @@ -38,9 +39,8 @@ test("bifrost route: returns 503 + fallback header when BIFROST_BASE_URL is unse
delete process.env.BIFROST_STREAMING_ENABLED;

// Dynamic import after env is set so the module reads the empty value.
const { POST } = await import(
"../../../../src/app/api/v1/relay/chat/completions/bifrost/route.ts"
);
const { POST } =
await import("../../../../src/app/api/v1/relay/chat/completions/bifrost/route.ts");

const req = new Request("http://localhost/api/v1/relay/chat/completions/bifrost", {
method: "POST",
Expand Down Expand Up @@ -117,3 +117,106 @@ test("bifrost route: OPTIONS responds with CORS headers", async () => {
"missing Access-Control-Allow-Headers header"
);
});

test("bifrost route: streaming usage and timeout cleanup wait for stream close", async (t) => {
process.env.BIFROST_BASE_URL = "http://bifrost.test.local:8080";
process.env.BIFROST_TIMEOUT_MS = "60000";
delete process.env.BIFROST_API_KEY;
delete process.env.OMNIROUTE_BIFROST_KEY;
delete process.env.BIFROST_STREAMING_ENABLED;

const core = await import("../../../../src/lib/db/core.ts");
const relayDb = await import("../../../../src/lib/db/relayProxies.ts");
core.resetDbInstance();

const rawToken = "relay_stream_regression_4595";
const token = {
id: "rl_stream_regression_4595",
rawToken,
tokenPrefix: "rl_stream",
};
const now = Math.floor(Date.now() / 1000);
const db = core.getDbInstance();
db.prepare(
`
INSERT INTO relay_tokens (id, name, token_hash, token_prefix, description, combo_id,
allowed_models, max_tokens_per_request, max_requests_per_minute, max_requests_per_day,
max_cost_per_day, enabled, created_at, updated_at, expires_at, metadata)
VALUES (?, ?, ?, ?, '', NULL, ?, 128000, 60, 10000, 0, 1, ?, ?, NULL, '{}')
`
).run(
token.id,
"stream-regression",
createHash("sha256").update(rawToken).digest("hex"),
token.tokenPrefix,
JSON.stringify(["gpt-4"]),
now,
now
);

const originalFetch = globalThis.fetch;
const originalClearTimeout = globalThis.clearTimeout;
const clearedTimeouts: unknown[] = [];

t.after(() => {
globalThis.fetch = originalFetch;
globalThis.clearTimeout = originalClearTimeout;
restoreEnv();
core.resetDbInstance();
});

let streamController: ReadableStreamDefaultController<Uint8Array> | undefined;
const upstreamBody = new ReadableStream<Uint8Array>({
start(controller) {
streamController = controller;
},
});

globalThis.fetch = (async (_input: RequestInfo | URL, _init?: RequestInit) => {
return new Response(upstreamBody, {
status: 200,
headers: { "Content-Type": "text/event-stream" },
});
}) as typeof fetch;

const { POST } = await import(
`../../../../src/app/api/v1/relay/chat/completions/bifrost/route.ts?case=${Date.now()}-${Math.random()}`
);

globalThis.clearTimeout = ((timeoutId: unknown) => {
clearedTimeouts.push(timeoutId);
return originalClearTimeout(timeoutId as NodeJS.Timeout);
}) as typeof clearTimeout;

const req = new Request("http://localhost/api/v1/relay/chat/completions/bifrost", {
method: "POST",
headers: {
authorization: `Bearer ${token.rawToken}`,
"content-type": "application/json",
"x-request-id": "stream-regression-4595",
},
body: JSON.stringify({
model: "gpt-4",
messages: [{ role: "user", content: "hi" }],
stream: true,
}),
});

const res = await POST(req);
assert.equal(res.status, 200);
assert.equal(relayDb.getRelayLogs(token.id).length, 0);
assert.equal(clearedTimeouts.length, 0);

streamController?.enqueue(encoder.encode('data: {"delta":"hi"}\n\n'));
assert.equal(relayDb.getRelayLogs(token.id).length, 0);
assert.equal(clearedTimeouts.length, 0);

streamController?.close();
assert.equal(await res.text(), 'data: {"delta":"hi"}\n\n');

const logs = relayDb.getRelayLogs(token.id);
assert.equal(logs.length, 1);
assert.equal(logs[0]?.status, "success");
assert.equal(logs[0]?.status_code, 200);
assert.equal(clearedTimeouts.length, 1);
Comment on lines +208 to +221

@gemini-code-assist gemini-code-assist Bot Jun 22, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Since the connection timeout should be cleared immediately after the initial response headers are received (when POST resolves), we should update the test assertions to expect clearedTimeouts.length to be 1 as soon as the response is returned, rather than waiting for the stream to close.

  assert.equal(clearedTimeouts.length, 1);

  streamController?.enqueue(encoder.encode('data: {\"delta\":\"hi\"}\\n\\n'));
  assert.equal(relayDb.getRelayLogs(token.id).length, 0);
  assert.equal(clearedTimeouts.length, 1);

  streamController?.close();
  assert.equal(await res.text(), 'data: {\"delta\":\"hi\"}\\n\\n');

  const logs = relayDb.getRelayLogs(token.id);
  assert.equal(logs.length, 1);
  assert.equal(logs[0]?.status, \"success\");
  assert.equal(logs[0]?.status_code, 200);
  assert.equal(clearedTimeouts.length, 1);

});
Loading