mirror of
https://github.com/diegosouzapw/OmniRoute.git
synced 2026-09-13 18:32:12 +03:00
fix(combo): stop quota-weighted routing onto out-of-credit connections (#13006)
Both gaps are real and they compound: `executeTargetAttempt` already classified the 402 through `isQuotaExhaustionResponse` and then dropped it, so the only writer into the quota cache was the 429 path in `chat.ts`. A snapshot reading `remaining=1%, is_exhausted=0` five hours stale is then exactly what quota-weighted routing will keep picking. --- Validated in one consolidated worktree cut from `release/v3.8.51`, boarded with the other 19 PRs of this batch. Two in-batch conflicts, both additive and resolved by keeping each side: the `ENVIRONMENT.md` table (#13035 + #13011) and the `chatHelpers.ts` import block (#12975 on the tip + #13017). - `typecheck:core` clean; `check:dashboard-typecheck` OK (206 pre-existing, within baseline); `check:changelog-integrity` OK; `check:docs-counts` migrations ✓ - complexity 2816 / baseline 3218 and cognitive-complexity 1271 / baseline 1437 — both under baseline - 531 of 532 focused assertions green across the batch's 46 test files - `check-file-size` rebaselined for the batch's real growth (annotation `_rebaseline_2026_09_11_mergebatch_v3851_houminxi`, landed on #13038), attributed per PR The single red is **not this batch**: `tests/unit/combo/quota-weighted-strategy.test.ts` → "A/B isolation: 7 hard-empty + 2 at 0.5% + 1 at 40%, floor=1" asserts an order between two connections of identical weight and flakes on the pure tip too — 2 failures in 4 runs at `origin/release/v3.8.51` with nothing from this batch applied. ⚠️ base-red inherited: #12732 — `Docs Gates`, `Merge integrity`, `No new ESLint warnings`, `Unit Tests fast-path` and `Fast Quality Gates` reproduce on the pure tip (provider count 356 vs the 358 the modules define, SKILL.md drift, and `open-sse/utils/stream.ts` at 3115 > frozen 3098, untouched here). Thanks @HouMinXi — the live evidence on these (X500 logs, `storage.sqlite` state, real `/v1/models` probes, the 36-minute outage write-up) is what let a 20-PR batch be reviewed as a unit.
This commit is contained in:
@@ -0,0 +1 @@
|
||||
- **fix(combo):** quota-weighted routing stops drawing on an out-of-credit connection — a 402 now invalidates the stored quota snapshot instead of leaving its stale remaining percentage in place, and a snapshot older than 10 minutes no longer counts as confident headroom for the primary pool ([#12972](https://github.com/diegosouzapw/OmniRoute/pull/12972)) — thanks @HouMinXi
|
||||
@@ -0,0 +1 @@
|
||||
- **fix(combo):** quota-aware expansion drops banned, inactive, missing, and wrong-provider connections before quota fetch or model dispatch; pins and allowlists stay selectors, not a bypass. Antigravity automatic exhaustion now requires a reported zero remaining, so a positive balance below 1% stays eligible.
|
||||
@@ -78,6 +78,7 @@ import {
|
||||
isQuotaExhaustionResponse,
|
||||
recordQuotaExhaustionClassification,
|
||||
} from "./quotaExhaustion.ts";
|
||||
import { markAccountExhaustedFromCredits } from "../../../src/domain/quotaCache.ts";
|
||||
import { classifyComboOutcome, redactConnectionLabel } from "./comboErrorAggregation.ts";
|
||||
import { readConnectionForCooldownGate } from "./executeTargetGates.ts";
|
||||
import {
|
||||
@@ -994,6 +995,12 @@ export async function executeTargetAttempt(opts: {
|
||||
|
||||
const quotaExhausted = await isQuotaExhaustionResponse(result, provider, rawModel, profile);
|
||||
recordQuotaExhaustionClassification(result, quotaExhausted);
|
||||
// Balance exhaustion is upstream truth about credits, and it outranks the
|
||||
// stored snapshot — which can be hours stale and still claim headroom. Mark
|
||||
// it so the next quota-weighted draw stops picking this connection.
|
||||
if (quotaExhausted && result.status === 402 && targetWithConnection.connectionId && provider) {
|
||||
markAccountExhaustedFromCredits(targetWithConnection.connectionId, provider);
|
||||
}
|
||||
state.observeFailure(quotaExhausted, target.executionKey);
|
||||
|
||||
// Check if this is a transient error worth retrying on same model.
|
||||
|
||||
@@ -1,18 +1,17 @@
|
||||
/**
|
||||
* Stateful + async reset-aware / reset-window quota strategies for combo routing.
|
||||
*
|
||||
* Holds the two mutable module-level caches that back reset-aware routing
|
||||
* (`resetAwareConnectionCache` for per-provider active connections and
|
||||
* `resetAwareQuotaCache` for per-connection quota snapshots), plus the helpers
|
||||
* Holds the per-connection quota snapshot cache and helpers
|
||||
* that read/write them and the strategy orderers. Extracted byte-identically
|
||||
* from combo.ts (QG v2 Fase 9 T5 D7b) — the larger, stateful half of the
|
||||
* reset-aware quota block. The pure scoring/window-math half lives in
|
||||
* ./quotaScoring.ts and is imported here.
|
||||
*
|
||||
* State cohesion: `resetAwareConnectionCache`, `resetAwareQuotaCache`, and
|
||||
* State cohesion: `resetAwareQuotaCache` and
|
||||
* `MAX_RESET_AWARE_CACHE` MUST remain single instances defined once here,
|
||||
* alongside their only readers/writers (getQuotaAwareConnectionsForTarget,
|
||||
* fetchResetAwareQuotaWithCache) — never duplicate a Map.
|
||||
* alongside their only readers/writers (`fetchResetAwareQuotaWithCache`).
|
||||
* Connection lists go through `getCachedProviderConnections` (5s TTL,
|
||||
* invalidated on connection writes). Do not add a second connection cache.
|
||||
*
|
||||
* Cross-module state: the tie-band round-robin in orderTargetsByResetAwareQuota
|
||||
* and orderTargetsByResetWindow shares the same rrCounters Map from ./rrState.ts
|
||||
@@ -50,18 +49,28 @@ import { rankByHeadroom, type HeadroomSaturation } from "./headroomRanking.ts";
|
||||
import { getInflight, incrementInflight } from "./quotaShareInflight.ts";
|
||||
import { preferAntigravityConnectionsWithStoredProject } from "../antigravityProjectPersist.ts";
|
||||
import { getQuotaFetchScope } from "../antigravityQuotaFamily.ts";
|
||||
import { isQuotaExhaustedForRequest } from "../../../src/domain/quotaCache.ts";
|
||||
import {
|
||||
getQuotaSnapshotFetchedAt,
|
||||
getQuotaWeightedRemainingPercent,
|
||||
isQuotaExhaustedForRequest,
|
||||
} from "../../../src/domain/quotaCache.ts";
|
||||
|
||||
/**
|
||||
* How long a stored quota snapshot stays good enough to be counted as confident
|
||||
* headroom by the quota-weighted A pool.
|
||||
*
|
||||
* Matches the background refresh cadence for active accounts (quotaCache's
|
||||
* ACTIVE_TTL_MS), doubled to absorb one missed refresh tick. Past that the
|
||||
* snapshot says "unknown", not "empty": the connection drops to the B pool and
|
||||
* is still routed to when nothing fresher has room.
|
||||
*/
|
||||
export const QUOTA_WEIGHTED_MAX_SNAPSHOT_AGE_MS = 10 * 60 * 1000;
|
||||
|
||||
const RESET_AWARE_CONNECTION_CACHE_TTL_MS = 30_000;
|
||||
const RESET_AWARE_QUOTA_FETCH_CONCURRENCY = 5;
|
||||
const HEADROOM_SATURATION_FETCH_CONCURRENCY = 5;
|
||||
|
||||
const MAX_RESET_AWARE_CACHE = 200;
|
||||
|
||||
const resetAwareConnectionCache = new Map<
|
||||
string,
|
||||
{ fetchedAt: number; connections: Array<Record<string, unknown>> }
|
||||
>();
|
||||
const resetAwareQuotaCache = new Map<
|
||||
string,
|
||||
{ fetchedAt: number; quota: unknown; refreshPromise: Promise<unknown> | null }
|
||||
@@ -77,12 +86,6 @@ async function getQuotaAwareConnectionsForTarget(
|
||||
const provider = getResetAwareProvider(target);
|
||||
if (!provider || !getQuotaFetcher(provider)) return [];
|
||||
if (!connectionCache.has(provider)) {
|
||||
const cached = resetAwareConnectionCache.get(provider);
|
||||
if (cached && Date.now() - cached.fetchedAt < RESET_AWARE_CONNECTION_CACHE_TTL_MS) {
|
||||
connectionCache.set(provider, cached.connections);
|
||||
return cached.connections;
|
||||
}
|
||||
|
||||
if (!connectionLoadPromises.has(provider)) {
|
||||
connectionLoadPromises.set(
|
||||
provider,
|
||||
@@ -90,22 +93,17 @@ async function getQuotaAwareConnectionsForTarget(
|
||||
try {
|
||||
const connections = await getCachedProviderConnections({ provider, isActive: true });
|
||||
let activeConnections = Array.isArray(connections)
|
||||
? (connections as Array<Record<string, unknown>>)
|
||||
? (connections as Array<Record<string, unknown>>).filter(
|
||||
(connection) =>
|
||||
connection.isActive !== false &&
|
||||
String(connection.testStatus || "")
|
||||
.trim()
|
||||
.toLowerCase() !== "banned"
|
||||
)
|
||||
: [];
|
||||
if (provider === "antigravity" || provider === "agy") {
|
||||
activeConnections = preferAntigravityConnectionsWithStoredProject(activeConnections);
|
||||
}
|
||||
if (
|
||||
!resetAwareConnectionCache.has(provider) &&
|
||||
resetAwareConnectionCache.size >= MAX_RESET_AWARE_CACHE
|
||||
) {
|
||||
const oldest = resetAwareConnectionCache.keys().next().value;
|
||||
if (oldest !== undefined) resetAwareConnectionCache.delete(oldest);
|
||||
}
|
||||
resetAwareConnectionCache.set(provider, {
|
||||
connections: activeConnections,
|
||||
fetchedAt: Date.now(),
|
||||
});
|
||||
return activeConnections;
|
||||
} catch (error) {
|
||||
log.warn?.("COMBO", "Reset-aware failed to load quota-aware connections.", {
|
||||
@@ -212,6 +210,8 @@ export async function expandTargetsByQuotaAwareConnections(
|
||||
apiKeyAllowedConnectionIds
|
||||
);
|
||||
if (connectionIds.length === 0) {
|
||||
const provider = getResetAwareProvider(target);
|
||||
if (provider && getQuotaFetcher(provider)) continue;
|
||||
if (
|
||||
unrestrictedConnectionIds.length > 0 &&
|
||||
normalizeConnectionIds(apiKeyAllowedConnectionIds)
|
||||
@@ -225,6 +225,7 @@ export async function expandTargetsByQuotaAwareConnections(
|
||||
for (const connectionId of connectionIds) {
|
||||
const provider = getResetAwareProvider(target);
|
||||
const connection = connectionById.get(connectionId);
|
||||
if (provider && getQuotaFetcher(provider) && connection?.provider !== provider) continue;
|
||||
if (
|
||||
connection &&
|
||||
typeof connection.rateLimitedUntil === "string" &&
|
||||
@@ -750,14 +751,15 @@ function sortByScoreThenIndex(a: QuotaWeightedScored, b: QuotaWeightedScored): n
|
||||
return a.index - b.index;
|
||||
}
|
||||
|
||||
function resolveQuotaWeightedFloor(configSource: Record<string, unknown> | null | undefined): number {
|
||||
function resolveQuotaWeightedFloor(
|
||||
configSource: Record<string, unknown> | null | undefined
|
||||
): number {
|
||||
// Number(null) and Number("") are both 0, so an unset or blank key would
|
||||
// switch the floor off instead of taking the default. Only a value that is
|
||||
// actually a number, or a non-empty numeric string, gets to move it.
|
||||
const configured = configSource?.quotaWeightedFloorPercent;
|
||||
const raw =
|
||||
typeof configured === "number" ||
|
||||
(typeof configured === "string" && configured.trim() !== "")
|
||||
typeof configured === "number" || (typeof configured === "string" && configured.trim() !== "")
|
||||
? Number(configured)
|
||||
: Number.NaN;
|
||||
return Number.isFinite(raw) ? Math.max(0, Math.min(100, raw)) : 1;
|
||||
@@ -798,12 +800,28 @@ export async function orderTargetsByQuotaWeighted(
|
||||
}),
|
||||
});
|
||||
|
||||
const eligible = scoredTargets.filter((entry) => entry.remainingPercent > 0);
|
||||
// The live snapshot outranks the freshly-scored fetch on two counts: a 402
|
||||
// recorded against this connection zeroes it, and an observation older than
|
||||
// the staleness bound is not confident enough to sit in the A pool.
|
||||
const now = Date.now();
|
||||
const withSnapshot = scoredTargets.map((entry) => {
|
||||
const connectionId = entry.target.connectionId ?? "";
|
||||
const marked = connectionId ? getQuotaWeightedRemainingPercent(connectionId) : null;
|
||||
const fetchedAt = connectionId ? getQuotaSnapshotFetchedAt(connectionId) : null;
|
||||
return {
|
||||
...entry,
|
||||
remainingPercent: marked === 0 ? 0 : entry.remainingPercent,
|
||||
stale: fetchedAt !== null && now - fetchedAt > QUOTA_WEIGHTED_MAX_SNAPSHOT_AGE_MS,
|
||||
};
|
||||
});
|
||||
|
||||
const eligible = withSnapshot.filter((entry) => entry.remainingPercent > 0);
|
||||
const floor = resolveQuotaWeightedFloor(configSource);
|
||||
const poolA =
|
||||
floor === 0 ? eligible : eligible.filter((entry) => entry.remainingPercent > floor);
|
||||
const poolB =
|
||||
floor === 0 ? [] : eligible.filter((entry) => entry.remainingPercent > 0 && entry.remainingPercent <= floor);
|
||||
const hasRoom = (entry: (typeof eligible)[number]) =>
|
||||
floor === 0 ? true : entry.remainingPercent > floor;
|
||||
// A holds only connections we both believe have room AND observed recently.
|
||||
const poolA = eligible.filter((entry) => hasRoom(entry) && !entry.stale);
|
||||
const poolB = eligible.filter((entry) => !hasRoom(entry) || entry.stale);
|
||||
const selected = poolA.length > 0 ? poolA : poolB;
|
||||
if (selected.length === 0) return [];
|
||||
|
||||
|
||||
@@ -304,8 +304,8 @@ function isAntigravityQuotaExhausted(
|
||||
matchingWindows.length > 0 &&
|
||||
matchingWindows.every(
|
||||
(windowName) =>
|
||||
getQuotaWindowStatus(connectionId, windowName, DEFAULT_QUOTA_THRESHOLD_PERCENT)
|
||||
?.reachedThreshold
|
||||
// Automatic exhaustion is not the operator's optional usage cutoff.
|
||||
getQuotaWindowStatus(connectionId, windowName, 100)?.reachedThreshold
|
||||
)
|
||||
);
|
||||
}
|
||||
@@ -683,6 +683,45 @@ export function getQuotaWindowObservation(
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Mark an account as out of credits from a 402-class response.
|
||||
*
|
||||
* Upstream refusing the request for balance is authoritative: it outranks
|
||||
* whatever remaining percentage the last snapshot happened to hold, which may
|
||||
* be hours old. Without this, a connection that answered 402 keeps its stale
|
||||
* non-zero remaining and the next quota-weighted draw can pick it again.
|
||||
*
|
||||
* The entry is kept (never deactivated or deleted) — credits come back, and a
|
||||
* later successful refresh or window reset clears the flag through the same
|
||||
* paths that clear a 429 mark.
|
||||
*/
|
||||
export function markAccountExhaustedFromCredits(connectionId: string, provider: string) {
|
||||
markAccountExhaustedFrom429(connectionId, provider);
|
||||
}
|
||||
|
||||
/**
|
||||
* Remaining headroom the quota-weighted strategy should credit this connection
|
||||
* with, as a percentage. Returns 0 once the connection is known exhausted so a
|
||||
* 402-marked account cannot be weighted back into the draw.
|
||||
*/
|
||||
export function getQuotaWeightedRemainingPercent(connectionId: string): number | null {
|
||||
const entry = getState().cache.get(connectionId) || hydrateQuotaCacheFromSnapshots(connectionId);
|
||||
if (!entry) return null;
|
||||
if (isAccountQuotaExhausted(connectionId)) return 0;
|
||||
|
||||
const remaining = Object.values(entry.quotas)
|
||||
.filter((quota) => quota.fractionReported !== false)
|
||||
.map((quota) => clampPercent(quota.remainingPercentage));
|
||||
if (remaining.length === 0) return null;
|
||||
return Math.min(...remaining);
|
||||
}
|
||||
|
||||
/** Epoch-ms of the observation backing this connection's snapshot, if any. */
|
||||
export function getQuotaSnapshotFetchedAt(connectionId: string): number | null {
|
||||
const entry = getState().cache.get(connectionId) || hydrateQuotaCacheFromSnapshots(connectionId);
|
||||
return entry ? entry.fetchedAt : null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Mark an account as quota-exhausted from a 429 response (no quota data available).
|
||||
* Uses 5-minute fixed TTL since we don't know the actual resetAt.
|
||||
|
||||
@@ -71,6 +71,7 @@
|
||||
"tests/unit/alibaba-free-tier-exhaustion.test.ts",
|
||||
"tests/unit/anthropic-thinking-signature-recovery.test.ts",
|
||||
"tests/unit/agy-family-not-connection-cooldown.test.ts",
|
||||
"tests/unit/agy-quota-exhaustion-threshold.test.ts",
|
||||
"tests/unit/antigravity-429-quota-cooldown.test.ts",
|
||||
"tests/unit/antigravity-429-quota-tdd.test.ts",
|
||||
"tests/unit/antigravity-prefer-stored-project.test.ts",
|
||||
@@ -230,6 +231,8 @@
|
||||
"tests/unit/combo/combo-exhausted-skip.test.ts",
|
||||
"tests/unit/combo/combo-target-timeout-standards.test.ts",
|
||||
"tests/unit/combo/effective-max-concurrency.test.ts",
|
||||
"tests/unit/combo/quota-connection-eligibility.test.ts",
|
||||
"tests/unit/combo/quota-weighted-stale-402.test.ts",
|
||||
"tests/unit/combo/quota-weighted-strategy.test.ts",
|
||||
"tests/unit/combo/recovery-hint.test.ts",
|
||||
"tests/unit/combo/reset-window-strategy-9330.test.ts",
|
||||
|
||||
83
tests/unit/agy-quota-exhaustion-threshold.test.ts
Normal file
83
tests/unit/agy-quota-exhaustion-threshold.test.ts
Normal file
@@ -0,0 +1,83 @@
|
||||
import test, { after, beforeEach } 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 previousDataDir = process.env.DATA_DIR;
|
||||
const dataDir = fs.mkdtempSync(path.join(os.tmpdir(), "agy-quota-threshold-"));
|
||||
process.env.DATA_DIR = dataDir;
|
||||
|
||||
const core = await import("../../src/lib/db/core.ts");
|
||||
const cache = await import("../../src/domain/quotaCache.ts");
|
||||
const { evaluateQuotaLimitPolicy } = await import("../../src/sse/services/auth.ts");
|
||||
const { toProviderConnection } = await import("../../src/lib/db/providers/lazyConnectionView.ts");
|
||||
|
||||
beforeEach(() => cache.__clearForTests());
|
||||
after(() => {
|
||||
cache.__clearForTests();
|
||||
core.resetDbInstance();
|
||||
if (previousDataDir === undefined) delete process.env.DATA_DIR;
|
||||
else process.env.DATA_DIR = previousDataDir;
|
||||
fs.rmSync(dataDir, { recursive: true, force: true });
|
||||
});
|
||||
|
||||
function seed(provider: string, remaining: number, fractionReported = true) {
|
||||
const resetAt = new Date(Date.now() + 86_400_000).toISOString();
|
||||
cache.setQuotaCache("threshold-account", provider, {
|
||||
"gemini-3.8-flash-high": { remainingPercentage: remaining, resetAt, fractionReported },
|
||||
gemini_weekly: { remainingPercentage: remaining, resetAt, fractionReported },
|
||||
"claude-opus-4-6-thinking": { remainingPercentage: 0, resetAt },
|
||||
claude_gpt_weekly: { remainingPercentage: 0, resetAt },
|
||||
});
|
||||
}
|
||||
|
||||
for (const provider of ["agy", "antigravity"]) {
|
||||
for (const remaining of [0.01, 0.94, 1, 1.01]) {
|
||||
test(`${provider}: positive ${remaining}% is not automatic exhaustion`, () => {
|
||||
seed(provider, remaining);
|
||||
assert.equal(
|
||||
cache.isQuotaExhaustedForRequest("threshold-account", provider, "gemini-3.8-flash-high"),
|
||||
false
|
||||
);
|
||||
assert.equal(
|
||||
cache.isQuotaExhaustedForRequest("threshold-account", provider, "claude-opus-4-6-thinking"),
|
||||
true
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
test(`${provider}: reported zero remains exhausted`, () => {
|
||||
seed(provider, 0);
|
||||
assert.equal(
|
||||
cache.isQuotaExhaustedForRequest("threshold-account", provider, "gemini-3.8-flash-high"),
|
||||
true
|
||||
);
|
||||
});
|
||||
|
||||
test(`${provider}: unreported zero remains unknown`, () => {
|
||||
seed(provider, 0, false);
|
||||
assert.equal(
|
||||
cache.isQuotaExhaustedForRequest("threshold-account", provider, "gemini-3.8-flash-high"),
|
||||
false
|
||||
);
|
||||
});
|
||||
|
||||
test(`${provider}: explicit 99% usage policy still blocks low remaining quota`, () => {
|
||||
seed(provider, 0.94);
|
||||
const decision = evaluateQuotaLimitPolicy(
|
||||
provider,
|
||||
toProviderConnection({
|
||||
id: "threshold-account",
|
||||
provider,
|
||||
isActive: true,
|
||||
providerSpecificData: {
|
||||
limitPolicy: { enabled: true, thresholdPercent: 99, windows: ["gemini_weekly"] },
|
||||
},
|
||||
}),
|
||||
"gemini-3.8-flash-high"
|
||||
);
|
||||
assert.equal(decision.blocked, true);
|
||||
assert.equal(decision.reasons.length, 1);
|
||||
});
|
||||
}
|
||||
@@ -137,7 +137,7 @@ test("isQuotaExhaustedForRequest scopes gemini exhaustion to the requested model
|
||||
);
|
||||
});
|
||||
|
||||
test("isQuotaExhaustedForRequest treats near-zero remaining as exhausted at default threshold", () => {
|
||||
test("isQuotaExhaustedForRequest keeps reported positive remaining available", () => {
|
||||
const connectionId = "conn-near-zero-test";
|
||||
quotaCache.setQuotaCache(connectionId, "antigravity", {
|
||||
"gemini-3.7-flash-medium": { remainingPercentage: 0.00000167, resetAt: null },
|
||||
@@ -149,8 +149,8 @@ test("isQuotaExhaustedForRequest treats near-zero remaining as exhausted at defa
|
||||
"antigravity",
|
||||
"antigravity/gemini-3.7-flash-medium"
|
||||
),
|
||||
true,
|
||||
"effectively-zero remaining should count as exhausted"
|
||||
false,
|
||||
"positive quota is not exhaustion; explicit usage cutoffs are evaluated separately"
|
||||
);
|
||||
});
|
||||
|
||||
|
||||
267
tests/unit/combo/quota-connection-eligibility.test.ts
Normal file
267
tests/unit/combo/quota-connection-eligibility.test.ts
Normal file
@@ -0,0 +1,267 @@
|
||||
import test, { after } from "node:test";
|
||||
import assert from "node:assert/strict";
|
||||
import fs from "node:fs";
|
||||
import os from "node:os";
|
||||
import path from "node:path";
|
||||
import { randomUUID } from "node:crypto";
|
||||
import http from "node:http";
|
||||
|
||||
const dataDir = fs.mkdtempSync(path.join(os.tmpdir(), "quota-eligibility-"));
|
||||
process.env.DATA_DIR = dataDir;
|
||||
const db = await import("../../../src/lib/db/providers.ts");
|
||||
const core = await import("../../../src/lib/db/core.ts");
|
||||
const { registerQuotaFetcher } = await import("../../../open-sse/services/quotaPreflight.ts");
|
||||
const { orderTargetsByResetAwareQuota, orderTargetsByQuotaWeighted } =
|
||||
await import("../../../open-sse/services/combo/quotaStrategies.ts");
|
||||
const { handleComboChat } = await import("../../../open-sse/services/combo.ts");
|
||||
const log = { info() {}, warn() {}, debug() {}, error() {} };
|
||||
const quota = { used: 20, total: 100, percentUsed: 0.2, limitReached: false };
|
||||
|
||||
after(() => {
|
||||
core.resetDbInstance();
|
||||
fs.rmSync(dataDir, { recursive: true, force: true });
|
||||
});
|
||||
|
||||
function target(provider: string, connectionId: string | null, allowedConnectionIds?: string[]) {
|
||||
return {
|
||||
kind: "model" as const,
|
||||
stepId: randomUUID(),
|
||||
executionKey: randomUUID(),
|
||||
modelStr: `${provider}/test-model`,
|
||||
provider,
|
||||
providerId: provider,
|
||||
connectionId,
|
||||
allowedConnectionIds,
|
||||
weight: 1,
|
||||
label: null,
|
||||
};
|
||||
}
|
||||
|
||||
async function fixture() {
|
||||
const provider = `eligibility-${randomUUID()}`;
|
||||
const rows = [];
|
||||
for (const [name, isActive, testStatus] of [
|
||||
["healthy", true, "active"],
|
||||
["disabled", false, "error"],
|
||||
["banned", true, "banned"],
|
||||
["transient", true, "error"],
|
||||
] as const) {
|
||||
rows.push(
|
||||
await db.createProviderConnection({
|
||||
provider,
|
||||
name,
|
||||
isActive,
|
||||
testStatus,
|
||||
authType: "apikey",
|
||||
})
|
||||
);
|
||||
}
|
||||
const [healthy, disabled, banned, transient] = rows;
|
||||
const fetched: string[] = [];
|
||||
registerQuotaFetcher(provider, async (id) => {
|
||||
fetched.push(id);
|
||||
return quota;
|
||||
});
|
||||
return { provider, healthy, disabled, banned, transient, fetched };
|
||||
}
|
||||
|
||||
for (const [strategy, order] of [
|
||||
["reset-aware", orderTargetsByResetAwareQuota],
|
||||
["quota-weighted", orderTargetsByQuotaWeighted],
|
||||
] as const) {
|
||||
for (const mode of ["pinned", "allowlisted", "expanded"] as const) {
|
||||
test(`${strategy}: ${mode} excludes ineligible IDs before quota workers`, async () => {
|
||||
const f = await fixture();
|
||||
const ids = [f.healthy.id, f.disabled.id, f.banned.id, f.transient.id, randomUUID()];
|
||||
const targets =
|
||||
mode === "pinned"
|
||||
? ids.map((id) => target(f.provider, id))
|
||||
: [target(f.provider, null, mode === "allowlisted" ? ids : undefined)];
|
||||
const ordered = await order(targets, randomUUID(), {}, log, ids);
|
||||
const eligible = [f.healthy.id, f.transient.id].sort();
|
||||
assert.deepEqual(
|
||||
[...f.fetched].sort(),
|
||||
eligible,
|
||||
"no quota calls for disabled, banned, or missing IDs"
|
||||
);
|
||||
assert.deepEqual(ordered.map((t) => t.connectionId).sort(), eligible);
|
||||
});
|
||||
}
|
||||
test(`${strategy}: API-key allowlist cannot admit disabled or banned pins`, async () => {
|
||||
const f = await fixture();
|
||||
const ids = [f.disabled.id, f.banned.id, f.healthy.id];
|
||||
const ordered = await order(
|
||||
[...ids, f.transient.id].map((id) => target(f.provider, id)),
|
||||
randomUUID(),
|
||||
{},
|
||||
log,
|
||||
ids
|
||||
);
|
||||
assert.deepEqual(f.fetched, [f.healthy.id]);
|
||||
assert.deepEqual(
|
||||
ordered.map((t) => t.connectionId),
|
||||
[f.healthy.id]
|
||||
);
|
||||
});
|
||||
test(`${strategy}: API-key allowlist that matches no eligible row does not fall back`, async () => {
|
||||
const f = await fixture();
|
||||
const emptyPool = [f.disabled.id, f.banned.id];
|
||||
const ordered = await order([target(f.provider, null)], randomUUID(), {}, log, emptyPool);
|
||||
assert.deepEqual(ordered, []);
|
||||
assert.deepEqual(f.fetched, []);
|
||||
});
|
||||
test(`${strategy}: a pin cannot borrow another provider's eligible row`, async () => {
|
||||
const first = await fixture();
|
||||
const second = await fixture();
|
||||
const ordered = await order(
|
||||
[target(second.provider, second.healthy.id), target(first.provider, second.healthy.id)],
|
||||
randomUUID(),
|
||||
{},
|
||||
log
|
||||
);
|
||||
assert.deepEqual(first.fetched, []);
|
||||
assert.deepEqual(second.fetched, [second.healthy.id]);
|
||||
assert.equal(ordered.length, 1);
|
||||
assert.equal(ordered[0].provider, second.provider);
|
||||
});
|
||||
test(`${strategy}: disabling all connections prevents provider fallback`, async () => {
|
||||
const f = await fixture();
|
||||
await db.updateProviderConnection(f.healthy.id, { isActive: false });
|
||||
await db.updateProviderConnection(f.banned.id, { isActive: false });
|
||||
await db.updateProviderConnection(f.transient.id, { isActive: false });
|
||||
const ordered = await order([target(f.provider, null)], randomUUID(), {}, log);
|
||||
assert.deepEqual(ordered, []);
|
||||
assert.deepEqual(f.fetched, []);
|
||||
});
|
||||
test(`${strategy}: retry reloads eligibility after disable and ban`, async () => {
|
||||
const f = await fixture();
|
||||
const targets = [f.healthy, f.transient].map((r) => target(f.provider, r.id));
|
||||
await order(targets, randomUUID(), {}, log);
|
||||
await db.updateProviderConnection(f.healthy.id, { isActive: false });
|
||||
await db.updateProviderConnection(f.transient.id, { testStatus: "banned" });
|
||||
f.fetched.length = 0;
|
||||
const ordered = await order(targets, randomUUID(), {}, log);
|
||||
assert.deepEqual(ordered, [], "cached active rows must not re-enter retry pool");
|
||||
assert.deepEqual(f.fetched, []);
|
||||
});
|
||||
}
|
||||
|
||||
test(
|
||||
"real combo routing sends only eligible pins to local HTTP upstream",
|
||||
{ timeout: 15_000 },
|
||||
async (t) => {
|
||||
const f = await fixture();
|
||||
const dispatched: string[] = [];
|
||||
const received: string[] = [];
|
||||
const server = http.createServer((req, res) => {
|
||||
received.push(String(req.headers["x-connection-id"]));
|
||||
res.setHeader("Content-Type", "application/json");
|
||||
res.end(
|
||||
JSON.stringify({ choices: [{ message: { role: "assistant", content: "healthy reply" } }] })
|
||||
);
|
||||
});
|
||||
t.after(() => server.closeAllConnections());
|
||||
await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", resolve));
|
||||
try {
|
||||
const address = server.address() as { port: number };
|
||||
const response = await handleComboChat({
|
||||
body: { stream: false, messages: [{ role: "user", content: "test" }] },
|
||||
combo: {
|
||||
name: randomUUID(),
|
||||
strategy: "reset-aware",
|
||||
config: { disableSessionStickiness: true, maxRetries: 0 },
|
||||
models: [f.disabled, f.banned, f.healthy].map((r) => ({
|
||||
model: `${f.provider}/test-model`,
|
||||
providerId: f.provider,
|
||||
connectionId: r.id,
|
||||
})),
|
||||
},
|
||||
settings: {},
|
||||
allCombos: [],
|
||||
log,
|
||||
handleSingleModel: async (_body, _model, options) => {
|
||||
assert.ok(options && "connectionId" in options && options.connectionId);
|
||||
dispatched.push(options.connectionId);
|
||||
const upstream = await fetch(`http://127.0.0.1:${address.port}`, {
|
||||
headers: { "x-connection-id": options.connectionId },
|
||||
});
|
||||
return new Response(upstream.body, {
|
||||
status: upstream.status,
|
||||
headers: upstream.headers,
|
||||
});
|
||||
},
|
||||
});
|
||||
assert.equal(response.status, 200);
|
||||
assert.equal((await response.json()).choices[0].message.content, "healthy reply");
|
||||
assert.deepEqual(f.fetched, [f.healthy.id]);
|
||||
assert.deepEqual(dispatched, [f.healthy.id]);
|
||||
assert.deepEqual(received, [f.healthy.id]);
|
||||
} finally {
|
||||
server.closeAllConnections();
|
||||
await new Promise<void>((resolve) => server.close(() => resolve()));
|
||||
}
|
||||
}
|
||||
);
|
||||
|
||||
test(
|
||||
"real combo retries local upstream 503 without dispatching ineligible accounts",
|
||||
{ timeout: 15_000 },
|
||||
async (t) => {
|
||||
const f = await fixture();
|
||||
const dispatched: string[] = [];
|
||||
const received: string[] = [];
|
||||
const server = http.createServer((req, res) => {
|
||||
received.push(String(req.headers["x-connection-id"]));
|
||||
res.setHeader("Content-Type", "application/json");
|
||||
if (received.length === 1) {
|
||||
res.statusCode = 503;
|
||||
res.end(JSON.stringify({ error: { message: "upstream temporarily unavailable" } }));
|
||||
return;
|
||||
}
|
||||
res.end(
|
||||
JSON.stringify({ choices: [{ message: { role: "assistant", content: "healthy reply" } }] })
|
||||
);
|
||||
});
|
||||
t.after(() => server.closeAllConnections());
|
||||
await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", resolve));
|
||||
try {
|
||||
const address = server.address() as { port: number };
|
||||
const response = await handleComboChat({
|
||||
body: { stream: false, messages: [{ role: "user", content: "test" }] },
|
||||
combo: {
|
||||
name: randomUUID(),
|
||||
strategy: "reset-aware",
|
||||
config: { disableSessionStickiness: true, maxRetries: 0 },
|
||||
models: [f.disabled, f.banned, f.healthy, f.transient].map((r) => ({
|
||||
model: `${f.provider}/test-model`,
|
||||
providerId: f.provider,
|
||||
connectionId: r.id,
|
||||
})),
|
||||
},
|
||||
settings: {},
|
||||
allCombos: [],
|
||||
log,
|
||||
handleSingleModel: async (_body, _model, options) => {
|
||||
assert.ok(options && "connectionId" in options && options.connectionId);
|
||||
dispatched.push(options.connectionId);
|
||||
const upstream = await fetch(`http://127.0.0.1:${address.port}`, {
|
||||
headers: { "x-connection-id": options.connectionId },
|
||||
});
|
||||
return new Response(upstream.body, {
|
||||
status: upstream.status,
|
||||
headers: upstream.headers,
|
||||
});
|
||||
},
|
||||
});
|
||||
assert.equal(response.status, 200);
|
||||
assert.equal((await response.json()).choices[0].message.content, "healthy reply");
|
||||
const eligible = [f.healthy.id, f.transient.id].sort();
|
||||
assert.deepEqual([...f.fetched].sort(), eligible);
|
||||
assert.deepEqual([...dispatched].sort(), eligible);
|
||||
assert.deepEqual([...received].sort(), eligible);
|
||||
} finally {
|
||||
server.closeAllConnections();
|
||||
await new Promise<void>((resolve) => server.close(() => resolve()));
|
||||
}
|
||||
}
|
||||
);
|
||||
439
tests/unit/combo/quota-weighted-stale-402.test.ts
Normal file
439
tests/unit/combo/quota-weighted-stale-402.test.ts
Normal file
@@ -0,0 +1,439 @@
|
||||
/**
|
||||
* Two ways an out-of-credit connection kept drawing quota-weighted traffic:
|
||||
*
|
||||
* 1. A 402 from upstream left the stored quota snapshot untouched, so the very
|
||||
* next weighted draw still saw the old non-zero remaining and could pick the
|
||||
* same dead connection again.
|
||||
* 2. A snapshot refreshed hours ago counted as confident headroom. Live incident:
|
||||
* remaining=1%, is_exhausted=0, last refreshed 5h earlier, upstream answered
|
||||
* 402 "Grok Build usage balance exhausted".
|
||||
*
|
||||
* Staleness means "unknown", not "dead": a stale connection drops out of the A
|
||||
* pool but stays reachable through B, and is never deactivated.
|
||||
*/
|
||||
import test, { after, afterEach } from "node:test";
|
||||
import assert from "node:assert/strict";
|
||||
import fs from "node:fs";
|
||||
import os from "node:os";
|
||||
import path from "node:path";
|
||||
import { randomUUID } from "node:crypto";
|
||||
|
||||
const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-qw-stale-402-"));
|
||||
const ORIGINAL_DATA_DIR = process.env.DATA_DIR;
|
||||
process.env.DATA_DIR = TEST_DATA_DIR;
|
||||
|
||||
const dbCore = await import("../../../src/lib/db/core.ts");
|
||||
const db = await import("../../../src/lib/db/providers.ts");
|
||||
const quotaCache = await import("../../../src/domain/quotaCache.ts");
|
||||
const { registerQuotaFetcher } = await import("../../../open-sse/services/quotaPreflight.ts");
|
||||
const { orderTargetsByQuotaWeighted, QUOTA_WEIGHTED_MAX_SNAPSHOT_AGE_MS } =
|
||||
await import("../../../open-sse/services/combo/quotaStrategies.ts");
|
||||
const { resetAllCircuitBreakers } = await import("../../../src/shared/utils/circuitBreaker.ts");
|
||||
const { _clearInflightForTest } =
|
||||
await import("../../../open-sse/services/combo/quotaShareInflight.ts");
|
||||
const { _setSecureRandomFloatSource } = await import("../../../src/shared/utils/secureRandom.ts");
|
||||
|
||||
after(() => {
|
||||
dbCore.resetDbInstance();
|
||||
fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true, maxRetries: 5, retryDelay: 100 });
|
||||
if (ORIGINAL_DATA_DIR === undefined) delete process.env.DATA_DIR;
|
||||
else process.env.DATA_DIR = ORIGINAL_DATA_DIR;
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
_setSecureRandomFloatSource(null);
|
||||
quotaCache.__clearForTests();
|
||||
resetAllCircuitBreakers();
|
||||
_clearInflightForTest();
|
||||
});
|
||||
|
||||
const CLOCK_BASE = Date.now();
|
||||
const iso = (ms = 86_400_000) => new Date(CLOCK_BASE + ms).toISOString();
|
||||
|
||||
function quotaAt(percentUsed: number, extra: Record<string, unknown> = {}) {
|
||||
return {
|
||||
used: percentUsed * 100,
|
||||
total: 100,
|
||||
percentUsed,
|
||||
resetAt: iso(7 * 86_400_000),
|
||||
window5h: { percentUsed, resetAt: iso(5 * 3600_000) },
|
||||
window7d: { percentUsed, resetAt: iso(7 * 86_400_000) },
|
||||
limitReached: false,
|
||||
...extra,
|
||||
};
|
||||
}
|
||||
|
||||
function makeTarget(provider: string, connectionId: string, model = "gemini-3.8-flash-high") {
|
||||
return {
|
||||
kind: "model" as const,
|
||||
stepId: `step-${connectionId}`,
|
||||
executionKey: `${provider}/${model}@${connectionId}`,
|
||||
modelStr: `${provider}/${model}`,
|
||||
provider,
|
||||
providerId: provider,
|
||||
connectionId,
|
||||
weight: 1,
|
||||
label: null,
|
||||
};
|
||||
}
|
||||
|
||||
async function seedConnection(provider: string, name: string) {
|
||||
const row = await db.createProviderConnection({
|
||||
provider,
|
||||
name,
|
||||
isActive: true,
|
||||
testStatus: "active",
|
||||
authType: "apikey",
|
||||
});
|
||||
return String(row.id);
|
||||
}
|
||||
|
||||
// ── Hole 1: a 402 must invalidate the snapshot ──────────────────────────────
|
||||
|
||||
test("markAccountExhaustedFromCredits: 402 flips the snapshot to exhausted", () => {
|
||||
const id = `credit-${randomUUID()}`;
|
||||
quotaCache.setQuotaCache(id, "grok-cli", {
|
||||
session: { remainingPercentage: 1, resetAt: iso() },
|
||||
});
|
||||
assert.equal(quotaCache.isAccountQuotaExhausted(id), false, "precondition: has headroom");
|
||||
|
||||
quotaCache.markAccountExhaustedFromCredits(id, "grok-cli");
|
||||
|
||||
assert.equal(quotaCache.isAccountQuotaExhausted(id), true);
|
||||
const entry = quotaCache.getQuotaCache(id);
|
||||
assert.equal(entry?.exhausted, true);
|
||||
assert.equal(
|
||||
quotaCache.getQuotaWeightedRemainingPercent(id),
|
||||
0,
|
||||
"a credit-exhausted connection reports no remaining headroom"
|
||||
);
|
||||
});
|
||||
|
||||
test("a 402-marked connection loses the weighted draw to a healthy peer", async () => {
|
||||
const provider = "agy";
|
||||
const dead = await seedConnection(provider, `dead-${randomUUID()}`);
|
||||
const healthy = await seedConnection(provider, `ok-${randomUUID()}`);
|
||||
// Upstream still reports headroom for the dead account — the stale snapshot
|
||||
// that caused the incident. Only the 402 mark tells the truth.
|
||||
registerQuotaFetcher(provider, async () => quotaAt(0.6));
|
||||
|
||||
quotaCache.setQuotaCache(dead, provider, { session: { remainingPercentage: 1, resetAt: iso() } });
|
||||
quotaCache.markAccountExhaustedFromCredits(dead, provider);
|
||||
|
||||
_setSecureRandomFloatSource(() => 0);
|
||||
const ordered = await orderTargetsByQuotaWeighted(
|
||||
[makeTarget(provider, dead), makeTarget(provider, healthy)],
|
||||
"credit-402",
|
||||
{ quotaWeightedFloorPercent: 1 },
|
||||
{ warn() {} },
|
||||
null
|
||||
);
|
||||
|
||||
assert.equal(ordered[0]?.connectionId, healthy, "402'd connection must not lead the order");
|
||||
});
|
||||
|
||||
test("a 402 mark never deactivates or deletes the connection", () => {
|
||||
const id = `keep-${randomUUID()}`;
|
||||
quotaCache.setQuotaCache(id, "grok-cli", {
|
||||
session: { remainingPercentage: 40, resetAt: iso() },
|
||||
});
|
||||
quotaCache.markAccountExhaustedFromCredits(id, "grok-cli");
|
||||
|
||||
const entry = quotaCache.getQuotaCache(id);
|
||||
assert.ok(entry, "the cache entry survives — a 402 is a credit state, not a dead key");
|
||||
assert.equal(entry?.connectionId, id);
|
||||
assert.equal(entry?.provider, "grok-cli");
|
||||
});
|
||||
|
||||
test("a successful quota refresh clears the 402 mark", () => {
|
||||
const id = `refresh-${randomUUID()}`;
|
||||
quotaCache.markAccountExhaustedFromCredits(id, "grok-cli");
|
||||
assert.equal(quotaCache.isAccountQuotaExhausted(id), true);
|
||||
|
||||
quotaCache.setQuotaCache(id, "grok-cli", {
|
||||
session: { remainingPercentage: 55, resetAt: iso() },
|
||||
});
|
||||
|
||||
assert.equal(
|
||||
quotaCache.isAccountQuotaExhausted(id),
|
||||
false,
|
||||
"upstream saying there is headroom again outranks the earlier 402"
|
||||
);
|
||||
});
|
||||
|
||||
// ── Hole 2: snapshot staleness is bounded ───────────────────────────────────
|
||||
|
||||
test("QUOTA_WEIGHTED_MAX_SNAPSHOT_AGE_MS is exported and shorter than the incident gap", () => {
|
||||
assert.equal(typeof QUOTA_WEIGHTED_MAX_SNAPSHOT_AGE_MS, "number");
|
||||
assert.ok(QUOTA_WEIGHTED_MAX_SNAPSHOT_AGE_MS > 0);
|
||||
assert.ok(
|
||||
QUOTA_WEIGHTED_MAX_SNAPSHOT_AGE_MS < 5 * 60 * 60 * 1000,
|
||||
"the 5h-old snapshot from the incident must not count as confident headroom"
|
||||
);
|
||||
});
|
||||
|
||||
test("a stale snapshot yields the A pool to a freshly-observed peer", async () => {
|
||||
const provider = "agy";
|
||||
const stale = await seedConnection(provider, `stale-${randomUUID()}`);
|
||||
const fresh = await seedConnection(provider, `fresh-${randomUUID()}`);
|
||||
registerQuotaFetcher(provider, async () => quotaAt(0.6));
|
||||
|
||||
quotaCache.setQuotaCache(stale, provider, {
|
||||
session: { remainingPercentage: 90, resetAt: iso() },
|
||||
});
|
||||
const staleEntry = quotaCache.getQuotaCache(stale);
|
||||
assert.ok(staleEntry);
|
||||
// Age the snapshot past the bound. Higher remaining than the fresh peer, so a
|
||||
// pass that ignored staleness would rank it first.
|
||||
staleEntry.fetchedAt = Date.now() - QUOTA_WEIGHTED_MAX_SNAPSHOT_AGE_MS - 60_000;
|
||||
|
||||
quotaCache.setQuotaCache(fresh, provider, {
|
||||
session: { remainingPercentage: 40, resetAt: iso() },
|
||||
});
|
||||
|
||||
_setSecureRandomFloatSource(() => 0);
|
||||
const ordered = await orderTargetsByQuotaWeighted(
|
||||
[makeTarget(provider, stale), makeTarget(provider, fresh)],
|
||||
"stale-vs-fresh",
|
||||
{ quotaWeightedFloorPercent: 1 },
|
||||
{ warn() {} },
|
||||
null
|
||||
);
|
||||
|
||||
assert.equal(ordered[0]?.connectionId, fresh, "a fresh observation outranks a stale one");
|
||||
assert.ok(
|
||||
ordered.some((t) => t.connectionId === stale),
|
||||
"stale means unknown, not dead — it stays reachable behind the fresh peer"
|
||||
);
|
||||
});
|
||||
|
||||
test("a snapshot exactly at the age bound still counts as fresh", async () => {
|
||||
const provider = "agy";
|
||||
const atBound = await seedConnection(provider, `at-bound-${randomUUID()}`);
|
||||
const younger = await seedConnection(provider, `younger-${randomUUID()}`);
|
||||
registerQuotaFetcher(provider, async () => quotaAt(0.6));
|
||||
|
||||
quotaCache.setQuotaCache(atBound, provider, {
|
||||
session: { remainingPercentage: 90, resetAt: iso() },
|
||||
});
|
||||
const boundEntry = quotaCache.getQuotaCache(atBound);
|
||||
assert.ok(boundEntry);
|
||||
// A second inside the bound, not past it. The staleness test is strictly
|
||||
// greater, so this snapshot keeps its A-pool seat and its higher remaining
|
||||
// wins. The second of slack absorbs the clock advancing during the await.
|
||||
boundEntry.fetchedAt = Date.now() - QUOTA_WEIGHTED_MAX_SNAPSHOT_AGE_MS + 1_000;
|
||||
|
||||
quotaCache.setQuotaCache(younger, provider, {
|
||||
session: { remainingPercentage: 40, resetAt: iso() },
|
||||
});
|
||||
|
||||
_setSecureRandomFloatSource(() => 0);
|
||||
const ordered = await orderTargetsByQuotaWeighted(
|
||||
[makeTarget(provider, atBound), makeTarget(provider, younger)],
|
||||
"at-bound",
|
||||
{ quotaWeightedFloorPercent: 1 },
|
||||
{ warn() {} },
|
||||
null
|
||||
);
|
||||
|
||||
assert.equal(
|
||||
ordered[0]?.connectionId,
|
||||
atBound,
|
||||
"a snapshot at exactly the bound has not aged out yet"
|
||||
);
|
||||
});
|
||||
|
||||
test("an all-stale set still routes rather than returning nothing", async () => {
|
||||
const provider = "agy";
|
||||
const a = await seedConnection(provider, `stale-a-${randomUUID()}`);
|
||||
const b = await seedConnection(provider, `stale-b-${randomUUID()}`);
|
||||
registerQuotaFetcher(provider, async () => quotaAt(0.6));
|
||||
|
||||
for (const id of [a, b]) {
|
||||
quotaCache.setQuotaCache(id, provider, {
|
||||
session: { remainingPercentage: 80, resetAt: iso() },
|
||||
});
|
||||
const entry = quotaCache.getQuotaCache(id);
|
||||
assert.ok(entry);
|
||||
entry.fetchedAt = Date.now() - QUOTA_WEIGHTED_MAX_SNAPSHOT_AGE_MS - 60_000;
|
||||
}
|
||||
|
||||
_setSecureRandomFloatSource(() => 0);
|
||||
const ordered = await orderTargetsByQuotaWeighted(
|
||||
[makeTarget(provider, a), makeTarget(provider, b)],
|
||||
"all-stale",
|
||||
{ quotaWeightedFloorPercent: 1 },
|
||||
{ warn() {} },
|
||||
null
|
||||
);
|
||||
|
||||
assert.equal(ordered.length, 2, "staleness must not empty the routing set");
|
||||
});
|
||||
|
||||
test("a fresh snapshot is unaffected by the staleness bound", async () => {
|
||||
const provider = "agy";
|
||||
const high = await seedConnection(provider, `high-${randomUUID()}`);
|
||||
const low = await seedConnection(provider, `low-${randomUUID()}`);
|
||||
registerQuotaFetcher(provider, async () => quotaAt(0.6));
|
||||
|
||||
quotaCache.setQuotaCache(high, provider, {
|
||||
session: { remainingPercentage: 90, resetAt: iso() },
|
||||
});
|
||||
quotaCache.setQuotaCache(low, provider, {
|
||||
session: { remainingPercentage: 20, resetAt: iso() },
|
||||
});
|
||||
|
||||
_setSecureRandomFloatSource(() => 0);
|
||||
const ordered = await orderTargetsByQuotaWeighted(
|
||||
[makeTarget(provider, low), makeTarget(provider, high)],
|
||||
"both-fresh",
|
||||
{ quotaWeightedFloorPercent: 1 },
|
||||
{ warn() {} },
|
||||
null
|
||||
);
|
||||
|
||||
assert.equal(ordered.length, 2);
|
||||
assert.ok(
|
||||
ordered.some((t) => t.connectionId === high),
|
||||
"both fresh connections remain eligible"
|
||||
);
|
||||
});
|
||||
|
||||
// ── The 402 mark is wired into the attempt path, not just available ─────────
|
||||
//
|
||||
// Calling the helper directly cannot prove the call site exists: with the
|
||||
// executeTargetAttempt hook deleted, every direct-call assertion above still
|
||||
// passes. This drives a real 402 through the attempt loop instead.
|
||||
|
||||
function credits402(): Response {
|
||||
return new Response(
|
||||
JSON.stringify({ error: { message: "Grok Build usage balance exhausted" } }),
|
||||
{ status: 402, headers: { "content-type": "application/json" } }
|
||||
);
|
||||
}
|
||||
|
||||
function attemptState(target: Record<string, unknown>) {
|
||||
return {
|
||||
orderedTargets: [target],
|
||||
fallbackCount: 0,
|
||||
recordedAttempts: 0,
|
||||
comboErrors: [],
|
||||
lastError: null,
|
||||
lastStatus: null,
|
||||
earliestRetryAfter: null,
|
||||
comboExpired: false,
|
||||
exhaustedProviders: new Set(),
|
||||
exhaustedConnections: new Set(),
|
||||
transientRateLimitedProviders: new Set(),
|
||||
abortControllers: new Map([[0, new AbortController()]]),
|
||||
dispatchedTargets: new Set(),
|
||||
targetFailureTrust: new Map(),
|
||||
comboAttemptOrder: [],
|
||||
skippedForCircuitOpen: false,
|
||||
earliestCircuitOpenRetryMs: 0,
|
||||
globalAttempts: 0,
|
||||
observedFailure: false,
|
||||
allObservedFailuresQuota: true,
|
||||
observeFailure() {},
|
||||
};
|
||||
}
|
||||
|
||||
function attemptDeps(response: () => Response) {
|
||||
return {
|
||||
strategy: "quota-weighted",
|
||||
combo: { name: "t", models: [] },
|
||||
config: {},
|
||||
log: { info() {}, warn() {}, debug() {}, error() {} },
|
||||
settings: null,
|
||||
resilienceSettings: { providerCooldown: { enabled: false } },
|
||||
sticky: { targets: [], messageHash: null, stuck: false },
|
||||
effectiveSessionId: null,
|
||||
preScreenMap: new Map(),
|
||||
quotaCutoffResetWindowConfig: {},
|
||||
maxRetries: 0,
|
||||
traceInvocationId: "inv-402",
|
||||
clientRequestedStream: false,
|
||||
handleSingleModelWithTimeout: async () => response(),
|
||||
body: { messages: [{ role: "user", content: "hi" }] },
|
||||
startTime: Date.now(),
|
||||
releaseStickyPinOnFailure() {},
|
||||
clearStaleLKGP() {},
|
||||
};
|
||||
}
|
||||
|
||||
test("a 402 through the attempt path marks the connection exhausted", async () => {
|
||||
const { executeTargetAttempt } =
|
||||
await import("../../../open-sse/services/combo/executeTargetAttempt.ts");
|
||||
const connectionId = `attempt-${randomUUID()}`;
|
||||
quotaCache.setQuotaCache(connectionId, "grok-cli", {
|
||||
session: { remainingPercentage: 1, resetAt: iso() },
|
||||
});
|
||||
assert.equal(
|
||||
quotaCache.isAccountQuotaExhausted(connectionId),
|
||||
false,
|
||||
"precondition: the stale snapshot still claims headroom"
|
||||
);
|
||||
|
||||
const target = {
|
||||
kind: "model" as const,
|
||||
stepId: "s1",
|
||||
executionKey: `grok-cli/grok@${connectionId}`,
|
||||
modelStr: "grok-cli/grok",
|
||||
provider: "grok-cli",
|
||||
providerId: null,
|
||||
connectionId,
|
||||
weight: 1,
|
||||
label: null,
|
||||
};
|
||||
|
||||
await executeTargetAttempt({
|
||||
index: 0,
|
||||
state: attemptState(target) as never,
|
||||
deps: attemptDeps(credits402) as never,
|
||||
targetForAttempt: target as never,
|
||||
profile: {},
|
||||
protectedPriorityTarget: false,
|
||||
});
|
||||
|
||||
assert.equal(
|
||||
quotaCache.isAccountQuotaExhausted(connectionId),
|
||||
true,
|
||||
"the 402 must invalidate the snapshot from inside the attempt path"
|
||||
);
|
||||
});
|
||||
|
||||
test("a non-credit failure through the attempt path leaves the snapshot alone", async () => {
|
||||
const { executeTargetAttempt } =
|
||||
await import("../../../open-sse/services/combo/executeTargetAttempt.ts");
|
||||
const connectionId = `attempt-500-${randomUUID()}`;
|
||||
quotaCache.setQuotaCache(connectionId, "grok-cli", {
|
||||
session: { remainingPercentage: 60, resetAt: iso() },
|
||||
});
|
||||
|
||||
const target = {
|
||||
kind: "model" as const,
|
||||
stepId: "s1",
|
||||
executionKey: `grok-cli/grok@${connectionId}`,
|
||||
modelStr: "grok-cli/grok",
|
||||
provider: "grok-cli",
|
||||
providerId: null,
|
||||
connectionId,
|
||||
weight: 1,
|
||||
label: null,
|
||||
};
|
||||
|
||||
await executeTargetAttempt({
|
||||
index: 0,
|
||||
state: attemptState(target) as never,
|
||||
deps: attemptDeps(() => new Response("boom", { status: 500 })) as never,
|
||||
targetForAttempt: target as never,
|
||||
profile: {},
|
||||
protectedPriorityTarget: false,
|
||||
});
|
||||
|
||||
assert.equal(
|
||||
quotaCache.isAccountQuotaExhausted(connectionId),
|
||||
false,
|
||||
"a 500 is not a credit signal — headroom must survive it"
|
||||
);
|
||||
});
|
||||
@@ -2,6 +2,7 @@
|
||||
* quota-weighted: skip empty accounts, weighted-draw the rest.
|
||||
* Spec: _tasks/superpowers/specs/2026-09-04-quota-weighted-routing-design.md
|
||||
*/
|
||||
import { resolveProviderId } from "../../../src/shared/constants/providers.ts";
|
||||
import test, { after, afterEach } from "node:test";
|
||||
import assert from "node:assert/strict";
|
||||
import fs from "node:fs";
|
||||
@@ -14,16 +15,15 @@ const ORIGINAL_DATA_DIR = process.env.DATA_DIR;
|
||||
process.env.DATA_DIR = TEST_DATA_DIR;
|
||||
|
||||
const dbCore = await import("../../../src/lib/db/core.ts");
|
||||
const { invalidateDbCache } = await import("../../../src/lib/db/readCache.ts");
|
||||
const quotaCache = await import("../../../src/domain/quotaCache.ts");
|
||||
const { getResetAwareRemainingPercent, resolveResetAwareConfig, scoreResetAwareQuota } =
|
||||
await import("../../../open-sse/services/combo/quotaScoring.ts");
|
||||
const { registerQuotaFetcher } = await import("../../../open-sse/services/quotaPreflight.ts");
|
||||
const { convertUsageToQuotaInfo } = await import("../../../open-sse/services/genericQuotaFetcher.ts");
|
||||
const {
|
||||
expandTargetsByQuotaAwareConnections,
|
||||
orderTargetsByQuotaWeighted,
|
||||
pickWeightedIndex,
|
||||
} = await import("../../../open-sse/services/combo/quotaStrategies.ts");
|
||||
const { convertUsageToQuotaInfo } =
|
||||
await import("../../../open-sse/services/genericQuotaFetcher.ts");
|
||||
const { expandTargetsByQuotaAwareConnections, orderTargetsByQuotaWeighted, pickWeightedIndex } =
|
||||
await import("../../../open-sse/services/combo/quotaStrategies.ts");
|
||||
const { getCircuitBreaker, resetAllCircuitBreakers } =
|
||||
await import("../../../src/shared/utils/circuitBreaker.ts");
|
||||
const { applyStrategyOrdering } =
|
||||
@@ -44,9 +44,7 @@ const { HANDLED_COMBO_STRATEGIES } =
|
||||
await import("../../../open-sse/services/combo/strategyDispatch.ts");
|
||||
const { comboStrategySchema } = await import("../../../src/shared/validation/schemas.ts");
|
||||
const { _setSecureRandomFloatSource } = await import("../../../src/shared/utils/secureRandom.ts");
|
||||
const { getQuotaFetchScope } = await import(
|
||||
"../../../open-sse/services/antigravityQuotaFamily.ts"
|
||||
);
|
||||
const { getQuotaFetchScope } = await import("../../../open-sse/services/antigravityQuotaFamily.ts");
|
||||
|
||||
after(() => {
|
||||
dbCore.resetDbInstance();
|
||||
@@ -88,7 +86,18 @@ function quotaAt(percentUsed: number, extra: Record<string, unknown> = {}) {
|
||||
};
|
||||
}
|
||||
|
||||
function seedConnection(provider: string, connectionId: string) {
|
||||
dbCore
|
||||
.getDbInstance()
|
||||
.prepare(
|
||||
"INSERT OR IGNORE INTO provider_connections (id, provider, is_active, test_status, created_at, updated_at) VALUES (?, ?, 1, 'active', '2026-09-09T00:00:00Z', '2026-09-09T00:00:00Z')"
|
||||
)
|
||||
.run(connectionId, resolveProviderId(provider));
|
||||
invalidateDbCache("connections");
|
||||
}
|
||||
|
||||
function makeTarget(provider: string, connectionId: string, model = "gemini-3.8-flash-high") {
|
||||
seedConnection(provider, connectionId);
|
||||
return {
|
||||
kind: "model" as const,
|
||||
stepId: `step-${connectionId}`,
|
||||
@@ -133,7 +142,7 @@ test("getResetAwareRemainingPercent: missing windows fall back to overall percen
|
||||
assert.equal(getResetAwareRemainingPercent({ percentUsed: 0.7 }), 30);
|
||||
});
|
||||
|
||||
test("dual: default expand drops 0.5% agy via 99% kick; skipExhaustionFilter keeps it", async () => {
|
||||
test("dual: default expansion and skipExhaustionFilter preserve positive quota", async () => {
|
||||
const provider = "agy";
|
||||
const low = `low-${randomUUID()}`;
|
||||
const healthy = `ok-${randomUUID()}`;
|
||||
@@ -152,8 +161,8 @@ test("dual: default expand drops 0.5% agy via 99% kick; skipExhaustionFilter kee
|
||||
);
|
||||
assert.equal(
|
||||
dropped.expandedTargets.some((t) => t.connectionId === low),
|
||||
false,
|
||||
"0.5% remaining must be treated as exhausted by the 99% dashboard kick"
|
||||
true,
|
||||
"positive remaining quota must not trigger automatic exhaustion"
|
||||
);
|
||||
assert.equal(
|
||||
dropped.expandedTargets.some((t) => t.connectionId === healthy),
|
||||
@@ -211,7 +220,10 @@ test("A/B isolation: 7 hard-empty + 2 at 0.5% + 1 at 40%, floor=1", async () =>
|
||||
low
|
||||
);
|
||||
for (const id of dead) {
|
||||
assert.equal(ordered.some((t) => t.connectionId === id), false);
|
||||
assert.equal(
|
||||
ordered.some((t) => t.connectionId === id),
|
||||
false
|
||||
);
|
||||
}
|
||||
});
|
||||
|
||||
@@ -231,8 +243,16 @@ test("7 empty + 3 healthy → length 3, no hard-empty", async () => {
|
||||
null
|
||||
);
|
||||
assert.equal(ordered.length, 3);
|
||||
for (const id of dead) assert.equal(ordered.some((t) => t.connectionId === id), false);
|
||||
for (const id of ok) assert.equal(ordered.some((t) => t.connectionId === id), true);
|
||||
for (const id of dead)
|
||||
assert.equal(
|
||||
ordered.some((t) => t.connectionId === id),
|
||||
false
|
||||
);
|
||||
for (const id of ok)
|
||||
assert.equal(
|
||||
ordered.some((t) => t.connectionId === id),
|
||||
true
|
||||
);
|
||||
});
|
||||
|
||||
test("pickWeightedIndex skips non-positive weights", () => {
|
||||
@@ -580,6 +600,7 @@ test("applyStrategyOrdering(quota-weighted) uses the orderer", async () => {
|
||||
const pipelineLog = { info() {}, warn() {}, error() {}, debug() {} };
|
||||
|
||||
function pinComboModels(provider, model, connectionIds) {
|
||||
connectionIds.forEach((id) => seedConnection(provider, id));
|
||||
return connectionIds.map((connectionId, index) => ({
|
||||
kind: "model",
|
||||
provider,
|
||||
@@ -857,8 +878,16 @@ test("three hard-empty of ten never win the first draw", async () => {
|
||||
);
|
||||
assert.equal(ordered.length, 7);
|
||||
assert.equal(dead.includes(ordered[0]?.connectionId ?? ""), false);
|
||||
for (const id of dead) assert.equal(ordered.some((t) => t.connectionId === id), false);
|
||||
for (const id of ok) assert.equal(ordered.some((t) => t.connectionId === id), true);
|
||||
for (const id of dead)
|
||||
assert.equal(
|
||||
ordered.some((t) => t.connectionId === id),
|
||||
false
|
||||
);
|
||||
for (const id of ok)
|
||||
assert.equal(
|
||||
ordered.some((t) => t.connectionId === id),
|
||||
true
|
||||
);
|
||||
});
|
||||
|
||||
test("quota-weighted Gemini keeps the account when only Claude weekly is empty", async () => {
|
||||
@@ -1059,7 +1088,11 @@ test("quota-share sticky pin transfers the inflight slot to the pinned account",
|
||||
if ("earlyResponse" in result) return;
|
||||
assert.equal(result.sticky.stuck, true);
|
||||
assert.equal(result.orderedTargets[0]?.connectionId, pinned);
|
||||
assert.equal(getInflight(drawn), 0, "drawn account must drop the slot after stickiness moves [0]");
|
||||
assert.equal(
|
||||
getInflight(drawn),
|
||||
0,
|
||||
"drawn account must drop the slot after stickiness moves [0]"
|
||||
);
|
||||
assert.equal(getInflight(pinned), 1, "pinned account must hold the transferred slot");
|
||||
result.quotaShareRelease?.();
|
||||
assert.equal(getInflight(pinned), 0);
|
||||
|
||||
@@ -200,8 +200,19 @@ test("#9330 canonically named windows keep their existing resolution (no regress
|
||||
test("#9330 orderTargetsByResetWindow dispatches the soonest-resetting account first", async () => {
|
||||
const antigravity = `agy-9330-${randomUUID()}`;
|
||||
const codex = `codex-9330-${randomUUID()}`;
|
||||
const antigravityConnection = `agy-conn-${randomUUID()}`;
|
||||
const codexConnection = `codex-conn-${randomUUID()}`;
|
||||
const { createProviderConnection } = await import("../../../src/lib/db/providers.ts");
|
||||
const { id: antigravityConnection } = (await createProviderConnection({
|
||||
provider: antigravity,
|
||||
authType: "oauth",
|
||||
isActive: true,
|
||||
testStatus: "active",
|
||||
})) as { id: string };
|
||||
const { id: codexConnection } = (await createProviderConnection({
|
||||
provider: codex,
|
||||
authType: "oauth",
|
||||
isActive: true,
|
||||
testStatus: "active",
|
||||
})) as { id: string };
|
||||
|
||||
registerQuotaFetcher(antigravity, async () => antigravityQuotaFresh);
|
||||
registerQuotaFetcher(codex, async () => codexQuota26Days);
|
||||
|
||||
Reference in New Issue
Block a user