diff --git a/changelog.d/fixes/7350-rerank-proxy-pinning.md b/changelog.d/fixes/7350-rerank-proxy-pinning.md new file mode 100644 index 0000000000..29bd56856b --- /dev/null +++ b/changelog.d/fixes/7350-rerank-proxy-pinning.md @@ -0,0 +1 @@ +- **fix(rerank):** Honor the connection's pinned proxy on rerank calls, so a provider that geo-blocks the host IP (Voyage AI) works on the same connection where chat and embeddings already did ([#7350](https://github.com/diegosouzapw/OmniRoute/issues/7350)) — thanks @kamenkadmitry diff --git a/open-sse/handlers/rerank.ts b/open-sse/handlers/rerank.ts index fb37f08267..719459d555 100644 --- a/open-sse/handlers/rerank.ts +++ b/open-sse/handlers/rerank.ts @@ -12,6 +12,9 @@ import { attachOmniRouteMetaHeaders } from "@/domain/omnirouteResponseMeta"; import { calculateModalCost } from "@/lib/usage/costCalculator"; import { generateRequestId } from "@/shared/utils/requestId"; import { saveCallLog } from "@/lib/usageDb"; +import { resolveProxyForConnection } from "@/lib/db/settings"; +import { runWithProxyContext } from "../utils/proxyFetch.ts"; +import * as log from "@/sse/utils/logger"; /** * Build authorization header for a rerank provider @@ -161,6 +164,7 @@ function buildAuthHeader(providerConfig, token) { * @param {number} [options.top_n] - Number of top results to return * @param {boolean} [options.return_documents] - Whether to include document text in results * @param {Object} options.credentials - Provider credentials { apiKey, accessToken } + * @param {string} [options.connectionId] - Connection ID for per-connection proxy resolution * @returns {Response} */ /** @returns {Promise} */ @@ -171,6 +175,7 @@ export async function handleRerank({ top_n, return_documents, credentials, + connectionId = null, }) { const startTime = Date.now(); if (!model) return errorResponse(400, "model is required"); @@ -210,8 +215,20 @@ export async function handleRerank({ ? `${providerConfig.baseUrl}/${modelId}` : providerConfig.baseUrl; - try { - const res = await fetch(rerankUrl, { + // Resolve per-connection proxy so rerank honors the same proxy pinning as chat + // and embeddings (#7350). Without this, rerank requests egress directly and fail + // when the provider blocks certain IP ranges (e.g. Voyage AI from Russian IPs). + let proxyInfo: Awaited> | null = null; + if (connectionId) { + try { + proxyInfo = await resolveProxyForConnection(connectionId); + } catch (err) { + log.error("RERANK", `Proxy resolution failed for connection ${connectionId}: ${err}`); + } + } + + const doFetch = () => + fetch(rerankUrl, { method: "POST", headers: { "Content-Type": "application/json", @@ -220,6 +237,11 @@ export async function handleRerank({ body: JSON.stringify(requestBody), }); + try { + const res = connectionId + ? await runWithProxyContext(proxyInfo?.proxy || null, doFetch) + : await doFetch(); + if (!res.ok) { const errData = await res.json().catch(() => ({})); return errorResponse( diff --git a/src/app/api/v1/rerank/route.ts b/src/app/api/v1/rerank/route.ts index 7130af0d20..bf9da386eb 100644 --- a/src/app/api/v1/rerank/route.ts +++ b/src/app/api/v1/rerank/route.ts @@ -120,6 +120,7 @@ async function postHandler(request, context) { top_n: body.top_n, return_documents: body.return_documents, credentials, + connectionId: (credentials as { connectionId?: string } | null)?.connectionId || null, }); if (response?.ok) { await clearRecoveredProviderState(credentials); diff --git a/src/lib/api/modelTestRunner.ts b/src/lib/api/modelTestRunner.ts index f042b89341..d519313b66 100644 --- a/src/lib/api/modelTestRunner.ts +++ b/src/lib/api/modelTestRunner.ts @@ -297,7 +297,8 @@ export async function runSingleModelTest( const runInner = async (signal: AbortSignal): Promise => { if (isEmbedding) { return handleValidatedEmbeddingRequestBody( - testBody as Record & { model: string } + testBody as Record & { model: string }, + { connectionId: connectionId || undefined } ); } if (isRerank) { diff --git a/tests/unit/rerank-proxy-pinning-7350.test.ts b/tests/unit/rerank-proxy-pinning-7350.test.ts new file mode 100644 index 0000000000..9a699c8f77 --- /dev/null +++ b/tests/unit/rerank-proxy-pinning-7350.test.ts @@ -0,0 +1,125 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +/** + * #7350 — rerank egressed directly instead of honoring the connection's proxy, so a + * provider that geo-blocks the host's IP (Voyage AI from some ranges) failed even when + * the connection had a working proxy pinned, while chat and embeddings on the SAME + * connection worked. `handleRerank` now takes a `connectionId`, resolves that + * connection's proxy and wraps the upstream fetch in `runWithProxyContext`. + * + * These tests drive the real DB + resolution cascade (no module mocking) and assert on + * the proxy agent that actually reaches undici, which is what "the request egressed + * through the proxy" concretely means here. + */ + +const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-rerank-proxy-")); +process.env.DATA_DIR = TEST_DATA_DIR; +process.env.API_KEY_SECRET = "test-secret"; + +const core = await import("../../src/lib/db/core.ts"); +const providersDb = await import("../../src/lib/db/providers.ts"); +const proxiesDb = await import("../../src/lib/db/proxies.ts"); +const proxyFetch = await import("../../open-sse/utils/proxyFetch.ts"); +const { handleRerank } = await import("../../open-sse/handlers/rerank.ts"); + +const originalFetch = globalThis.fetch; + +/** Captures the proxy URL visible inside the dispatch context at fetch time. */ +function stubFetch(seen: { proxyUrl: string | null | undefined }[]) { + globalThis.fetch = (async () => { + seen.push({ proxyUrl: proxyFetch.getCurrentProxyUrlForTests?.() ?? undefined }); + return new Response( + JSON.stringify({ data: [{ index: 0, relevance_score: 0.9 }], model: "rerank-2" }), + { status: 200, headers: { "Content-Type": "application/json" } } + ); + }) as typeof fetch; +} + +test.after(() => { + globalThis.fetch = originalFetch; + core.resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); +}); + +test("#7350 handleRerank routes the upstream call through the connection's pinned proxy", async () => { + core.resetDbInstance(); + const conn = await providersDb.createProviderConnection({ + provider: "voyage", + authType: "apikey", + name: "voyage-proxied", + apiKey: "pa-test-key", + }); + const proxy = await proxiesDb.createProxy({ + name: "Rerank Egress Proxy", + type: "http", + host: "rerank-egress.local", + port: 8080, + }); + await proxiesDb.assignProxyToScope("account", (conn as { id: string }).id, proxy.id); + + // The stub only ever answers a DIRECT call: runWithProxyContext dispatches through + // undici with the pinned proxy agent instead, so a pinned-but-unreachable proxy is + // observable as "the stub was bypassed and the request did not succeed". That + // difference IS the wiring — before #7350 this call egressed directly and got a 200. + const seen: { proxyUrl: string | null | undefined }[] = []; + stubFetch(seen); + + const res = (await handleRerank({ + model: "voyage/rerank-2", + query: "q", + documents: ["a", "b"], + credentials: { apiKey: "pa-test-key" }, + connectionId: (conn as { id: string }).id, + })) as Response; + + assert.equal(seen.length, 0, "a pinned proxy must bypass the direct-egress path entirely"); + assert.notEqual( + res.status, + 200, + "the unreachable pinned proxy must surface as a failure rather than silently egressing direct" + ); +}); + +test("#7350 an unresolvable connectionId degrades to a direct call instead of failing the request", async () => { + core.resetDbInstance(); + const seen: { proxyUrl: string | null | undefined }[] = []; + stubFetch(seen); + + const res = await handleRerank({ + model: "voyage/rerank-2", + query: "q", + documents: ["a"], + credentials: { apiKey: "pa-test-key" }, + connectionId: "connection-that-does-not-exist", + }); + + assert.equal(seen.length, 1, "proxy resolution failure must not swallow the upstream call"); + assert.equal( + (res as Response).status, + 200, + "a failed proxy lookup is logged and skipped, never turned into a request error" + ); +}); + +test("#7350 omitting connectionId keeps the previous direct-egress behavior", async () => { + core.resetDbInstance(); + const seen: { proxyUrl: string | null | undefined }[] = []; + stubFetch(seen); + + await handleRerank({ + model: "voyage/rerank-2", + query: "q", + documents: ["a"], + credentials: { apiKey: "pa-test-key" }, + }); + + assert.equal(seen.length, 1); + assert.ok( + !seen[0].proxyUrl, + `no connectionId must mean no proxy context, saw: ${seen[0].proxyUrl}` + ); +});