diff --git a/src/lib/usage/providerLimits.ts b/src/lib/usage/providerLimits.ts index d3de997e75..cf02868f2d 100644 --- a/src/lib/usage/providerLimits.ts +++ b/src/lib/usage/providerLimits.ts @@ -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 { + 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 { 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, diff --git a/tests/unit/provider-limits-recovery.test.ts b/tests/unit/provider-limits-recovery.test.ts new file mode 100644 index 0000000000..f719d19090 --- /dev/null +++ b/tests/unit/provider-limits-recovery.test.ts @@ -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) { + 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"); +});