From 92c511550d76e3e27f9fcddfa7019b3144054754 Mon Sep 17 00:00:00 2001 From: Jan Leon Date: Tue, 30 Jun 2026 02:58:46 +0200 Subject: [PATCH] chore(duplication): share relay stream finalizer (#5497) --- .../relay/chat/completions/bifrost/route.ts | 49 ++----------------- .../api/v1/relay/chat/completions/route.ts | 39 +-------------- .../relay/chat/completions/streamFinalizer.ts | 37 ++++++++++++++ .../api/v1/relay-stream-finalizer.test.ts | 46 +++++++++++++++++ 4 files changed, 89 insertions(+), 82 deletions(-) create mode 100644 src/app/api/v1/relay/chat/completions/streamFinalizer.ts create mode 100644 tests/unit/api/v1/relay-stream-finalizer.test.ts diff --git a/src/app/api/v1/relay/chat/completions/bifrost/route.ts b/src/app/api/v1/relay/chat/completions/bifrost/route.ts index 936cbc7a64..6d1c29b0a4 100644 --- a/src/app/api/v1/relay/chat/completions/bifrost/route.ts +++ b/src/app/api/v1/relay/chat/completions/bifrost/route.ts @@ -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 { z } from "zod"; import { @@ -44,6 +40,7 @@ import { hashToken, sanitizeForensicHeader, } from "../relaySecurity"; +import { finalizeReadableStream } from "../streamFinalizer"; // Minimal request-shape validation (Rule #7). `.passthrough()` keeps every other // OpenAI chat-completion field intact (temperature, tools, response_format, …) — @@ -76,44 +73,6 @@ export async function OPTIONS() { return handleCorsOptions(); } -function finalizeReadableStream( - body: ReadableStream, - onFinalize: (error?: unknown) => void -): ReadableStream { - const reader = body.getReader(); - let finalized = false; - - const finalizeOnce = (error?: unknown) => { - if (finalized) return; - finalized = true; - onFinalize(error); - }; - - return new ReadableStream({ - async pull(controller) { - try { - const { done, value } = await reader.read(); - if (done) { - finalizeOnce(); - controller.close(); - return; - } - controller.enqueue(value); - } catch (error) { - finalizeOnce(error); - controller.error(error); - } - }, - async cancel(reason) { - try { - await reader.cancel(reason); - } finally { - finalizeOnce(reason); - } - }, - }); -} - export async function POST(request: Request) { const startTime = Date.now(); const clientIp = getClientIp(request); @@ -249,7 +208,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 } ); } diff --git a/src/app/api/v1/relay/chat/completions/route.ts b/src/app/api/v1/relay/chat/completions/route.ts index 32bae3ef08..27154bc574 100644 --- a/src/app/api/v1/relay/chat/completions/route.ts +++ b/src/app/api/v1/relay/chat/completions/route.ts @@ -24,6 +24,7 @@ import { shouldTryBifrost, type BifrostRoutingConfig, } from "./routingBackend"; +import { finalizeReadableStream } from "./streamFinalizer"; import type { RelayToken } from "@/lib/db/relayProxies"; const JSON_CORS_HEADERS = { ...CORS_HEADERS, "Content-Type": "application/json" } as const; @@ -51,44 +52,6 @@ function recordUsage( }); } -function finalizeReadableStream( - body: ReadableStream, - onFinalize: (error?: unknown) => void -): ReadableStream { - const reader = body.getReader(); - let finalized = false; - - const finalizeOnce = (error?: unknown) => { - if (finalized) return; - finalized = true; - onFinalize(error); - }; - - return new ReadableStream({ - async pull(controller) { - try { - const { done, value } = await reader.read(); - if (done) { - finalizeOnce(); - controller.close(); - return; - } - controller.enqueue(value); - } catch (error) { - finalizeOnce(error); - controller.error(error); - } - }, - async cancel(reason) { - try { - await reader.cancel(reason); - } finally { - finalizeOnce(reason); - } - }, - }); -} - async function forwardToBifrost( request: Request, body: unknown, diff --git a/src/app/api/v1/relay/chat/completions/streamFinalizer.ts b/src/app/api/v1/relay/chat/completions/streamFinalizer.ts new file mode 100644 index 0000000000..4eb6aeada0 --- /dev/null +++ b/src/app/api/v1/relay/chat/completions/streamFinalizer.ts @@ -0,0 +1,37 @@ +export function finalizeReadableStream( + body: ReadableStream, + onFinalize: (error?: unknown) => void +): ReadableStream { + const reader = body.getReader(); + let finalized = false; + + const finalizeOnce = (error?: unknown) => { + if (finalized) return; + finalized = true; + onFinalize(error); + }; + + return new ReadableStream({ + async pull(controller) { + try { + const { done, value } = await reader.read(); + if (done) { + finalizeOnce(); + controller.close(); + return; + } + controller.enqueue(value); + } catch (error) { + finalizeOnce(error); + controller.error(error); + } + }, + async cancel(reason) { + try { + await reader.cancel(reason); + } finally { + finalizeOnce(reason); + } + }, + }); +} diff --git a/tests/unit/api/v1/relay-stream-finalizer.test.ts b/tests/unit/api/v1/relay-stream-finalizer.test.ts new file mode 100644 index 0000000000..3be9fbeb14 --- /dev/null +++ b/tests/unit/api/v1/relay-stream-finalizer.test.ts @@ -0,0 +1,46 @@ +import test from "node:test"; +import assert from "node:assert/strict"; + +import { finalizeReadableStream } from "../../../../src/app/api/v1/relay/chat/completions/streamFinalizer.ts"; + +test("finalizeReadableStream finalizes once after the wrapped stream completes", async () => { + const finalized: unknown[] = []; + const stream = finalizeReadableStream( + new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode("hello")); + controller.close(); + }, + }), + (error) => finalized.push(error) + ); + + assert.equal(await new Response(stream).text(), "hello"); + assert.deepEqual(finalized, [undefined]); +}); + +test("finalizeReadableStream finalizes once when the consumer cancels", async () => { + const finalized: unknown[] = []; + let cancelReason: unknown; + const stream = finalizeReadableStream( + new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode("chunk")); + }, + cancel(reason) { + cancelReason = reason; + }, + }), + (error) => finalized.push(error) + ); + + const reader = stream.getReader(); + const first = await reader.read(); + assert.equal(new TextDecoder().decode(first.value), "chunk"); + + await reader.cancel("client disconnected"); + await reader.cancel("second cancel"); + + assert.equal(cancelReason, "client disconnected"); + assert.deepEqual(finalized, ["client disconnected"]); +});