mirror of
https://github.com/diegosouzapw/OmniRoute.git
synced 2026-08-01 21:02:12 +03:00
feat(quota): force quota-exclusive keys onto pool connection in account selection (Phase A4)
This commit is contained in:
@@ -20,6 +20,27 @@ export interface QuotaKeyScope {
|
||||
providers: string[];
|
||||
}
|
||||
|
||||
/**
|
||||
* Constrain an existing connection allow-list to the connections belonging to a
|
||||
* quota key's pool scope.
|
||||
*
|
||||
* Semantics mirror `intersectAllowedConnectionIds` in chat.ts:
|
||||
* - Empty `quotaConnectionIds` (non-quota key) → return `existing` unchanged.
|
||||
* - Empty / null `existing` (no prior constraint) → return `quotaConnectionIds`.
|
||||
* - Both non-empty → intersection.
|
||||
* - Disjoint sets → empty array (no eligible connection).
|
||||
*
|
||||
* This is a pure, synchronous function — easy to unit-test without DB setup.
|
||||
*/
|
||||
export function constrainConnectionsToQuota(
|
||||
existing: string[],
|
||||
quotaConnectionIds: string[]
|
||||
): string[] {
|
||||
if (quotaConnectionIds.length === 0) return existing;
|
||||
if (existing.length === 0) return quotaConnectionIds;
|
||||
return existing.filter((id) => quotaConnectionIds.includes(id));
|
||||
}
|
||||
|
||||
/**
|
||||
* Given the `allowedQuotas` field of an API key (array of quota-pool IDs),
|
||||
* returns the set of connection IDs and provider slugs that the key is
|
||||
|
||||
@@ -96,6 +96,10 @@ import {
|
||||
resolveCooldownAwareRetrySettings,
|
||||
waitForCooldownAwareRetry,
|
||||
} from "../services/cooldownAwareRetry";
|
||||
import {
|
||||
constrainConnectionsToQuota,
|
||||
resolveQuotaKeyScope,
|
||||
} from "../../lib/quota/quotaKey";
|
||||
|
||||
registerCodexQuotaFetcher();
|
||||
|
||||
@@ -460,11 +464,20 @@ export async function handleChat(request: any, clientRawRequest: any = null) {
|
||||
const resolvedModel = modelInfo.model || modelString;
|
||||
const hasForcedConnection =
|
||||
typeof target?.connectionId === "string" && target.connectionId.trim().length > 0;
|
||||
const allowedConnections = intersectAllowedConnectionIds(
|
||||
let allowedConnections = intersectAllowedConnectionIds(
|
||||
apiKeyInfo?.allowedConnections ?? null,
|
||||
target?.allowedConnectionIds ?? null
|
||||
);
|
||||
|
||||
// A4: quota-exclusive keys must only use the pool's connection(s).
|
||||
if (apiKeyInfo?.allowedQuotas && apiKeyInfo.allowedQuotas.length > 0) {
|
||||
const quotaScope = await resolveQuotaKeyScope(apiKeyInfo.allowedQuotas);
|
||||
allowedConnections = constrainConnectionsToQuota(
|
||||
allowedConnections ?? [],
|
||||
quotaScope.connectionIds
|
||||
);
|
||||
}
|
||||
|
||||
if (Array.isArray(allowedConnections) && allowedConnections.length === 0) {
|
||||
return false;
|
||||
}
|
||||
@@ -731,10 +744,20 @@ async function handleSingleModelChat(
|
||||
const hasForcedConnection =
|
||||
typeof runtimeOptions.forcedConnectionId === "string" &&
|
||||
runtimeOptions.forcedConnectionId.trim().length > 0;
|
||||
const effectiveAllowedConnections = intersectAllowedConnectionIds(
|
||||
let effectiveAllowedConnections = intersectAllowedConnectionIds(
|
||||
apiKeyInfo?.allowedConnections ?? null,
|
||||
runtimeOptions.allowedConnectionIds ?? null
|
||||
);
|
||||
|
||||
// A4: quota-exclusive keys must only use the pool's connection(s).
|
||||
if (apiKeyInfo?.allowedQuotas && apiKeyInfo.allowedQuotas.length > 0) {
|
||||
const quotaScope = await resolveQuotaKeyScope(apiKeyInfo.allowedQuotas);
|
||||
effectiveAllowedConnections = constrainConnectionsToQuota(
|
||||
effectiveAllowedConnections ?? [],
|
||||
quotaScope.connectionIds
|
||||
);
|
||||
}
|
||||
|
||||
const bypassReason = forceLiveComboTest
|
||||
? "combo live test"
|
||||
: hasForcedConnection
|
||||
|
||||
38
tests/unit/quota-constrain-connections.test.ts
Normal file
38
tests/unit/quota-constrain-connections.test.ts
Normal file
@@ -0,0 +1,38 @@
|
||||
/**
|
||||
* tests/unit/quota-constrain-connections.test.ts
|
||||
*
|
||||
* TDD coverage for src/lib/quota/quotaKey.ts::constrainConnectionsToQuota.
|
||||
* Pure-function tests — no DB setup needed.
|
||||
*
|
||||
* Cases:
|
||||
* 1. Non-quota key (empty quotaConnectionIds) → existing returned unchanged.
|
||||
* 2. No prior constraint (empty existing) + quota [c1,c2] → [c1,c2].
|
||||
* 3. Existing [c1,c2,c3] ∩ quota [c2] → [c2].
|
||||
* 4. Disjoint sets → [].
|
||||
*/
|
||||
|
||||
import { test } from "node:test";
|
||||
import assert from "node:assert/strict";
|
||||
|
||||
import { constrainConnectionsToQuota } from "../../src/lib/quota/quotaKey.ts";
|
||||
|
||||
test("non-quota key: empty quotaConnectionIds returns existing unchanged", () => {
|
||||
const existing = ["c1", "c2", "c3"];
|
||||
const result = constrainConnectionsToQuota(existing, []);
|
||||
assert.deepStrictEqual(result, existing);
|
||||
});
|
||||
|
||||
test("no prior constraint: empty existing with quota [c1,c2] returns [c1,c2]", () => {
|
||||
const result = constrainConnectionsToQuota([], ["c1", "c2"]);
|
||||
assert.deepStrictEqual(result, ["c1", "c2"]);
|
||||
});
|
||||
|
||||
test("intersection: existing [c1,c2,c3] ∩ quota [c2] returns [c2]", () => {
|
||||
const result = constrainConnectionsToQuota(["c1", "c2", "c3"], ["c2"]);
|
||||
assert.deepStrictEqual(result, ["c2"]);
|
||||
});
|
||||
|
||||
test("disjoint sets: no common elements returns empty array", () => {
|
||||
const result = constrainConnectionsToQuota(["c1", "c2"], ["c3", "c4"]);
|
||||
assert.deepStrictEqual(result, []);
|
||||
});
|
||||
Reference in New Issue
Block a user