feat(quota): force quota-exclusive keys onto pool connection in account selection (Phase A4)

This commit is contained in:
diegosouzapw
2026-05-30 21:00:36 -03:00
parent 8316c618b2
commit a921300a53
3 changed files with 84 additions and 2 deletions

View File

@@ -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

View File

@@ -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

View 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, []);
});