From 00b7b71bd370a841dc84c78a7b06fdbb9936a13d Mon Sep 17 00:00:00 2001 From: "Bob.Hou" Date: Fri, 28 Aug 2026 11:22:18 -0400 Subject: [PATCH] feat(routing): expand connection-aware quota prefilter across combo strategies (#11682) (#11850) Expands per-connection quota-aware pre-filtering across all 20 combo strategies so exhausted accounts are filtered before strategy resolution instead of causing avoidable upstream errors. Closes #11682. 15/15 + 131/131 + 457/457 (vitest) focused tests passing. Thanks! --- open-sse/services/combo.ts | 20 +- .../combo/connectionAwareExpansion.ts | 158 ++++ open-sse/services/combo/dispatchPrelude.ts | 61 +- open-sse/services/combo/quotaStrategies.ts | 8 +- open-sse/services/combo/targetResolution.ts | 18 +- open-sse/services/combo/types.ts | 2 +- open-sse/services/comboConfig.ts | 3 + src/lib/db/settings.ts | 2 + src/shared/validation/schemas/combo.ts | 3 + src/shared/validation/settingsSchemas.ts | 2 + .../combo/connection-aware-expansion.test.ts | 675 ++++++++++++++++++ 11 files changed, 944 insertions(+), 8 deletions(-) create mode 100644 open-sse/services/combo/connectionAwareExpansion.ts create mode 100644 tests/unit/combo/connection-aware-expansion.test.ts diff --git a/open-sse/services/combo.ts b/open-sse/services/combo.ts index 1f73cbcc9f..38c210ee5b 100644 --- a/open-sse/services/combo.ts +++ b/open-sse/services/combo.ts @@ -168,6 +168,7 @@ import { recordStickyWeightedSuccess, resolveComboStickyRoundRobinLimit, } from "./combo/rrState.ts"; +import { expandTargetsForAllStrategies } from "./combo/connectionAwareExpansion.ts"; import { validateResponseQuality, releaseQualityClone, @@ -839,6 +840,8 @@ async function handleComboChatInner({ combo, config, strategy, + settings, + apiKeyAllowedConnections, allCombos, handleSingleModelWithTimeout, log, @@ -889,6 +892,7 @@ async function handleComboChatInner({ settings, allCombos, signal, + apiKeyAllowedConnections, hiddenModelsByProvider, clientManagedResponsesContext, deferContextOverflowWhenCompressible, @@ -2876,6 +2880,7 @@ async function handleRoundRobinCombo({ settings, allCombos, signal, + apiKeyAllowedConnections = null, nesting = null, hiddenModelsByProvider = getHiddenModelsByProvider(), clientManagedResponsesContext, @@ -2932,12 +2937,25 @@ async function handleRoundRobinCombo({ } : allCombos; - const orderedTargets = resolveComboTargets( + let orderedTargets = resolveComboTargets( rrExpandedCombo, rrExpandedAllCombos, clampComboDepth(config.maxComboDepth), hiddenModelsByProvider ); + // Connection-aware expansion is opt-in. RR runs outside + // resolveComboTargetPipeline, so it wires the same stage here. Rotation + // granularity becomes model x connection: rrStartIndex takes mod over the + // expanded list, so each account occupies its own rotation slot. + orderedTargets = await expandTargetsForAllStrategies({ + strategy: "round-robin", + targets: orderedTargets, + comboName: combo.name, + config: combo.config, + settings: settings as Record | null | undefined, + log, + apiKeyAllowedConnectionIds: apiKeyAllowedConnections, + }); const tagFilteredTargets = await applyRequestTagRouting(orderedTargets, body, log); const evalRankedTargets = orderTargetsByEvalScores(tagFilteredTargets, config.evalRouting, log); // Align with the main/auto paths: combo config OR top-level settings (#8488 / #8494). diff --git a/open-sse/services/combo/connectionAwareExpansion.ts b/open-sse/services/combo/connectionAwareExpansion.ts new file mode 100644 index 0000000000..7cdfeab382 --- /dev/null +++ b/open-sse/services/combo/connectionAwareExpansion.ts @@ -0,0 +1,158 @@ +/** + * Connection-aware expansion -- shared pipeline stage for all combo strategies. + * + * Group-A strategies (reset-aware / reset-window / headroom / quota-share) and + * `auto` already expand targets to concrete per-connection candidates inside + * their own ordering (applyStrategyOrdering / buildAutoCandidates). The other + * 15 strategies ("group B": priority, weighted, round-robin, random, p2c, + * least-used, cost-optimized, lkgp, fill-first, strict-random, + * context-optimized, cache-optimized, context-relay, fusion, pipeline) kept a + * provider-level view, so an exhausted multi-account provider kept getting + * picked -- the per-target exhaustion gates in combo.ts were dead code for + * them (target.connectionId === null). + * + * This module promotes the A-group expander to a shared stage, gated behind + * the opt-in config key `connectionAwareExpansion` (combo config or settings + * fallback, default false). When off, targets pass through byte-identical. + * + */ + +import type { ResolvedComboTarget, ComboLogger } from "./types.ts"; +import { expandTargetsByQuotaAwareConnections } from "./quotaStrategies.ts"; + +/** + * The 15 group-B strategies that gain per-connection awareness through this + * stage. A-group strategies (reset-aware, reset-window, headroom, quota-share) + * and `auto` are deliberately absent -- they expand inside their own ordering + * and a second pass here would only burn another connection-list read + * (idempotent but wasteful). + */ +export const CONNECTION_AWARE_EXPANSION_GROUP_B: readonly string[] = [ + "priority", + "weighted", + "round-robin", + "random", + "p2c", + "least-used", + "cost-optimized", + "lkgp", + "fill-first", + "strict-random", + "context-optimized", + "cache-optimized", + "context-relay", + "fusion", + "pipeline", +] as const; + +const GROUP_B_SET = new Set(CONNECTION_AWARE_EXPANSION_GROUP_B); + +/** + * Upper bound on per-target expansion. Bounds pool blowup when a provider has + * many accounts: a 10-step combo over an 8-connection provider + * stays at most 80 candidates instead of unbounded growth. + */ +export const DEFAULT_CONNECTION_AWARE_EXPANSION_MAX_PER_TARGET = 8; + +export interface ExpandTargetsForAllStrategiesArgs { + strategy: string; + targets: ResolvedComboTarget[]; + comboName: string; + /** Resolved combo config (cascade output) or raw combo config. */ + config: Record | null | undefined; + /** Global settings layer; consulted when the combo config layer is unset. */ + settings?: Record | null | undefined; + log: ComboLogger; + /** API-key allowedConnections scope, intersected with the expansion. */ + apiKeyAllowedConnectionIds?: string[] | null; + /** + * Test-only injection point replacing the underlying expander, used to + * exercise the fail-open path deterministically (T10). Production callers + * leave it unset so the real expander runs. + */ + __testExpander?: ( + targets: ResolvedComboTarget[], + comboName: string, + log: ComboLogger, + apiKeyAllowedConnectionIds: string[] | null + ) => Promise<{ expandedTargets: ResolvedComboTarget[] }>; +} + +/** + * Should the connection-aware expansion stage run for this request? + * + * true iff the strategy is a group-B strategy AND the resolved config (combo + * config layer, falling back to the settings layer) enables + * `connectionAwareExpansion`. Both layers default to false; flipping the + * switch must never change a combo that did not ask for it. + */ +export function shouldApplyConnectionAwareExpansion( + strategy: string, + config: Record | null | undefined, + settings?: Record | null | undefined +): boolean { + if (!GROUP_B_SET.has(strategy)) return false; + if (config?.connectionAwareExpansion === true) return true; + if (config?.connectionAwareExpansion === false) return false; + return settings?.connectionAwareExpansion === true; +} + +function clampMaxPerTarget(value: unknown): number { + const numeric = Number(value); + if (!Number.isFinite(numeric) || numeric < 1) { + return DEFAULT_CONNECTION_AWARE_EXPANSION_MAX_PER_TARGET; + } + return Math.min(Math.floor(numeric), 64); +} + +/** + * Expand group-B targets into per-connection candidates using the A-group + * expander. Fail-open everywhere: when anything in the expansion path throws, + * the original targets come back untouched (spec section 3.1; the stage must never + * take a combo down). + */ +export async function expandTargetsForAllStrategies( + args: ExpandTargetsForAllStrategiesArgs +): Promise { + const { strategy, targets, comboName, config, settings, log } = args; + if (targets.length === 0) return targets; + if (!shouldApplyConnectionAwareExpansion(strategy, config, settings)) return targets; + + const maxPerTarget = clampMaxPerTarget(config?.connectionAwareExpansionMaxPerTarget); + + try { + const expander = + args.__testExpander ?? + ((t: ResolvedComboTarget[], name: string, l: ComboLogger, allowed: string[] | null) => + expandTargetsByQuotaAwareConnections(t, name, l, allowed)); + const { expandedTargets } = await expander( + targets, + comboName, + log, + args.apiKeyAllowedConnectionIds ?? null + ); + + // Cap per ORIGINAL target: group entries by stepId and truncate each + // group to maxPerTarget. (The expander preserves original order within a + // target, so truncation keeps the first N connections in priority order.) + const counts = new Map(); + const capped: ResolvedComboTarget[] = []; + for (const target of expandedTargets) { + const key = target.stepId ?? target.executionKey; + const seen = counts.get(key) ?? 0; + if (seen >= maxPerTarget) continue; + counts.set(key, seen + 1); + capped.push(target); + } + return capped; + } catch (error) { + // Fail-open (spec section 3.1): expansion is a best-effort pre-filter, never a + // hard dependency. Auth-layer gates remain the backstop. + log.warn?.("COMBO", "Connection-aware expansion failed; passing targets through", { + comboName, + strategy, + err: error, + }); + return targets; + } +} diff --git a/open-sse/services/combo/dispatchPrelude.ts b/open-sse/services/combo/dispatchPrelude.ts index 604d5d228a..b39f592b51 100644 --- a/open-sse/services/combo/dispatchPrelude.ts +++ b/open-sse/services/combo/dispatchPrelude.ts @@ -30,6 +30,7 @@ import { } from "./comboStructure.ts"; import { isComboModelVisible } from "./comboVisibility.ts"; import { buildFusionHandleSingleModel, extractFusionPanelSpec } from "./fusionPanel.ts"; +import { expandTargetsForAllStrategies } from "./connectionAwareExpansion.ts"; import { expandComboSystemPromptIfPresent, resolveTargetFingerprint, @@ -438,12 +439,25 @@ export async function tryFusionDispatch(args: { } if (strategy !== "fusion") return null; - const allResolvedFusionTargets = resolveComboTargets( + let allResolvedFusionTargets = resolveComboTargets( combo, args.allCombos, clampComboDepth(config.maxComboDepth), args.hiddenModelsByProvider ); + // Connection-aware expansion is opt-in. The fusion panel itself is + // keyed by model string below (`resolvedByModelStr` keeps ONE target per + // modelStr -- the first healthy connection), so the panel size is unchanged; + // only each member's connectionId becomes a vetted, non-exhausted account. + allResolvedFusionTargets = await expandTargetsForAllStrategies({ + strategy, + targets: allResolvedFusionTargets, + comboName: combo.name, + config: combo.config, + settings: args.settings as Record | null | undefined, + log, + apiKeyAllowedConnectionIds: args.apiKeyAllowedConnections ?? null, + }); // #3378 (ported from upstream decolua/9router): every non-fusion combo // strategy runs candidates through filterTargetsByRequestCompatibility before // dispatch, which excludes a target whose vision support cannot be *confirmed* @@ -490,8 +504,19 @@ export async function tryFusionDispatch(args: { for (const target of resolvedFusionTargets) { if (!resolvedByModelStr.has(target.modelStr)) resolvedByModelStr.set(target.modelStr, target); } + // Deduplicate targets by stepId / modelStr so panel size does not inflate + // when models expand across multiple connections. + const seenPanelKeys = new Set(); + const distinctTargets: typeof resolvedFusionTargets = []; + for (const target of resolvedFusionTargets) { + const key = target.stepId ?? target.modelStr; + if (!seenPanelKeys.has(key)) { + seenPanelKeys.add(key); + distinctTargets.push(target); + } + } const { panel: fusionPanel, comboRefUnits } = extractFusionPanelSpec( - resolvedFusionTargets.map((target) => target.modelStr), + distinctTargets.map((target) => target.modelStr), combo.name, null ); @@ -539,6 +564,8 @@ export async function tryPipelineDispatch(args: { combo: ComboLike; config: ComboSetupConfig; strategy: string; + settings?: Record | null; + apiKeyAllowedConnections?: string[] | null; allCombos?: ComboCollectionLike; handleSingleModelWithTimeout: HandleSingleModel; log: ComboLogger; @@ -549,18 +576,44 @@ export async function tryPipelineDispatch(args: { combo, config, strategy, + settings, + apiKeyAllowedConnections, allCombos, handleSingleModelWithTimeout, log, hiddenModelsByProvider, } = args; if (strategy !== "pipeline") return null; - const pipelineSteps: PipelineStep[] = resolveComboTargets( + const resolvedTargets = resolveComboTargets( combo, allCombos, clampComboDepth(config.maxComboDepth), hiddenModelsByProvider - ).map((target) => ({ target, prompt: target.prompt })); + ); + const expanded = await expandTargetsForAllStrategies({ + strategy, + targets: resolvedTargets, + comboName: combo.name, + config: combo.config, + settings: settings as Record | null | undefined, + log, + apiKeyAllowedConnectionIds: apiKeyAllowedConnections ?? null, + }); + // Pipeline: each stage is one step. If a step expanded to multiple connections, + // keep the first healthy connection for that stage. + const seenSteps = new Set(); + const pipelineTargets: typeof expanded = []; + for (const target of expanded) { + const key = target.stepId ?? target.modelStr; + if (!seenSteps.has(key)) { + seenSteps.add(key); + pipelineTargets.push(target); + } + } + const pipelineSteps: PipelineStep[] = pipelineTargets.map((target) => ({ + target, + prompt: target.prompt, + })); return handlePipelineChat({ body, steps: pipelineSteps, diff --git a/open-sse/services/combo/quotaStrategies.ts b/open-sse/services/combo/quotaStrategies.ts index ad47e5d2df..c4cb4c5b52 100644 --- a/open-sse/services/combo/quotaStrategies.ts +++ b/open-sse/services/combo/quotaStrategies.ts @@ -163,7 +163,13 @@ function getTargetConnectionIds( return connectionIds; } -async function expandTargetsByQuotaAwareConnections( +/** + * Exported for the connection-aware expansion pipeline stage + * (connectionAwareExpansion.ts) so all 20 combo strategies can share the + * A-group per-connection expander without duplicating its logic. The + * function body is unchanged; only the visibility is widened. + */ +export async function expandTargetsByQuotaAwareConnections( targets: ResolvedComboTarget[], comboName: string, log: { warn?: (...args: unknown[]) => void }, diff --git a/open-sse/services/combo/targetResolution.ts b/open-sse/services/combo/targetResolution.ts index 5022f2caa7..33ef333ba7 100644 --- a/open-sse/services/combo/targetResolution.ts +++ b/open-sse/services/combo/targetResolution.ts @@ -41,6 +41,7 @@ import { errorResponseWithComboDiagnostics } from "../../utils/error.ts"; import { getCircuitBreaker } from "../../../src/shared/utils/circuitBreaker"; import type { ResilienceSettings } from "../../../src/lib/resilience/settings"; import { applyStrategyOrdering } from "./applyStrategyOrdering.ts"; +import { expandTargetsForAllStrategies } from "./connectionAwareExpansion.ts"; import { clampComboDepth } from "./comboPredicates.ts"; import { describeCapabilityFilterExhaustion, @@ -695,7 +696,7 @@ async function applyPromptCacheStage( export async function resolveComboTargetPipeline( deps: ResolveComboTargetPipelineDeps ): Promise { - const { body, combo, strategy, config, allCombos, log, isModelAvailable } = deps; + const { body, combo, strategy, config, allCombos, log, isModelAvailable, settings } = deps; const { expandedCombo, expandedAllCombos } = await expandComboWildcards(combo, allCombos); const stickyWeightedLimit = clampStickyWeightedTargetLimit( @@ -720,6 +721,21 @@ export async function resolveComboTargetPipeline( orderedTargets = await applyRequestTagRouting(orderedTargets, body, log); + // Connection-aware expansion for group-B strategies is opt-in. Runs + // BEFORE orderByStrategy so every downstream consumer (strategy ordering, + // continuity/stickiness, prompt-cache stage) sees per-connection targets. + // Stickiness is applied later inside applyContinuityFilters, so its pin key + // naturally matches the expanded connectionId targets. + orderedTargets = await expandTargetsForAllStrategies({ + strategy, + targets: orderedTargets, + comboName: combo.name, + config: combo.config, + settings: settings as Record | null | undefined, + log, + apiKeyAllowedConnectionIds: deps.apiKeyAllowedConnections, + }); + logTargetPoolSize(strategy, allCombos, orderedTargets, stickyWeightedKey, log); const pipelineResponse = await dispatchSmartPipeline( diff --git a/open-sse/services/combo/types.ts b/open-sse/services/combo/types.ts index 83a7f26693..15800b25b5 100644 --- a/open-sse/services/combo/types.ts +++ b/open-sse/services/combo/types.ts @@ -143,7 +143,7 @@ export type HandleComboChatOptions = { requestHeaders?: Headers | Record | null; }; -export type HandleRoundRobinOptions = Omit; +export type HandleRoundRobinOptions = HandleComboChatOptions; export type HistoricalLatencyStatsEntry = { totalRequests?: number; diff --git a/open-sse/services/comboConfig.ts b/open-sse/services/comboConfig.ts index 9fa87f1568..ee67218894 100644 --- a/open-sse/services/comboConfig.ts +++ b/open-sse/services/comboConfig.ts @@ -194,6 +194,9 @@ const DEFAULT_COMBO_CONFIG = { latencyWeight: 0.15, cacheTtlMs: 60000, }, + // Connection-aware expansion for group-B combo strategies is opt-in. + connectionAwareExpansion: false, + connectionAwareExpansionMaxPerTarget: 8, // Context window requirements for combo target filtering/sorting (undefined by // default — declared here so resolveComboSetupConfig's inferred return type // includes the key; combo.ts reads config.contextRequirements). diff --git a/src/lib/db/settings.ts b/src/lib/db/settings.ts index b9393db3c8..15bfdf6981 100644 --- a/src/lib/db/settings.ts +++ b/src/lib/db/settings.ts @@ -147,6 +147,8 @@ export async function getSettings() { tailscaleUrl: "", stickyRoundRobinLimit: 3, disableSessionStickiness: false, + // Global connection-aware expansion fallback for group-B combo strategies is opt-in. + connectionAwareExpansion: false, promptCacheAffinityEnabled: true, comboStrategy: "fallback", comboStickyRoundRobinLimit: null, // null = inherit stickyRoundRobinLimit (a literal default here shadows the documented batched-rotation default of 3 — #6678 regression caught by the v3.8.47 release CI) diff --git a/src/shared/validation/schemas/combo.ts b/src/shared/validation/schemas/combo.ts index 988a4ad2db..b6587705e1 100644 --- a/src/shared/validation/schemas/combo.ts +++ b/src/shared/validation/schemas/combo.ts @@ -224,6 +224,9 @@ export const comboRuntimeConfigSchema = z resetWindowTieBandMs: z.coerce.number().int().min(0).max(86_400_000).optional(), resetWindowQuotaCacheTtlMs: z.coerce.number().int().min(0).max(300_000).optional(), resetWindowQuotaCacheMaxStaleMs: z.coerce.number().int().min(0).max(3_600_000).optional(), + // Connection-aware expansion for group-B combo strategies is opt-in. + connectionAwareExpansion: z.boolean().optional(), + connectionAwareExpansionMaxPerTarget: z.coerce.number().int().min(1).max(64).optional(), shadowRouting: shadowRoutingSchema.optional(), evalRouting: evalRoutingSchema.optional(), // Fusion strategy (open-sse/services/fusion.ts): the panel is the combo's diff --git a/src/shared/validation/settingsSchemas.ts b/src/shared/validation/settingsSchemas.ts index dd5a65e15a..c50dcc76b8 100644 --- a/src/shared/validation/settingsSchemas.ts +++ b/src/shared/validation/settingsSchemas.ts @@ -324,6 +324,8 @@ export const updateSettingsSchema = z.object({ }), // #6168: global session-stickiness opt-out (per-combo config overrides this). disableSessionStickiness: z.boolean().optional(), + // Global connection-aware expansion fallback for group-B combo strategies is opt-in. + connectionAwareExpansion: z.boolean().optional(), /** Keep eligible combo targets close to the provider-side prompt cache. */ promptCacheAffinityEnabled: z.boolean().optional(), /** diff --git a/tests/unit/combo/connection-aware-expansion.test.ts b/tests/unit/combo/connection-aware-expansion.test.ts new file mode 100644 index 0000000000..6dbd2a3f4a --- /dev/null +++ b/tests/unit/combo/connection-aware-expansion.test.ts @@ -0,0 +1,675 @@ +/** + * Connection-aware expansion for all combo strategies. + * + * Group-B strategies (priority / weighted / round-robin / random / p2c / + * least-used / cost-optimized / lkgp / fill-first / strict-random / + * context-optimized / cache-optimized / context-relay / fusion / pipeline) + * historically resolved targets WITHOUT a per-connection view: exhausted + * accounts kept getting picked. This suite pins the new opt-in pipeline + * stage `expandTargetsForAllStrategies` -- reusing the A-group expander + * (`expandTargetsByQuotaAwareConnections`) -- behind the + * `connectionAwareExpansion` config key (default false). + * + * T1-T10 mirror spec section 6 verbatim. + */ +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"; +import { randomUUID } from "node:crypto"; + +const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-conn-aware-")); +process.env.DATA_DIR = TEST_DATA_DIR; + +const coreDb = await import("../../../src/lib/db/core.ts"); +const providersDb = await import("../../../src/lib/db/providers.ts"); +const quotaCache = await import("../../../src/domain/quotaCache.ts"); +const { registerQuotaFetcher } = await import("../../../open-sse/services/quotaPreflight.ts"); +const { + expandTargetsForAllStrategies, + shouldApplyConnectionAwareExpansion, + CONNECTION_AWARE_EXPANSION_GROUP_B, +} = await import("../../../open-sse/services/combo/connectionAwareExpansion.ts"); +const { handleComboChat } = await import("../../../open-sse/services/combo.ts"); +const { lockExactModel } = await import("../../../open-sse/services/accountFallback.ts"); +const { tryPipelineDispatch } = await import("../../../open-sse/services/combo/dispatchPrelude.ts"); +const { resolveComboSetupConfig } = await import("../../../open-sse/services/comboConfig.ts"); + +const noopLog = { warn: () => {}, info: () => {}, debug: () => {}, error: () => {} }; + +function makeTarget(overrides: Record = {}) { + return { + kind: "model" as const, + stepId: "step-1", + executionKey: "step-1", + modelStr: "test-provider/model-x", + provider: "test-provider", + providerId: null, + connectionId: null, + weight: 1, + label: null, + ...overrides, + }; +} + +test.after(() => { + coreDb.resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); +}); + +// Gate: strategy + config resolution + +test("T0a: group B strategies are the 15 non-quota-aware strategies", () => { + const groupA = new Set(["reset-aware", "reset-window", "headroom", "quota-share", "auto"]); + for (const strategy of CONNECTION_AWARE_EXPANSION_GROUP_B) { + assert.ok(!groupA.has(strategy), `group B must not contain A-group strategy ${strategy}`); + } + assert.equal(CONNECTION_AWARE_EXPANSION_GROUP_B.length, 15); +}); + +test("T0b: gate is closed by default and for A-group strategies", () => { + assert.equal( + shouldApplyConnectionAwareExpansion("priority", {}), + false, + "default off for group B" + ); + assert.equal( + shouldApplyConnectionAwareExpansion("priority", { connectionAwareExpansion: true }), + true, + "on for group B when enabled" + ); + assert.equal( + shouldApplyConnectionAwareExpansion("priority", null, { connectionAwareExpansion: true }), + true, + "fallback to settings when combo config is unset" + ); + assert.equal( + shouldApplyConnectionAwareExpansion( + "priority", + { connectionAwareExpansion: false }, + { connectionAwareExpansion: true } + ), + false, + "combo config overrides settings" + ); + assert.equal( + shouldApplyConnectionAwareExpansion("reset-aware", { + connectionAwareExpansion: true, + }), + false, + "A-group strategies never double-expand" + ); +}); + +// T1: expansion filters exhausted, tags healthy @connectionId + +test("T1: B-group target expands to healthy connections only (handleSingleModel receives healthy connectionId)", async () => { + const provider = "cae-t1-" + randomUUID(); + registerQuotaFetcher(provider, async () => ({ used: 0, total: 100, percentUsed: 0 })); + const healthy = await providersDb.createProviderConnection({ + provider, + authType: "apikey", + name: "healthy", + apiKey: "key-" + randomUUID(), + isActive: true, + }); + const limited = await providersDb.createProviderConnection({ + provider, + authType: "apikey", + name: "limited", + apiKey: "key-" + randomUUID(), + isActive: true, + rateLimitedUntil: new Date(Date.now() + 60_000).toISOString(), + }); + + const target = makeTarget({ + modelStr: `${provider}/model-x`, + provider, + }); + const out = await expandTargetsForAllStrategies({ + strategy: "priority", + targets: [target], + comboName: "cae-t1", + config: { connectionAwareExpansion: true }, + log: noopLog, + }); + + assert.equal(out.length, 1, "one expanded target (healthy connection only)"); + assert.equal(out[0].connectionId, healthy.id); + assert.ok(out[0].executionKey.includes("@" + healthy.id)); + assert.ok( + !out.some((t) => t.connectionId === limited.id), + "rate-limited connection must be filtered out" + ); + + // End-to-end through handleComboChat + let receivedConnectionId: string | null = null; + const res = await handleComboChat({ + body: { messages: [{ role: "user", content: "hi" }] }, + combo: { + name: "cae-t1-combo-" + randomUUID(), + strategy: "priority", + models: [`${provider}/model-x`], + config: { connectionAwareExpansion: true }, + }, + handleSingleModel: async (_b, _m, t) => { + receivedConnectionId = + (t as { connectionId?: string | null } | undefined)?.connectionId ?? null; + return new Response(JSON.stringify({ choices: [] }), { status: 200 }); + }, + isModelAvailable: async () => true, + log: noopLog, + settings: null, + allCombos: null, + }); + assert.equal(res.status, 200); + assert.equal( + receivedConnectionId, + healthy.id, + "handleSingleModel must receive the healthy connectionId" + ); +}); + +// T2: switch off = byte-identical passthrough + +test("T2: switch off keeps targets byte-identical (no @ suffix)", async () => { + const provider = "cae-t2-" + randomUUID(); + registerQuotaFetcher(provider, async () => ({ used: 0, total: 100, percentUsed: 0 })); + await providersDb.createProviderConnection({ + provider, + authType: "apikey", + name: "only", + apiKey: "key-" + randomUUID(), + isActive: true, + }); + + const target = makeTarget({ modelStr: `${provider}/model-x`, provider }); + const out = await expandTargetsForAllStrategies({ + strategy: "priority", + targets: [target], + comboName: "cae-t2", + config: { connectionAwareExpansion: false }, + log: noopLog, + }); + + assert.equal(out.length, 1); + assert.equal(out[0].connectionId, null, "no connection pinning when off"); + assert.equal(out[0].executionKey, "step-1", "no @connectionId suffix when off"); +}); + +test("T2b: global setting enables expansion when combo config is unset", async () => { + const provider = "cae-t2b-" + randomUUID(); + registerQuotaFetcher(provider, async () => ({ used: 0, total: 100, percentUsed: 0 })); + const healthy = await providersDb.createProviderConnection({ + provider, + authType: "apikey", + name: "healthy", + isActive: true, + }); + + let receivedConnectionId: string | null = null; + const res = await handleComboChat({ + body: { messages: [{ role: "user", content: "hi" }] }, + combo: { + name: "cae-t2b-combo-" + randomUUID(), + strategy: "priority", + models: [`${provider}/model-x`], + }, + handleSingleModel: async (_b, _m, target) => { + receivedConnectionId = + (target as { connectionId?: string | null } | undefined)?.connectionId ?? null; + return new Response(JSON.stringify({ choices: [] }), { status: 200 }); + }, + isModelAvailable: async () => true, + log: noopLog, + settings: { connectionAwareExpansion: true }, + allCombos: null, + }); + + assert.equal(res.status, 200); + assert.equal(receivedConnectionId, healthy.id); +}); + +// T3: provider without a quota fetcher is not expanded + +test("T3: no-fetcher provider passes through untouched", async () => { + const provider = "cae-t3-nofetcher-" + randomUUID(); + // NOTE: no registerQuotaFetcher call for this provider. + await providersDb.createProviderConnection({ + provider, + authType: "apikey", + name: "only", + apiKey: "key-" + randomUUID(), + isActive: true, + }); + + const target = makeTarget({ modelStr: `${provider}/model-x`, provider }); + const out = await expandTargetsForAllStrategies({ + strategy: "random", + targets: [target], + comboName: "cae-t3", + config: { connectionAwareExpansion: true }, + log: noopLog, + }); + + assert.equal(out.length, 1); + assert.equal(out[0].connectionId, null); + assert.equal(out[0].executionKey, "step-1"); +}); + +// T4: pinned connection that is exhausted drops the target + +test("T4: pinned exhausted connection drops the whole target", async () => { + const provider = "cae-t4-" + randomUUID(); + registerQuotaFetcher(provider, async () => ({ used: 0, total: 100, percentUsed: 0 })); + const pinned = await providersDb.createProviderConnection({ + provider, + authType: "apikey", + name: "pinned", + apiKey: "key-" + randomUUID(), + isActive: true, + rateLimitedUntil: new Date(Date.now() + 120_000).toISOString(), + }); + + const target = makeTarget({ + modelStr: `${provider}/model-x`, + provider, + connectionId: pinned.id, + }); + const out = await expandTargetsForAllStrategies({ + strategy: "priority", + targets: [target], + comboName: "cae-t4", + config: { connectionAwareExpansion: true }, + log: noopLog, + }); + + assert.equal(out.length, 0, "pinned exhausted target must be dropped entirely"); +}); + +// T5: allowedConnectionIds intersects with the active pool + +test("T5: step allowlist intersects with active connections", async () => { + const provider = "cae-t5-" + randomUUID(); + registerQuotaFetcher(provider, async () => ({ used: 0, total: 100, percentUsed: 0 })); + const connA = await providersDb.createProviderConnection({ + provider, + authType: "apikey", + name: "a", + apiKey: "key-" + randomUUID(), + isActive: true, + }); + const connB = await providersDb.createProviderConnection({ + provider, + authType: "apikey", + name: "b", + apiKey: "key-" + randomUUID(), + isActive: true, + }); + + const target = makeTarget({ + modelStr: `${provider}/model-x`, + provider, + allowedConnectionIds: [connA.id], + }); + const out = await expandTargetsForAllStrategies({ + strategy: "round-robin", + targets: [target], + comboName: "cae-t5", + config: { connectionAwareExpansion: true }, + log: noopLog, + }); + + assert.equal(out.length, 1); + assert.equal(out[0].connectionId, connA.id); + assert.ok(!out.some((t) => t.connectionId === connB.id)); +}); + +// T6: RR rotation granularity becomes model x connection + +test("T6: RR rotation over expanded targets spans connections across consecutive requests", async () => { + const provider = "cae-t6-" + randomUUID(); + registerQuotaFetcher(provider, async () => ({ used: 0, total: 100, percentUsed: 0 })); + const ids: string[] = []; + for (let i = 0; i < 3; i++) { + const conn = await providersDb.createProviderConnection({ + provider, + authType: "apikey", + name: "conn" + i, + apiKey: "key-" + randomUUID(), + isActive: true, + }); + ids.push(conn.id); + } + + const comboName = "cae-t6-combo-" + randomUUID(); + const combo = { + name: comboName, + strategy: "round-robin", + models: [`${provider}/model-x`], + config: { connectionAwareExpansion: true }, + }; + + const dispatchedConnections: string[] = []; + for (let i = 0; i < 3; i++) { + const res = await handleComboChat({ + body: { messages: [{ role: "user", content: `req-${i}` }] }, + combo, + handleSingleModel: async (_b, _m, t) => { + dispatchedConnections.push( + (t as { connectionId?: string | null } | undefined)?.connectionId || "" + ); + return new Response(JSON.stringify({ choices: [] }), { status: 200 }); + }, + isModelAvailable: async () => true, + log: noopLog, + settings: null, + allCombos: null, + }); + assert.equal(res.status, 200); + } + + assert.equal(dispatchedConnections.length, 3); + const distinct = new Set(dispatchedConnections); + assert.equal(distinct.size, 3, "consecutive RR requests must hit 3 distinct connections"); + for (const id of ids) { + assert.ok(distinct.has(id), `connection ${id} must be hit`); + } +}); + +// T7: main loop model lockout gate skips locked connection + +test("T7: main loop model lockout gate (combo.ts:1286) skips locked target", async () => { + const provider = "cae-t7-" + randomUUID(); + registerQuotaFetcher(provider, async () => ({ used: 0, total: 100, percentUsed: 0 })); + const connLocked = await providersDb.createProviderConnection({ + provider, + authType: "apikey", + name: "locked", + apiKey: "key-" + randomUUID(), + isActive: true, + }); + const connHealthy = await providersDb.createProviderConnection({ + provider, + authType: "apikey", + name: "healthy", + apiKey: "key-" + randomUUID(), + isActive: true, + }); + + // Lock model-x on connLocked + lockExactModel(provider, connLocked.id, "model-x", "cooldown", 60_000); + + const logs: string[] = []; + const captureLog = { + warn: () => {}, + info: (_cat: string, msg: string) => { + logs.push(msg); + }, + debug: () => {}, + error: () => {}, + }; + + const dispatchedConnections: string[] = []; + const res = await handleComboChat({ + body: { messages: [{ role: "user", content: "test" }] }, + combo: { + name: "cae-t7-combo-" + randomUUID(), + strategy: "priority", + models: [`${provider}/model-x`], + config: { connectionAwareExpansion: true }, + }, + handleSingleModel: async (_b, _m, t) => { + dispatchedConnections.push( + (t as { connectionId?: string | null } | undefined)?.connectionId || "" + ); + return new Response(JSON.stringify({ choices: [] }), { status: 200 }); + }, + isModelAvailable: async () => true, + log: captureLog, + settings: null, + allCombos: null, + }); + + assert.equal(res.status, 200); + assert.deepEqual( + dispatchedConnections, + [connHealthy.id], + "only the unlocked healthy connection is dispatched" + ); + assert.ok( + logs.some((msg) => msg.includes("model locked by resilience")), + "must log that model is locked" + ); +}); + +// T8: fusion panel sizing is not inflated by expansion + +test("T8: fusion panel does not inflate; members carry vetted connectionId", async () => { + const provider1 = "cae-t8-p1-" + randomUUID(); + const provider2 = "cae-t8-p2-" + randomUUID(); + registerQuotaFetcher(provider1, async () => ({ used: 0, total: 100, percentUsed: 0 })); + registerQuotaFetcher(provider2, async () => ({ used: 0, total: 100, percentUsed: 0 })); + + const p1Conns = []; + for (let i = 0; i < 3; i++) { + p1Conns.push( + await providersDb.createProviderConnection({ + provider: provider1, + authType: "apikey", + name: "p1-c" + i, + apiKey: "key-" + randomUUID(), + isActive: true, + }) + ); + } + const p2Conns = []; + for (let i = 0; i < 3; i++) { + p2Conns.push( + await providersDb.createProviderConnection({ + provider: provider2, + authType: "apikey", + name: "p2-c" + i, + apiKey: "key-" + randomUUID(), + isActive: true, + }) + ); + } + + const combo = { + name: "cae-t8-fusion-" + randomUUID(), + strategy: "fusion", + models: [`${provider1}/model-a`, `${provider2}/model-b`], + config: { connectionAwareExpansion: true, judgeModel: "judge-p/judge-m" }, + }; + + const dispatchedPanelTargets: Array<{ modelStr: string; connectionId?: string | null }> = []; + const res = await handleComboChat({ + body: { messages: [{ role: "user", content: "fuse" }] }, + combo, + handleSingleModel: async (_b, modelStr, t) => { + dispatchedPanelTargets.push({ + modelStr, + connectionId: (t as { connectionId?: string | null } | undefined)?.connectionId, + }); + return new Response( + JSON.stringify({ + choices: [{ message: { content: `response-from-${modelStr}` } }], + }), + { status: 200 } + ); + }, + isModelAvailable: async () => true, + log: noopLog, + settings: null, + allCombos: null, + }); + + assert.equal(res.status, 200); + // Panel size should equal 2 (one per model), plus 1 judge synthesis call (or panel members) + // Each panel member dispatched should have a vetted connectionId + const p1Calls = dispatchedPanelTargets.filter((c) => c.modelStr === `${provider1}/model-a`); + const p2Calls = dispatchedPanelTargets.filter((c) => c.modelStr === `${provider2}/model-b`); + assert.equal(p1Calls.length, 1, "model-a should be called once in panel"); + assert.equal(p2Calls.length, 1, "model-b should be called once in panel"); + assert.equal( + p1Calls[0].connectionId, + p1Conns[0].id, + "model-a should carry first healthy connection" + ); + assert.equal( + p2Calls[0].connectionId, + p2Conns[0].id, + "model-b should carry first healthy connection" + ); +}); + +// T9: quotaCache snapshot drives filtering + +test("T9: isQuotaExhaustedForRequest=true connection is filtered", async () => { + const provider = "cae-t9-" + randomUUID(); + registerQuotaFetcher(provider, async () => ({ used: 0, total: 100, percentUsed: 0 })); + const exhausted = await providersDb.createProviderConnection({ + provider, + authType: "apikey", + name: "exhausted", + apiKey: "key-" + randomUUID(), + isActive: true, + }); + const healthy = await providersDb.createProviderConnection({ + provider, + authType: "apikey", + name: "healthy", + apiKey: "key-" + randomUUID(), + isActive: true, + }); + // Standard provider: aggregate exhausted=true with a far-future reset. + quotaCache.setQuotaCache(exhausted.id, provider, { + default: { remainingPercentage: 0, resetAt: new Date(Date.now() + 86_400_000).toISOString() }, + }); + + const target = makeTarget({ modelStr: `${provider}/model-x`, provider }); + const out = await expandTargetsForAllStrategies({ + strategy: "priority", + targets: [target], + comboName: "cae-t9", + config: { connectionAwareExpansion: true }, + log: noopLog, + }); + + assert.equal(out.length, 1); + assert.equal(out[0].connectionId, healthy.id); +}); + +// T10: fail-open on connection-list load error + +test("T10: load failure returns the original targets (fail-open)", async () => { + const provider = "cae-t10-" + randomUUID(); + registerQuotaFetcher(provider, async () => ({ used: 0, total: 100, percentUsed: 0 })); + const target = makeTarget({ modelStr: `${provider}/model-x`, provider }); + const out = await expandTargetsForAllStrategies({ + strategy: "priority", + targets: [target], + comboName: "cae-t10", + config: { connectionAwareExpansion: true }, + log: noopLog, + __testExpander: async () => { + throw new Error("synthetic load failure"); + }, + }); + assert.ok(out[0] === target, "fail-open must return the ORIGINAL target objects"); + + assert.equal(out.length, 1); + assert.equal(out[0].connectionId, null); + assert.equal(out[0].executionKey, "step-1"); +}); + +// API-key allowedConnections intersects + +test("apiKey allowedConnections intersects the expanded pool", async () => { + const provider = "cae-apikey-" + randomUUID(); + registerQuotaFetcher(provider, async () => ({ used: 0, total: 100, percentUsed: 0 })); + const connA = await providersDb.createProviderConnection({ + provider, + authType: "apikey", + name: "a", + apiKey: "key-" + randomUUID(), + isActive: true, + }); + const connB = await providersDb.createProviderConnection({ + provider, + authType: "apikey", + name: "b", + apiKey: "key-" + randomUUID(), + isActive: true, + }); + + const target = makeTarget({ modelStr: `${provider}/model-x`, provider }); + const out = await expandTargetsForAllStrategies({ + strategy: "priority", + targets: [target], + comboName: "cae-apikey", + config: { connectionAwareExpansion: true }, + log: noopLog, + apiKeyAllowedConnectionIds: [connA.id], + }); + + assert.equal(out.length, 1); + assert.equal(out[0].connectionId, connA.id); + assert.ok(!out.some((t) => t.connectionId === connB.id)); +}); + +// Pipeline dispatch + +test("Pipeline strategy: each stage picks the first healthy connection", async () => { + const provider1 = "cae-pipe-p1-" + randomUUID(); + const provider2 = "cae-pipe-p2-" + randomUUID(); + registerQuotaFetcher(provider1, async () => ({ used: 0, total: 100, percentUsed: 0 })); + registerQuotaFetcher(provider2, async () => ({ used: 0, total: 100, percentUsed: 0 })); + + const p1Conn = await providersDb.createProviderConnection({ + provider: provider1, + authType: "apikey", + name: "p1-c0", + apiKey: "key-" + randomUUID(), + isActive: true, + }); + const p2Conn = await providersDb.createProviderConnection({ + provider: provider2, + authType: "apikey", + name: "p2-c0", + apiKey: "key-" + randomUUID(), + isActive: true, + }); + + const combo = { + name: "cae-pipe-combo-" + randomUUID(), + strategy: "pipeline", + models: [`${provider1}/model-1`, `${provider2}/model-2`], + config: { connectionAwareExpansion: true }, + }; + + const executedStages: string[] = []; + const res = await tryPipelineDispatch({ + body: { messages: [{ role: "user", content: "pipe" }] }, + combo, + config: resolveComboSetupConfig(combo, null), + strategy: "pipeline", + handleSingleModelWithTimeout: async (_b, modelStr, target) => { + executedStages.push( + `${modelStr}@${(target as { connectionId?: string | null } | undefined)?.connectionId}` + ); + return new Response(JSON.stringify({ choices: [{ message: { content: "stage-out" } }] }), { + status: 200, + }); + }, + log: noopLog, + }); + + assert.ok(res !== null); + assert.equal(res.status, 200); + assert.equal(executedStages.length, 2); + assert.equal(executedStages[0], `${provider1}/model-1@${p1Conn.id}`); + assert.equal(executedStages[1], `${provider2}/model-2@${p2Conn.id}`); +});