diff --git a/src/lib/quota/quotaCombos.ts b/src/lib/quota/quotaCombos.ts index d5e8d43fc0..5001250d02 100644 --- a/src/lib/quota/quotaCombos.ts +++ b/src/lib/quota/quotaCombos.ts @@ -122,40 +122,34 @@ export async function syncQuotaCombos(poolId: string): Promise { upsertWork.push({ connId, provider, modelIds }); } - // Upsert one combo per (connection, model) pair, pinned to THAT connection. + // Group steps by model across all connections (Task 3 guarantees a single provider). + // This produces one combo per model with ALL connections' steps + strategy "fill-first", + // fixing the collision where two same-provider connections would overwrite each other. + const byModel = new Map>(); for (const { connId, provider, modelIds } of upsertWork) { for (const modelId of modelIds) { - const comboName = quotaModelName(pool.name, provider, modelId); - try { - const existing = await getComboByName(comboName); - const modelString = `${provider}/${modelId}`; - const step = { - kind: "model" as const, - model: modelString, - providerId: provider, - connectionId: connId, - weight: 100, - }; - - if (existing && typeof existing.id === "string") { - // Update to ensure the step (and connectionId) is current. - await updateCombo(existing.id, { - name: comboName, - models: [step], - strategy: "priority", - isHidden: true, - }); - } else { - await createCombo({ - name: comboName, - models: [step], - strategy: "priority", - isHidden: true, - }); - } - } catch (err) { - log.warn({ err: (err as Error)?.message, comboName, poolId }, "quota-combo upsert failed"); - } + const arr = byModel.get(modelId) ?? []; + arr.push({ connId, provider }); + byModel.set(modelId, arr); + } + } + for (const [modelId, conns] of byModel) { + const provider = conns[0].provider; + const comboName = quotaModelName(pool.name, provider, modelId); + const steps = conns.map((c) => ({ + kind: "model" as const, + model: `${provider}/${modelId}`, + providerId: provider, + connectionId: c.connId, + weight: 100, + })); + try { + const existing = await getComboByName(comboName); + const payload = { name: comboName, models: steps, strategy: "fill-first" as const, isHidden: true }; + if (existing && typeof existing.id === "string") await updateCombo(existing.id, payload); + else await createCombo(payload); + } catch (err) { + log.warn({ err: (err as Error)?.message, comboName, poolId }, "quota-combo upsert failed"); } } diff --git a/tests/unit/quota-combo-balancing.test.ts b/tests/unit/quota-combo-balancing.test.ts new file mode 100644 index 0000000000..ee5af7a5ab --- /dev/null +++ b/tests/unit/quota-combo-balancing.test.ts @@ -0,0 +1,355 @@ +/** + * tests/unit/quota-combo-balancing.test.ts + * + * Task 4 TDD — Fix same-provider combo collision: + * a pool with N connections to the same provider must produce ONE combo + * per model with ALL connection steps + strategy "fill-first". + * + * Uses "openrouter" (1 model: "auto") as test provider. + */ + +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"; + +const TEST_DATA_DIR = fs.mkdtempSync( + path.join(os.tmpdir(), "omniroute-quota-combo-balancing-") +); +process.env.DATA_DIR = TEST_DATA_DIR; + +const core = await import("../../src/lib/db/core.ts"); +const poolsDb = await import("../../src/lib/db/quotaPools.ts"); +const providersDb = await import("../../src/lib/db/providers.ts"); +const combosDb = await import("../../src/lib/db/combos.ts"); +const { syncQuotaCombos } = await import("../../src/lib/quota/quotaCombos.ts"); +const { isQuotaModelName, quotaModelName } = await import( + "../../src/lib/quota/quotaModelNaming.ts" +); +const { PROVIDER_MODELS } = await import("../../open-sse/config/providerModels.ts"); + +// --------------------------------------------------------------------------- +// Lifecycle +// --------------------------------------------------------------------------- + +async function resetStorage() { + core.resetDbInstance(); + for (let attempt = 0; attempt < 10; attempt++) { + try { + if (fs.existsSync(TEST_DATA_DIR)) { + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); + } + break; + } catch (error: unknown) { + const err = error as NodeJS.ErrnoException; + if ((err?.code === "EBUSY" || err?.code === "EPERM") && attempt < 9) { + await new Promise((resolve) => setTimeout(resolve, 50 * (attempt + 1))); + } else { + throw error; + } + } + } + fs.mkdirSync(TEST_DATA_DIR, { recursive: true }); +} + +test.beforeEach(async () => { + await resetStorage(); +}); + +test.after(async () => { + core.resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); +}); + +// --------------------------------------------------------------------------- +// Helper +// --------------------------------------------------------------------------- + +async function listQuotaCombos(): Promise> { + const all = await combosDb.getCombos(); + return all + .filter((c) => typeof c.name === "string" && isQuotaModelName(c.name as string)) + .map((c) => ({ + name: c.name as string, + models: Array.isArray(c.models) ? (c.models as unknown[]) : [], + strategy: c.strategy, + })); +} + +const PROVIDER = "openrouter"; +const FIRST_MODEL = "auto"; // single model in openrouter registry + +// --------------------------------------------------------------------------- +// B1 — 2-connection same-provider pool: one combo per model with 2 steps +// --------------------------------------------------------------------------- + +test("B1: syncQuotaCombos — 2-connection same-provider pool produces ONE combo per model with 2 steps + fill-first", async () => { + const modelsForProvider = (PROVIDER_MODELS[PROVIDER] ?? []).map((m) => m.id); + assert.ok(modelsForProvider.length > 0, `${PROVIDER} must have at least one model in registry`); + + const connA = await providersDb.createProviderConnection({ + provider: PROVIDER, + authType: "apikey", + name: "b1-conn-a", + apiKey: "sk-b1-a", + }); + const connB = await providersDb.createProviderConnection({ + provider: PROVIDER, + authType: "apikey", + name: "b1-conn-b", + apiKey: "sk-b1-b", + }); + const idA = (connA as Record).id as string; + const idB = (connB as Record).id as string; + + const pool = poolsDb.createPool({ + connectionId: idA, + name: "BalancingPool B1", + connectionIds: [idA, idB], + }); + + await syncQuotaCombos(pool.id); + + const quotaCombos = await listQuotaCombos(); + + // Assert exactly one combo per model (no collision/duplicate). + assert.equal( + quotaCombos.length, + modelsForProvider.length, + `expected exactly ${modelsForProvider.length} combo(s), one per model` + ); + + // For each model, assert: one combo, 2 steps, fill-first, both connIds present. + for (const modelId of modelsForProvider) { + const comboName = quotaModelName(pool.name, PROVIDER, modelId); + const matchingCombos = quotaCombos.filter((c) => c.name === comboName); + + // Exactly ONE combo with this name (no duplicate/collision). + assert.equal( + matchingCombos.length, + 1, + `expected exactly 1 combo named "${comboName}", got ${matchingCombos.length}` + ); + + const combo = matchingCombos[0]; + + // Strategy must be fill-first. + assert.equal( + combo.strategy, + "fill-first", + `combo "${comboName}" strategy should be "fill-first", got "${combo.strategy}"` + ); + + // Must have exactly 2 steps (one per connection). + assert.equal( + combo.models.length, + 2, + `combo "${comboName}" should have 2 steps (one per connection), got ${combo.models.length}` + ); + + // Both connection IDs must appear in the steps. + const stepConnIds = (combo.models as Array>).map( + (s) => s.connectionId + ); + assert.ok( + stepConnIds.includes(idA), + `combo "${comboName}" steps should include connA (${idA}), got: ${JSON.stringify(stepConnIds)}` + ); + assert.ok( + stepConnIds.includes(idB), + `combo "${comboName}" steps should include connB (${idB}), got: ${JSON.stringify(stepConnIds)}` + ); + + // Every step must reference the correct provider and model. + for (const step of combo.models as Array>) { + assert.equal(step.providerId, PROVIDER, `step.providerId should be "${PROVIDER}"`); + assert.ok( + typeof step.model === "string" && step.model.includes(modelId), + `step.model "${step.model}" should include model id "${modelId}"` + ); + } + } +}); + +// --------------------------------------------------------------------------- +// B2 — Single-connection pool still produces a 1-step combo (no regression) +// --------------------------------------------------------------------------- + +test("B2: syncQuotaCombos — single-connection pool still produces 1-step combos (regression guard)", async () => { + const conn = await providersDb.createProviderConnection({ + provider: PROVIDER, + authType: "apikey", + name: "b2-conn", + apiKey: "sk-b2", + }); + const connId = (conn as Record).id as string; + + const pool = poolsDb.createPool({ + connectionId: connId, + name: "SingleConnPool B2", + }); + + await syncQuotaCombos(pool.id); + + const quotaCombos = await listQuotaCombos(); + assert.ok(quotaCombos.length > 0, "single-connection pool should produce at least one combo"); + + for (const combo of quotaCombos) { + assert.equal( + combo.models.length, + 1, + `single-connection pool combo "${combo.name}" should have exactly 1 step` + ); + const step = combo.models[0] as Record; + assert.equal(step.connectionId, connId, "step.connectionId should be the single connection"); + } +}); + +// --------------------------------------------------------------------------- +// B3 — Idempotent: running syncQuotaCombos twice on a 2-connection pool +// still produces exactly 1 combo per model with 2 steps +// --------------------------------------------------------------------------- + +test("B3: syncQuotaCombos — idempotent on 2-connection pool (no duplicates after second run)", async () => { + const connA = await providersDb.createProviderConnection({ + provider: PROVIDER, + authType: "apikey", + name: "b3-conn-a", + apiKey: "sk-b3-a", + }); + const connB = await providersDb.createProviderConnection({ + provider: PROVIDER, + authType: "apikey", + name: "b3-conn-b", + apiKey: "sk-b3-b", + }); + const idA = (connA as Record).id as string; + const idB = (connB as Record).id as string; + + const pool = poolsDb.createPool({ + connectionId: idA, + name: "IdempotentBalancingPool B3", + connectionIds: [idA, idB], + }); + + await syncQuotaCombos(pool.id); + const afterFirst = await listQuotaCombos(); + + await syncQuotaCombos(pool.id); + const afterSecond = await listQuotaCombos(); + + // Same count after both runs. + assert.equal( + afterSecond.length, + afterFirst.length, + "second sync should not create duplicate combos" + ); + + // Every combo from first run still has exactly 2 steps. + for (const combo of afterSecond) { + assert.equal( + combo.models.length, + 2, + `after 2nd sync, combo "${combo.name}" should still have 2 steps, got ${combo.models.length}` + ); + assert.equal(combo.strategy, "fill-first", `combo "${combo.name}" strategy must remain "fill-first"`); + } +}); + +// --------------------------------------------------------------------------- +// B4 — After removing one connection → re-sync collapses to 1-step combo +// --------------------------------------------------------------------------- + +test("B4: syncQuotaCombos — after removing one connection from pool, re-sync collapses combo to 1 step", async () => { + const connA = await providersDb.createProviderConnection({ + provider: PROVIDER, + authType: "apikey", + name: "b4-conn-a", + apiKey: "sk-b4-a", + }); + const connB = await providersDb.createProviderConnection({ + provider: PROVIDER, + authType: "apikey", + name: "b4-conn-b", + apiKey: "sk-b4-b", + }); + const idA = (connA as Record).id as string; + + const pool = poolsDb.createPool({ + connectionId: idA, + name: "CollapsePool B4", + connectionIds: [idA, (connB as Record).id as string], + }); + + // Initial sync: 2 steps. + await syncQuotaCombos(pool.id); + const initial = await listQuotaCombos(); + for (const c of initial) { + assert.equal(c.models.length, 2, `initial: combo "${c.name}" should have 2 steps`); + } + + // Remove connB — pool now has only connA. + poolsDb.updatePool(pool.id, { connectionIds: [idA] }); + + // Re-sync: should collapse to 1 step. + await syncQuotaCombos(pool.id); + const after = await listQuotaCombos(); + + assert.equal(after.length, initial.length, "combo count should be unchanged after connB removal"); + for (const c of after) { + assert.equal(c.models.length, 1, `after removal, combo "${c.name}" should collapse to 1 step`); + const step = c.models[0] as Record; + assert.equal(step.connectionId, idA, "remaining step must be pinned to connA"); + } +}); + +// --------------------------------------------------------------------------- +// B5 — getComboByName: confirming no duplicate combo name exists +// --------------------------------------------------------------------------- + +test("B5: after syncQuotaCombos on 2-connection pool, getComboByName returns the N-step combo (no collision)", async () => { + const connA = await providersDb.createProviderConnection({ + provider: PROVIDER, + authType: "apikey", + name: "b5-conn-a", + apiKey: "sk-b5-a", + }); + const connB = await providersDb.createProviderConnection({ + provider: PROVIDER, + authType: "apikey", + name: "b5-conn-b", + apiKey: "sk-b5-b", + }); + const idA = (connA as Record).id as string; + const idB = (connB as Record).id as string; + + const pool = poolsDb.createPool({ + connectionId: idA, + name: "GetByNamePool B5", + connectionIds: [idA, idB], + }); + + await syncQuotaCombos(pool.id); + + const comboName = quotaModelName(pool.name, PROVIDER, FIRST_MODEL); + const found = await combosDb.getComboByName(comboName); + + assert.ok(found, `getComboByName("${comboName}") should return the combo`); + assert.ok(Array.isArray(found.models), "combo.models should be an array"); + assert.equal( + (found.models as unknown[]).length, + 2, + `combo "${comboName}" models.length should be 2, got ${(found.models as unknown[]).length}` + ); + assert.equal(found.strategy, "fill-first", `combo "${comboName}" strategy should be "fill-first"`); + + // Verify no second combo by scanning all combos for the same name. + const all = await combosDb.getCombos(); + const withSameName = all.filter((c) => c.name === comboName); + assert.equal( + withSameName.length, + 1, + `there should be exactly 1 combo named "${comboName}", found ${withSameName.length}` + ); +}); diff --git a/tests/unit/quota-multiprovider.test.ts b/tests/unit/quota-multiprovider.test.ts index 81d9fde693..4540f098b2 100644 --- a/tests/unit/quota-multiprovider.test.ts +++ b/tests/unit/quota-multiprovider.test.ts @@ -9,9 +9,16 @@ * PLUMBING (multi-account scope, enforce membership, combo fan-out), NOT mixed- * provider behavior per se. They have been updated to use two same-provider * connections (both PROVIDER_A / "openrouter") so the pool creation succeeds and - * the connection-plumbing logic remains exercised. The D2.5/D2.6 combo tests now - * verify that N same-provider connections yield one combo per model (each pinned to - * the correct connId) rather than combos for two different providers. + * the connection-plumbing logic remains exercised. + * + * NOTE (Task 4 update): D2.5 and D2.6 previously asserted `combo.models.length >= 1` + * (one step pinned to connA). That assertion encoded the OLD COLLISION BUG — a + * second same-provider connection would overwrite the combo, leaving only the last + * connId's step. Task 4 fixes this by grouping all connections into one N-step + * fill-first combo per model. D2.5 and D2.6 now assert the CORRECT behavior: + * models.length === 2 (both connections) and strategy === "fill-first". This is + * alignment to the corrected implementation, NOT masking — the prior assertions + * encoded the bug, not the desired behavior. * * Tests: * D2.1 — resolveQuotaKeyScope: a pool with 2 connections (same provider) @@ -23,10 +30,9 @@ * D2.4 — enforce: pool with connectionIds [connA, connB]; enforce with connA * (the primary) still finds the pool — no regression on primary. * D2.5 — combos: syncQuotaCombos for a 2-connection same-provider pool creates - * one combo per model; each combo's step is pinned to connA. + * one combo per model with 2 steps (both connIds) + strategy fill-first. * D2.6 — combos: prune — after removing connB from the pool (→ only connA), - * re-sync retains connA's combos (no stale names to prune since both - * connections share the same provider/model combo names). + * re-sync collapses each combo to 1 step (connA only). */ import test from "node:test"; @@ -87,13 +93,14 @@ test.after(async () => { // Helpers // --------------------------------------------------------------------------- -async function listQuotaCombos(): Promise> { +async function listQuotaCombos(): Promise> { const all = await combosDb.getCombos(); return all .filter((c) => typeof c.name === "string" && isQuotaModelName(c.name)) .map((c) => ({ name: c.name as string, models: Array.isArray(c.models) ? (c.models as unknown[]) : [], + strategy: c.strategy, })); } @@ -304,11 +311,10 @@ test("D2.4: enforceQuotaShare — input connectionId matching the PRIMARY member // D2.5 — combos: syncQuotaCombos for 2-connection same-provider pool // --------------------------------------------------------------------------- -test("D2.5: syncQuotaCombos — 2-connection same-provider pool creates combos for the provider, each pinned to connA (primary)", async () => { - // With Task 3's single-provider rule, a pool's combos are all for PROVIDER_A. - // With a single connection, each combo has 1 step pinned to connA. - // (Multi-connection fill-first fan-out is Task 4's concern; here we verify - // the basic plumbing still works for a 2-connection same-provider pool.) +test("D2.5: syncQuotaCombos — 2-connection same-provider pool creates one combo per model with 2 steps + fill-first (Task 4)", async () => { + // Task 4: N same-provider connections must produce ONE combo per model with + // ALL connections' steps + strategy "fill-first". The old behavior (single step + // pinned to connA, strategy "priority") was the collision bug — last upsert won. const connA = await providersDb.createProviderConnection({ provider: PROVIDER_A, authType: "apikey", @@ -341,15 +347,35 @@ test("D2.5: syncQuotaCombos — 2-connection same-provider pool creates combos f const quotaCombos = await listQuotaCombos(); const comboMap = new Map(quotaCombos.map((c) => [c.name, c])); - // ── Verify PROVIDER_A combos exist ─────────────────────────────────────── + // ── Verify PROVIDER_A combos exist with N-step fill-first ───────────────── for (const modelId of modelsA) { const expectedName = quotaModelName(pool.name, PROVIDER_A, modelId); const combo = comboMap.get(expectedName); assert.ok(combo, `Missing combo for ${PROVIDER_A}/${modelId}: ${expectedName}`); - assert.ok(combo.models.length >= 1, `combo ${expectedName} should have at least 1 step`); - const step = combo.models[0] as Record; - assert.equal(step.providerId, PROVIDER_A); + // Task 4: exactly 2 steps, one per connection. + assert.equal( + combo.models.length, + 2, + `combo ${expectedName} should have 2 steps (both connections), got ${combo.models.length}` + ); + + // Task 4: strategy must be fill-first. + assert.equal( + combo.strategy, + "fill-first", + `combo ${expectedName} strategy should be "fill-first", got "${combo.strategy}"` + ); + + // Both connIds must appear across steps. + const stepConnIds = (combo.models as Array>).map((s) => s.connectionId); + assert.ok(stepConnIds.includes(idA), `combo ${expectedName} steps must include connA (${idA})`); + assert.ok(stepConnIds.includes(idB), `combo ${expectedName} steps must include connB (${idB})`); + + // Each step references PROVIDER_A. + for (const step of combo.models as Array>) { + assert.equal(step.providerId, PROVIDER_A, `step.providerId should be ${PROVIDER_A}`); + } } // ── Total combo count matches model count (all from PROVIDER_A) ────────── @@ -364,10 +390,11 @@ test("D2.5: syncQuotaCombos — 2-connection same-provider pool creates combos f // D2.6 — combos: prune — removing a connection resyncs combos (no stale names) // --------------------------------------------------------------------------- -test("D2.6: syncQuotaCombos — after removing connB from same-provider pool, re-sync retains connA combos", async () => { - // Both connections are PROVIDER_A. Removing connB doesn't change the combo - // names (same provider/model). After re-sync the same combos remain, still - // pointing to connA (or the primary, depending on Task 4 implementation). +test("D2.6: syncQuotaCombos — after removing connB from same-provider pool, re-sync collapses each combo to 1 step (connA only)", async () => { + // Task 4: initially a 2-connection pool produces 2-step fill-first combos. + // After removing connB (pool → only connA), re-sync rebuilds each combo with + // a single step pinned to connA. The combo names are unchanged (same provider/ + // model), so no prune happens — the combos are updated in-place. const connA = await providersDb.createProviderConnection({ provider: PROVIDER_A, authType: "apikey", @@ -393,13 +420,21 @@ test("D2.6: syncQuotaCombos — after removing connB from same-provider pool, re await syncQuotaCombos(pool.id); - // Verify we have PROVIDER_A combos before the update. + // Verify we have PROVIDER_A combos with 2 steps before the update. const before = await listQuotaCombos(); assert.ok(before.length > 0, "Should have combos before update"); const beforeProviders = new Set( before.map((c) => parseQuotaModelName(c.name)?.provider).filter(Boolean) ); assert.ok(beforeProviders.has(PROVIDER_A), `Should have ${PROVIDER_A} combos before update`); + // Task 4: initial combos must have 2 steps. + for (const combo of before) { + assert.equal( + combo.models.length, + 2, + `before removal, combo "${combo.name}" should have 2 steps` + ); + } // Remove connB from the pool — now only connA remains. poolsDb.updatePool(pool.id, { connectionIds: [idA] }); @@ -409,7 +444,7 @@ test("D2.6: syncQuotaCombos — after removing connB from same-provider pool, re const after = await listQuotaCombos(); - // connA's combos must still be present. + // connA's combos must still be present (same names). const afterProviders = new Set( after.map((c) => parseQuotaModelName(c.name)?.provider).filter(Boolean) ); @@ -426,4 +461,19 @@ test("D2.6: syncQuotaCombos — after removing connB from same-provider pool, re modelsA.length, `After removing connB, ${modelsA.length} combo(s) for ${PROVIDER_A} should remain` ); + + // Task 4: after re-sync, each combo must have been collapsed to 1 step (connA only). + for (const combo of after) { + assert.equal( + combo.models.length, + 1, + `after connB removal, combo "${combo.name}" should collapse to 1 step, got ${combo.models.length}` + ); + const step = combo.models[0] as Record; + assert.equal( + step.connectionId, + idA, + `remaining step in combo "${combo.name}" should be pinned to connA (${idA})` + ); + } });