mirror of
https://github.com/diegosouzapw/OmniRoute.git
synced 2026-07-09 17:28:51 +00:00
fix(provider-limits): clear transient rate-limit state when quota recovers (#6128)
Clear transient rate-limit state when quota recovers. Integrated into release/v3.8.44.
This commit is contained in:
parent
8cded07233
commit
716041223e
2 changed files with 295 additions and 4 deletions
|
|
@ -14,6 +14,7 @@ import {
|
|||
import { syncToCloud } from "@/lib/cloudSync";
|
||||
import { setQuotaCache } from "@/domain/quotaCache";
|
||||
import { buildClaudeExtraUsageConnectionUpdate } from "@/lib/providers/claudeExtraUsage";
|
||||
import { clearRecoveredProviderState } from "@/sse/services/auth";
|
||||
import { getMachineId } from "@/shared/utils/machine";
|
||||
import { USAGE_SUPPORTED_PROVIDERS } from "@/shared/constants/providers";
|
||||
import { getExecutor } from "@omniroute/open-sse/executors/index.ts";
|
||||
|
|
@ -392,9 +393,87 @@ export function quotaPathShouldMarkExpired(
|
|||
return true;
|
||||
}
|
||||
|
||||
async function syncExpiredStatusIfNeeded(connection: ProviderConnectionLike, usage: JsonRecord) {
|
||||
const TERMINAL_STATUSES_FOR_QUOTA_RECOVERY = new Set([
|
||||
"credits_exhausted",
|
||||
"banned",
|
||||
"expired",
|
||||
"deactivated",
|
||||
]);
|
||||
|
||||
function isTerminalStatusForQuotaRecovery(testStatus: string | null | undefined): boolean {
|
||||
if (!testStatus) return false;
|
||||
return TERMINAL_STATUSES_FOR_QUOTA_RECOVERY.has(testStatus);
|
||||
}
|
||||
|
||||
export function hasUsableQuota(usage: JsonRecord): boolean {
|
||||
const quotas = usage?.quotas;
|
||||
if (!isRecord(quotas)) return false;
|
||||
for (const value of Object.values(quotas)) {
|
||||
if (!isRecord(value)) continue;
|
||||
if (value.unlimited === true) return true;
|
||||
const remaining =
|
||||
typeof value.remaining === "number"
|
||||
? value.remaining
|
||||
: typeof value.remainingPercentage === "number"
|
||||
? value.remainingPercentage
|
||||
: null;
|
||||
if (remaining !== null && remaining > 0) return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
export async function maybeClearRecoveredQuotaState(
|
||||
connection: ProviderConnectionLike,
|
||||
usage: JsonRecord
|
||||
): Promise<ProviderConnectionLike> {
|
||||
if (!hasUsableQuota(usage)) return connection;
|
||||
if (isTerminalStatusForQuotaRecovery(connection.testStatus)) return connection;
|
||||
|
||||
const hasTransientState =
|
||||
connection.testStatus === "unavailable" ||
|
||||
Boolean(connection.rateLimitedUntil) ||
|
||||
Boolean(connection.lastError) ||
|
||||
Boolean(connection.errorCode) ||
|
||||
Boolean(connection.lastErrorType) ||
|
||||
Boolean(connection.lastErrorSource) ||
|
||||
(connection.backoffLevel ?? 0) > 0;
|
||||
|
||||
if (!hasTransientState) return connection;
|
||||
|
||||
try {
|
||||
await clearRecoveredProviderState({
|
||||
connectionId: connection.id,
|
||||
testStatus: connection.testStatus,
|
||||
lastError: connection.lastError ?? null,
|
||||
rateLimitedUntil: connection.rateLimitedUntil ?? null,
|
||||
errorCode: connection.errorCode ?? null,
|
||||
lastErrorType: connection.lastErrorType ?? null,
|
||||
lastErrorSource: connection.lastErrorSource ?? null,
|
||||
});
|
||||
} catch (dbError) {
|
||||
console.warn("[ProviderLimits] Failed to clear recovered quota state:", dbError);
|
||||
return connection;
|
||||
}
|
||||
|
||||
return {
|
||||
...connection,
|
||||
testStatus: "active",
|
||||
lastError: null,
|
||||
lastErrorAt: null,
|
||||
lastErrorType: null,
|
||||
lastErrorSource: null,
|
||||
errorCode: null,
|
||||
rateLimitedUntil: null,
|
||||
backoffLevel: 0,
|
||||
};
|
||||
}
|
||||
|
||||
async function syncExpiredStatusIfNeeded(
|
||||
connection: ProviderConnectionLike,
|
||||
usage: JsonRecord
|
||||
): Promise<ProviderConnectionLike> {
|
||||
if (!quotaPathShouldMarkExpired(connection.provider, usage.message, connection.testStatus)) {
|
||||
return;
|
||||
return connection;
|
||||
}
|
||||
|
||||
try {
|
||||
|
|
@ -405,7 +484,14 @@ async function syncExpiredStatusIfNeeded(connection: ProviderConnectionLike, usa
|
|||
});
|
||||
} catch (dbError) {
|
||||
console.error("[ProviderLimits] Failed to sync expired status to DB:", dbError);
|
||||
return connection;
|
||||
}
|
||||
|
||||
return {
|
||||
...connection,
|
||||
testStatus: "expired",
|
||||
lastErrorType: "token_expired",
|
||||
};
|
||||
}
|
||||
|
||||
async function syncClaudeExtraUsageStateIfNeeded(
|
||||
|
|
@ -626,10 +712,11 @@ async function fetchLiveProviderLimitsWithOptions(
|
|||
if (isRecord(usage.quotas)) {
|
||||
setQuotaCache(connectionId, connection.provider, usage.quotas);
|
||||
}
|
||||
await syncExpiredStatusIfNeeded(connection, usage);
|
||||
connection = await syncExpiredStatusIfNeeded(connection, usage);
|
||||
connection = await syncClaudeExtraUsageStateIfNeeded(connection, usage);
|
||||
connection = await syncClaudeBootstrapIfNeeded(connection, usage);
|
||||
connection = await syncAntigravitySubscriptionIfNeeded(connection, usage);
|
||||
connection = await maybeClearRecoveredQuotaState(connection, usage);
|
||||
return { connection, usage };
|
||||
}
|
||||
|
||||
|
|
@ -740,10 +827,11 @@ async function fetchLiveProviderLimitsWithOptions(
|
|||
if (isRecord(result.usage.quotas)) {
|
||||
setQuotaCache(connectionId, connection.provider, result.usage.quotas);
|
||||
}
|
||||
await syncExpiredStatusIfNeeded(connection, result.usage);
|
||||
connection = await syncExpiredStatusIfNeeded(connection, result.usage);
|
||||
connection = await syncClaudeExtraUsageStateIfNeeded(connection, result.usage);
|
||||
connection = await syncClaudeBootstrapIfNeeded(connection, result.usage);
|
||||
connection = await syncAntigravitySubscriptionIfNeeded(connection, result.usage);
|
||||
connection = await maybeClearRecoveredQuotaState(connection, result.usage);
|
||||
|
||||
return {
|
||||
connection,
|
||||
|
|
|
|||
203
tests/unit/provider-limits-recovery.test.ts
Normal file
203
tests/unit/provider-limits-recovery.test.ts
Normal file
|
|
@ -0,0 +1,203 @@
|
|||
import test from "node:test";
|
||||
import assert from "node:assert/strict";
|
||||
import fs from "node:fs";
|
||||
import os from "node:os";
|
||||
import path from "node:path";
|
||||
|
||||
const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-provider-limits-recovery-"));
|
||||
process.env.DATA_DIR = TEST_DATA_DIR;
|
||||
process.env.API_KEY_SECRET = "test-provider-limits-recovery-secret";
|
||||
|
||||
const core = await import("../../src/lib/db/core.ts");
|
||||
const providersDb = await import("../../src/lib/db/providers.ts");
|
||||
const providerLimits = await import("../../src/lib/usage/providerLimits.ts");
|
||||
|
||||
const originalFetch = globalThis.fetch;
|
||||
|
||||
async function resetStorage() {
|
||||
core.resetDbInstance();
|
||||
fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true });
|
||||
fs.mkdirSync(TEST_DATA_DIR, { recursive: true });
|
||||
}
|
||||
|
||||
async function withMockedFetch(fetchImpl: typeof fetch, fn: () => Promise<void>) {
|
||||
const previousFetch = globalThis.fetch;
|
||||
globalThis.fetch = fetchImpl;
|
||||
try {
|
||||
await fn();
|
||||
} finally {
|
||||
globalThis.fetch = previousFetch;
|
||||
}
|
||||
}
|
||||
|
||||
async function createGlmConnectionWithTransientCooldown() {
|
||||
return providersDb.createProviderConnection({
|
||||
provider: "glm",
|
||||
authType: "apikey",
|
||||
name: `GLM Recovery ${Date.now()}`,
|
||||
apiKey: "glm-test-key",
|
||||
testStatus: "unavailable",
|
||||
rateLimitedUntil: new Date(Date.now() + 60_000).toISOString(),
|
||||
lastError: "rate limit exceeded",
|
||||
lastErrorType: "rate_limited",
|
||||
lastErrorSource: "executor",
|
||||
errorCode: 429,
|
||||
backoffLevel: 2,
|
||||
});
|
||||
}
|
||||
|
||||
function glmQuotaResponse() {
|
||||
// Mirrors open-sse/services/usage/glm.ts: TOKENS_LIMIT window with remaining.
|
||||
return new Response(
|
||||
JSON.stringify({
|
||||
code: 200,
|
||||
success: true,
|
||||
data: {
|
||||
planName: "max",
|
||||
limits: [
|
||||
{
|
||||
type: "TOKENS_LIMIT",
|
||||
unit: 3,
|
||||
number: 5,
|
||||
percentage: 13,
|
||||
nextResetTime: Math.floor(Date.now() / 1000) + 3 * 3600,
|
||||
models: [],
|
||||
},
|
||||
],
|
||||
},
|
||||
}),
|
||||
{ status: 200, headers: { "content-type": "application/json" } }
|
||||
);
|
||||
}
|
||||
|
||||
test.beforeEach(async () => {
|
||||
globalThis.fetch = originalFetch;
|
||||
await resetStorage();
|
||||
});
|
||||
|
||||
test.after(async () => {
|
||||
globalThis.fetch = originalFetch;
|
||||
core.resetDbInstance();
|
||||
fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true });
|
||||
});
|
||||
|
||||
test("successful GLM quota refresh clears transient rate-limit state", async () => {
|
||||
const connection = await createGlmConnectionWithTransientCooldown();
|
||||
const connectionId = (connection as { id: string }).id;
|
||||
|
||||
await withMockedFetch(
|
||||
(() => glmQuotaResponse()) as typeof fetch,
|
||||
async () => {
|
||||
await providerLimits.fetchAndPersistProviderLimits(connectionId, "manual");
|
||||
}
|
||||
);
|
||||
|
||||
const updated = (await providersDb.getProviderConnectionById(connectionId)) as Record<
|
||||
string,
|
||||
unknown
|
||||
>;
|
||||
assert.equal(updated.testStatus, "active", "testStatus should be reset to active");
|
||||
assert.equal(updated.rateLimitedUntil, undefined, "rateLimitedUntil should be cleared");
|
||||
assert.equal(updated.errorCode, undefined, "errorCode should be cleared");
|
||||
assert.equal(updated.lastErrorType, undefined, "lastErrorType should be cleared");
|
||||
assert.equal(updated.backoffLevel, 0, "backoffLevel should be reset to 0");
|
||||
});
|
||||
|
||||
async function createGlmConnectionWithStatus(status: string) {
|
||||
return providersDb.createProviderConnection({
|
||||
provider: "glm",
|
||||
authType: "apikey",
|
||||
name: "GLM " + status + " " + Date.now(),
|
||||
apiKey: "glm-test-key",
|
||||
testStatus: status,
|
||||
lastError: "permanent failure",
|
||||
lastErrorType: "permanent",
|
||||
errorCode: 403,
|
||||
backoffLevel: 1,
|
||||
});
|
||||
}
|
||||
|
||||
test("successful quota refresh does not clear terminal credits_exhausted status", async () => {
|
||||
const connection = await createGlmConnectionWithStatus("credits_exhausted");
|
||||
const connectionId = (connection as { id: string }).id;
|
||||
|
||||
await withMockedFetch(
|
||||
(() => glmQuotaResponse()) as typeof fetch,
|
||||
async () => {
|
||||
await providerLimits.fetchAndPersistProviderLimits(connectionId, "manual");
|
||||
}
|
||||
);
|
||||
|
||||
const updated = (await providersDb.getProviderConnectionById(connectionId)) as Record<
|
||||
string,
|
||||
unknown
|
||||
>;
|
||||
assert.equal(updated.testStatus, "credits_exhausted");
|
||||
assert.equal(updated.lastErrorType, "permanent");
|
||||
});
|
||||
|
||||
test("successful quota refresh does not clear terminal banned status", async () => {
|
||||
const connection = await createGlmConnectionWithStatus("banned");
|
||||
const connectionId = (connection as { id: string }).id;
|
||||
|
||||
await withMockedFetch(
|
||||
(() => glmQuotaResponse()) as typeof fetch,
|
||||
async () => {
|
||||
await providerLimits.fetchAndPersistProviderLimits(connectionId, "manual");
|
||||
}
|
||||
);
|
||||
|
||||
const updated = (await providersDb.getProviderConnectionById(connectionId)) as Record<
|
||||
string,
|
||||
unknown
|
||||
>;
|
||||
assert.equal(updated.testStatus, "banned");
|
||||
});
|
||||
|
||||
test("successful quota refresh does not clear terminal expired status", async () => {
|
||||
const connection = await createGlmConnectionWithStatus("expired");
|
||||
const connectionId = (connection as { id: string }).id;
|
||||
|
||||
await withMockedFetch(
|
||||
(() => glmQuotaResponse()) as typeof fetch,
|
||||
async () => {
|
||||
await providerLimits.fetchAndPersistProviderLimits(connectionId, "manual");
|
||||
}
|
||||
);
|
||||
|
||||
const updated = (await providersDb.getProviderConnectionById(connectionId)) as Record<
|
||||
string,
|
||||
unknown
|
||||
>;
|
||||
assert.equal(updated.testStatus, "expired");
|
||||
});
|
||||
|
||||
test("error-only quota response does not clear transient state", async () => {
|
||||
const connection = await createGlmConnectionWithTransientCooldown();
|
||||
const connectionId = (connection as { id: string }).id;
|
||||
|
||||
await withMockedFetch(
|
||||
(() =>
|
||||
new Response(JSON.stringify({ message: "GLM quota API error (429)" }), {
|
||||
status: 429,
|
||||
headers: { "content-type": "application/json" },
|
||||
})) as typeof fetch,
|
||||
async () => {
|
||||
// The live GLM usage path throws on a 429 (it does not return an error
|
||||
// envelope), so the fetch rejects. The transient-state assertions below then
|
||||
// confirm the throw happened BEFORE maybeClearRecoveredQuotaState — i.e. an
|
||||
// errored refresh never clears the connection's cooldown.
|
||||
await assert.rejects(
|
||||
() => providerLimits.fetchAndPersistProviderLimits(connectionId, "manual"),
|
||||
/429/
|
||||
);
|
||||
}
|
||||
);
|
||||
|
||||
const updated = (await providersDb.getProviderConnectionById(connectionId)) as Record<
|
||||
string,
|
||||
unknown
|
||||
>;
|
||||
assert.equal(updated.testStatus, "unavailable", "transient state should not be cleared on error");
|
||||
assert.equal(updated.lastErrorType, "rate_limited");
|
||||
});
|
||||
Loading…
Add table
Add a link
Reference in a new issue