Skip to content

fix usage-limit cooldown persistence #354

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
ndycode wants to merge 1 commit into from
Closed

fix usage-limit cooldown persistence #354

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 29 additions & 15 deletions index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1865,13 +1865,17 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
quotaKey,
rateLimit.retryAfterMs,
);
const cooldownMs = Math.max(
delayMs,
rateLimit.retryAfterMs,
);
preemptiveQuotaScheduler.markRateLimited(
quotaScheduleKey,
delayMs,
cooldownMs,
);
const waitLabel = formatWaitTime(delayMs);
const waitLabel = formatWaitTime(cooldownMs);

if (delayMs <= RATE_LIMIT_SHORT_RETRY_THRESHOLD_MS) {
if (cooldownMs <= RATE_LIMIT_SHORT_RETRY_THRESHOLD_MS) {
if (
accountManager.shouldShowAccountToast(
account.index,
Expand All @@ -1887,14 +1891,14 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
}

await sleep(
addJitter(Math.max(MIN_BACKOFF_MS, delayMs), 0.2),
addJitter(Math.max(MIN_BACKOFF_MS, cooldownMs), 0.2),
);
continue;
}

accountManager.markRateLimitedWithReason(
account,
delayMs,
cooldownMs,
modelFamily,
parseRateLimitReason(rateLimit.code),
model,
Expand Down Expand Up @@ -2142,19 +2146,29 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
);
}
if (!fallbackResponse.ok) {
try {
await fallbackResponse.body?.cancel();
} catch {
// Best effort cleanup before trying next fallback account.
}
if (fallbackResponse.status === 429) {
const { response: handledFallbackResponse, rateLimit: fallbackRateLimit } =
await handleErrorResponse(fallbackResponse);
if (handledFallbackResponse.status === 429) {
const retryAfterMs =
parseRetryAfterHintMs(
fallbackResponse.headers,
) ?? 60_000;
fallbackRateLimit?.retryAfterMs ?? 60_000;
const fallbackQuotaKey =
model ? `${modelFamily}:${model}` : modelFamily;
const { delayMs } = getRateLimitBackoff(
fallbackAccount.index,
fallbackQuotaKey,
retryAfterMs,
);
const cooldownMs = Math.max(
delayMs,
retryAfterMs,
);
preemptiveQuotaScheduler.markRateLimited(
`${fallbackEntitlementAccountKey}:${model ?? modelFamily}`,
cooldownMs,
);
accountManager.markRateLimitedWithReason(
fallbackAccount,
retryAfterMs,
cooldownMs,
modelFamily,
"quota",
model,
Expand Down
6 changes: 4 additions & 2 deletions lib/accounts.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -783,15 +783,17 @@ export class AccountManager {

const baseKey = getQuotaKey(family);
if (!model || reason === "quota" || reason === "unknown") {
account.rateLimitResetTimes[baseKey] = resetAt;
const currentResetAt = account.rateLimitResetTimes[baseKey] ?? 0;
account.rateLimitResetTimes[baseKey] = Math.max(currentResetAt, resetAt);
}

if (
model &&
(reason === "tokens" || reason === "concurrent" || reason === "unknown")
) {
const modelKey = getQuotaKey(family, model);
account.rateLimitResetTimes[modelKey] = resetAt;
const currentResetAt = account.rateLimitResetTimes[modelKey] ?? 0;
account.rateLimitResetTimes[modelKey] = Math.max(currentResetAt, resetAt);
}

account.lastRateLimitReason = reason;
Expand Down
10 changes: 8 additions & 2 deletions lib/preemptive-quota-scheduler.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -212,14 +212,20 @@ export class PreemptiveQuotaScheduler {
markRateLimited(key: string, retryAfterMs: number, now = Date.now()): void {
if (!key) return;
const waitMs = Number.isFinite(retryAfterMs) ? Math.max(0, Math.floor(retryAfterMs)) : 0;
const existing = this.snapshots.get(key);
const existingResetAtMs = Math.max(
existing?.primary.resetAtMs ?? 0,
existing?.secondary.resetAtMs ?? 0,
);
const nextResetAtMs = Math.max(existingResetAtMs, now + waitMs);
this.snapshots.set(key, {
status: 429,
primary: {
usedPercent: 100,
resetAtMs: now + waitMs,
resetAtMs: nextResetAtMs,
},
secondary: {},
updatedAt: now,
updatedAt: Math.max(existing?.updatedAt ?? 0, now),
});
}

Expand Down
212 changes: 150 additions & 62 deletions lib/request/fetch-helpers.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,11 @@ const CHATGPT_CODEX_UNSUPPORTED_MODEL_PATTERN =
/model is not supported when using codex with a chatgpt account/i;
const NORMALIZED_UNSUPPORTED_MODEL_PATTERN =
/the model ['"]([^'"]+)['"] is not currently available for this chatgpt account/i;
const MAX_RETRY_DELAY_MS = 5 * 60 * 1000;
const MAX_RATE_LIMIT_DELAY_MS = 7 * 24 * 60 * 60 * 1000;
const RETRY_AFTER_DURATION_PATTERN =
/\b(?:try|retry)\s+again\s+in\s+(\d+)\s*(second|minute|hour|day)s?\b/i;
const RETRY_AFTER_CLOCK_TIME_PATTERN =
/\b(?:try|retry)\s+again\s+at\s+(\d{1,2})(?::(\d{2}))?\s*(am|pm)\b/i;
const CREATE_CODEX_HEADERS_PARAM_KEYS = new Set(["init", "accountId", "accessToken", "opts"]);
const DEFAULT_PROXY_PORTS: Record<string, number> = {
"http:": 80,
Expand Down Expand Up @@ -1010,23 +1014,16 @@ function extractRateLimitInfoFromBody(
const isStatusRateLimit =
response.status === HTTP_STATUS.TOO_MANY_REQUESTS;
const parsed = parseRateLimitBody(bodyText);

const haystack = `${parsed?.code ?? ""} ${bodyText}`.toLowerCase();

// Entitlement errors should not be treated as rate limits
if (isEntitlementError(parsed?.code ?? "", bodyText)) {
return undefined;
}

const isRateLimit =
isStatusRateLimit ||
/usage_limit_reached|rate_limit_exceeded|rate_limit|usage limit/i.test(
haystack,
);
if (!isRateLimit) return undefined;

if (!isStatusRateLimit) return undefined;

const retryAfterMs =
parseRetryAfterMs(response, parsed) ?? 60000;
parseRetryAfterMs(response, bodyText, parsed) ?? 60000;

return { retryAfterMs, code: parsed?.code };
}
Expand Down Expand Up @@ -1197,8 +1194,24 @@ function ensureJsonErrorResponse(response: Response, payload: ErrorPayload): Res
});
}

function parseResetTimestampMs(rawValue: string): number | null {
const trimmed = rawValue.trim();
if (trimmed.length === 0) return null;

if (/^\d+$/.test(trimmed)) {
const parsed = Number.parseInt(trimmed, 10);
if (Number.isFinite(parsed) && parsed > 0) {
return parsed < 10_000_000_000 ? parsed * 1000 : parsed;
}
}

const parsedDate = Date.parse(trimmed);
return Number.isFinite(parsedDate) ? parsedDate : null;
}

function parseRetryAfterMs(
response: Response,
bodyText: string,
parsedBody?: { resetsAt?: number; retryAfterMs?: number; retryAfterSeconds?: number },
): number | null {
if (parsedBody?.retryAfterMs !== undefined) {
Expand All @@ -1211,69 +1224,144 @@ function parseRetryAfterMs(
if (normalized !== null) return normalized;
}

const retryAfterMsHeader = response.headers.get("retry-after-ms");
if (retryAfterMsHeader) {
const parsed = Number.parseInt(retryAfterMsHeader, 10);
if (!Number.isNaN(parsed) && parsed > 0) {
return parsed;
}
}
const retryAfterMsHeader = response.headers.get("retry-after-ms");
if (retryAfterMsHeader) {
const parsed = Number.parseInt(retryAfterMsHeader, 10);
const normalized = normalizeRetryAfterMs(parsed);
if (normalized !== null) {
return normalized;
}
}

const retryAfterHeader = response.headers.get("retry-after");
if (retryAfterHeader) {
const parsed = Number.parseInt(retryAfterHeader, 10);
if (!Number.isNaN(parsed) && parsed > 0) {
return parsed * 1000;
}
}
const retryAfterHeader = response.headers.get("retry-after");
if (retryAfterHeader) {
const parsed = Number.parseInt(retryAfterHeader, 10);
const normalized = normalizeRetryAfterSeconds(parsed);
if (normalized !== null) {
return normalized;
}
}

const resetAtHeaders = [
"x-codex-primary-reset-at",
"x-codex-secondary-reset-at",
"x-ratelimit-reset",
];
const now = Date.now();
const resetCandidates: number[] = [];
for (const header of resetAtHeaders) {
const value = response.headers.get(header);
if (!value) continue;
const parsed = Number.parseInt(value, 10);
if (!Number.isNaN(parsed) && parsed > 0) {
const timestamp =
parsed < 10_000_000_000 ? parsed * 1000 : parsed;
const delta = timestamp - now;
if (delta > 0) resetCandidates.push(delta);
}
}
const resetAfterSecondsHeaders = [
"x-codex-primary-reset-after-seconds",
"x-codex-secondary-reset-after-seconds",
];
const resetCandidates: number[] = [];
for (const header of resetAfterSecondsHeaders) {
const value = response.headers.get(header);
if (!value) continue;
const parsed = Number.parseInt(value, 10);
const normalized = normalizeRetryAfterSeconds(parsed);
if (normalized !== null) {
resetCandidates.push(normalized);
}
}

if (parsedBody?.resetsAt) {
const timestamp =
parsedBody.resetsAt < 10_000_000_000
? parsedBody.resetsAt * 1000
: parsedBody.resetsAt;
const delta = timestamp - now;
if (delta > 0) resetCandidates.push(delta);
}
const resetAtHeaders = [
"x-codex-primary-reset-at",
"x-codex-secondary-reset-at",
"x-ratelimit-reset",
];
const now = Date.now();
for (const header of resetAtHeaders) {
const value = response.headers.get(header);
if (!value) continue;
const timestamp = parseResetTimestampMs(value);
if (timestamp === null) continue;
const delta = normalizeRetryAfterMs(timestamp - now);
if (delta !== null) resetCandidates.push(delta);
}

if (resetCandidates.length > 0) {
return Math.min(...resetCandidates);
}
if (parsedBody?.resetsAt) {
const timestamp =
parsedBody.resetsAt < 10_000_000_000
? parsedBody.resetsAt * 1000
: parsedBody.resetsAt;
const delta = normalizeRetryAfterMs(timestamp - now);
if (delta !== null) resetCandidates.push(delta);
}

if (resetCandidates.length > 0) {
return Math.max(...resetCandidates);
}

const naturalLanguageRetryAfterMs = parseRetryAfterTextMs(bodyText, now);
if (naturalLanguageRetryAfterMs !== null) {
return naturalLanguageRetryAfterMs;
}

return null;
return null;
}

function normalizeRetryAfterMs(value: number): number | null {
function parseRetryAfterTextMs(bodyText: string, now: number): number | null {
if (!bodyText) return null;

const durationMatch = bodyText.match(RETRY_AFTER_DURATION_PATTERN);
if (durationMatch) {
const amount = Number.parseInt(durationMatch[1] ?? "", 10);
const unit = (durationMatch[2] ?? "").toLowerCase();
if (Number.isFinite(amount) && amount > 0) {
const multiplier =
unit === "second"
? 1000
: unit === "minute"
? 60_000
: unit === "hour"
? 60 * 60_000
: unit === "day"
? 24 * 60 * 60_000
: 0;
if (multiplier > 0) {
return clampRateLimitDelayMs(amount * multiplier);
}
}
}

const clockTimeMatch = bodyText.match(RETRY_AFTER_CLOCK_TIME_PATTERN);
if (!clockTimeMatch) return null;

const hour12 = Number.parseInt(clockTimeMatch[1] ?? "", 10);
const minute = Number.parseInt(clockTimeMatch[2] ?? "0", 10);
const meridiem = (clockTimeMatch[3] ?? "").toLowerCase();
if (
!Number.isFinite(hour12) ||
hour12 < 1 ||
hour12 > 12 ||
!Number.isFinite(minute) ||
minute < 0 ||
minute > 59 ||
(meridiem !== "am" && meridiem !== "pm")
) {
return null;
}

const target = new Date(now);
let hour24 = hour12 % 12;
if (meridiem === "pm") {
hour24 += 12;
}
target.setHours(hour24, minute, 0, 0);
if (target.getTime() <= now) {
target.setDate(target.getDate() + 1);
}

return clampRateLimitDelayMs(target.getTime() - now);
}

function clampRateLimitDelayMs(value: number): number | null {
if (!Number.isFinite(value)) return null;
const ms = Math.floor(value);
if (ms <= 0) return null;
return Math.min(ms, MAX_RETRY_DELAY_MS);
const normalized = Math.floor(value);
if (normalized <= 0) return null;
return Math.min(normalized, MAX_RATE_LIMIT_DELAY_MS);
}

function normalizeRetryAfterMs(value: number): number | null {
return clampRateLimitDelayMs(value);
}

function normalizeRetryAfterSeconds(value: number): number | null {
if (!Number.isFinite(value)) return null;
const ms = Math.floor(value * 1000);
if (ms <= 0) return null;
return Math.min(ms, MAX_RETRY_DELAY_MS);
return clampRateLimitDelayMs(value * 1000);
}

function toNumber(value: unknown): number | undefined {
Expand Down
Loading