Handle codex OAuth usage-limit cooldowns (#316)

* Handle codex oauth usage-limit cooldowns * Generalize short-window 429 limit cooldowns * Fan out short-window cooldowns by credential
2026-03-30 13:39:07 +08:00
parent d47c3c4a29
commit 71be4cb63a
2 changed files with 462 additions and 11 deletions
@@ -175,6 +175,373 @@ describe('TokenRouter runtime cache', () => {
    expect(thirdCooldownMs).toBeLessThanOrEqual(35_000);
  });

+  it('uses codex oauth reset hints for usage-limit cooldowns', async () => {
+    const site = await db.insert(schema.sites).values({
+      name: 'codex-oauth-site',
+      url: 'https://codex-oauth.example.com',
+      platform: 'codex',
+      status: 'active',
+    }).returning().get();
+
+    const account = await db.insert(schema.accounts).values({
+      siteId: site.id,
+      username: 'codex-oauth-user',
+      accessToken: 'codex-access-token',
+      status: 'active',
+      oauthProvider: 'codex',
+      extraConfig: JSON.stringify({
+        credentialMode: 'session',
+        oauth: {
+          provider: 'codex',
+        },
+      }),
+    }).returning().get();
+
+    const route = await db.insert(schema.tokenRoutes).values({
+      modelPattern: 'gpt-5.4',
+      routingStrategy: 'weighted',
+      enabled: true,
+    }).returning().get();
+
+    const channel = await db.insert(schema.routeChannels).values({
+      routeId: route.id,
+      accountId: account.id,
+      priority: 0,
+      weight: 10,
+      enabled: true,
+    }).returning().get();
+
+    const router = new TokenRouter();
+    const startedAt = Date.now();
+    await router.recordFailure(channel.id, {
+      status: 429,
+      errorText: JSON.stringify({
+        error: {
+          type: 'usage_limit_reached',
+          resets_in_seconds: 600,
+          message: 'quota exceeded',
+        },
+      }),
+      modelName: 'gpt-5.4',
+    });
+
+    const record = await db.select().from(schema.routeChannels)
+      .where(eq(schema.routeChannels.id, channel.id))
+      .get();
+    const cooldownMs = Date.parse(String(record?.cooldownUntil || '')) - startedAt;
+
+    expect(cooldownMs).toBeGreaterThanOrEqual(9 * 60 * 1000);
+    expect(cooldownMs).toBeLessThanOrEqual(11 * 60 * 1000);
+  });
+
+  it('falls back to a short cooldown for codex oauth usage-limit failures without reset hints', async () => {
+    const site = await db.insert(schema.sites).values({
+      name: 'codex-oauth-fallback-site',
+      url: 'https://codex-oauth-fallback.example.com',
+      platform: 'codex',
+      status: 'active',
+    }).returning().get();
+
+    const account = await db.insert(schema.accounts).values({
+      siteId: site.id,
+      username: 'codex-oauth-fallback-user',
+      accessToken: 'codex-fallback-access-token',
+      status: 'active',
+      oauthProvider: 'codex',
+      extraConfig: JSON.stringify({
+        credentialMode: 'session',
+        oauth: {
+          provider: 'codex',
+        },
+      }),
+    }).returning().get();
+
+    const route = await db.insert(schema.tokenRoutes).values({
+      modelPattern: 'gpt-5.4',
+      routingStrategy: 'weighted',
+      enabled: true,
+    }).returning().get();
+
+    const channel = await db.insert(schema.routeChannels).values({
+      routeId: route.id,
+      accountId: account.id,
+      priority: 0,
+      weight: 10,
+      enabled: true,
+    }).returning().get();
+
+    const router = new TokenRouter();
+    const startedAt = Date.now();
+    await router.recordFailure(channel.id, {
+      status: 429,
+      errorText: JSON.stringify({
+        error: {
+          type: 'usage_limit_reached',
+          message: 'The usage limit has been reached',
+        },
+      }),
+      modelName: 'gpt-5.4',
+    });
+
+    const record = await db.select().from(schema.routeChannels)
+      .where(eq(schema.routeChannels.id, channel.id))
+      .get();
+    const cooldownMs = Date.parse(String(record?.cooldownUntil || '')) - startedAt;
+
+    expect(cooldownMs).toBeGreaterThanOrEqual(4 * 60 * 1000);
+    expect(cooldownMs).toBeLessThanOrEqual(6 * 60 * 1000);
+  });
+
+  it('reuses stored codex oauth reset hints when the latest 429 omits reset timing', async () => {
+    const site = await db.insert(schema.sites).values({
+      name: 'codex-oauth-stored-reset-site',
+      url: 'https://codex-oauth-stored-reset.example.com',
+      platform: 'codex',
+      status: 'active',
+    }).returning().get();
+
+    const resetAt = new Date(Date.now() + 8 * 60 * 1000).toISOString();
+    const account = await db.insert(schema.accounts).values({
+      siteId: site.id,
+      username: 'codex-oauth-stored-reset-user',
+      accessToken: 'codex-stored-reset-access-token',
+      status: 'active',
+      oauthProvider: 'codex',
+      extraConfig: JSON.stringify({
+        credentialMode: 'session',
+        oauth: {
+          provider: 'codex',
+          quota: {
+            status: 'supported',
+            source: 'reverse_engineered',
+            lastLimitResetAt: resetAt,
+            providerMessage: 'current codex oauth signals do not expose stable 5h/7d remaining values',
+            windows: {
+              fiveHour: {
+                supported: false,
+                message: 'official 5h quota window is not exposed by current codex oauth artifacts',
+              },
+              sevenDay: {
+                supported: false,
+                message: 'official 7d quota window is not exposed by current codex oauth artifacts',
+              },
+            },
+          },
+        },
+      }),
+    }).returning().get();
+
+    const route = await db.insert(schema.tokenRoutes).values({
+      modelPattern: 'gpt-5.4',
+      routingStrategy: 'weighted',
+      enabled: true,
+    }).returning().get();
+
+    const channel = await db.insert(schema.routeChannels).values({
+      routeId: route.id,
+      accountId: account.id,
+      priority: 0,
+      weight: 10,
+      enabled: true,
+    }).returning().get();
+
+    const router = new TokenRouter();
+    await router.recordFailure(channel.id, {
+      status: 429,
+      errorText: JSON.stringify({
+        error: {
+          type: 'usage_limit_reached',
+          message: 'The usage limit has been reached',
+        },
+      }),
+      modelName: 'gpt-5.4',
+    });
+
+    const record = await db.select().from(schema.routeChannels)
+      .where(eq(schema.routeChannels.id, channel.id))
+      .get();
+
+    expect(record?.cooldownUntil).toBe(resetAt);
+  });
+
+  it('treats non-oauth 429 limit failures as short-window cooldowns', async () => {
+    const site = await db.insert(schema.sites).values({
+      name: 'limit-site',
+      url: 'https://limit-site.example.com',
+      platform: 'new-api',
+      status: 'active',
+    }).returning().get();
+
+    const account = await db.insert(schema.accounts).values({
+      siteId: site.id,
+      username: 'limit-user',
+      accessToken: 'limit-access-token',
+      apiToken: 'limit-api-token',
+      status: 'active',
+    }).returning().get();
+
+    const route = await db.insert(schema.tokenRoutes).values({
+      modelPattern: 'gpt-4o-mini',
+      routingStrategy: 'weighted',
+      enabled: true,
+    }).returning().get();
+
+    const channel = await db.insert(schema.routeChannels).values({
+      routeId: route.id,
+      accountId: account.id,
+      priority: 0,
+      weight: 10,
+      enabled: true,
+    }).returning().get();
+
+    const router = new TokenRouter();
+    const startedAt = Date.now();
+    await router.recordFailure(channel.id, {
+      status: 429,
+      errorText: JSON.stringify({
+        error: {
+          message: 'rate limit exceeded',
+        },
+      }),
+      modelName: 'gpt-4o-mini',
+    });
+
+    const record = await db.select().from(schema.routeChannels)
+      .where(eq(schema.routeChannels.id, channel.id))
+      .get();
+    const cooldownMs = Date.parse(String(record?.cooldownUntil || '')) - startedAt;
+
+    expect(cooldownMs).toBeGreaterThanOrEqual(4 * 60 * 1000);
+    expect(cooldownMs).toBeLessThanOrEqual(6 * 60 * 1000);
+    expect(record?.failCount).toBe(0);
+  });
+
+  it('keeps generic 429 backoff when the error is not limit-related', async () => {
+    const site = await db.insert(schema.sites).values({
+      name: 'generic-429-site',
+      url: 'https://generic-429.example.com',
+      platform: 'new-api',
+      status: 'active',
+    }).returning().get();
+
+    const account = await db.insert(schema.accounts).values({
+      siteId: site.id,
+      username: 'generic-429-user',
+      accessToken: 'generic-429-access-token',
+      apiToken: 'generic-429-api-token',
+      status: 'active',
+    }).returning().get();
+
+    const route = await db.insert(schema.tokenRoutes).values({
+      modelPattern: 'gpt-4o-mini',
+      routingStrategy: 'weighted',
+      enabled: true,
+    }).returning().get();
+
+    const channel = await db.insert(schema.routeChannels).values({
+      routeId: route.id,
+      accountId: account.id,
+      priority: 0,
+      weight: 10,
+      enabled: true,
+    }).returning().get();
+
+    const router = new TokenRouter();
+    const startedAt = Date.now();
+    await router.recordFailure(channel.id, {
+      status: 429,
+      errorText: JSON.stringify({
+        error: {
+          message: 'upstream overloaded',
+        },
+      }),
+      modelName: 'gpt-4o-mini',
+    });
+
+    const record = await db.select().from(schema.routeChannels)
+      .where(eq(schema.routeChannels.id, channel.id))
+      .get();
+    const cooldownMs = Date.parse(String(record?.cooldownUntil || '')) - startedAt;
+
+    expect(cooldownMs).toBeGreaterThanOrEqual(10_000);
+    expect(cooldownMs).toBeLessThanOrEqual(20_000);
+    expect(record?.failCount).toBe(1);
+  });
+
+  it('applies short-window cooldown to sibling channels that share the same account-level credential', async () => {
+    const site = await db.insert(schema.sites).values({
+      name: 'shared-credential-site',
+      url: 'https://shared-credential.example.com',
+      platform: 'codex',
+      status: 'active',
+    }).returning().get();
+
+    const account = await db.insert(schema.accounts).values({
+      siteId: site.id,
+      username: 'shared-credential-user',
+      accessToken: 'shared-credential-access-token',
+      status: 'active',
+      oauthProvider: 'codex',
+      extraConfig: JSON.stringify({
+        credentialMode: 'session',
+        oauth: {
+          provider: 'codex',
+        },
+      }),
+    }).returning().get();
+
+    const primaryRoute = await db.insert(schema.tokenRoutes).values({
+      modelPattern: 'gpt-5.4',
+      routingStrategy: 'weighted',
+      enabled: true,
+    }).returning().get();
+
+    const siblingRoute = await db.insert(schema.tokenRoutes).values({
+      modelPattern: 'gpt-4o-mini',
+      routingStrategy: 'weighted',
+      enabled: true,
+    }).returning().get();
+
+    const primaryChannel = await db.insert(schema.routeChannels).values({
+      routeId: primaryRoute.id,
+      accountId: account.id,
+      priority: 0,
+      weight: 10,
+      enabled: true,
+    }).returning().get();
+
+    const siblingChannel = await db.insert(schema.routeChannels).values({
+      routeId: siblingRoute.id,
+      accountId: account.id,
+      priority: 0,
+      weight: 10,
+      enabled: true,
+    }).returning().get();
+
+    const router = new TokenRouter();
+    const initialSiblingSelection = await router.selectChannel('gpt-4o-mini');
+    expect(initialSiblingSelection?.channel.id).toBe(siblingChannel.id);
+
+    await router.recordFailure(primaryChannel.id, {
+      status: 429,
+      errorText: JSON.stringify({
+        error: {
+          type: 'usage_limit_reached',
+          message: 'The usage limit has been reached',
+        },
+      }),
+      modelName: 'gpt-5.4',
+    });
+
+    const cooledSibling = await db.select().from(schema.routeChannels)
+      .where(eq(schema.routeChannels.id, siblingChannel.id))
+      .get();
+
+    expect(cooledSibling?.cooldownUntil).toBeTruthy();
+    expect(cooledSibling?.failCount).toBe(0);
+    expect(await router.selectChannel('gpt-4o-mini')).toBeNull();
+  });
+
  it('round robins across all available channels regardless of priority', async () => {
    const site = await db.insert(schema.sites).values({
      name: 'round-robin-site',
@@ -1,4 +1,4 @@
-import { eq, inArray } from 'drizzle-orm';
+import { and, eq, inArray, isNull } from 'drizzle-orm';
 import { db, schema } from '../db/index.js';
 import { upsertSetting } from '../db/upsertSetting.js';
 import { config } from '../config.js';
@@ -11,6 +11,7 @@ import {
 import { type DownstreamRoutingPolicy, EMPTY_DOWNSTREAM_ROUTING_POLICY } from './downstreamPolicyTypes.js';
 import { isUsableAccountToken } from './accountTokenService.js';
 import { getOauthInfoFromAccount } from './oauth/oauthAccount.js';
+import { parseCodexQuotaResetHint } from './oauth/quota.js';
 import {
  isExactTokenRouteModelPattern,
  isTokenRouteRegexPattern,
@@ -70,6 +71,7 @@ type SiteRuntimeHealthState = {
 };

 const FAILURE_BACKOFF_BASE_SEC = 15;
+const SHORT_WINDOW_LIMIT_COOLDOWN_MS = 5 * 60 * 1000;
 const MIN_EFFECTIVE_UNIT_COST = 1e-6;
 const ROUND_ROBIN_FAILURE_THRESHOLD = 3;
 const ROUND_ROBIN_COOLDOWN_LEVELS_SEC = [0, 10 * 60, 60 * 60, 24 * 60 * 60] as const;
@@ -145,6 +147,14 @@ const SITE_TRANSIENT_FAILURE_PATTERNS: RegExp[] = [
  /econnrefused/i,
 ];

+const USAGE_LIMIT_RATE_LIMIT_PATTERNS: RegExp[] = [
+  /usage_limit_reached/i,
+  /usage\s+limit\s+has\s+been\s+reached/i,
+  /quota\s+exceeded/i,
+  /rate\s+limit/i,
+  /\blimit\b/i,
+];
+
 type SiteRuntimeHealthPersistencePayload = {
  version: 1;
  savedAtMs: number;
@@ -203,6 +213,12 @@ function matchesAnyPattern(patterns: RegExp[], input?: string | null): boolean {
  return patterns.some((pattern) => pattern.test(text));
 }

+function isUsageLimitRateLimitFailure(context: SiteRuntimeFailureContext = {}): boolean {
+  const status = typeof context.status === 'number' ? context.status : 0;
+  if (status !== 429) return false;
+  return matchesAnyPattern(USAGE_LIMIT_RATE_LIMIT_PATTERNS, context.errorText);
+}
+
 function clampNumber(value: number, min: number, max: number): number {
  return Math.max(min, Math.min(max, value));
 }
@@ -230,6 +246,10 @@ function resolveSiteRuntimeFailurePenalty(context: SiteRuntimeFailureContext = {
  const status = typeof context.status === 'number' ? context.status : 0;
  const errorText = (context.errorText || '').trim();

+  if (isUsageLimitRateLimitFailure({ status, errorText })) {
+    return 0.4;
+  }
+
  if (status >= 500 || matchesAnyPattern(SITE_TRANSIENT_FAILURE_PATTERNS, errorText)) {
    return 2.5;
  }
@@ -264,9 +284,63 @@ function resolveSiteRuntimeFailurePenalty(context: SiteRuntimeFailureContext = {
 function isTransientSiteRuntimeFailure(context: SiteRuntimeFailureContext = {}): boolean {
  const status = typeof context.status === 'number' ? context.status : 0;
  const errorText = (context.errorText || '').trim();
+  if (isUsageLimitRateLimitFailure({ status, errorText })) {
+    return false;
+  }
  return status >= 500 || status === 429 || matchesAnyPattern(SITE_TRANSIENT_FAILURE_PATTERNS, errorText);
 }

+function resolveShortWindowLimitCooldown(
+  account: typeof schema.accounts.$inferSelect,
+  context: SiteRuntimeFailureContext = {},
+  nowMs = Date.now(),
+): string | null {
+  const status = typeof context.status === 'number' ? context.status : 0;
+  const errorText = (context.errorText || '').trim();
+  if (!isUsageLimitRateLimitFailure({ status, errorText })) return null;
+
+  const resetHint = parseCodexQuotaResetHint(status, errorText, nowMs);
+  if (resetHint) {
+    const hintMs = Date.parse(resetHint.resetAt);
+    if (Number.isFinite(hintMs) && hintMs > nowMs) {
+      return new Date(hintMs).toISOString();
+    }
+  }
+
+  const oauth = getOauthInfoFromAccount(account);
+  const storedResetAt = oauth?.quota?.lastLimitResetAt;
+  if (oauth?.provider === 'codex' && storedResetAt) {
+    const storedMs = Date.parse(storedResetAt);
+    if (Number.isFinite(storedMs) && storedMs > nowMs) {
+      return new Date(storedMs).toISOString();
+    }
+  }
+
+  return new Date(nowMs + SHORT_WINDOW_LIMIT_COOLDOWN_MS).toISOString();
+}
+
+async function loadCredentialScopedChannelIds(
+  channel: typeof schema.routeChannels.$inferSelect,
+  accountId: number,
+): Promise<number[]> {
+  if (typeof channel.tokenId === 'number' && channel.tokenId > 0) {
+    const rows = await db.select({ id: schema.routeChannels.id })
+      .from(schema.routeChannels)
+      .where(eq(schema.routeChannels.tokenId, channel.tokenId))
+      .all();
+    return rows.map((row) => row.id);
+  }
+
+  const rows = await db.select({ id: schema.routeChannels.id })
+    .from(schema.routeChannels)
+    .where(and(
+      eq(schema.routeChannels.accountId, accountId),
+      isNull(schema.routeChannels.tokenId),
+    ))
+    .all();
+  return rows.map((row) => row.id);
+}
+
 function getDecayedSiteRuntimePenalty(state: SiteRuntimeHealthState, nowMs: number): number {
  if (!Number.isFinite(state.penaltyScore) || state.penaltyScore <= 0) return 0;
  const elapsedMs = Math.max(0, nowMs - state.lastUpdatedAtMs);
@@ -1709,16 +1783,24 @@ export class TokenRouter {
    const route = row.token_routes;
    const nowMs = Date.now();
    const nowIso = new Date(nowMs).toISOString();
-    const failCount = (ch.failCount ?? 0) + 1;
    const normalizedContext: SiteRuntimeFailureContext = typeof context === 'string'
      ? { modelName: context }
      : (context ?? {});
+    const shortWindowLimitCooldownUntil = resolveShortWindowLimitCooldown(account, normalizedContext, nowMs);
+    const failCount = shortWindowLimitCooldownUntil ? 0 : ((ch.failCount ?? 0) + 1);
    const routeStrategy = resolveRouteStrategy(route);
+    const affectedChannelIds = shortWindowLimitCooldownUntil
+      ? await loadCredentialScopedChannelIds(ch, account.id)
+      : [channelId];
    let cooldownUntil: string | null = null;
    let consecutiveFailCount = Math.max(0, ch.consecutiveFailCount ?? 0) + 1;
    let cooldownLevel = Math.max(0, ch.cooldownLevel ?? 0);

-    if (routeStrategy === 'round_robin') {
+    if (shortWindowLimitCooldownUntil) {
+      cooldownUntil = shortWindowLimitCooldownUntil;
+      consecutiveFailCount = 0;
+      cooldownLevel = 0;
+    } else if (routeStrategy === 'round_robin') {
      if (consecutiveFailCount >= ROUND_ROBIN_FAILURE_THRESHOLD) {
        cooldownLevel = Math.min(cooldownLevel + 1, ROUND_ROBIN_COOLDOWN_LEVELS_SEC.length - 1);
        const cooldownSec = resolveRoundRobinCooldownSec(cooldownLevel);
@@ -1738,15 +1820,17 @@ export class TokenRouter {
      consecutiveFailCount,
      cooldownLevel,
      cooldownUntil,
-    }).where(eq(schema.routeChannels.id, channelId)).run();
+    }).where(inArray(schema.routeChannels.id, affectedChannelIds)).run();

-    patchCachedChannel(channelId, (channel) => {
-      channel.failCount = failCount;
-      channel.lastFailAt = nowIso;
-      channel.cooldownUntil = cooldownUntil;
-      channel.consecutiveFailCount = consecutiveFailCount;
-      channel.cooldownLevel = cooldownLevel;
-    });
+    for (const affectedChannelId of affectedChannelIds) {
+      patchCachedChannel(affectedChannelId, (channel) => {
+        channel.failCount = failCount;
+        channel.lastFailAt = nowIso;
+        channel.cooldownUntil = cooldownUntil;
+        channel.consecutiveFailCount = consecutiveFailCount;
+        channel.cooldownLevel = cooldownLevel;
+      });
+    }

    recordSiteRuntimeFailure(account.siteId, normalizedContext, nowMs);
  }