feat(quota): auto-sync quotaShared-* combos on pool allocation changes (Phase B2)

Mints one combo per model of the pool's provider when a quota pool is
created/updated/reallocated, and prunes stale quota combos on deletion.
This commit is contained in:
diegosouzapw
2026-05-30 21:23:08 -03:00
parent 6214ea6768
commit 49f6092099
3 changed files with 560 additions and 2 deletions

View File

@@ -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<void> {
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<void> {
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);
}
/**

View File

@@ -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<string, unknown> | 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<void> {
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<ReturnType<typeof getCombos>> = [];
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<void> {
// 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<ReturnType<typeof getCombos>> = [];
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");
}
}
}

View File

@@ -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<Array<{ name: string; models: unknown[] }>> {
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<string, unknown>).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<string, unknown>).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<string, unknown>;
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<string, unknown>).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<string, unknown>).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<string, unknown>).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<string, unknown>).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"
);
});