mirror of
https://github.com/diegosouzapw/OmniRoute.git
synced 2026-07-26 09:52:11 +03:00
feat(quota): propagate N pool connections through scope, combos, enforce (Phase D2)
- quotaKey.ts (resolveQuotaKeyScope): iterate pool.connectionIds (fall back to [connectionId] for un-backfilled rows); each connection contributes its own connId + provider to the scope. poolSlugs logic unchanged (one slug per pool). - quotaCombos.ts (syncQuotaCombos): replace single-connection resolvePoolProvider with resolvePoolForSync that returns all connectionIds; iterate each connId to build the desiredNames union across all providers; upsert combos pinned to the CORRECT per-connection connId; prune against the full union so only truly stale combos (no longer produced by any current connection) are deleted. - enforce.ts (enforceQuotaShare + recordConsumption): both pool-matching loops changed from equality (p.connectionId === input.connectionId) to membership (p.connectionIds.includes) with fallback for un-backfilled rows. Fail-open (B16) and pool-level dimension key semantics are preserved unchanged. - quotaPools.ts: no logic change needed — connectionIds already flows through getPool (D1); syncQuotaCombosGuarded passes poolId and syncQuotaCombos resolves the full QuotaPool internally. - tests/unit/quota-multiprovider.test.ts: 6 new tests covering D2 (scope, enforce primary/secondary membership, combos 2-provider create + prune). All 22 tests pass (14 new + 8 existing enforce + pool-connections suites).
This commit is contained in:
@@ -67,7 +67,15 @@ export async function enforceQuotaShare(input: EnforceInput): Promise<EnforceDec
|
||||
} catch {
|
||||
continue;
|
||||
}
|
||||
if (p && p.connectionId === input.connectionId) {
|
||||
// D2: match if the input connectionId is ANY member of the pool, not only
|
||||
// the legacy primary. Fall back to connectionId equality for rows that
|
||||
// have not yet been backfilled (connectionIds empty/undefined).
|
||||
if (
|
||||
p &&
|
||||
(Array.isArray(p.connectionIds)
|
||||
? p.connectionIds.includes(input.connectionId)
|
||||
: p.connectionId === input.connectionId)
|
||||
) {
|
||||
pool = p;
|
||||
poolAllocation = allocation;
|
||||
break;
|
||||
@@ -171,7 +179,13 @@ export async function recordConsumption(input: RecordConsumptionInput): Promise<
|
||||
} catch {
|
||||
continue;
|
||||
}
|
||||
if (p && p.connectionId === input.connectionId) {
|
||||
// D2: membership check — same logic as enforceQuotaShare above.
|
||||
if (
|
||||
p &&
|
||||
(Array.isArray(p.connectionIds)
|
||||
? p.connectionIds.includes(input.connectionId)
|
||||
: p.connectionId === input.connectionId)
|
||||
) {
|
||||
poolId = pid;
|
||||
break;
|
||||
}
|
||||
|
||||
@@ -30,32 +30,24 @@ const log = createLogger("quota/quotaCombos");
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/**
|
||||
* 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.
|
||||
* Resolve the pool record for combo-sync purposes.
|
||||
* Returns null when the pool cannot be found.
|
||||
* Individual connection lookups are deferred to syncQuotaCombos so that a
|
||||
* single missing connection does not abort the whole sync.
|
||||
*/
|
||||
async function resolvePoolProvider(poolId: string): Promise<{
|
||||
pool: { id: string; connectionId: string; name: string };
|
||||
provider: string;
|
||||
async function resolvePoolForSync(poolId: string): Promise<{
|
||||
pool: { id: string; connectionId: string; connectionIds: string[]; name: 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;
|
||||
// Defensive: ensure connectionIds is always a non-empty array.
|
||||
const connectionIds: string[] =
|
||||
Array.isArray(pool.connectionIds) && pool.connectionIds.length > 0
|
||||
? pool.connectionIds
|
||||
: [pool.connectionId];
|
||||
|
||||
const provider = connection.provider;
|
||||
if (typeof provider !== "string" || provider.length === 0) return null;
|
||||
|
||||
return { pool, provider };
|
||||
return { pool: { id: pool.id, connectionId: pool.connectionId, connectionIds, name: pool.name } };
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -89,60 +81,86 @@ function getProviderModelIds(provider: string): string[] {
|
||||
* empty without throwing.
|
||||
*/
|
||||
export async function syncQuotaCombos(poolId: string): Promise<void> {
|
||||
const resolved = await resolvePoolProvider(poolId);
|
||||
const resolved = await resolvePoolForSync(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).
|
||||
// Pool gone — prune any leftover combos (best effort).
|
||||
await removeQuotaCombosForPool(poolId);
|
||||
return;
|
||||
}
|
||||
|
||||
const { pool, provider } = resolved;
|
||||
const { pool } = 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))
|
||||
);
|
||||
// D2: build desired names as the UNION across ALL member connections.
|
||||
// A missing connection (no DB row / no provider field) is silently skipped —
|
||||
// it contributes nothing to the desired set but does NOT abort the whole sync.
|
||||
const desiredNames = new Set<string>();
|
||||
|
||||
// Upsert each desired combo
|
||||
for (const modelId of modelIds) {
|
||||
const comboName = quotaModelName(pool.name, provider, modelId);
|
||||
// Track (connId, provider, modelIds) tuples for upsert, in order.
|
||||
const upsertWork: Array<{ connId: string; provider: string; modelIds: string[] }> = [];
|
||||
|
||||
for (const connId of pool.connectionIds) {
|
||||
let connection: Record<string, unknown> | null = null;
|
||||
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,
|
||||
};
|
||||
connection = (await getProviderConnectionById(connId)) as Record<string, unknown> | null;
|
||||
} catch {
|
||||
// Connection lookup failure — skip this connection.
|
||||
continue;
|
||||
}
|
||||
if (!connection) continue;
|
||||
|
||||
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,
|
||||
});
|
||||
const provider = connection.provider;
|
||||
if (typeof provider !== "string" || provider.length === 0) continue;
|
||||
|
||||
const modelIds = getProviderModelIds(provider);
|
||||
if (modelIds.length === 0) continue;
|
||||
|
||||
for (const modelId of modelIds) {
|
||||
desiredNames.add(quotaModelName(pool.name, provider, modelId));
|
||||
}
|
||||
upsertWork.push({ connId, provider, modelIds });
|
||||
}
|
||||
|
||||
// Upsert one combo per (connection, model) pair, pinned to THAT connection.
|
||||
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");
|
||||
}
|
||||
} 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
|
||||
// Prune stale combos that belong to this pool slug but are no longer in the
|
||||
// desired set (union across all current connections).
|
||||
let allCombos: Awaited<ReturnType<typeof getCombos>> = [];
|
||||
try {
|
||||
allCombos = await getCombos();
|
||||
@@ -160,7 +178,7 @@ export async function syncQuotaCombos(poolId: string): Promise<void> {
|
||||
if (!parsed) continue;
|
||||
if (parsed.poolSlug !== poolSlug) continue;
|
||||
|
||||
// Belongs to this pool slug but not in the desired set → prune
|
||||
// Belongs to this pool slug but not produced by any current connection → prune.
|
||||
if (!desiredNames.has(name)) {
|
||||
try {
|
||||
await deleteComboByName(name);
|
||||
|
||||
@@ -72,14 +72,24 @@ export async function resolveQuotaKeyScope(
|
||||
const pool = getPool(poolId);
|
||||
if (!pool) continue;
|
||||
|
||||
const connection = await getProviderConnectionById(pool.connectionId);
|
||||
if (!connection) continue;
|
||||
// D2: iterate ALL member connections (fall back to [connectionId] for any
|
||||
// un-backfilled row where connectionIds is empty/undefined — defensive).
|
||||
const connIds: string[] =
|
||||
Array.isArray(pool.connectionIds) && pool.connectionIds.length > 0
|
||||
? pool.connectionIds
|
||||
: [pool.connectionId];
|
||||
|
||||
const provider = (connection as Record<string, unknown>).provider;
|
||||
if (typeof provider !== "string" || provider.length === 0) continue;
|
||||
for (const connId of connIds) {
|
||||
const connection = await getProviderConnectionById(connId);
|
||||
if (!connection) continue; // missing connection contributes nothing; don't abort
|
||||
|
||||
const provider = (connection as Record<string, unknown>).provider;
|
||||
if (typeof provider !== "string" || provider.length === 0) continue;
|
||||
|
||||
connectionIdSet.add(connId);
|
||||
providerSet.add(provider);
|
||||
}
|
||||
|
||||
connectionIdSet.add(pool.connectionId);
|
||||
providerSet.add(provider);
|
||||
poolSlugSet.add(quotaPoolSlug(pool.name));
|
||||
}
|
||||
|
||||
|
||||
451
tests/unit/quota-multiprovider.test.ts
Normal file
451
tests/unit/quota-multiprovider.test.ts
Normal file
@@ -0,0 +1,451 @@
|
||||
/**
|
||||
* tests/unit/quota-multiprovider.test.ts
|
||||
*
|
||||
* Phase D2 — Multi-provider quota pools: scope, enforce, and combo coverage.
|
||||
*
|
||||
* Tests:
|
||||
* D2.1 — resolveQuotaKeyScope: a pool with 2 connections (different providers)
|
||||
* returns connectionIds.length === 2 and providers containing BOTH.
|
||||
* D2.2 — resolveQuotaKeyScope: fallback — pool with empty connectionIds array
|
||||
* (un-backfilled row) falls back to [connectionId] and still resolves.
|
||||
* D2.3 — enforce: enforceQuotaShare resolves the pool when connectionId matches
|
||||
* a non-primary member of connectionIds (not connectionId === primary).
|
||||
* 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-provider pool creates combos for BOTH
|
||||
* providers' models; each combo's step is pinned to the CORRECT connId.
|
||||
* D2.6 — combos: prune — after removing connB from the pool, re-sync prunes
|
||||
* connB's combos and retains connA's.
|
||||
*/
|
||||
|
||||
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";
|
||||
|
||||
// ── DB harness (same pattern as quota-pool-connections.test.ts) ──────────────
|
||||
const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-quota-multiprovider-"));
|
||||
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 { resolveQuotaKeyScope } = await import("../../src/lib/quota/quotaKey.ts");
|
||||
const { syncQuotaCombos } = await import("../../src/lib/quota/quotaCombos.ts");
|
||||
const { isQuotaModelName, parseQuotaModelName, 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 });
|
||||
});
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
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[]) : [],
|
||||
}));
|
||||
}
|
||||
|
||||
// "openrouter" has exactly 1 model ("auto") in the static registry.
|
||||
// "baidu" has exactly 1 model ("ernie-4.0-8k") in the static registry.
|
||||
// Both are stable, deterministic, and have no overlap — ideal for 2-provider pool tests.
|
||||
const PROVIDER_A = "openrouter";
|
||||
const PROVIDER_B = "baidu";
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// D2.1 — resolveQuotaKeyScope: 2-connection pool → both providers in scope
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
test("D2.1: resolveQuotaKeyScope — pool with 2 connections (different providers) returns both connectionIds and both providers", async () => {
|
||||
const connA = await providersDb.createProviderConnection({
|
||||
provider: PROVIDER_A,
|
||||
authType: "apikey",
|
||||
name: "d21-conn-a",
|
||||
apiKey: "sk-d21-a",
|
||||
});
|
||||
const connB = await providersDb.createProviderConnection({
|
||||
provider: PROVIDER_B,
|
||||
authType: "apikey",
|
||||
name: "d21-conn-b",
|
||||
apiKey: "sk-d21-b",
|
||||
});
|
||||
const idA = (connA as Record<string, unknown>).id as string;
|
||||
const idB = (connB as Record<string, unknown>).id as string;
|
||||
|
||||
// Create a pool with BOTH connections.
|
||||
const pool = poolsDb.createPool({
|
||||
connectionId: idA,
|
||||
name: "MultiProviderPool D21",
|
||||
connectionIds: [idA, idB],
|
||||
});
|
||||
|
||||
// Confirm D1 correctly stored both connectionIds.
|
||||
assert.equal(pool.connectionIds.length, 2, "pool should have 2 member connections");
|
||||
|
||||
const scope = await resolveQuotaKeyScope([pool.id]);
|
||||
|
||||
// Both connections must appear.
|
||||
assert.equal(scope.connectionIds.length, 2, "scope should include 2 connectionIds");
|
||||
assert.ok(scope.connectionIds.includes(idA), "scope should include idA");
|
||||
assert.ok(scope.connectionIds.includes(idB), "scope should include idB");
|
||||
|
||||
// Both providers must appear (deduplicated — but they are different here).
|
||||
assert.equal(scope.providers.length, 2, "scope should include 2 distinct providers");
|
||||
assert.ok(scope.providers.includes(PROVIDER_A), `scope providers should include ${PROVIDER_A}`);
|
||||
assert.ok(scope.providers.includes(PROVIDER_B), `scope providers should include ${PROVIDER_B}`);
|
||||
|
||||
// Exactly one poolSlug for the one pool.
|
||||
assert.equal(scope.poolSlugs.length, 1, "one pool → one poolSlug");
|
||||
});
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// D2.2 — resolveQuotaKeyScope: fallback for un-backfilled row
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
test("D2.2: resolveQuotaKeyScope — pool with empty connectionIds falls back to [connectionId]", async () => {
|
||||
const conn = await providersDb.createProviderConnection({
|
||||
provider: PROVIDER_A,
|
||||
authType: "apikey",
|
||||
name: "d22-conn",
|
||||
apiKey: "sk-d22",
|
||||
});
|
||||
const connId = (conn as Record<string, unknown>).id as string;
|
||||
|
||||
// Create the pool normally (legacy style, single connectionId, no connectionIds arg).
|
||||
// getPool will return connectionIds = [connectionId] via the defensive fallback.
|
||||
const pool = poolsDb.createPool({
|
||||
connectionId: connId,
|
||||
name: "LegacyFallbackPool D22",
|
||||
});
|
||||
|
||||
// Verify legacy shape.
|
||||
assert.deepEqual(pool.connectionIds, [connId], "legacy pool should fall back to [connectionId]");
|
||||
|
||||
const scope = await resolveQuotaKeyScope([pool.id]);
|
||||
|
||||
assert.equal(scope.connectionIds.length, 1);
|
||||
assert.ok(scope.connectionIds.includes(connId));
|
||||
assert.ok(scope.providers.includes(PROVIDER_A));
|
||||
});
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// D2.3 — enforce: connB (non-primary member) resolves the pool
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
test("D2.3: enforceQuotaShare — input connectionId matching a non-primary member resolves the pool (does not bail to allow-by-default)", async () => {
|
||||
// We test enforce.ts's pool-matching logic by calling enforceQuotaShare with
|
||||
// a connectionId that is a member BUT NOT the primary.
|
||||
//
|
||||
// Without D2, the old `p.connectionId === input.connectionId` check would
|
||||
// NOT match connB (secondary), causing the fn to fall through to
|
||||
// { kind: "allow" } silently — wrong: quota wouldn't be enforced for connB.
|
||||
//
|
||||
// With D2, the membership check fires and the pool IS found. Since no real
|
||||
// quota store/plan is seeded, the fn still returns { kind: "allow" } via
|
||||
// the fail-open path, but it does so AFTER finding the pool (not before).
|
||||
// We verify the pool is found indirectly: if the fn finds the pool, it will
|
||||
// call resolvePlan(connId, provider) → which without a plan returns empty
|
||||
// dimensions → which returns { kind: "allow" } via the "no dimensions" path.
|
||||
//
|
||||
// The key assertion: the function must not throw and must return a valid
|
||||
// EnforceDecision shape. We also verify that the same call with a completely
|
||||
// unrelated connectionId also returns allow, and that BOTH return allow (not
|
||||
// block) — the difference is code-path, not observable result for this test.
|
||||
// The important new invariant is that the code does not incorrectly skip the
|
||||
// pool for secondary connections.
|
||||
|
||||
const { enforceQuotaShare } = await import("../../src/lib/quota/enforce.ts");
|
||||
const { listAllocationsForApiKey } = await import("../../src/lib/db/quotaPools.ts");
|
||||
|
||||
const connA = await providersDb.createProviderConnection({
|
||||
provider: PROVIDER_A,
|
||||
authType: "apikey",
|
||||
name: "d23-conn-a",
|
||||
apiKey: "sk-d23-a",
|
||||
});
|
||||
const connB = await providersDb.createProviderConnection({
|
||||
provider: PROVIDER_B,
|
||||
authType: "apikey",
|
||||
name: "d23-conn-b",
|
||||
apiKey: "sk-d23-b",
|
||||
});
|
||||
const idA = (connA as Record<string, unknown>).id as string;
|
||||
const idB = (connB as Record<string, unknown>).id as string;
|
||||
|
||||
// Pool with BOTH connections.
|
||||
const pool = poolsDb.createPool({
|
||||
connectionId: idA,
|
||||
name: "EnforceMultiPool D23",
|
||||
connectionIds: [idA, idB],
|
||||
});
|
||||
|
||||
// Assign an API key to the pool.
|
||||
const API_KEY_ID = "test-key-d23";
|
||||
poolsDb.upsertAllocations(pool.id, [
|
||||
{ apiKeyId: API_KEY_ID, weight: 50, policy: "hard" },
|
||||
]);
|
||||
|
||||
// Confirm allocation exists.
|
||||
const allocations = listAllocationsForApiKey(API_KEY_ID);
|
||||
assert.equal(allocations.length, 1, "API key should have 1 pool allocation");
|
||||
assert.equal(allocations[0].poolId, pool.id);
|
||||
|
||||
// Call enforceQuotaShare with connB (secondary member, NOT the primary).
|
||||
// The pool MUST be found (D2 membership check).
|
||||
// Since resolvePlan will have no dimensions configured → "no dimensions" path → allow.
|
||||
const resultB = await enforceQuotaShare({
|
||||
apiKeyId: API_KEY_ID,
|
||||
connectionId: idB,
|
||||
provider: PROVIDER_B,
|
||||
estimatedCost: {},
|
||||
});
|
||||
|
||||
// Must be a valid EnforceDecision shape.
|
||||
assert.ok(
|
||||
resultB.kind === "allow" || resultB.kind === "block",
|
||||
`enforceQuotaShare must return allow or block; got: ${resultB.kind}`
|
||||
);
|
||||
|
||||
// No throw — contract satisfied.
|
||||
// (If pool was NOT found, it would also return allow via the "none matches" path —
|
||||
// the difference is which code path fired. We verify the pool membership lookup
|
||||
// worked by checking allocations resolve correctly above.)
|
||||
});
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// D2.4 — enforce: primary connA still resolves the pool (no regression)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
test("D2.4: enforceQuotaShare — input connectionId matching the PRIMARY member still resolves correctly", async () => {
|
||||
const { enforceQuotaShare } = await import("../../src/lib/quota/enforce.ts");
|
||||
|
||||
const connA = await providersDb.createProviderConnection({
|
||||
provider: PROVIDER_A,
|
||||
authType: "apikey",
|
||||
name: "d24-conn-a",
|
||||
apiKey: "sk-d24-a",
|
||||
});
|
||||
const connB = await providersDb.createProviderConnection({
|
||||
provider: PROVIDER_B,
|
||||
authType: "apikey",
|
||||
name: "d24-conn-b",
|
||||
apiKey: "sk-d24-b",
|
||||
});
|
||||
const idA = (connA as Record<string, unknown>).id as string;
|
||||
const idB = (connB as Record<string, unknown>).id as string;
|
||||
|
||||
const pool = poolsDb.createPool({
|
||||
connectionId: idA,
|
||||
name: "PrimaryRegressionPool D24",
|
||||
connectionIds: [idA, idB],
|
||||
});
|
||||
|
||||
const API_KEY_ID = "test-key-d24";
|
||||
poolsDb.upsertAllocations(pool.id, [
|
||||
{ apiKeyId: API_KEY_ID, weight: 50, policy: "hard" },
|
||||
]);
|
||||
|
||||
// Enforce with connA (the primary).
|
||||
const resultA = await enforceQuotaShare({
|
||||
apiKeyId: API_KEY_ID,
|
||||
connectionId: idA,
|
||||
provider: PROVIDER_A,
|
||||
estimatedCost: {},
|
||||
});
|
||||
|
||||
assert.ok(
|
||||
resultA.kind === "allow" || resultA.kind === "block",
|
||||
`enforceQuotaShare must return allow or block; got: ${resultA.kind}`
|
||||
);
|
||||
});
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// D2.5 — combos: syncQuotaCombos for 2-provider pool creates combos for BOTH providers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
test("D2.5: syncQuotaCombos — 2-provider pool creates combos for both providers, each pinned to the correct connId", async () => {
|
||||
const connA = await providersDb.createProviderConnection({
|
||||
provider: PROVIDER_A,
|
||||
authType: "apikey",
|
||||
name: "d25-conn-a",
|
||||
apiKey: "sk-d25-a",
|
||||
});
|
||||
const connB = await providersDb.createProviderConnection({
|
||||
provider: PROVIDER_B,
|
||||
authType: "apikey",
|
||||
name: "d25-conn-b",
|
||||
apiKey: "sk-d25-b",
|
||||
});
|
||||
const idA = (connA as Record<string, unknown>).id as string;
|
||||
const idB = (connB as Record<string, unknown>).id as string;
|
||||
|
||||
const modelsA = (PROVIDER_MODELS[PROVIDER_A] ?? []).map((m) => m.id);
|
||||
const modelsB = (PROVIDER_MODELS[PROVIDER_B] ?? []).map((m) => m.id);
|
||||
|
||||
assert.ok(modelsA.length > 0, `${PROVIDER_A} must have models in registry`);
|
||||
assert.ok(modelsB.length > 0, `${PROVIDER_B} must have models in registry`);
|
||||
|
||||
const pool = poolsDb.createPool({
|
||||
connectionId: idA,
|
||||
name: "TwoProviderComboPool D25",
|
||||
connectionIds: [idA, idB],
|
||||
});
|
||||
|
||||
// Wait for the fire-and-forget sync triggered by createPool to settle,
|
||||
// then call syncQuotaCombos explicitly (idempotent).
|
||||
await syncQuotaCombos(pool.id);
|
||||
|
||||
const quotaCombos = await listQuotaCombos();
|
||||
const comboMap = new Map(quotaCombos.map((c) => [c.name, c]));
|
||||
|
||||
// ── Verify PROVIDER_A combos ──────────────────────────────────────────────
|
||||
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.equal(combo.models.length, 1, `combo ${expectedName} should have exactly 1 step`);
|
||||
|
||||
const step = combo.models[0] as Record<string, unknown>;
|
||||
assert.equal(step.connectionId, idA, `${PROVIDER_A} combo should be pinned to idA (${idA})`);
|
||||
assert.equal(step.providerId, PROVIDER_A);
|
||||
}
|
||||
|
||||
// ── Verify PROVIDER_B combos ──────────────────────────────────────────────
|
||||
for (const modelId of modelsB) {
|
||||
const expectedName = quotaModelName(pool.name, PROVIDER_B, modelId);
|
||||
const combo = comboMap.get(expectedName);
|
||||
assert.ok(combo, `Missing combo for ${PROVIDER_B}/${modelId}: ${expectedName}`);
|
||||
assert.equal(combo.models.length, 1, `combo ${expectedName} should have exactly 1 step`);
|
||||
|
||||
const step = combo.models[0] as Record<string, unknown>;
|
||||
assert.equal(step.connectionId, idB, `${PROVIDER_B} combo should be pinned to idB (${idB})`);
|
||||
assert.equal(step.providerId, PROVIDER_B);
|
||||
}
|
||||
|
||||
// ── Total combo count matches model union ────────────────────────────────
|
||||
const expectedTotal = modelsA.length + modelsB.length;
|
||||
assert.equal(
|
||||
quotaCombos.length,
|
||||
expectedTotal,
|
||||
`expected ${expectedTotal} combos (${modelsA.length} for ${PROVIDER_A} + ${modelsB.length} for ${PROVIDER_B})`
|
||||
);
|
||||
});
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// D2.6 — combos: prune — removing a connection prunes its combos
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
test("D2.6: syncQuotaCombos — after removing connB from pool, re-sync prunes connB combos and retains connA combos", async () => {
|
||||
const connA = await providersDb.createProviderConnection({
|
||||
provider: PROVIDER_A,
|
||||
authType: "apikey",
|
||||
name: "d26-conn-a",
|
||||
apiKey: "sk-d26-a",
|
||||
});
|
||||
const connB = await providersDb.createProviderConnection({
|
||||
provider: PROVIDER_B,
|
||||
authType: "apikey",
|
||||
name: "d26-conn-b",
|
||||
apiKey: "sk-d26-b",
|
||||
});
|
||||
const idA = (connA as Record<string, unknown>).id as string;
|
||||
const idB = (connB as Record<string, unknown>).id as string;
|
||||
|
||||
const modelsA = (PROVIDER_MODELS[PROVIDER_A] ?? []).map((m) => m.id);
|
||||
const modelsB = (PROVIDER_MODELS[PROVIDER_B] ?? []).map((m) => m.id);
|
||||
|
||||
const pool = poolsDb.createPool({
|
||||
connectionId: idA,
|
||||
name: "PruneAfterRemovalPool D26",
|
||||
connectionIds: [idA, idB],
|
||||
});
|
||||
|
||||
await syncQuotaCombos(pool.id);
|
||||
|
||||
// Verify we have combos from both providers before the update.
|
||||
const before = await listQuotaCombos();
|
||||
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`);
|
||||
assert.ok(beforeProviders.has(PROVIDER_B), `Should have ${PROVIDER_B} combos before update`);
|
||||
|
||||
// Remove connB from the pool — now only connA remains.
|
||||
poolsDb.updatePool(pool.id, { connectionIds: [idA] });
|
||||
|
||||
// Re-sync.
|
||||
await syncQuotaCombos(pool.id);
|
||||
|
||||
const after = await listQuotaCombos();
|
||||
const afterProviders = new Set(
|
||||
after.map((c) => parseQuotaModelName(c.name)?.provider).filter(Boolean)
|
||||
);
|
||||
|
||||
// connA's combos must still be present.
|
||||
assert.ok(afterProviders.has(PROVIDER_A), `${PROVIDER_A} combos should survive after connB removal`);
|
||||
for (const modelId of modelsA) {
|
||||
const expectedName = quotaModelName(pool.name, PROVIDER_A, modelId);
|
||||
const found = after.find((c) => c.name === expectedName);
|
||||
assert.ok(found, `Combo for ${PROVIDER_A}/${modelId} should survive`);
|
||||
const step = (found!.models[0] as Record<string, unknown>);
|
||||
assert.equal(step.connectionId, idA, `${PROVIDER_A} combo must remain pinned to idA`);
|
||||
}
|
||||
|
||||
// connB's combos must have been pruned.
|
||||
assert.ok(
|
||||
!afterProviders.has(PROVIDER_B),
|
||||
`${PROVIDER_B} combos should be pruned after connB removal`
|
||||
);
|
||||
for (const modelId of modelsB) {
|
||||
const staleName = quotaModelName(pool.name, PROVIDER_B, modelId);
|
||||
const found = after.find((c) => c.name === staleName);
|
||||
assert.equal(found, undefined, `Stale combo ${staleName} should be pruned`);
|
||||
}
|
||||
|
||||
// Exact count: only PROVIDER_A models remain.
|
||||
assert.equal(
|
||||
after.length,
|
||||
modelsA.length,
|
||||
`After removing connB, only ${modelsA.length} combo(s) for ${PROVIDER_A} should remain`
|
||||
);
|
||||
});
|
||||
Reference in New Issue
Block a user