From c77215f20498e2643cc4febe0ab3fe738dbab3f3 Mon Sep 17 00:00:00 2001 From: PizzaV <103120356+pizzav-xyz@users.noreply.github.com> Date: Tue, 9 Jun 2026 16:14:27 +0200 Subject: [PATCH] fix(combo): parallel pre-screen + circuit-breaker fast-exit for priority combos (#3169) Integrated into release/v3.8.18 --- open-sse/services/combo.ts | 73 +++++++- open-sse/services/comboConfig.ts | 6 + tests/unit/combo-prescreen.test.ts | 205 +++++++++++++++++++++ tests/unit/combo-provider-cooldown.test.ts | 52 ++++++ 4 files changed, 334 insertions(+), 2 deletions(-) create mode 100644 tests/unit/combo-prescreen.test.ts diff --git a/open-sse/services/combo.ts b/open-sse/services/combo.ts index 0ec3a8485c..15193ee2ce 100644 --- a/open-sse/services/combo.ts +++ b/open-sse/services/combo.ts @@ -13,6 +13,7 @@ import { recordProviderFailure, isProviderFailureCode, isProviderExhaustedReason, + type ProviderProfile, } from "./accountFallback.ts"; import { FETCH_TIMEOUT_MS, RateLimitReason } from "../config/constants.ts"; import { errorResponse, unavailableResponse } from "../utils/error.ts"; @@ -27,6 +28,7 @@ import { resolveComboConfig, getDefaultComboConfig, resolveComboTargetTimeoutMs, + PRE_SCREEN_CONCURRENCY, } from "./comboConfig.ts"; import { maybeGenerateHandoff, @@ -1833,6 +1835,42 @@ async function fetchResetAwareQuotaWithCache({ return refresh(); } +type PreScreenResult = { profile: ProviderProfile | null; available: boolean }; + +export async function preScreenTargets( + targets: ResolvedComboTarget[], + isModelAvailable?: ((model: string, target: ResolvedComboTarget) => Promise) | null +): Promise> { + if (targets.length === 0) { + return new Map(); + } + + const results = await mapWithConcurrency( + targets, + PRE_SCREEN_CONCURRENCY, + async (target): Promise<{ key: string; result: PreScreenResult }> => { + const profile = await getRuntimeProviderProfile(target.provider).catch(() => null); + + const breaker = getCircuitBreaker(target.provider); + if (breaker.getStatus().state === "OPEN") { + return { key: target.executionKey, result: { profile, available: false } }; + } + + let available = true; + if (isModelAvailable) { + available = await isModelAvailable(target.modelStr, target).catch(() => true); + } + return { key: target.executionKey, result: { profile, available } }; + } + ); + + const map = new Map(); + for (const { key, result } of results) { + map.set(key, result); + } + return map; +} + async function orderTargetsByResetAwareQuota( targets: ResolvedComboTarget[], comboName: string, @@ -3217,6 +3255,15 @@ export async function handleComboChat({ orderedTargets = orderTargetsByEvalScores(orderedTargets, config.evalRouting, log); orderedTargets = filterTargetsByRequestCompatibility(orderedTargets, body, log); + // Parallel pre-screen: check provider profiles and model availability for all targets + // Only runs for priority strategy where sequential checking causes latency + const preScreenMap = + strategy === "priority" + ? await preScreenTargets(orderedTargets, isModelAvailable).catch( + () => new Map() + ) + : new Map(); + if (orderedTargets.length === 0) { return comboModelNotFoundResponse("Combo has no executable targets"); } @@ -3285,7 +3332,18 @@ export async function handleComboChat({ const target = orderedTargets[i]; const modelStr = target.modelStr; const provider = target.provider; - const profile = await getRuntimeProviderProfile(provider); + + const cb = getCircuitBreaker(provider); + if (cb.getStatus().state === "OPEN") { + log.info("COMBO", `Skipping ${modelStr} — circuit breaker OPEN for ${provider}`); + if (i > 0) fallbackCount++; + return null; + } + + // Use pre-screened profile if available, otherwise fetch on demand + const preScreenEntry = preScreenMap.get(target.executionKey); + const profile = preScreenEntry?.profile ?? (await getRuntimeProviderProfile(provider)); + const allowRateLimitedConnection = Boolean(provider && provider !== "unknown") && transientRateLimitedProviders.has(provider); @@ -3307,7 +3365,18 @@ export async function handleComboChat({ return null; } - // Pre-check: skip models where no credentials are available (excluded, rate-limited, or unavailable) + // Pre-screen may have already determined this target unavailable (e.g. + // circuit-breaker OPEN at resolve time). Skip immediately in that case. + // For targets pre-screened as "available" we still call isModelAvailable + // below because connection cooldowns (rateLimitedUntil) can change + // mid-request after a same-provider failure — the pre-screen snapshot is + // stale by the time we reach the 2nd/3rd same-provider target. + const preCheckedAvailable = preScreenEntry?.available ?? null; + if (preCheckedAvailable === false) { + log.info("COMBO", `Skipping ${modelStr} — pre-screen marked unavailable`); + if (i > 0) fallbackCount++; + return null; + } if (isModelAvailable) { const available = await isModelAvailable(modelStr, targetForAttempt); if (!available) { diff --git a/open-sse/services/comboConfig.ts b/open-sse/services/comboConfig.ts index 64ef751d00..e8c9c33a52 100644 --- a/open-sse/services/comboConfig.ts +++ b/open-sse/services/comboConfig.ts @@ -7,6 +7,12 @@ import { MAX_TIMER_TIMEOUT_MS } from "../../src/shared/utils/runtimeTimeouts.ts"; +/** + * Maximum number of concurrent pre-screen checks (provider profile + availability) + * when running parallel pre-screening for priority strategy combos. + */ +export const PRE_SCREEN_CONCURRENCY = 5; + const DEFAULT_COMBO_CONFIG = { strategy: "priority", maxRetries: 1, diff --git a/tests/unit/combo-prescreen.test.ts b/tests/unit/combo-prescreen.test.ts new file mode 100644 index 0000000000..aa5008ea1f --- /dev/null +++ b/tests/unit/combo-prescreen.test.ts @@ -0,0 +1,205 @@ +import test, { after } from "node:test"; +import assert from "node:assert/strict"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-combo-prescreen-")); +const ORIGINAL_DATA_DIR = process.env.DATA_DIR; +process.env.DATA_DIR = TEST_DATA_DIR; + +const dbCore = await import("../../src/lib/db/core.ts"); +const { handleComboChat } = await import("../../open-sse/services/combo.ts"); +const combosDb = await import("../../src/lib/db/combos.ts"); + +after(() => { + dbCore.resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); + if (ORIGINAL_DATA_DIR === undefined) { + delete process.env.DATA_DIR; + } else { + process.env.DATA_DIR = ORIGINAL_DATA_DIR; + } +}); + +function okResponse(model: string) { + return Response.json({ choices: [{ message: { role: "assistant", content: model } }] }); +} + +function makeLog() { + return { + info() {}, + warn() {}, + debug() {}, + error() {}, + }; +} + +const reqBody = { + model: "prescreen-test", + messages: [{ role: "user", content: "hi" }], + stream: false, +}; + +test("pre-screen: all targets checked in parallel", async () => { + const checkOrder: string[] = []; + + const combo = await combosDb.createCombo({ + name: "prescreen-parallel", + strategy: "priority", + models: ["p1/m1", "p2/m2", "p3/m3"], + }); + + const response = await handleComboChat({ + body: { ...reqBody, model: combo.name }, + combo, + allCombos: [combo], + isModelAvailable: async (modelStr: string) => { + checkOrder.push(modelStr); + return true; + }, + relayOptions: undefined, + signal: undefined, + settings: {}, + log: makeLog(), + handleSingleModel: async (_body: unknown, modelStr: string) => { + return okResponse(modelStr); + }, + }); + + assert.equal(response.status, 200); + assert.ok(checkOrder.length >= 3, "pre-screen must check all 3 targets concurrently"); +}); + +test("pre-screen: unavailable targets are skipped", async () => { + const calls: string[] = []; + const availability = new Map([ + ["p1/m1", false], + ["p2/m2", true], + ["p3/m3", true], + ]); + + const combo = await combosDb.createCombo({ + name: "prescreen-skip", + strategy: "priority", + models: ["p1/m1", "p2/m2", "p3/m3"], + }); + + const response = await handleComboChat({ + body: { ...reqBody, model: combo.name }, + combo, + allCombos: [combo], + isModelAvailable: async (modelStr: string) => { + return availability.get(modelStr) ?? true; + }, + relayOptions: undefined, + signal: undefined, + settings: {}, + log: makeLog(), + handleSingleModel: async (_body: unknown, modelStr: string) => { + calls.push(modelStr); + return okResponse(modelStr); + }, + }); + + assert.equal(response.status, 200); + assert.equal(calls.length, 1); + assert.equal(calls[0], "p2/m2"); +}); + +test("pre-screen: failure treated as unknown availability", async () => { + const calls: string[] = []; + let checkCount = 0; + + const combo = await combosDb.createCombo({ + name: "prescreen-failure", + strategy: "priority", + models: ["p1/m1", "p2/m2"], + }); + + const response = await handleComboChat({ + body: { ...reqBody, model: combo.name }, + combo, + allCombos: [combo], + isModelAvailable: async (modelStr: string) => { + checkCount++; + if (checkCount === 1) { + throw new Error("DB connection failed"); + } + return true; + }, + relayOptions: undefined, + signal: undefined, + settings: {}, + log: makeLog(), + handleSingleModel: async (_body: unknown, modelStr: string) => { + calls.push(modelStr); + return okResponse(modelStr); + }, + }); + + assert.equal(response.status, 200); + // m1 pre-screen failed (treated as available), m2 pre-screen succeeded + // Both should be tried, m1 first since it's available + assert.equal(calls.length, 1); + assert.equal(calls[0], "p1/m1"); +}); + +test("pre-screen: only runs for priority strategy", async () => { + const checkOrder: string[] = []; + + const combo = await combosDb.createCombo({ + name: "prescreen-strategy", + strategy: "round-robin", + models: ["p1/m1", "p2/m2"], + }); + + const response = await handleComboChat({ + body: { ...reqBody, model: combo.name }, + combo, + allCombos: [combo], + isModelAvailable: async (modelStr: string) => { + checkOrder.push(modelStr); + return true; + }, + relayOptions: undefined, + signal: undefined, + settings: {}, + log: makeLog(), + handleSingleModel: async (_body: unknown, modelStr: string) => { + return okResponse(modelStr); + }, + }); + + assert.equal(response.status, 200); + assert.equal(checkOrder.length >= 1, true); +}); + +test("pre-screen: backward compatible with all targets available", async () => { + const calls: string[] = []; + + const combo = await combosDb.createCombo({ + name: "prescreen-compat", + strategy: "priority", + models: ["p1/m1", "p2/m2", "p3/m3"], + }); + + const response = await handleComboChat({ + body: { ...reqBody, model: combo.name }, + combo, + allCombos: [combo], + isModelAvailable: async () => true, + relayOptions: undefined, + signal: undefined, + settings: {}, + log: makeLog(), + handleSingleModel: async (_body: unknown, modelStr: string) => { + calls.push(modelStr); + return okResponse(modelStr); + }, + }); + + assert.equal(response.status, 200); + assert.equal(calls.length, 1); + assert.equal(calls[0], "p1/m1"); +}); diff --git a/tests/unit/combo-provider-cooldown.test.ts b/tests/unit/combo-provider-cooldown.test.ts index 046e15aef8..daaceec9b0 100644 --- a/tests/unit/combo-provider-cooldown.test.ts +++ b/tests/unit/combo-provider-cooldown.test.ts @@ -14,6 +14,12 @@ const { seedConnection, settingsDb, } = harness; +const { preScreenTargets } = await import( + "../../open-sse/services/combo.ts" +); +const { getCircuitBreaker } = await import( + "../../src/shared/utils/circuitBreaker.ts" +); test.beforeEach(async () => { await resetStorage(); @@ -105,3 +111,49 @@ test("combo failover skips the cooled provider target on the next request", asyn assert.equal(openaiCalls, 1); assert.equal(claudeCalls, 2); }); + +test("pre-screen marks target unavailable when circuit breaker is OPEN", async () => { + const breaker = getCircuitBreaker("openai", { failureThreshold: 1, resetTimeout: 60_000 }); + try { + await breaker.execute(async () => { + throw new Error("simulated failure"); + }); + } catch { + // expected + } + + const targets = [ + { + kind: "model" as const, + stepId: "step-1", + executionKey: "openai/gpt-4o", + modelStr: "openai/gpt-4o", + provider: "openai", + providerId: "conn-1", + connectionId: "conn-1", + weight: 1, + label: null, + }, + { + kind: "model" as const, + stepId: "step-2", + executionKey: "claude/claude-3-5-sonnet-20241022", + modelStr: "claude/claude-3-5-sonnet-20241022", + provider: "claude", + providerId: "conn-2", + connectionId: "conn-2", + weight: 1, + label: null, + }, + ]; + + const results = await preScreenTargets(targets as any); + + const openaiResult = results.get("openai/gpt-4o"); + assert.ok(openaiResult, "openai target should have a pre-screen result"); + assert.equal(openaiResult.available, false, "open-circuit-breaker target should be unavailable"); + + const claudeResult = results.get("claude/claude-3-5-sonnet-20241022"); + assert.ok(claudeResult, "claude target should have a pre-screen result"); + assert.equal(claudeResult.available, true, "closed-circuit-breaker target should be available"); +});