diff --git a/src/lib/db/quotaPools.ts b/src/lib/db/quotaPools.ts index 1f76bd6205..c4f70ec52b 100644 --- a/src/lib/db/quotaPools.ts +++ b/src/lib/db/quotaPools.ts @@ -9,6 +9,28 @@ */ import { getDbInstance } from "./core"; +// Phase B2: auto-mint/prune quotaShared-* combos when pool allocations change. +// Imported lazily (dynamic import in the hook) to avoid circular-dependency +// risk between db/ and quota/ modules. The import is fire-and-forget; combo +// failures never break pool CRUD. +async function syncQuotaCombosGuarded(poolId: string): Promise { + try { + const { syncQuotaCombos } = await import("@/lib/quota/quotaCombos"); + await syncQuotaCombos(poolId); + } catch (err) { + // Guard: combo-sync failure must never break pool CRUD callers. + console.warn("[quota-pools] syncQuotaCombos failed (non-fatal):", (err as Error)?.message); + } +} + +async function removeQuotaCombosGuarded(poolId: string): Promise { + try { + const { removeQuotaCombosForPool } = await import("@/lib/quota/quotaCombos"); + await removeQuotaCombosForPool(poolId); + } catch (err) { + console.warn("[quota-pools] removeQuotaCombosForPool failed (non-fatal):", (err as Error)?.message); + } +} // --------------------------------------------------------------------------- // Local type shapes (aligned with src/lib/quota/dimensions.ts — merged by F7) @@ -160,10 +182,15 @@ export function createPool(input: PoolCreate): QuotaPool { upsertAllocations(id, input.allocations); } - return rowToPool( + const result = rowToPool( { id, connection_id: input.connectionId, name: input.name, created_at: now }, getAllocations(id) ); + + // Phase B2: fire-and-forget combo sync; failures are logged but never thrown. + void syncQuotaCombosGuarded(id); + + return result; } /** @@ -185,7 +212,12 @@ export function updatePool(id: string, input: PoolUpdate): QuotaPool | null { upsertAllocations(id, input.allocations); } - return rowToPool(existing, getAllocations(id)); + const result = rowToPool(existing, getAllocations(id)); + + // Phase B2: fire-and-forget combo sync; failures are logged but never thrown. + void syncQuotaCombosGuarded(id); + + return result; } /** @@ -193,6 +225,10 @@ export function updatePool(id: string, input: PoolUpdate): QuotaPool | null { * Returns true if a row was deleted, false if not found. */ export function deletePool(id: string): boolean { + // Phase B2: remove quota combos BEFORE deleting the pool row so that + // removeQuotaCombosForPool can still resolve the pool name → slug. + void removeQuotaCombosGuarded(id); + const result = getDb().prepare("DELETE FROM quota_pools WHERE id = ?").run(id); return result.changes > 0; } @@ -221,6 +257,9 @@ export function upsertAllocations(poolId: string, allocations: PoolAllocation[]) } }); doUpsert(); + + // Phase B2: fire-and-forget combo sync; failures are logged but never thrown. + void syncQuotaCombosGuarded(poolId); } /** diff --git a/src/lib/quota/quotaCombos.ts b/src/lib/quota/quotaCombos.ts new file mode 100644 index 0000000000..3bd4004a2d --- /dev/null +++ b/src/lib/quota/quotaCombos.ts @@ -0,0 +1,212 @@ +/** + * quota/quotaCombos.ts — Auto-mint / prune `quotaShared-*` virtual combo models + * when a quota pool gains or loses allocations (Phase B2). + * + * Each combo routes to a single {provider, model} target and is pinned to the + * pool's connectionId via ComboModelStep.connectionId (supported by the combo + * target schema). Phase B4 wires resolution — this module only keeps the combo + * rows in sync with the pool's provider model list. + * + * Guard: combo-sync failures never propagate to pool CRUD callers. + */ + +import { getPool } from "@/lib/db/quotaPools"; +import { getProviderConnectionById } from "@/lib/db/providers"; +import { + getCombos, + createCombo, + deleteComboByName, + getComboByName, + updateCombo, +} from "@/lib/db/combos"; +import { PROVIDER_MODELS } from "@omniroute/open-sse/config/providerModels"; +import { quotaModelName, parseQuotaModelName, isQuotaModelName, quotaPoolSlug } from "./quotaModelNaming"; +import { createLogger } from "@/shared/utils/logger"; + +const log = createLogger("quota/quotaCombos"); + +// --------------------------------------------------------------------------- +// Internal helpers +// --------------------------------------------------------------------------- + +/** + * Resolve the provider slug for a pool's connection. + * Returns null when the pool or connection cannot be found, or when the + * provider field is missing/empty. + */ +async function resolvePoolProvider(poolId: string): Promise<{ + pool: { id: string; connectionId: string; name: string }; + provider: string; +} | null> { + const pool = getPool(poolId); + if (!pool) return null; + + let connection: Record | null = null; + try { + connection = (await getProviderConnectionById(pool.connectionId)) as Record< + string, + unknown + > | null; + } catch { + return null; + } + if (!connection) return null; + + const provider = connection.provider; + if (typeof provider !== "string" || provider.length === 0) return null; + + return { pool, provider }; +} + +/** + * Return the list of model IDs for a provider from the static registry. + * Empty array when the provider is unknown or has no registered models. + */ +function getProviderModelIds(provider: string): string[] { + const models = PROVIDER_MODELS[provider]; + if (!Array.isArray(models) || models.length === 0) return []; + return models + .map((m) => (typeof m === "object" && m !== null && typeof (m as { id?: unknown }).id === "string" ? (m as { id: string }).id : null)) + .filter((id): id is string => id !== null && id.length > 0); +} + +// --------------------------------------------------------------------------- +// Public API +// --------------------------------------------------------------------------- + +/** + * Synchronise `quotaShared-*` combos for a pool: + * + * 1. Resolve pool → connection → provider. + * 2. For each model in PROVIDER_MODELS[provider], upsert a combo named + * `quotaModelName(pool.name, provider, model)` with a single model-step + * pinned to the pool's connectionId. + * 3. Prune stale quota combos for this pool slug that are no longer in the + * desired set. + * + * Idempotent: running twice produces no changes on the second call. + * Defensive: missing pool, missing connection, or empty model list → prune to + * empty without throwing. + */ +export async function syncQuotaCombos(poolId: string): Promise { + const resolved = await resolvePoolProvider(poolId); + + if (!resolved) { + // Pool or connection gone — prune any leftover combos if we can find the + // pool slug from poolId (best effort: we won't have the name, so skip). + await removeQuotaCombosForPool(poolId); + return; + } + + const { pool, provider } = resolved; + const poolSlug = quotaPoolSlug(pool.name); + const modelIds = getProviderModelIds(provider); + + // Build the set of desired combo names + const desiredNames = new Set( + modelIds.map((modelId) => quotaModelName(pool.name, provider, modelId)) + ); + + // Upsert each desired combo + 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: pool.connectionId, + weight: 100, + }; + + if (existing && typeof existing.id === "string") { + // Update to ensure connectionId / step 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"); + } + } + + // Prune stale combos that belong to this pool slug but are no longer desired + let allCombos: Awaited> = []; + try { + allCombos = await getCombos(); + } catch (err) { + log.warn({ err: (err as Error)?.message, poolId }, "quota-combo prune: getCombos failed"); + return; + } + + for (const combo of allCombos) { + const name = typeof combo.name === "string" ? combo.name : null; + if (!name) continue; + if (!isQuotaModelName(name)) continue; + + const parsed = parseQuotaModelName(name); + if (!parsed) continue; + if (parsed.poolSlug !== poolSlug) continue; + + // Belongs to this pool slug but not in the desired set → prune + if (!desiredNames.has(name)) { + try { + await deleteComboByName(name); + } catch (err) { + log.warn({ err: (err as Error)?.message, comboName: name, poolId }, "quota-combo prune failed"); + } + } + } +} + +/** + * Delete ALL `quotaShared-*` combos that belong to the given pool. + * + * Used on pool deletion. Because the pool may already be gone from the DB when + * this is called, we look up the pool name first; if missing, we fall back to + * scanning all quota combos and deleting those whose parsed slug matches the + * pool's last-known slug (best-effort via poolId as slug). + */ +export async function removeQuotaCombosForPool(poolId: string): Promise { + // Try to get the pool's name to compute the canonical slug + const pool = getPool(poolId); + const slug = pool ? quotaPoolSlug(pool.name) : null; + + let allCombos: Awaited> = []; + try { + allCombos = await getCombos(); + } catch (err) { + log.warn({ err: (err as Error)?.message, poolId }, "removeQuotaCombosForPool: getCombos failed"); + return; + } + + for (const combo of allCombos) { + const name = typeof combo.name === "string" ? combo.name : null; + if (!name) continue; + if (!isQuotaModelName(name)) continue; + + const parsed = parseQuotaModelName(name); + if (!parsed) continue; + + // Match by slug when we have a pool name; otherwise no match possible + if (slug !== null && parsed.poolSlug !== slug) continue; + + try { + await deleteComboByName(name); + } catch (err) { + log.warn({ err: (err as Error)?.message, comboName: name, poolId }, "quota-combo remove failed"); + } + } +} diff --git a/tests/unit/quota-combos-sync.test.ts b/tests/unit/quota-combos-sync.test.ts new file mode 100644 index 0000000000..a63db3d4ca --- /dev/null +++ b/tests/unit/quota-combos-sync.test.ts @@ -0,0 +1,307 @@ +/** + * tests/unit/quota-combos-sync.test.ts + * + * TDD coverage for src/lib/quota/quotaCombos.ts::syncQuotaCombos and + * src/lib/quota/quotaCombos.ts::removeQuotaCombosForPool (Phase B2). + * + * Uses "glm" as the test provider because it has a small, stable model list + * in the static registry (10 models). Mirrors the seeding pattern from + * quota-key-resolve.test.ts. + */ + +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-combos-sync-")); +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, removeQuotaCombosForPool } = await import( + "../../src/lib/quota/quotaCombos.ts" +); +const { quotaModelName, isQuotaModelName, parseQuotaModelName, quotaPoolSlug } = await import( + "../../src/lib/quota/quotaModelNaming.ts" +); +const { PROVIDER_MODELS } = await import("../../open-sse/config/providerModels.ts"); + +// --------------------------------------------------------------------------- +// Test lifecycle helpers +// --------------------------------------------------------------------------- + +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 to list all quota combos from the DB +// --------------------------------------------------------------------------- + +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[]) : [], + })); +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +test("syncQuotaCombos: creates one combo per glm model with correct name and target", async () => { + // Seed a glm connection + const conn = await providersDb.createProviderConnection({ + provider: "glm", + authType: "apikey", + name: "quota-combos-sync-glm", + apiKey: "sk-test-glm-quota-b2", + }); + const connId = (conn as Record).id as string; + assert.ok(connId, "connection should have an id"); + + const pool = poolsDb.createPool({ connectionId: connId, name: "TestGlmPool" }); + + await syncQuotaCombos(pool.id); + + const glmModels = PROVIDER_MODELS["glm"] ?? []; + assert.ok(glmModels.length > 0, "glm should have models in registry"); + + const quotaCombos = await listQuotaCombos(); + const quotaComboNames = new Set(quotaCombos.map((c) => c.name)); + + // Every glm model should have a combo + for (const model of glmModels) { + const expectedName = quotaModelName("TestGlmPool", "glm", model.id); + assert.ok( + quotaComboNames.has(expectedName), + `Missing combo for model ${model.id}: ${expectedName}` + ); + } + + // No extra quota combos for other pools + for (const c of quotaCombos) { + const parsed = parseQuotaModelName(c.name); + assert.ok(parsed, `Could not parse quota model name: ${c.name}`); + assert.equal(parsed?.poolSlug, quotaPoolSlug("TestGlmPool")); + } +}); + +test("syncQuotaCombos: each combo has a single step with provider=glm and connectionId pinned", async () => { + const conn = await providersDb.createProviderConnection({ + provider: "glm", + authType: "apikey", + name: "quota-combos-step-check", + apiKey: "sk-test-glm-step", + }); + const connId = (conn as Record).id as string; + const pool = poolsDb.createPool({ connectionId: connId, name: "StepCheckPool" }); + + await syncQuotaCombos(pool.id); + + const quotaCombos = await listQuotaCombos(); + assert.ok(quotaCombos.length > 0, "expected at least one quota combo"); + + for (const c of quotaCombos) { + const parsed = parseQuotaModelName(c.name); + assert.ok(parsed, `unparseable combo name: ${c.name}`); + assert.equal(parsed?.provider, "glm"); + + assert.equal(c.models.length, 1, `combo ${c.name} should have exactly 1 step`); + const step = c.models[0] as Record; + assert.equal(step.kind, "model"); + + // Model string includes the provider prefix + const modelStr = typeof step.model === "string" ? step.model : ""; + assert.ok( + modelStr.startsWith("glm/") || modelStr === parsed.model, + `step.model "${modelStr}" should contain the model id "${parsed.model}"` + ); + + // connectionId is pinned to the pool's connection + assert.equal( + step.connectionId, + connId, + `step.connectionId should be pinned to pool connection ${connId}` + ); + } +}); + +test("syncQuotaCombos: idempotent — calling twice produces no duplicates", async () => { + const conn = await providersDb.createProviderConnection({ + provider: "glm", + authType: "apikey", + name: "quota-combos-idempotent", + apiKey: "sk-test-glm-idem", + }); + const connId = (conn as Record).id as string; + const pool = poolsDb.createPool({ connectionId: connId, name: "IdempotentPool" }); + + await syncQuotaCombos(pool.id); + const afterFirst = await listQuotaCombos(); + const firstCount = afterFirst.length; + + await syncQuotaCombos(pool.id); + const afterSecond = await listQuotaCombos(); + + assert.equal(afterSecond.length, firstCount, "second sync must not create duplicate combos"); + + // All names should be identical sets + const firstNames = new Set(afterFirst.map((c) => c.name)); + const secondNames = new Set(afterSecond.map((c) => c.name)); + for (const name of firstNames) { + assert.ok(secondNames.has(name), `Name disappeared after second sync: ${name}`); + } +}); + +test("syncQuotaCombos: prunes stale combos for same pool slug", async () => { + // Seed two separate connections and pools, both named to produce different slugs + const conn = await providersDb.createProviderConnection({ + provider: "glm", + authType: "apikey", + name: "quota-combos-prune-conn", + apiKey: "sk-test-glm-prune", + }); + const connId = (conn as Record).id as string; + const pool = poolsDb.createPool({ connectionId: connId, name: "PrunePool" }); + + await syncQuotaCombos(pool.id); + + const afterInitial = await listQuotaCombos(); + const initialCount = afterInitial.length; + assert.ok(initialCount > 0, "should have combos after initial sync"); + + // Manually insert a stale combo with the same pool slug but a nonexistent model + const staleComboName = `quotaShared-${quotaPoolSlug("PrunePool")}-glm/fake-model-stale`; + await combosDb.createCombo({ + name: staleComboName, + models: [{ kind: "model", model: "glm/fake-model-stale", providerId: "glm", weight: 100 }], + strategy: "priority", + isHidden: true, + }); + + // Verify the stale combo exists + const stale = await combosDb.getComboByName(staleComboName); + assert.ok(stale, "stale combo should exist before prune"); + + // Re-sync — should prune the stale combo + await syncQuotaCombos(pool.id); + + const pruned = await combosDb.getComboByName(staleComboName); + assert.equal(pruned, null, "stale combo should be pruned after re-sync"); + + // Desired combos should still be present + const afterPrune = await listQuotaCombos(); + assert.equal( + afterPrune.length, + initialCount, + "combo count should return to initial after pruning stale" + ); +}); + +test("removeQuotaCombosForPool: removes all quota combos for the pool", async () => { + const conn = await providersDb.createProviderConnection({ + provider: "glm", + authType: "apikey", + name: "quota-combos-remove", + apiKey: "sk-test-glm-remove", + }); + const connId = (conn as Record).id as string; + const pool = poolsDb.createPool({ connectionId: connId, name: "RemovePool" }); + + await syncQuotaCombos(pool.id); + + const before = await listQuotaCombos(); + assert.ok(before.length > 0, "expected combos to remove"); + + await removeQuotaCombosForPool(pool.id); + + const after = await listQuotaCombos(); + assert.equal(after.length, 0, "all quota combos should be removed"); +}); + +test("syncQuotaCombos: does not affect quota combos for a different pool slug", async () => { + const conn = await providersDb.createProviderConnection({ + provider: "glm", + authType: "apikey", + name: "quota-combos-isolation", + apiKey: "sk-test-glm-isolation", + }); + const connId = (conn as Record).id as string; + + const poolA = poolsDb.createPool({ connectionId: connId, name: "PoolAlpha" }); + const poolB = poolsDb.createPool({ connectionId: connId, name: "PoolBeta" }); + + await syncQuotaCombos(poolA.id); + await syncQuotaCombos(poolB.id); + + const all = await listQuotaCombos(); + const slugA = quotaPoolSlug("PoolAlpha"); + const slugB = quotaPoolSlug("PoolBeta"); + + const forA = all.filter((c) => parseQuotaModelName(c.name)?.poolSlug === slugA); + const forB = all.filter((c) => parseQuotaModelName(c.name)?.poolSlug === slugB); + + assert.ok(forA.length > 0, "PoolAlpha should have combos"); + assert.ok(forB.length > 0, "PoolBeta should have combos"); + + // Removing PoolA's combos should not touch PoolB's + await removeQuotaCombosForPool(poolA.id); + + const remaining = await listQuotaCombos(); + const remainingForA = remaining.filter((c) => parseQuotaModelName(c.name)?.poolSlug === slugA); + const remainingForB = remaining.filter((c) => parseQuotaModelName(c.name)?.poolSlug === slugB); + + assert.equal(remainingForA.length, 0, "PoolAlpha combos should all be removed"); + assert.equal(remainingForB.length, forB.length, "PoolBeta combos should be untouched"); +}); + +test("syncQuotaCombos: unknown pool id — no throw, prunes nothing (no combos exist)", async () => { + // Should not throw + await assert.doesNotReject( + () => syncQuotaCombos("nonexistent-pool-id"), + "syncQuotaCombos with unknown poolId should not throw" + ); + + const quotaCombos = await listQuotaCombos(); + assert.equal(quotaCombos.length, 0, "no combos should exist"); +}); + +test("removeQuotaCombosForPool: unknown pool id — no throw", async () => { + await assert.doesNotReject( + () => removeQuotaCombosForPool("nonexistent-pool-id"), + "removeQuotaCombosForPool with unknown poolId should not throw" + ); +});