pjandro/OmniRoute
1
0
Fork 0
mirror of https://github.com/diegosouzapw/OmniRoute.git synced 2026-08-04 06:12:10 +03:00
Code Issues Packages Projects Releases Wiki Activity

chore(duplication): share relay stream finalizer (#5497)

Browse Source
This commit is contained in:
Jan Leon
2026-06-30 02:58:46 +02:00
committed by GitHub
parent b9fd69ecf3
commit 92c511550d
4 changed files with 89 additions and 82 deletions
Download Patch File Download Diff File Expand all files Collapse all files

49
src/app/api/v1/relay/chat/completions/bifrost/route.ts
View File

@@ -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<Uint8Array>,
onFinalize: (error?: unknown) => void
): ReadableStream<Uint8Array> {
const reader = body.getReader();
let finalized = false;
const finalizeOnce = (error?: unknown) => {
if (finalized) return;
finalized = true;
onFinalize(error);
};
return new ReadableStream<Uint8Array>({
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 }
);
}

39
src/app/api/v1/relay/chat/completions/route.ts
View File

@@ -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<Uint8Array>,
onFinalize: (error?: unknown) => void
): ReadableStream<Uint8Array> {
const reader = body.getReader();
let finalized = false;
const finalizeOnce = (error?: unknown) => {
if (finalized) return;
finalized = true;
onFinalize(error);
};
return new ReadableStream<Uint8Array>({
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,

37
src/app/api/v1/relay/chat/completions/streamFinalizer.ts Normal file
View File

@@ -0,0 +1,37 @@
export function finalizeReadableStream(
body: ReadableStream<Uint8Array>,
onFinalize: (error?: unknown) => void
): ReadableStream<Uint8Array> {
const reader = body.getReader();
let finalized = false;
const finalizeOnce = (error?: unknown) => {
if (finalized) return;
finalized = true;
onFinalize(error);
};
return new ReadableStream<Uint8Array>({
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);
}
},
});
}

46
tests/unit/api/v1/relay-stream-finalizer.test.ts Normal file
View File

@@ -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<Uint8Array>({
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<Uint8Array>({
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"]);
});