From 0ab90f63e50960e95baea54f49dc52c111f88390 Mon Sep 17 00:00:00 2001 From: Jeff Scott Ward Date: Sat, 11 Jul 2026 20:34:57 -0400 Subject: [PATCH 1/2] fix(auth): rotate through quota-limited accounts Keep ordinary 401 retries bounded while replay-safe quota failures walk every distinct eligible credential. Anchor blocks to the failed credential and stop on cycles, aborts, or 64 attempts. --- packages/ai/CHANGELOG.md | 1 + packages/ai/src/auth-retry.ts | 244 +++++++++++--- packages/ai/src/auth-storage.ts | 110 +++++-- packages/ai/src/error/auth-classify.ts | 3 +- packages/ai/src/stream.ts | 49 ++- packages/ai/test/auth-retry.test.ts | 307 +++++++++++++++++- .../auth-storage-broker-no-sentinel.test.ts | 1 + .../auth-storage-force-refresh-rotate.test.ts | 155 +++++++++ packages/ai/test/rate-limit-utils.test.ts | 10 + packages/ai/test/stream-auth-retry.test.ts | 124 ++++++- packages/coding-agent/CHANGELOG.md | 1 + .../src/config/api-key-resolver.ts | 3 +- .../test/agent-session-retry-cap.test.ts | 94 ++++-- .../test/auth-storage-rotation.test.ts | 69 +++- 14 files changed, 1008 insertions(+), 163 deletions(-) diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index f72bb7642b9..b212eeb8ab6 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -12,6 +12,7 @@ ### Fixed - Fixed OAuth credential resolution returning "No API key found" when every plan-eligible OpenAI Codex account was rate-limit blocked and the only unblocked account failed the model's plan gate: resolution now runs a last-resort ladder that first yields a plan-fitting account regardless of usage blocks (so callers get real usage-limit retry semantics), then tries every account with the plan filter dropped before reporting no credential +- Fixed provider-agnostic replay-safe usage/account-quota failures to rotate through every distinct eligible credential instead of stopping after the fixed a/b/c ladder, while preserving transient-429 backoff, cycle/abort guards, exact failed-credential targeting despite stale session stickiness, and a finite safety ceiling. ## [16.4.5] - 2026-07-11 diff --git a/packages/ai/src/auth-retry.ts b/packages/ai/src/auth-retry.ts index 22a8f8e552e..4a140610aa8 100644 --- a/packages/ai/src/auth-retry.ts +++ b/packages/ai/src/auth-retry.ts @@ -2,12 +2,13 @@ import type { OAuthAccess } from "./auth-storage"; import * as AIError from "./error"; import { isAuthRetryableError } from "./error/auth-classify"; import { isUsageLimit } from "./error/flags"; +import { isUsageLimitOutcome } from "./error/rate-limit"; /** * Context passed to an {@link ApiKeyResolver} on each resolution attempt. * - * The `error`/`lastChance` pair drives the central a/b/c retry policy shared by - * the streaming ({@link streamSimple}) and non-streaming ({@link withAuth}) + * The `error`/`lastChance` pair preserves the legacy a/b/c resolver contract + * shared by streaming ({@link streamSimple}) and non-streaming ({@link withAuth}) * drivers: * - `error === undefined` → **initial resolve** (no force-refresh; cheap, may * return a locally-cached not-yet-expired token). @@ -16,11 +17,13 @@ import { isUsageLimit } from "./error/flags"; * - `error !== undefined && lastChance` → **step (c): switch account** * (invalidate/usage-limit the current credential and rotate to a sibling). * - * The resolver returns the bearer to send, or `undefined` to stop retrying and - * surface the last error to the caller. + * Current drivers preserve that bounded a/b/c sequence for ordinary 401/auth + * failures. Usage/account-limit failures skip refresh and may repeat step (c) + * until the resolver returns `undefined`, cycles, or hits + * {@link AUTH_RETRY_MAX_ATTEMPTS}. */ export interface ApiKeyResolveContext { - /** True on the final retry step — the resolver should rotate to a sibling credential. */ + /** True when the resolver should rotate to a sibling credential. */ lastChance: boolean; /** The auth error that triggered this re-resolution, or `undefined` on the initial resolve. */ error: unknown; @@ -77,13 +80,22 @@ export function seedApiKeyResolver(seed: string | undefined, resolver: ApiKeyRes export { isAuthRetryableError }; /** - * The ordered `lastChance` values for the retry steps after the initial - * attempt fails: `false` → step (b) refresh-same, `true` → step (c) switch. - * Shared by {@link withAuth} and the streaming retry driver so both run the - * same policy. + * Legacy bounded a/b/c retry sequence retained for public compatibility: + * `false` → refresh-same, `true` → rotate/switch. Current drivers consume it + * once for ordinary 401/auth failures; usage/account-limit failures may repeat + * sibling rotation until a termination guard fires. */ export const AUTH_RETRY_STEPS: readonly boolean[] = [false, true]; +export const AUTH_RETRY_MAX_ATTEMPTS = 64; + +function isDirectCredentialRotationError(error: unknown): boolean { + if (isUsageLimit(error)) return true; + const status = AIError.status(error); + const message = error instanceof Error ? error.message : typeof error === "string" ? error : undefined; + return isUsageLimitOutcome(status, message); +} + /** Resolve a single retry step, swallowing resolver failures into `undefined`. */ export async function resolveRetryKey( resolver: ApiKeyResolver, @@ -93,22 +105,102 @@ export async function resolveRetryKey( previousKey?: string, ): Promise { try { - const rotateSibling = lastChance || (!lastChance && isUsageLimit(error)); + const rotateSibling = lastChance || (!lastChance && isDirectCredentialRotationError(error)); return (await resolver({ lastChance: rotateSibling, error, signal, previousKey })) || undefined; } catch { return undefined; } } +export interface AuthRetryKeyState { + /** Bearer strings already sent during this logical operation. */ + attemptedKeys: Set; + /** Bearer used by the most recent failed attempt. */ + lastKey: string; + /** Whether the current credential already consumed its 401 refresh-same retry. */ + refreshedCurrent: boolean; + /** Whether the legacy non-usage auth path already switched to one sibling. */ + legacyAuthSwitchUsed: boolean; + /** Total outbound attempts accepted for this logical operation, including the initial request. */ + attempts: number; +} + +export function createAuthRetryKeyState(initialKey: string): AuthRetryKeyState { + return { + attemptedKeys: new Set([initialKey]), + lastKey: initialKey, + refreshedCurrent: false, + legacyAuthSwitchUsed: false, + attempts: 1, + }; +} + +function acceptRetryKey(state: AuthRetryKeyState, key: string, refreshedCurrent: boolean): string | undefined { + if (state.attemptedKeys.has(key) || state.attempts >= AUTH_RETRY_MAX_ATTEMPTS) return undefined; + state.attemptedKeys.add(key); + state.attempts += 1; + state.lastKey = key; + state.refreshedCurrent = refreshedCurrent; + return key; +} + +export async function resolveNextAuthRetryKey( + state: AuthRetryKeyState, + resolver: ApiKeyResolver, + error: unknown, + signal?: AbortSignal, +): Promise { + if (signal?.aborted) return undefined; + if (state.attempts >= AUTH_RETRY_MAX_ATTEMPTS) return undefined; + const directRotation = isDirectCredentialRotationError(error); + if (!directRotation) { + if (state.legacyAuthSwitchUsed) return undefined; + if (!state.refreshedCurrent) { + const refreshed = await resolveRetryKey(resolver, false, error, signal, state.lastKey); + state.refreshedCurrent = true; + if (signal?.aborted) return undefined; + if (refreshed !== undefined) { + const accepted = acceptRetryKey(state, refreshed, true); + if (accepted !== undefined) return accepted; + } + } + } + + if (signal?.aborted) return undefined; + const rotated = await resolveRetryKey(resolver, true, error, signal, state.lastKey); + if (signal?.aborted || rotated === undefined) return undefined; + const accepted = acceptRetryKey(state, rotated, !directRotation); + if (accepted !== undefined && !directRotation) state.legacyAuthSwitchUsed = true; + return accepted; +} + +function oauthCredentialIdentity(access: OAuthAccess): string { + return access.credentialId !== undefined ? `credential:${access.credentialId}` : `bearer:${access.accessToken}`; +} + +async function runOAuthAttempt( + access: OAuthAccess, + attempt: (access: OAuthAccess) => Promise, + isAuthError: (error: unknown) => boolean, +): Promise<{ ok: true; result: T } | { ok: false; error: unknown }> { + try { + return { ok: true, result: await attempt(access) }; + } catch (error) { + if (!isAuthError(error)) throw error; + return { ok: false, error }; + } +} + /** * Runs an auth-protected operation through the central a/b/c retry policy. * * - A static string key (or any non-resolver) → a single `attempt` with no * retry (identical to the legacy static-key path). - * - A resolver → initial `attempt`, then on a retryable auth error up to two - * more attempts (refresh-same, then switch). A step is skipped when the - * resolver returns the same key it just tried or `undefined`; non-auth errors - * propagate immediately. + * - A resolver → initial `attempt`, then resolver-driven retries until the + * applicable policy is exhausted, the resolver declines or cycles, or the + * operation reaches {@link AUTH_RETRY_MAX_ATTEMPTS}. Ordinary 401/auth + * failures retain one refresh-same plus one sibling switch; usage/account + * limits rotate directly through distinct siblings. * * Used by non-streaming consumers (image generation, web search, completion * helpers). The streaming driver in `stream.ts` implements the same policy with @@ -129,21 +221,21 @@ export async function withAuth( const resolver = key; const signal = opts?.signal; - let lastKey = await resolveRetryKey(resolver, false, undefined, signal); - if (lastKey === undefined) throw missingKey(); + const initialKey = await resolveRetryKey(resolver, false, undefined, signal); + if (initialKey === undefined) throw missingKey(); + const state = createAuthRetryKeyState(initialKey); let lastError: unknown; try { - return await attempt(lastKey); + return await attempt(initialKey); } catch (error) { if (!isAuthError(error)) throw error; lastError = error; } - for (let i = 0; i < AUTH_RETRY_STEPS.length; i++) { - const nextKey = await resolveRetryKey(resolver, AUTH_RETRY_STEPS[i]!, lastError, signal, lastKey); - if (nextKey === undefined || nextKey === lastKey) continue; - lastKey = nextKey; + while (true) { + const nextKey = await resolveNextAuthRetryKey(state, resolver, lastError, signal); + if (nextKey === undefined) break; try { return await attempt(nextKey); } catch (error) { @@ -169,7 +261,7 @@ export interface OAuthAccessSource { rotateSessionCredential( provider: string, sessionId: string | undefined, - options?: { error?: unknown; signal?: AbortSignal }, + options?: { error?: unknown; signal?: AbortSignal; apiKey?: string; credentialId?: number }, ): Promise; } @@ -195,14 +287,18 @@ export interface WithOAuthAccessOptions { * `projectId`, `enterpriseUrl`) instead of bare API-key bytes. * * - initial → `getOAuthAccess` (or `opts.seed`). - * - step (b) → `getOAuthAccess` with `forceRefresh: true` (re-mint the SAME - * account; picks up peer/broker rotations). - * - step (c) → `rotateSessionCredential` then re-resolve (switch to a sibling). + * - 401/auth failure → one `getOAuthAccess` with `forceRefresh: true` for the + * current account, then sibling rotation. + * - usage-limit failure → `rotateSessionCredential` directly, without a + * force-refresh detour. * - * A step is skipped when it yields no access or the same `accessToken` that - * just failed; non-auth errors propagate immediately. Use this instead of - * hand-rolled `getOAuthAccess` + fetch flows so 401s and usage-limits rotate - * credentials instead of failing the call. + * A refresh-same step may retry a new bearer for the same credential identity; + * sibling rotation stops when it yields a credential identity + * (`credentialId ?? accessToken`) or bearer already attempted in this turn. + * All OAuth attempts share the {@link AUTH_RETRY_MAX_ATTEMPTS} ceiling. + * Non-auth errors propagate immediately. Use this instead of hand-rolled + * `getOAuthAccess` + fetch flows so 401s and usage-limits rotate credentials + * instead of failing the call. */ export async function withOAuthAccess( storage: OAuthAccessSource, @@ -221,34 +317,76 @@ export async function withOAuthAccess( ); } - const resolveStep = async (lastChance: boolean, error: unknown): Promise => { - try { - if (!lastChance) return await storage.getOAuthAccess(provider, sessionId, { forceRefresh: true, signal }); - await storage.rotateSessionCredential(provider, sessionId, { error, signal }); - return await storage.getOAuthAccess(provider, sessionId, { signal }); - } catch { - return undefined; - } - }; + const attemptedBearers = new Set([lastAccess.accessToken]); + const attemptedCredentialIdentities = new Set([oauthCredentialIdentity(lastAccess)]); + let attemptCount = 1; + let legacyAuthSwitchUsed = false; + let refreshedCurrent = false; + let attemptResult = await runOAuthAttempt(lastAccess, attempt, isAuthError); + if (attemptResult.ok) return attemptResult.result; - let lastError: unknown; - try { - return await attempt(lastAccess); - } catch (error) { - if (!isAuthError(error)) throw error; - lastError = error; - } + let lastError = attemptResult.error; + while (true) { + let next: OAuthAccess | undefined; + if (signal?.aborted || attemptCount >= AUTH_RETRY_MAX_ATTEMPTS) break; + const directRotation = isDirectCredentialRotationError(lastError); + if (!directRotation) { + if (legacyAuthSwitchUsed) break; + if (!refreshedCurrent) { + refreshedCurrent = true; + try { + next = await storage.getOAuthAccess(provider, sessionId, { forceRefresh: true, signal }); + } catch { + next = undefined; + } + if (signal?.aborted) break; + if (next) { + const bearer = next.accessToken; + if (!attemptedBearers.has(bearer) && attemptCount < AUTH_RETRY_MAX_ATTEMPTS) { + attemptedCredentialIdentities.add(oauthCredentialIdentity(next)); + attemptedBearers.add(bearer); + attemptCount += 1; + lastAccess = next; + attemptResult = await runOAuthAttempt(next, attempt, isAuthError); + if (attemptResult.ok) return attemptResult.result; + lastError = attemptResult.error; + continue; + } + } + } + } - for (const lastChance of AUTH_RETRY_STEPS) { - const next = await resolveStep(lastChance, lastError); - if (!next || next.accessToken === lastAccess.accessToken) continue; - lastAccess = next; + if (signal?.aborted || attemptCount >= AUTH_RETRY_MAX_ATTEMPTS) break; try { - return await attempt(next); - } catch (error) { - if (!isAuthError(error)) throw error; - lastError = error; + const rotated = await storage.rotateSessionCredential(provider, sessionId, { + error: lastError, + signal, + apiKey: lastAccess.accessToken, + credentialId: lastAccess.credentialId, + }); + if (!rotated) break; + next = await storage.getOAuthAccess(provider, sessionId, { signal }); + } catch { + next = undefined; } + if (signal?.aborted || !next) break; + const credentialIdentity = oauthCredentialIdentity(next); + if ( + attemptedCredentialIdentities.has(credentialIdentity) || + attemptedBearers.has(next.accessToken) || + attemptCount >= AUTH_RETRY_MAX_ATTEMPTS + ) { + break; + } + attemptedCredentialIdentities.add(credentialIdentity); + attemptedBearers.add(next.accessToken); + attemptCount += 1; + lastAccess = next; + refreshedCurrent = !directRotation; + if (!directRotation) legacyAuthSwitchUsed = true; + attemptResult = await runOAuthAttempt(next, attempt, isAuthError); + if (attemptResult.ok) return attemptResult.result; + lastError = attemptResult.error; } throw lastError; diff --git a/packages/ai/src/auth-storage.ts b/packages/ai/src/auth-storage.ts index 245c118c609..252d7ce6a4c 100644 --- a/packages/ai/src/auth-storage.ts +++ b/packages/ai/src/auth-storage.ts @@ -682,7 +682,7 @@ type AuthApiKeyOptions = { */ forceRefresh?: boolean; }; -type OAuthResolutionResult = { apiKey: string; credential: OAuthCredential }; +type OAuthResolutionResult = { apiKey: string; credential: OAuthCredential; credentialId?: number }; /** * Refreshed OAuth access plus identity metadata returned by @@ -3401,6 +3401,39 @@ export class AuthStorage { return results; } + async #resolveCredentialTarget( + provider: string, + sessionId: string | undefined, + options?: { credentialId?: number; apiKey?: string }, + ): Promise<{ type: AuthCredential["type"]; index: number; explicit: boolean } | undefined> { + const explicit = options?.credentialId !== undefined || options?.apiKey !== undefined; + if (explicit) { + const latestRows = this.#store.listAuthCredentials(provider); + this.#setStoredCredentials( + provider, + latestRows.map(row => ({ id: row.id, credential: row.credential })), + ); + } + if (options?.credentialId !== undefined) { + const stored = this.#getStoredCredentials(provider); + const index = stored.findIndex(entry => entry.id === options.credentialId); + const entry = index === -1 ? undefined : stored[index]; + if (entry) return { type: entry.credential.type, index, explicit: true }; + } + if (options?.apiKey !== undefined) { + const stored = this.#getStoredCredentials(provider); + for (let index = 0; index < stored.length; index++) { + const entry = stored[index]; + if (entry && (await this.#credentialMatchesApiKey(entry.credential, options.apiKey))) { + return { type: entry.credential.type, index, explicit: true }; + } + } + } + if (explicit) return undefined; + const sessionCredential = this.#getSessionCredential(provider, sessionId); + return sessionCredential ? { ...sessionCredential, explicit: false } : undefined; + } + /** * Marks the current session's credential as temporarily blocked due to usage limits. * Uses usage reports to determine accurate reset time when available. @@ -3411,20 +3444,19 @@ export class AuthStorage { async markUsageLimitReached( provider: string, sessionId: string | undefined, - options?: { retryAfterMs?: number; baseUrl?: string; modelId?: string; apiKey?: string; signal?: AbortSignal }, + options?: { + retryAfterMs?: number; + baseUrl?: string; + modelId?: string; + apiKey?: string; + credentialId?: number; + signal?: AbortSignal; + }, ): Promise { - let sessionCredential: { type: AuthCredential["type"]; index: number } | undefined; - if (options?.apiKey) { - const stored = this.#getStoredCredentials(provider); - for (let index = 0; index < stored.length; index++) { - const entry = stored[index]; - if (entry && (await this.#credentialMatchesApiKey(entry.credential, options.apiKey))) { - sessionCredential = { type: entry.credential.type, index }; - break; - } - } - } - sessionCredential ??= this.#getSessionCredential(provider, sessionId); + const sessionCredential = await this.#resolveCredentialTarget(provider, sessionId, { + credentialId: options?.credentialId, + apiKey: options?.apiKey, + }); if (!sessionCredential) return { switched: false }; const providerKey = this.#getProviderTypeKey(provider, sessionCredential.type); @@ -4209,7 +4241,7 @@ export class AuthStorage { } } this.#recordSessionCredential(provider, sessionId, "oauth", selection.index); - return { apiKey: result.apiKey, credential: updated }; + return { apiKey: result.apiKey, credential: updated, credentialId }; } catch (error) { const errorMsg = String(error); // Only remove credentials for definitive auth failures @@ -4434,9 +4466,10 @@ export class AuthStorage { } const resolved = await this.#resolveOAuthSelection(provider, sessionId, options); if (!resolved) return undefined; - const { credential } = resolved; + const { credential, credentialId } = resolved; return { accessToken: credential.access, + credentialId, accountId: credential.accountId, email: credential.email, projectId: credential.projectId, @@ -4922,36 +4955,45 @@ export class AuthStorage { } /** - * Rotate away from the session's current credential after a retryable auth - * error — step (c) of the auth-retry policy. Stateless: looks up the - * session-sticky credential (no API-key matching needed), applies the - * storage action for the error class, then clears the sticky so the next - * {@link AuthStorage.getApiKey} for this session picks a sibling. + * Rotate away from the credential that failed after a retryable auth error — + * step (c) of the auth-retry policy. Prefer the failed stored row id supplied + * in `options.credentialId`, then the failed bearer supplied in + * `options.apiKey`, so overlapping requests cannot redirect rotation through + * stale session stickiness. Fall back to the session-sticky credential only + * when neither explicit target is available. If an explicit target is supplied + * but no longer matches storage, return `false` without mutating sticky state. + * Apply the storage action for the error class, then let the next resolve + * select a sibling. * * - usage-limit / account-rate-limit error → {@link AuthStorage.markUsageLimitReached} * (temporary block via its own backoff — default plus server usage-report * reset; sticky left intact so the next resolve re-ranks around the block). * - otherwise (hard 401 / auth failure) → mark the credential suspect (or - * reload when no broker hook is wired) and block it, then drop the sticky. + * reload when no broker hook is wired) and block it, then drop matching + * sticky state. * * Returns whether another usable credential of the same type remains. */ async rotateSessionCredential( provider: string, sessionId: string | undefined, - options?: { error?: unknown; modelId?: string; apiKey?: string; signal?: AbortSignal }, + options?: { error?: unknown; modelId?: string; apiKey?: string; credentialId?: number; signal?: AbortSignal }, ): Promise { - const sessionCredential = this.#getSessionCredential(provider, sessionId); + const sessionCredential = await this.#resolveCredentialTarget(provider, sessionId, { + credentialId: options?.credentialId, + apiKey: options?.apiKey, + }); if (!sessionCredential) return false; const error = options?.error; const status = AIError.status(error); const message = error instanceof Error ? error.message : typeof error === "string" ? error : undefined; - if (isUsageLimitOutcome(status, message)) { + if (AIError.isUsageLimit(error) || isUsageLimitOutcome(status, message)) { return ( await this.markUsageLimitReached(provider, sessionId, { modelId: options?.modelId, apiKey: options?.apiKey, + credentialId: options?.credentialId, signal: options?.signal, }) ).switched; @@ -4967,7 +5009,13 @@ export class AuthStorage { !this.#isCredentialBlocked(provider, providerKey, index), ); const target = this.#getStoredCredentials(provider)[sessionCredential.index]; - this.#clearSessionCredential(provider, sessionId); + const sticky = this.#getSessionCredential(provider, sessionId); + if ( + !sessionCredential.explicit || + (sticky?.type === sessionCredential.type && sticky.index === sessionCredential.index) + ) { + this.#clearSessionCredential(provider, sessionId); + } this.#markCredentialBlocked( provider, providerKey, @@ -5005,12 +5053,18 @@ export class AuthStorage { */ resolver(provider: string, options?: { sessionId?: string; baseUrl?: string; modelId?: string }): ApiKeyResolver { const { sessionId, baseUrl, modelId } = options ?? {}; - return async ({ lastChance, error, signal }) => { + return async ({ lastChance, error, signal, previousKey }) => { if (error === undefined) { return this.getApiKey(provider, sessionId, { baseUrl, modelId, signal }); } if (lastChance) { - await this.rotateSessionCredential(provider, sessionId, { error, modelId, signal }); + const rotated = await this.rotateSessionCredential(provider, sessionId, { + error, + modelId, + signal, + apiKey: previousKey, + }); + if (!rotated) return undefined; return this.getApiKey(provider, sessionId, { baseUrl, modelId, signal }); } return this.getApiKey(provider, sessionId, { baseUrl, modelId, forceRefresh: true, signal }); diff --git a/packages/ai/src/error/auth-classify.ts b/packages/ai/src/error/auth-classify.ts index 2245007fa04..9234ada9bbf 100644 --- a/packages/ai/src/error/auth-classify.ts +++ b/packages/ai/src/error/auth-classify.ts @@ -1,5 +1,5 @@ import { extractHttpStatusFromError } from "@oh-my-pi/pi-utils"; -import { isOAuthExpiry } from "./flags"; +import { isOAuthExpiry, isUsageLimit } from "./flags"; import { isUsageLimitOutcome } from "./rate-limit"; /** @@ -21,6 +21,7 @@ export function isDefinitiveOAuthFailure(errorMsg: string): boolean { * upstream-backoff lane. */ export function isAuthRetryableError(error: unknown): boolean { + if (isUsageLimit(error)) return true; const httpStatus = extractHttpStatusFromError(error); if (httpStatus === 401) return true; const message = error instanceof Error ? error.message : typeof error === "string" ? error : undefined; diff --git a/packages/ai/src/stream.ts b/packages/ai/src/stream.ts index 7c66e6c9f92..98ad2458d40 100644 --- a/packages/ai/src/stream.ts +++ b/packages/ai/src/stream.ts @@ -17,7 +17,7 @@ import { CATALOG_PROVIDERS, type ProviderCatalogEntry } from "@oh-my-pi/pi-catal import { CODEX_BASE_URL } from "@oh-my-pi/pi-catalog/wire/codex"; import { $env, $pickenv, getConfigRootDir, isEnoent, logger, withExtraCaFetch } from "@oh-my-pi/pi-utils"; import { getCustomApi } from "./api-registry"; -import { AUTH_RETRY_STEPS, isApiKeyResolver, resolveRetryKey } from "./auth-retry"; +import { createAuthRetryKeyState, isApiKeyResolver, resolveNextAuthRetryKey } from "./auth-retry"; import * as AIError from "./error"; import { ProviderHttpError } from "./error"; import { isUsageLimitOutcome } from "./error/rate-limit"; @@ -978,17 +978,19 @@ function isRetryableUpstreamError(error: unknown, status: number | undefined, me // per-minute caps) classify as RATE_LIMIT_EXCEEDED in // `parseRateLimitReason` and stay in the provider's own backoff layer // instead of burning siblings. + if (AIError.isUsageLimit(error)) return true; if (status === 401) return true; - void error; return isUsageLimitOutcome(status, message); } function createAssistantAuthError(message: AssistantMessage): Error { const text = message.errorMessage ?? "Provider authentication failed"; const status = extractStatusFromAssistantError(message); - return status === undefined - ? new AIError.ProviderResponseError(text, { kind: "runtime" }) - : new ProviderHttpError(text, status); + const error = + status === undefined + ? new AIError.ProviderResponseError(text, { kind: "runtime" }) + : new ProviderHttpError(text, status); + return typeof message.errorId === "number" ? AIError.attach(error, message.errorId) : error; } function emitBufferedEvents(stream: AssistantMessageEventStream, events: AssistantMessageEvent[]): void { @@ -1012,12 +1014,12 @@ export function streamSimple( if (apiKeyResolver) { const outer = new AssistantMessageEventStream(); const signal = requestOptions?.signal; - // One inner attempt against a resolved string key. When - // `captureAuthFailure` is set, a retryable auth error that arrives before - // any replay-unsafe event is buffered and returned (so the caller can - // retry with a fresh key) instead of surfaced. The terminal attempt - // clears the flag and emits whatever it gets. - const runAttempt = async (apiKey: string, captureAuthFailure: boolean): Promise => { + // One inner attempt against a resolved string key. A retryable auth error + // that arrives before any replay-unsafe event is buffered and returned + // (so the caller can retry with a fresh key) instead of surfaced. Once any + // non-start event escapes, retry is no longer safe and the failure is + // emitted directly. + const runAttempt = async (apiKey: string): Promise => { const bufferedEvents: AssistantMessageEvent[] = []; let emittedReplayUnsafeEvent = false; const flushBuffered = (): void => { @@ -1034,7 +1036,6 @@ export function streamSimple( } if ( !emittedReplayUnsafeEvent && - captureAuthFailure && event.type === "error" && isRetryableUpstreamError( event.error, @@ -1054,7 +1055,6 @@ export function streamSimple( } catch (error) { if ( !emittedReplayUnsafeEvent && - captureAuthFailure && isRetryableUpstreamError( error, AIError.status(error), @@ -1096,27 +1096,16 @@ export function streamSimple( outer.fail(new AIError.MissingApiKeyError(model.provider)); return; } - let failure = await runAttempt(lastKey, true); + const retryState = createAuthRetryKeyState(lastKey); + let failure = await runAttempt(lastKey); if (!failure) return; - // a/b/c policy: refresh the same account (lastChance=false), then - // switch to a sibling (lastChance=true). A step is skipped when the - // resolver yields the same key it just tried or `undefined`; the - // final step's attempt clears the capture flag so it emits directly. - for (let step = 0; step < AUTH_RETRY_STEPS.length; step++) { + while (true) { // Caller aborted between attempts: don't mint a fresh token or fire // another doomed request — emit the captured failure instead. if (signal?.aborted) break; - const nextKey = await resolveRetryKey( - apiKeyResolver, - AUTH_RETRY_STEPS[step]!, - failure.error, - signal, - lastKey, - ); - if (nextKey === undefined || nextKey === lastKey) continue; - lastKey = nextKey; - const isLastStep = step === AUTH_RETRY_STEPS.length - 1; - const next = await runAttempt(nextKey, !isLastStep); + const nextKey = await resolveNextAuthRetryKey(retryState, apiKeyResolver, failure.error, signal); + if (nextKey === undefined) break; + const next = await runAttempt(nextKey); if (!next) return; failure = next; } diff --git a/packages/ai/test/auth-retry.test.ts b/packages/ai/test/auth-retry.test.ts index c851d8f3b13..84b0969dbb9 100644 --- a/packages/ai/test/auth-retry.test.ts +++ b/packages/ai/test/auth-retry.test.ts @@ -1,6 +1,14 @@ import { describe, expect, it } from "bun:test"; import type { ApiKeyResolveContext, OAuthAccess, OAuthAccessSource } from "@oh-my-pi/pi-ai"; -import { isApiKeyResolver, isAuthRetryableError, resolveApiKeyOnce, withAuth, withOAuthAccess } from "@oh-my-pi/pi-ai"; +import { + AUTH_RETRY_MAX_ATTEMPTS, + isApiKeyResolver, + isAuthRetryableError, + resolveApiKeyOnce, + withAuth, + withOAuthAccess, +} from "@oh-my-pi/pi-ai"; +import { ProviderHttpError } from "@oh-my-pi/pi-ai/error"; function authError(status = 401): Error & { status: number } { return Object.assign(new Error(`${status} authentication_error`), { status }); @@ -12,6 +20,10 @@ function usageLimitError(): Error & { status: number } { }); } +function opaque429Error(): Error & { status: number } { + return Object.assign(new Error(""), { status: 429 }); +} + describe("isApiKeyResolver / resolveApiKeyOnce", () => { it("narrows resolver vs static key and resolves the initial value", async () => { expect(isApiKeyResolver("static")).toBe(false); @@ -36,6 +48,12 @@ describe("isAuthRetryableError", () => { it("treats 401 and usage-limit phrasing as retryable, everything else as not", () => { expect(isAuthRetryableError(authError(401))).toBe(true); expect(isAuthRetryableError(usageLimitError())).toBe(true); + expect( + isAuthRetryableError(new ProviderHttpError("Generic provider failure", 429, { code: "insufficient_quota" })), + ).toBe(true); + expect( + isAuthRetryableError(new ProviderHttpError("Generic provider failure", 429, { code: "rate_limit_error" })), + ).toBe(false); // A 429 whose body names the *account's* rate limit is rotatable (switch // account), even though it isn't a 401 and isn't phrased "usage limit". expect( @@ -111,6 +129,58 @@ describe("withAuth", () => { ]); }); + it("does not exhaust every sibling on pure 401 auth failures", async () => { + const keys: string[] = []; + const contexts: ApiKeyResolveContext[] = []; + const pool = ["k0", "k1", "k2", "k3"]; + let resolveIndex = 0; + let lastError: unknown; + let caught: unknown; + + try { + await withAuth( + ctx => { + contexts.push(ctx); + return ctx.error === undefined ? pool[0] : pool[++resolveIndex]; + }, + async key => { + keys.push(key); + lastError = authError(); + throw lastError; + }, + ); + } catch (error) { + caught = error; + } + + expect(caught).toBe(lastError); + expect(keys).toEqual(["k0", "k1", "k2"]); + expect(contexts.map(ctx => ctx.lastChance)).toEqual([false, false, true]); + }); + + it("continues quota rotation when a refreshed 401 retry becomes a usage limit", async () => { + const keys: string[] = []; + const contexts: ApiKeyResolveContext[] = []; + const pool = ["k0", "k1", "k2", "k3"]; + let resolveIndex = 0; + const result = await withAuth( + ctx => { + contexts.push(ctx); + return ctx.error === undefined ? pool[0] : pool[++resolveIndex]; + }, + async key => { + keys.push(key); + if (key === "k3") return "success"; + if (key === "k0") throw authError(); + throw usageLimitError(); + }, + ); + + expect(result).toBe("success"); + expect(keys).toEqual(pool); + expect(contexts.map(ctx => ctx.lastChance)).toEqual([false, false, true, true]); + }); + it("switches accounts before refreshing the same account on usage limits", async () => { const keys: string[] = []; const contexts: ApiKeyResolveContext[] = []; @@ -133,6 +203,115 @@ describe("withAuth", () => { ]); }); + it("switches accounts before refreshing on opaque 429 usage outcomes", async () => { + const keys: string[] = []; + const contexts: ApiKeyResolveContext[] = []; + const result = await withAuth( + ctx => { + contexts.push(ctx); + return ctx.error === undefined ? "k0" : ctx.lastChance ? "k2" : "k1"; + }, + async key => { + keys.push(key); + if (key === "k2") return "success"; + throw opaque429Error(); + }, + ); + expect(result).toBe("success"); + expect(keys).toEqual(["k0", "k2"]); + expect(contexts.map(ctx => ({ lastChance: ctx.lastChance, hasError: ctx.error !== undefined }))).toEqual([ + { lastChance: false, hasError: false }, + { lastChance: true, hasError: true }, + ]); + }); + + it("rotates through every distinct sibling after consecutive usage limits", async () => { + const keys: string[] = []; + const contexts: ApiKeyResolveContext[] = []; + const pool = ["k0", "k1", "k2", "k3"]; + let nextSibling = 0; + const result = await withAuth( + ctx => { + contexts.push(ctx); + return ctx.error === undefined ? pool[0] : pool[++nextSibling]; + }, + async key => { + keys.push(key); + if (key === "k3") return "success"; + throw usageLimitError(); + }, + ); + + expect(result).toBe("success"); + expect(keys).toEqual(pool); + expect(contexts.map(ctx => ctx.lastChance)).toEqual([false, true, true, true]); + }); + + it("stops usage-limit rotation before retrying an already-attempted credential", async () => { + const keys: string[] = []; + const errors = [usageLimitError(), usageLimitError()]; + const resolved = ["k0", "k1", "k0"]; + let resolveIndex = 0; + let attemptIndex = 0; + + await expect( + withAuth( + () => resolved[resolveIndex++], + async key => { + keys.push(key); + throw errors[Math.min(attemptIndex++, errors.length - 1)]!; + }, + ), + ).rejects.toBe(errors[1]); + expect(keys).toEqual(["k0", "k1"]); + }); + + it("caps endlessly unique resolver retries", async () => { + const keys: string[] = []; + let resolveIndex = 0; + let lastError: unknown; + let caught: unknown; + + try { + await withAuth( + () => `k${resolveIndex++}`, + async key => { + keys.push(key); + lastError = usageLimitError(); + throw lastError; + }, + ); + } catch (error) { + caught = error; + } + + expect(caught).toBe(lastError); + expect(keys).toHaveLength(AUTH_RETRY_MAX_ATTEMPTS); + expect(resolveIndex).toBe(AUTH_RETRY_MAX_ATTEMPTS); + }); + + it("does not attempt a retry key resolved after abort", async () => { + const controller = new AbortController(); + const keys: string[] = []; + const original = usageLimitError(); + + await expect( + withAuth( + ctx => { + if (ctx.error === undefined) return "k0"; + controller.abort(); + return "k1"; + }, + async key => { + keys.push(key); + throw original; + }, + { signal: controller.signal }, + ), + ).rejects.toBe(original); + expect(keys).toEqual(["k0"]); + }); + it("stops retrying when the resolver returns undefined", async () => { const keys: string[] = []; const original = authError(); @@ -255,22 +434,57 @@ describe("withOAuthAccess", () => { expect(storage.calls).toEqual([{ forceRefresh: undefined }, { forceRefresh: true }]); }); - it("skips an unchanged force-refresh token and rotates to a sibling", async () => { + it("tries a refreshed bearer for the same credential id on 401 before rotating", async () => { const storage = fakeStorage({ - initial: access("dead"), - forced: access("dead"), - rotated: access("sibling"), + initial: access("stale", { credentialId: 7 }), + forced: access("fresh", { credentialId: 7 }), + rotated: access("sibling", { credentialId: 8 }), }); const attempts: string[] = []; const result = await withOAuthAccess(storage, "prov", async a => { attempts.push(a.accessToken); - if (a.accessToken === "dead") throw usageLimitError(); + if (a.accessToken === "stale") throw authError(); return "ok"; }); expect(result).toBe("ok"); - // "dead" must not be re-attempted after the no-op force refresh. - expect(attempts).toEqual(["dead", "sibling"]); - expect(storage.calls).toEqual([ + expect(attempts).toEqual(["stale", "fresh"]); + expect(storage.calls).toEqual([{ forceRefresh: undefined }, { forceRefresh: true }]); + }); + + it("does not exhaust every OAuth sibling on pure 401 auth failures", async () => { + const attempts: string[] = []; + const calls: Array<{ forceRefresh: boolean | undefined } | "rotate"> = []; + const rotated = [access("sibling-1", { credentialId: 2 }), access("sibling-2", { credentialId: 3 })]; + let rotateIndex = 0; + let lastError: unknown; + let caught: unknown; + const storage: OAuthAccessSource = { + async getOAuthAccess(_provider, _sessionId, options) { + calls.push({ forceRefresh: options?.forceRefresh }); + if (options?.forceRefresh) return access("fresh", { credentialId: 1 }); + if (rotateIndex > 0) return rotated[rotateIndex - 1]; + return access("stale", { credentialId: 1 }); + }, + async rotateSessionCredential() { + calls.push("rotate"); + rotateIndex += 1; + return true; + }, + }; + + try { + await withOAuthAccess(storage, "prov", async a => { + attempts.push(a.accessToken); + lastError = authError(); + throw lastError; + }); + } catch (error) { + caught = error; + } + + expect(caught).toBe(lastError); + expect(attempts).toEqual(["stale", "fresh", "sibling-1"]); + expect(calls).toEqual([ { forceRefresh: undefined }, { forceRefresh: true }, "rotate", @@ -278,6 +492,81 @@ describe("withOAuthAccess", () => { ]); }); + it("rotates directly to a sibling on usage limits", async () => { + const storage = fakeStorage({ + initial: access("dead"), + rotated: access("sibling"), + }); + const attempts: string[] = []; + const result = await withOAuthAccess(storage, "prov", async a => { + attempts.push(a.accessToken); + if (a.accessToken === "dead") throw usageLimitError(); + return "ok"; + }); + expect(result).toBe("ok"); + // Usage-limit failures burn/rotate the exhausted account directly; a + // force-refresh of the same account would duplicate the failed side effect. + expect(attempts).toEqual(["dead", "sibling"]); + expect(storage.calls).toEqual([{ forceRefresh: undefined }, "rotate", { forceRefresh: undefined }]); + }); + + it("passes the failed OAuth bearer to rotation", async () => { + const rotationTargets: Array<{ apiKey: string | undefined; credentialId: number | undefined }> = []; + const storage: OAuthAccessSource = { + async getOAuthAccess() { + return rotationTargets.length === 0 + ? access("dead", { credentialId: 17 }) + : access("sibling", { credentialId: 18 }); + }, + async rotateSessionCredential(_provider, _sessionId, options) { + rotationTargets.push({ apiKey: options?.apiKey, credentialId: options?.credentialId }); + return true; + }, + }; + const attempts: string[] = []; + const result = await withOAuthAccess(storage, "prov", async a => { + attempts.push(a.accessToken); + if (a.accessToken === "dead") throw usageLimitError(); + return "ok"; + }); + + expect(result).toBe("ok"); + expect(attempts).toEqual(["dead", "sibling"]); + expect(rotationTargets).toEqual([{ apiKey: "dead", credentialId: 17 }]); + }); + + it("caps endlessly unique OAuth rotation attempts", async () => { + const attempts: string[] = []; + let nextCredential = 0; + let rotateCalls = 0; + let lastError: unknown; + let caught: unknown; + const storage: OAuthAccessSource = { + async getOAuthAccess() { + const credentialId = nextCredential++; + return access(`token-${credentialId}`, { credentialId }); + }, + async rotateSessionCredential() { + rotateCalls += 1; + return true; + }, + }; + + try { + await withOAuthAccess(storage, "prov", async a => { + attempts.push(a.accessToken); + lastError = usageLimitError(); + throw lastError; + }); + } catch (error) { + caught = error; + } + + expect(caught).toBe(lastError); + expect(attempts).toHaveLength(AUTH_RETRY_MAX_ATTEMPTS); + expect(rotateCalls).toBe(AUTH_RETRY_MAX_ATTEMPTS - 1); + }); + it("propagates non-auth errors immediately and surfaces the last auth error when exhausted", async () => { const boom = new Error("syntax error"); await expect( diff --git a/packages/ai/test/auth-storage-broker-no-sentinel.test.ts b/packages/ai/test/auth-storage-broker-no-sentinel.test.ts index 36054716a9c..909ffe02531 100644 --- a/packages/ai/test/auth-storage-broker-no-sentinel.test.ts +++ b/packages/ai/test/auth-storage-broker-no-sentinel.test.ts @@ -70,6 +70,7 @@ describe("AuthStorage broker sentinel refresh", () => { expect(access).toEqual({ accessToken: "broker-access-rotated", + credentialId: expect.any(Number), accountId: "broker-account", email: "broker@example.com", projectId: "broker-project", diff --git a/packages/ai/test/auth-storage-force-refresh-rotate.test.ts b/packages/ai/test/auth-storage-force-refresh-rotate.test.ts index 94856c83e21..3a5218a260b 100644 --- a/packages/ai/test/auth-storage-force-refresh-rotate.test.ts +++ b/packages/ai/test/auth-storage-force-refresh-rotate.test.ts @@ -3,6 +3,7 @@ import * as fs from "node:fs/promises"; import * as os from "node:os"; import * as path from "node:path"; import { type AuthCredentialStore, AuthStorage, SqliteAuthCredentialStore } from "@oh-my-pi/pi-ai/auth-storage"; +import { ProviderHttpError } from "@oh-my-pi/pi-ai/error"; import { registerOAuthProvider, unregisterOAuthProviders } from "@oh-my-pi/pi-ai/registry/oauth"; import { removeWithRetries } from "../../utils/src/temp"; @@ -100,6 +101,28 @@ describe("AuthStorage forceRefresh + rotateSessionCredential", () => { expect(after).toBe("minted-access"); }); + test("getOAuthAccess includes a stable credentialId across cached and forced refresh resolves", async () => { + if (!authStorage) throw new Error("test setup failed"); + registerProvider(); + await authStorage.set(PROVIDER, [ + { type: "oauth", access: "cached-access", refresh: "cached-refresh", expires: farExpiry() }, + ]); + + const cached = await authStorage.getOAuthAccess(PROVIDER, "oauth-identity"); + expect(cached?.accessToken).toBe("cached-access"); + expect(typeof cached?.credentialId).toBe("number"); + const credentialId = cached?.credentialId; + if (credentialId === undefined) throw new Error("expected OAuth credential id"); + + const forced = await authStorage.getOAuthAccess(PROVIDER, "oauth-identity", { forceRefresh: true }); + expect(forced?.accessToken).toBe("minted-access"); + expect(forced?.credentialId).toBe(credentialId); + + const after = await authStorage.getOAuthAccess(PROVIDER, "oauth-identity"); + expect(after?.accessToken).toBe("minted-access"); + expect(after?.credentialId).toBe(credentialId); + }); + test("rotateSessionCredential(401) blocks + clears the sticky and rotates to a sibling", async () => { if (!authStorage) throw new Error("test setup failed"); registerProvider(); @@ -123,6 +146,117 @@ describe("AuthStorage forceRefresh + rotateSessionCredential", () => { expect(second).not.toBe(first); }); + test("resolver rotates the credential matching previousKey instead of a stale sticky", async () => { + if (!authStorage) throw new Error("test setup failed"); + await authStorage.set(PROVIDER, [ + { type: "api_key", key: "sticky-key" }, + { type: "api_key", key: "failed-key" }, + { type: "api_key", key: "survivor-key" }, + ]); + + const sessionId = "resolver-previous-key"; + const sticky = await authStorage.getApiKey(PROVIDER, sessionId); + if (!sticky) throw new Error("expected initial sticky credential"); + const failed = sticky === "failed-key" ? "sticky-key" : "failed-key"; + const resolver = authStorage.resolver(PROVIDER, { sessionId }); + + const retry = await resolver({ + lastChance: true, + error: authError(), + previousKey: failed, + }); + + expect(retry).toBe(sticky); + expect(retry).not.toBe(failed); + + const laterSelections = new Set(); + for (let index = 0; index < 6; index += 1) { + const selected = await authStorage.getApiKey(PROVIDER); + if (selected) laterSelections.add(selected); + } + expect(laterSelections.has(failed)).toBe(false); + expect(laterSelections.has(sticky)).toBe(true); + }); + + test("explicit missing rotation targets do not fall back to stale stickiness", async () => { + if (!authStorage || !store) throw new Error("test setup failed"); + await authStorage.set(PROVIDER, [ + { type: "api_key", key: "acc-A" }, + { type: "api_key", key: "acc-B" }, + { type: "api_key", key: "acc-C" }, + ]); + + const sessionId = "explicit-missing-target"; + const sticky = await authStorage.getApiKey(PROVIDER, sessionId); + if (!sticky) throw new Error("expected sticky credential"); + const maxCredentialId = Math.max(...store.listAuthCredentials(PROVIDER).map(row => row.id)); + const missingCredentialId = maxCredentialId + 1000; + + const rotated = await authStorage.rotateSessionCredential(PROVIDER, sessionId, { + error: authError(), + apiKey: "missing-or-changed-failed-bearer", + }); + expect(rotated).toBe(false); + expect(await authStorage.getApiKey(PROVIDER, sessionId)).toBe(sticky); + + const rotatedByMissingId = await authStorage.rotateSessionCredential(PROVIDER, sessionId, { + error: authError(), + credentialId: missingCredentialId, + }); + expect(rotatedByMissingId).toBe(false); + expect(await authStorage.getApiKey(PROVIDER, sessionId)).toBe(sticky); + + const marked = await authStorage.markUsageLimitReached(PROVIDER, sessionId, { + apiKey: "missing-or-changed-failed-bearer", + }); + expect(marked.switched).toBe(false); + expect(await authStorage.getApiKey(PROVIDER, sessionId)).toBe(sticky); + + const markedByMissingId = await authStorage.markUsageLimitReached(PROVIDER, sessionId, { + credentialId: missingCredentialId, + }); + expect(markedByMissingId.switched).toBe(false); + expect(await authStorage.getApiKey(PROVIDER, sessionId)).toBe(sticky); + }); + + test("credentialId rotation targets the failed row after bearer changes without clearing stale sticky", async () => { + if (!authStorage || !store) throw new Error("test setup failed"); + await authStorage.set(PROVIDER, [ + { type: "api_key", key: "acc-A" }, + { type: "api_key", key: "acc-B" }, + { type: "api_key", key: "acc-C" }, + ]); + + const sessionId = "credential-id-target"; + const sticky = await authStorage.getApiKey(PROVIDER, sessionId); + if (!sticky) throw new Error("expected sticky credential"); + const targetRow = store.listAuthCredentials(PROVIDER).find(row => { + const credential = row.credential; + return credential.type === "api_key" && credential.key !== sticky; + }); + if (targetRow?.credential.type !== "api_key") throw new Error("expected non-sticky target row"); + const oldKey = targetRow.credential.key; + const changedKey = `${oldKey}-rotated`; + store.updateAuthCredential(targetRow.id, { type: "api_key", key: changedKey }); + await authStorage.reload(); + + const rotated = await authStorage.rotateSessionCredential(PROVIDER, sessionId, { + error: authError(), + apiKey: oldKey, + credentialId: targetRow.id, + }); + expect(rotated).toBe(true); + expect(await authStorage.getApiKey(PROVIDER, sessionId)).toBe(sticky); + + const laterSelections = new Set(); + for (let index = 0; index < 6; index += 1) { + const selected = await authStorage.getApiKey(PROVIDER); + if (selected) laterSelections.add(selected); + } + expect(laterSelections.has(changedKey)).toBe(false); + expect(laterSelections.has(sticky)).toBe(true); + }); + test("rotateSessionCredential(usage-limit) delegates to markUsageLimitReached", async () => { if (!authStorage) throw new Error("test setup failed"); registerProvider(); @@ -150,6 +284,27 @@ describe("AuthStorage forceRefresh + rotateSessionCredential", () => { expect(second).not.toBe(first); }); + test("rotateSessionCredential treats structured usage codes as quota blocks despite generic messages", async () => { + if (!authStorage) throw new Error("test setup failed"); + registerProvider(); + await authStorage.set(PROVIDER, [ + { type: "oauth", access: "acc-A", refresh: "ref-A", expires: farExpiry() }, + { type: "oauth", access: "acc-B", refresh: "ref-B", expires: farExpiry() }, + ]); + + const first = await authStorage.getApiKey(PROVIDER, "machine-code-quota"); + const usageLimitSpy = vi.spyOn(authStorage, "markUsageLimitReached"); + const rotated = await authStorage.rotateSessionCredential(PROVIDER, "machine-code-quota", { + error: new ProviderHttpError("Generic provider failure", 401, { code: "insufficient_quota" }), + }); + + expect(rotated).toBe(true); + expect(usageLimitSpy).toHaveBeenCalledTimes(1); + expect(usageLimitSpy.mock.calls[0]?.[0]).toBe(PROVIDER); + expect(usageLimitSpy.mock.calls[0]?.[1]).toBe("machine-code-quota"); + expect(await authStorage.getApiKey(PROVIDER, "machine-code-quota")).not.toBe(first); + }); + test("rotateSessionCredential(xAI credits 403) blocks the exhausted account and rotates", async () => { if (!authStorage) throw new Error("test setup failed"); registerProvider(); diff --git a/packages/ai/test/rate-limit-utils.test.ts b/packages/ai/test/rate-limit-utils.test.ts index b2c9f57dd61..0b39b333b4f 100644 --- a/packages/ai/test/rate-limit-utils.test.ts +++ b/packages/ai/test/rate-limit-utils.test.ts @@ -1,4 +1,5 @@ import { describe, expect, it } from "bun:test"; +import { ProviderHttpError } from "@oh-my-pi/pi-ai/error"; import { isUsageLimit } from "@oh-my-pi/pi-ai/error/flags"; import { calculateRateLimitBackoffMs, @@ -149,6 +150,15 @@ describe("isUsageLimit", () => { expect(isUsageLimitStatus(429)).toBe(true); expect(isUsageLimitStatus(400)).toBe(false); }); + + it("detects structured provider usage codes without quota wording", () => { + expect(isUsageLimit(new ProviderHttpError("Generic provider failure", 429, { code: "insufficient_quota" }))).toBe( + true, + ); + expect(isUsageLimit(new ProviderHttpError("Generic provider failure", 429, { code: "rate_limit_error" }))).toBe( + false, + ); + }); }); describe("isUsageLimitOutcome", () => { diff --git a/packages/ai/test/stream-auth-retry.test.ts b/packages/ai/test/stream-auth-retry.test.ts index a3e7e5beea5..1ede17ae6ac 100644 --- a/packages/ai/test/stream-auth-retry.test.ts +++ b/packages/ai/test/stream-auth-retry.test.ts @@ -1,6 +1,8 @@ import { afterEach, describe, expect, it } from "bun:test"; import type { ApiKeyResolveContext } from "@oh-my-pi/pi-ai"; import { registerCustomApi, unregisterCustomApis } from "@oh-my-pi/pi-ai"; +import { ProviderHttpError } from "@oh-my-pi/pi-ai/error"; +import { classify } from "@oh-my-pi/pi-ai/error/flags"; import { streamSimple } from "@oh-my-pi/pi-ai/stream"; import type { Api, AssistantMessage, Context, Model, SimpleStreamOptions, Usage } from "@oh-my-pi/pi-ai/types"; import { AssistantMessageEventStream } from "@oh-my-pi/pi-ai/utils/event-stream"; @@ -32,8 +34,8 @@ function assistant(content: string[] = []): AssistantMessage { }; } -function assistantError(errorMessage: string, errorStatus?: number): AssistantMessage { - return { ...assistant(), stopReason: "error", errorMessage, errorStatus }; +function assistantError(errorMessage: string, errorStatus?: number, errorId?: number): AssistantMessage { + return { ...assistant(), stopReason: "error", errorMessage, errorStatus, errorId }; } function authError(): Error & { status: number } { @@ -357,6 +359,124 @@ describe("streamSimple resolver auth retry", () => { expect(keys).toEqual(["old-key", "new-key"]); }); + it("rotates on a machine-code-only usage error event before content", async () => { + const keys: unknown[] = []; + const contexts: ApiKeyResolveContext[] = []; + const errorId = classify(new ProviderHttpError("Generic provider failure", 429, { code: "insufficient_quota" })); + registerCustomApi( + API, + (_model: Model, _context: Context, options?: SimpleStreamOptions) => { + pushKey(keys, options); + const stream = new AssistantMessageEventStream(); + queueMicrotask(() => { + if (keys.length === 1) { + stream.push({ type: "start", partial: assistant() }); + stream.push({ + type: "error", + reason: "error", + error: assistantError("Generic provider failure", 429, errorId), + }); + return; + } + ok(stream); + }); + return stream; + }, + SOURCE_ID, + ); + + const stream = streamSimple(model(), context, { + apiKey: async ctx => { + contexts.push(ctx); + return ctx.error === undefined ? "old-key" : "new-key"; + }, + }); + for await (const _event of stream) { + // drain + } + + expect((await stream.result()).content).toEqual([{ type: "text", text: "ok" }]); + expect(keys).toEqual(["old-key", "new-key"]); + expect(contexts.map(ctx => ctx.lastChance)).toEqual([false, true]); + }); + + it("rotates through every distinct sibling while usage failures remain replay-safe", async () => { + const keys: unknown[] = []; + const eventTypes: string[] = []; + const contexts: ApiKeyResolveContext[] = []; + const pool = ["credential-A", "credential-B", "credential-C", "credential-D"]; + let nextSibling = 0; + registerCustomApi( + API, + (_model: Model, _context: Context, options?: SimpleStreamOptions) => { + pushKey(keys, options); + const stream = new AssistantMessageEventStream(); + queueMicrotask(() => { + if (options?.apiKey === "credential-D") { + ok(stream); + return; + } + stream.push({ type: "start", partial: assistant() }); + stream.push({ + type: "error", + reason: "error", + error: assistantError("You have hit your ChatGPT usage limit (pro plan). Try again later.", 429), + }); + }); + return stream; + }, + SOURCE_ID, + ); + + const stream = streamSimple(model(), context, { + apiKey: async ctx => { + contexts.push(ctx); + return ctx.error === undefined ? pool[0] : pool[++nextSibling]; + }, + }); + for await (const event of stream) { + eventTypes.push(event.type); + } + + expect((await stream.result()).content).toEqual([{ type: "text", text: "ok" }]); + expect(keys).toEqual(pool); + expect(contexts.map(ctx => ctx.lastChance)).toEqual([false, true, true, true]); + expect(eventTypes).toEqual(["start", "text_start", "text_delta", "text_end", "done"]); + }); + + it("stops replay-safe usage rotation when the resolver cycles to an attempted credential", async () => { + const keys: unknown[] = []; + const resolved = ["credential-A", "credential-B", "credential-A"]; + let resolveIndex = 0; + registerCustomApi( + API, + (_model: Model, _context: Context, options?: SimpleStreamOptions) => { + pushKey(keys, options); + const stream = new AssistantMessageEventStream(); + queueMicrotask(() => { + stream.push({ type: "start", partial: assistant() }); + stream.push({ + type: "error", + reason: "error", + error: assistantError("You have hit your ChatGPT usage limit (pro plan). Try again later.", 429), + }); + }); + return stream; + }, + SOURCE_ID, + ); + + const stream = streamSimple(model(), context, { + apiKey: async () => resolved[resolveIndex++], + }); + for await (const _event of stream) { + // drain + } + + expect((await stream.result()).stopReason).toBe("error"); + expect(keys).toEqual(["credential-A", "credential-B"]); + }); + it("rotates before emitting content for Codex quota payloads", async () => { const payloads: Array<{ message: string; status?: number }> = [ { message: "429", status: 429 }, diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 87cd5f7e25d..fbbc95b9d7e 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -30,6 +30,7 @@ - Fixed the Model Hub role-assignment strip hiding the selected chip once the row overflowed; the strip now scrolls horizontally, truncating passed chips behind a leading ellipsis so the selection (plus one chip of lookahead) stays visible. - Fixed mouse hover and clicks in the /models Roles view landing one row above the pointer (the row mapping subtracted the status row twice). - Fixed model search keeping the most-recently-used model on top of the results: match quality now ranks first (an exact `gpt-5.5` beats the active `gpt-5.6-sol`), with MRU order only breaking ties between equally good matches. +- Fixed ModelRegistry/AuthStorage resolvers to continue replay-safe usage/account-quota turns across every distinct eligible credential and return exhaustion when no sibling switches, instead of re-resolving the same failed credential. ## [16.4.5] - 2026-07-11 diff --git a/packages/coding-agent/src/config/api-key-resolver.ts b/packages/coding-agent/src/config/api-key-resolver.ts index d237302cb93..8440b14c227 100644 --- a/packages/coding-agent/src/config/api-key-resolver.ts +++ b/packages/coding-agent/src/config/api-key-resolver.ts @@ -59,12 +59,13 @@ export function createApiKeyResolver( // sibling exists we switch immediately; the precise no-sibling backoff // is owned by `markUsageLimitReached` (default + server usage-report // reset) and the outer whole-turn retry layer. - await registry.authStorage.rotateSessionCredential(provider, sessionId, { + const rotated = await registry.authStorage.rotateSessionCredential(provider, sessionId, { error, modelId, signal, apiKey: previousKey, }); + if (!rotated) return undefined; return registry.getApiKeyForProvider(provider, sessionId, { baseUrl, modelId }); } return registry.getApiKeyForProvider(provider, sessionId, { baseUrl, modelId, forceRefresh: true, signal }); diff --git a/packages/coding-agent/test/agent-session-retry-cap.test.ts b/packages/coding-agent/test/agent-session-retry-cap.test.ts index 1efa5c1bd80..28f9780f796 100644 --- a/packages/coding-agent/test/agent-session-retry-cap.test.ts +++ b/packages/coding-agent/test/agent-session-retry-cap.test.ts @@ -3,7 +3,8 @@ import * as path from "node:path"; import { scheduler } from "node:timers/promises"; import { Agent } from "@oh-my-pi/pi-agent-core"; import type { ApiKeyResolveContext, AssistantMessage, ToolCall } from "@oh-my-pi/pi-ai"; -import { createMockModel } from "@oh-my-pi/pi-ai/providers/mock"; +import { unregisterCustomApis } from "@oh-my-pi/pi-ai/api-registry"; +import { createMockModel, registerMockApi } from "@oh-my-pi/pi-ai/providers/mock"; import * as aiStream from "@oh-my-pi/pi-ai/stream"; import { AssistantMessageEventStream } from "@oh-my-pi/pi-ai/utils/event-stream"; import { getBundledModel } from "@oh-my-pi/pi-catalog/models"; @@ -17,6 +18,8 @@ import { TempDir } from "@oh-my-pi/pi-utils"; type AutoRetryEndEvent = Extract; type AutoRetryStartEvent = Extract; +const RETRY_CAP_MOCK_API_SOURCE = "agent-session-retry-cap-test"; + function lastAssistant(session: AgentSession): AssistantMessage { const message = session.agent.state.messages.at(-1); if (message?.role !== "assistant") { @@ -67,9 +70,10 @@ describe("AgentSession retry delay cap", () => { await session.dispose(); session = undefined; } + unregisterCustomApis(RETRY_CAP_MOCK_API_SOURCE); + vi.restoreAllMocks(); authStorage.close(); tempDir.removeSync(); - vi.restoreAllMocks(); }); it("bails immediately when retry-after exceeds retry.maxDelayMs", async () => { @@ -205,25 +209,45 @@ describe("AgentSession retry delay cap", () => { expect(last.content).toContainEqual({ type: "text", text: "recovered after stream read retry" }); }); - it("switches credentials instead of failing the delay cap for account rate limits", async () => { + it("rolls through four sibling credentials inside one AgentSession prompt before delay-cap retry", async () => { const model = getBundledModel("anthropic", "claude-sonnet-4-5"); - if (!model) { - throw new Error("Expected bundled Anthropic test model to exist"); + const fallbackModel = getBundledModel("openai", "gpt-5"); + if (!model || !fallbackModel) { + throw new Error("Expected bundled primary and fallback test models to exist"); } + const providerSessionId = "retry-four-credential-session-3"; + registerMockApi(RETRY_CAP_MOCK_API_SOURCE); authStorage.removeRuntimeApiKey("anthropic"); + authStorage.setRuntimeApiKey("openai", "openai-fallback-key"); await authStorage.set("anthropic", [ - { type: "api_key", key: "anthropic-key-1" }, - { type: "api_key", key: "anthropic-key-2" }, + { type: "api_key", key: "anthropic-key-A" }, + { type: "api_key", key: "anthropic-key-B" }, + { type: "api_key", key: "anthropic-key-C" }, + { type: "api_key", key: "anthropic-key-D" }, ]); const rateLimitError = '429 {"type":"error","error":{"type":"rate_limit_error","message":"This request would exceed your account\'s rate limit. Please try again later."}} retry-after-ms=11180000'; - const mock = createMockModel(); const requestedKeys: string[] = []; - let agent!: Agent; - agent = new Agent({ - getApiKey: model => modelRegistry.resolver(model, agent.sessionId), + const mock = createMockModel({ + id: model.id, + provider: model.provider, + handler: (_context, options) => { + const apiKey = typeof options?.apiKey === "string" ? options.apiKey : undefined; + if (!apiKey) { + throw new Error("Expected streamSimple to pass a resolved string API key"); + } + requestedKeys.push(apiKey); + return apiKey === "anthropic-key-D" + ? { content: ["recovered on fourth credential"], stopReason: "stop" } + : { throw: rateLimitError }; + }, + }); + const requestedModels: string[] = []; + const agent = new Agent({ + getApiKey: model => modelRegistry.resolver(model, providerSessionId), + sessionId: providerSessionId, initialState: { model, systemPrompt: ["Test"], @@ -231,22 +255,20 @@ describe("AgentSession retry delay cap", () => { messages: [], }, streamFn: (requestedModel, context, options) => { - const apiKey = resolveInitialApiKey(options?.apiKey); - requestedKeys.push(apiKey); - if (requestedKeys.length === 1) { - mock.push({ throw: rateLimitError }); - } else { - mock.push({ content: ["recovered after credential switch"] }); - } - return mock.stream(requestedModel, context, options); + requestedModels.push(`${requestedModel.provider}/${requestedModel.id}`); + return aiStream.streamSimple(mock.model, context, options); }, }); const settings = Settings.isolated({ "compaction.enabled": false, - "retry.baseDelayMs": 5, - "retry.maxDelayMs": 100, - "retry.maxRetries": 1, + "retry.baseDelayMs": 1, + "retry.maxDelayMs": 1, + "retry.maxRetries": 0, + "retry.modelFallback": true, + "retry.fallbackChains": { + default: [`${fallbackModel.provider}/${fallbackModel.id}`], + }, }); settings.setModelRole("default", `${model.provider}/${model.id}`); @@ -255,9 +277,9 @@ describe("AgentSession retry delay cap", () => { sessionManager: SessionManager.inMemory(), settings, modelRegistry, + providerSessionId, }); - const waitSpy = vi.spyOn(scheduler, "wait").mockResolvedValue(undefined); const retryStartEvents: AutoRetryStartEvent[] = []; const retryEndEvents: AutoRetryEndEvent[] = []; session.subscribe(event => { @@ -268,18 +290,26 @@ describe("AgentSession retry delay cap", () => { await session.prompt("Trigger account rate limit with long retry-after"); await session.waitForIdle(); - expect(requestedKeys).toHaveLength(2); - expect(new Set(requestedKeys).size).toBe(2); - expect(retryStartEvents).toHaveLength(1); - expect(retryStartEvents[0]).toMatchObject({ delayMs: 0 }); - expect(retryEndEvents).toHaveLength(1); - expect(retryEndEvents[0]).toMatchObject({ success: true, attempt: 1 }); - for (const call of waitSpy.mock.calls) { - expect(call[0]).toBeLessThanOrEqual(100); + expect(requestedModels).toEqual([`${model.provider}/${model.id}`]); + expect(requestedKeys).toEqual(["anthropic-key-A", "anthropic-key-B", "anthropic-key-C", "anthropic-key-D"]); + expect(new Set(requestedKeys).size).toBe(4); + expect(mock.calls).toHaveLength(4); + expect(retryStartEvents).toHaveLength(0); + expect(retryEndEvents).toHaveLength(0); + expect(session.model?.provider).toBe(model.provider); + expect(session.model?.id).toBe(model.id); + for (const call of mock.calls) { + expect(call.context.messages.filter(message => message.role === "user")).toHaveLength(1); } + expect(session.agent.state.messages.filter(message => message.role === "user")).toHaveLength(1); + expect( + session.agent.state.messages.some( + message => message.role === "custom" && "customType" in message && message.customType === "irc:incoming", + ), + ).toBe(false); const last = lastAssistant(session); expect(last.stopReason).toBe("stop"); - expect(last.content).toContainEqual({ type: "text", text: "recovered after credential switch" }); + expect(last.content).toContainEqual({ type: "text", text: "recovered on fourth credential" }); }); it("switches same-provider credentials before model fallback on ChatGPT usage limits", async () => { diff --git a/packages/coding-agent/test/auth-storage-rotation.test.ts b/packages/coding-agent/test/auth-storage-rotation.test.ts index bf7e3e6c560..162d3485133 100644 --- a/packages/coding-agent/test/auth-storage-rotation.test.ts +++ b/packages/coding-agent/test/auth-storage-rotation.test.ts @@ -2,9 +2,11 @@ import { afterEach, beforeEach, describe, expect, test, vi } from "bun:test"; import * as fs from "node:fs"; import * as os from "node:os"; import * as path from "node:path"; -import type { UsageProvider } from "@oh-my-pi/pi-ai"; +import { type UsageProvider, withAuth } from "@oh-my-pi/pi-ai"; import * as oauth from "@oh-my-pi/pi-ai/oauth"; import type { OAuthCredentials } from "@oh-my-pi/pi-ai/oauth/types"; +import { getBundledModel } from "@oh-my-pi/pi-catalog/models"; +import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry"; import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage"; import { removeSyncWithRetries, Snowflake } from "@oh-my-pi/pi-utils"; @@ -44,7 +46,7 @@ describe("AuthStorage account rotation", () => { // Stub the refresh path so AuthStorage doesn't hit a real OAuth endpoint // when the credential lands inside the 60s skew. Returning the credential - // unchanged preserves the test's deterministic accountId routing. + // unchanged preserves deterministic access-token routing. vi.spyOn(oauth, "refreshOAuthToken").mockImplementation(async (_provider, credential) => { return credential; }); @@ -52,7 +54,7 @@ describe("AuthStorage account rotation", () => { const credential = credentials["openai-codex"] as OAuthCredentials | undefined; if (!credential) return null; return { - apiKey: `api-${credential.accountId ?? "unknown"}`, + apiKey: credential.access, newCredentials: credential, }; }); @@ -86,14 +88,14 @@ describe("AuthStorage account rotation", () => { const sessionId = "issue-55-session"; const firstKey = await authStorage.getApiKey("openai-codex", sessionId); - expect(firstKey).toMatch(/^api-acct-/); + expect(firstKey).toMatch(/^access-/); usageExhausted = true; const { switched } = await authStorage.markUsageLimitReached("openai-codex", sessionId); expect(switched).toBe(true); const exhaustedFallbackKey = await authStorage.getApiKey("openai-codex", sessionId); - expect(exhaustedFallbackKey).toMatch(/^api-acct-/); + expect(exhaustedFallbackKey).toMatch(/^access-/); }); test("usage-limit rotation can match the failed bearer when session stickiness is missing", async () => { @@ -117,7 +119,7 @@ describe("AuthStorage account rotation", () => { const sessionId = "missing-sticky-session"; const result = await authStorage.markUsageLimitReached("openai-codex", sessionId, { apiKey: "access-1" }); expect(result.switched).toBe(true); - expect(await authStorage.getApiKey("openai-codex", sessionId)).toBe("api-acct-2"); + expect(await authStorage.getApiKey("openai-codex", sessionId)).toBe("access-2"); }); test("usage-limit rotation trusts the failed bearer over stale session stickiness", async () => { @@ -140,9 +142,62 @@ describe("AuthStorage account rotation", () => { const sessionId = "stale-sticky-session"; const stickyKey = await authStorage.getApiKey("openai-codex", sessionId); - const failedKey = stickyKey === "api-plus-acct" ? "k12-access" : "plus-access"; + const failedKey = stickyKey === "plus-access" ? "k12-access" : "plus-access"; const result = await authStorage.markUsageLimitReached("openai-codex", sessionId, { apiKey: failedKey }); expect(result.switched).toBe(true); expect(await authStorage.getApiKey("openai-codex", sessionId)).toBe(stickyKey); }); + + test("withAuth reaches a fourth healthy Codex OAuth sibling through ModelRegistry", async () => { + await authStorage.set("openai-codex", [ + { + type: "oauth", + access: "access-a", + refresh: "refresh-a", + expires: Date.now() + 60_000, + accountId: "acct-a", + }, + { + type: "oauth", + access: "access-b", + refresh: "refresh-b", + expires: Date.now() + 60_000, + accountId: "acct-b", + }, + { + type: "oauth", + access: "access-c", + refresh: "refresh-c", + expires: Date.now() + 60_000, + accountId: "acct-c", + }, + { + type: "oauth", + access: "access-d", + refresh: "refresh-d", + expires: Date.now() + 60_000, + accountId: "acct-d", + }, + ]); + + const model = getBundledModel("openai-codex", "gpt-5.5"); + if (!model) { + throw new Error("Expected bundled Codex test model to exist"); + } + + const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); + const attemptedKeys: string[] = []; + const result = await withAuth(modelRegistry.resolver(model, "codex-four-oauth-session"), async key => { + attemptedKeys.push(key); + if (key !== "access-d") { + throw new Error("You have hit your ChatGPT usage limit (pro plan). Try again later."); + } + return key; + }); + + expect(result).toBe("access-d"); + expect(attemptedKeys.at(-1)).toBe("access-d"); + expect([...attemptedKeys].sort()).toEqual(["access-a", "access-b", "access-c", "access-d"]); + expect(new Set(attemptedKeys).size).toBe(4); + }); }); From b929ed16448150baf031f64082c887f3c2ce78fe Mon Sep 17 00:00:00 2001 From: Jeff Scott Ward Date: Sun, 12 Jul 2026 00:35:26 -0400 Subject: [PATCH 2/2] fix: honor gateway quota usage limits --- packages/ai/CHANGELOG.md | 5 +- packages/ai/src/auth-gateway/server.ts | 4 +- packages/ai/src/auth-storage.ts | 124 +++++++--- .../test/auth-gateway-thinking-loop.test.ts | 80 ++++++- .../auth-storage-force-refresh-rotate.test.ts | 226 +++++++++++++++++- packages/coding-agent/CHANGELOG.md | 5 +- .../src/config/api-key-resolver.ts | 16 +- .../test/auth-storage-rotation.test.ts | 58 +++++ 8 files changed, 480 insertions(+), 38 deletions(-) diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index b212eeb8ab6..72ea5d81352 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Fixed provider-agnostic replay-safe usage/account-quota failures to rotate through every distinct eligible credential instead of stopping after the fixed a/b/c ladder, while preserving transient-429 backoff, cycle/abort guards, exact failed-credential targeting despite stale session stickiness, and a finite safety ceiling. + ## [16.4.6] - 2026-07-12 ### Added @@ -12,7 +16,6 @@ ### Fixed - Fixed OAuth credential resolution returning "No API key found" when every plan-eligible OpenAI Codex account was rate-limit blocked and the only unblocked account failed the model's plan gate: resolution now runs a last-resort ladder that first yields a plan-fitting account regardless of usage blocks (so callers get real usage-limit retry semantics), then tries every account with the plan filter dropped before reporting no credential -- Fixed provider-agnostic replay-safe usage/account-quota failures to rotate through every distinct eligible credential instead of stopping after the fixed a/b/c ladder, while preserving transient-429 backoff, cycle/abort guards, exact failed-credential targeting despite stale session stickiness, and a finite safety ceiling. ## [16.4.5] - 2026-07-11 diff --git a/packages/ai/src/auth-gateway/server.ts b/packages/ai/src/auth-gateway/server.ts index ae4d47415d6..faaf11e71cc 100644 --- a/packages/ai/src/auth-gateway/server.ts +++ b/packages/ai/src/auth-gateway/server.ts @@ -22,6 +22,7 @@ import { Effort } from "@oh-my-pi/pi-catalog/effort"; import { extractHttpStatusFromError, extractRetryHint, logger } from "@oh-my-pi/pi-utils"; import type { ApiKeyResolver } from "../auth-retry"; import type { AuthStorage } from "../auth-storage"; +import * as AIError from "../error"; import { classifyGatewayError } from "../error/gateway"; import { isUsageLimitOutcome } from "../error/rate-limit"; import * as anthropicMessages from "../providers/anthropic-messages-server"; @@ -227,7 +228,8 @@ async function refreshGatewayApiKeyAfterAuthError( peer: string, ): Promise { const message = error instanceof Error ? error.message : String(error); - if (isUsageLimitOutcome(extractHttpStatusFromError(error), message)) { + const status = extractHttpStatusFromError(error); + if (AIError.isUsageLimit(error) || isUsageLimitOutcome(status, message)) { const retryAfterMs = extractRetryHint(undefined, message); const { switched, retryAtMs } = await storage.markUsageLimitReached(provider, sessionId, { retryAfterMs, diff --git a/packages/ai/src/auth-storage.ts b/packages/ai/src/auth-storage.ts index 252d7ce6a4c..e92e3d1adca 100644 --- a/packages/ai/src/auth-storage.ts +++ b/packages/ai/src/auth-storage.ts @@ -8,6 +8,7 @@ * - `SqliteAuthCredentialStore`: concrete SQLite-backed implementation */ import { Database, type Statement } from "bun:sqlite"; +import { createHash } from "node:crypto"; import * as fs from "node:fs/promises"; import * as path from "node:path"; import { getAgentDbPath, logger } from "@oh-my-pi/pi-utils"; @@ -58,6 +59,12 @@ import { opencodeGoUsageProvider } from "./usage/opencode-go"; import { zaiUsageProvider } from "./usage/zai"; const USAGE_RANKING_METRIC_EPSILON = 1e-9; +const OAUTH_BEARER_FINGERPRINT_HISTORY_LIMIT = 8; + +/** SHA-256 bearer fingerprint, so superseded OAuth token bytes never enter the identity cache. */ +function fingerprintOAuthBearer(bearer: string): string { + return createHash("sha256").update(bearer).digest("base64url"); +} // ───────────────────────────────────────────────────────────────────────────── // Credential Types @@ -1076,6 +1083,8 @@ export class AuthStorage { #providerRoundRobinIndex: Map = new Map(); /** Tracks the last used credential per provider for a session (used for rate-limit switching). */ #sessionLastCredential: Map> = new Map(); + /** Recent bearer fingerprints resolved for each durable OAuth row; used only for delayed usage-limit attribution. */ + #oauthBearerFingerprints: Map> = new Map(); /** Maps provider:type -> credentialIndex -> blockedUntilMs for temporary backoff. */ #credentialBackoff: Map> = new Map(); /** Earliest time a freshly-set in-memory block may be cleared by live usage reconciliation. */ @@ -1329,6 +1338,16 @@ export class AuthStorage { #setStoredCredentials(provider: string, credentials: StoredCredential[]): void { const current = this.#data.get(provider) ?? []; if (storedCredentialArraysEqual(current, credentials)) return; + const trackedBearerFingerprints = this.#oauthBearerFingerprints.get(provider); + if (trackedBearerFingerprints) { + const activeOAuthIds = new Set( + credentials.filter(entry => entry.credential.type === "oauth").map(entry => entry.id), + ); + for (const credentialId of trackedBearerFingerprints.keys()) { + if (!activeOAuthIds.has(credentialId)) trackedBearerFingerprints.delete(credentialId); + } + if (trackedBearerFingerprints.size === 0) this.#oauthBearerFingerprints.delete(provider); + } if (credentials.length === 0) { this.#data.delete(provider); } else { @@ -1337,6 +1356,26 @@ export class AuthStorage { this.#bumpGeneration("credentials"); } + #recordOAuthBearerCredentialId(provider: string, bearer: string, credentialId: number | undefined): void { + if (credentialId === undefined) return; + const fingerprint = fingerprintOAuthBearer(bearer); + const byCredentialId = this.#oauthBearerFingerprints.get(provider) ?? new Map(); + const history = byCredentialId.get(credentialId) ?? []; + const nextHistory = history.filter(previous => previous !== fingerprint); + nextHistory.push(fingerprint); + if (nextHistory.length > OAUTH_BEARER_FINGERPRINT_HISTORY_LIMIT) nextHistory.shift(); + byCredentialId.set(credentialId, nextHistory); + this.#oauthBearerFingerprints.set(provider, byCredentialId); + } + + #findOAuthCredentialIdForBearer(provider: string, bearer: string): number | undefined { + const fingerprint = fingerprintOAuthBearer(bearer); + for (const [credentialId, history] of this.#oauthBearerFingerprints.get(provider) ?? []) { + if (history.includes(fingerprint)) return credentialId; + } + return undefined; + } + #resolveOAuthDedupeIdentityKey(provider: string, credential: OAuthCredential): string | null { return resolveCredentialIdentityKey(provider, credential); } @@ -3453,42 +3492,63 @@ export class AuthStorage { signal?: AbortSignal; }, ): Promise { - const sessionCredential = await this.#resolveCredentialTarget(provider, sessionId, { + let sessionCredential = await this.#resolveCredentialTarget(provider, sessionId, { credentialId: options?.credentialId, apiKey: options?.apiKey, }); + if (!sessionCredential && options?.credentialId === undefined && options?.apiKey !== undefined) { + // Account quota survives OAuth bearer rotation. Attribute a delayed + // usage-limit response through the durable row id captured when this + // exact bearer was resolved; never use this alias for hard auth errors. + const credentialId = this.#findOAuthCredentialIdForBearer(provider, options.apiKey); + const index = + credentialId === undefined + ? -1 + : this.#getStoredCredentials(provider).findIndex( + entry => entry.id === credentialId && entry.credential.type === "oauth", + ); + if (index >= 0) sessionCredential = { type: "oauth", index, explicit: true }; + } if (!sessionCredential) return { switched: false }; + const target = this.#getStoredCredentials(provider)[sessionCredential.index]; + if (!target || target.credential.type !== sessionCredential.type) return { switched: false }; + const credentialType = sessionCredential.type; + const targetCredentialId = target.id; - const providerKey = this.#getProviderTypeKey(provider, sessionCredential.type); + const providerKey = this.#getProviderTypeKey(provider, credentialType); const strategy = this.#rankingStrategyResolver?.(provider); const rankingContext: CredentialRankingContext = { modelId: options?.modelId }; const blockScope = strategy?.blockScope?.(rankingContext); const now = Date.now(); let blockedUntil = now + (options?.retryAfterMs ?? AuthStorage.#defaultBackoffMs); - if (sessionCredential.type === "oauth" && strategy) { - const credential = this.#getCredentialsForProvider(provider)[sessionCredential.index]; - if (credential?.type === "oauth") { - const report = await this.#getUsageReport(provider, credential, options); - if (report) { - const scopedLimits = this.#getScopedUsageLimits(strategy, report, rankingContext); - if (this.#isUsageLimitReached(scopedLimits)) { - const resetAtMs = this.#getUsageResetAtMs(scopedLimits, Date.now()); - if (resetAtMs && resetAtMs > blockedUntil) { - blockedUntil = resetAtMs; - } + if (credentialType === "oauth" && target.credential.type === "oauth" && strategy) { + const report = await this.#getUsageReport(provider, target.credential, options); + if (report) { + const scopedLimits = this.#getScopedUsageLimits(strategy, report, rankingContext); + if (this.#isUsageLimitReached(scopedLimits)) { + const resetAtMs = this.#getUsageResetAtMs(scopedLimits, Date.now()); + if (resetAtMs && resetAtMs > blockedUntil) { + blockedUntil = resetAtMs; } } } } - this.#markCredentialBlocked(provider, providerKey, sessionCredential.index, blockedUntil, blockScope); + // Usage lookup may refresh, disable, or remove a row. Re-resolve its + // durable id before applying positional in-memory and persisted blocks. + const targetIndex = this.#getStoredCredentials(provider).findIndex( + entry => entry.id === targetCredentialId && entry.credential.type === credentialType, + ); + if (targetIndex >= 0) { + this.#markCredentialBlocked(provider, providerKey, targetIndex, blockedUntil, blockScope); + } const remainingCredentials = this.#getCredentialsForProvider(provider) .map((credential, index) => ({ credential, index })) .filter( (entry): entry is { credential: AuthCredential; index: number } => - entry.credential.type === sessionCredential.type && entry.index !== sessionCredential.index, + entry.credential.type === credentialType && entry.index !== targetIndex, ); let retryAtMs: number | undefined; @@ -4240,6 +4300,7 @@ export class AuthStorage { } } } + this.#recordOAuthBearerCredentialId(provider, result.apiKey, credentialId); this.#recordSessionCredential(provider, sessionId, "oauth", selection.index); return { apiKey: result.apiKey, credential: updated, credentialId }; } catch (error) { @@ -4960,10 +5021,10 @@ export class AuthStorage { * in `options.credentialId`, then the failed bearer supplied in * `options.apiKey`, so overlapping requests cannot redirect rotation through * stale session stickiness. Fall back to the session-sticky credential only - * when neither explicit target is available. If an explicit target is supplied - * but no longer matches storage, return `false` without mutating sticky state. - * Apply the storage action for the error class, then let the next resolve - * select a sibling. + * when neither explicit target is available. For hard-auth errors, an explicit + * target that no longer matches storage returns `false` without mutation. + * Delayed usage-limit errors may instead recover the durable OAuth row from + * the bearer fingerprint recorded when the request resolved. * * - usage-limit / account-rate-limit error → {@link AuthStorage.markUsageLimitReached} * (temporary block via its own backoff — default plus server usage-report @@ -4979,12 +5040,6 @@ export class AuthStorage { sessionId: string | undefined, options?: { error?: unknown; modelId?: string; apiKey?: string; credentialId?: number; signal?: AbortSignal }, ): Promise { - const sessionCredential = await this.#resolveCredentialTarget(provider, sessionId, { - credentialId: options?.credentialId, - apiKey: options?.apiKey, - }); - if (!sessionCredential) return false; - const error = options?.error; const status = AIError.status(error); const message = error instanceof Error ? error.message : typeof error === "string" ? error : undefined; @@ -4999,6 +5054,12 @@ export class AuthStorage { ).switched; } + const sessionCredential = await this.#resolveCredentialTarget(provider, sessionId, { + credentialId: options?.credentialId, + apiKey: options?.apiKey, + }); + if (!sessionCredential) return false; + const providerKey = this.#getProviderTypeKey(provider, sessionCredential.type); // Snapshot sibling availability before mutating so a soft-deleting // suspect hook can't reindex the answer out from under us. @@ -5046,7 +5107,7 @@ export class AuthStorage { * * - initial (`error: undefined`) → resolve the session credential. * - step (b) `!lastChance` → force-refresh the SAME session-sticky credential. - * - step (c) `lastChance` → rotate to a sibling credential, then re-resolve. + * - step (c) `lastChance` → rotate to a sibling and re-resolve, unless quota exhaustion has no sibling. * * Used by web-search providers and other consumers that hold an AuthStorage * directly (no ModelRegistry in scope). @@ -5058,13 +5119,20 @@ export class AuthStorage { return this.getApiKey(provider, sessionId, { baseUrl, modelId, signal }); } if (lastChance) { - const rotated = await this.rotateSessionCredential(provider, sessionId, { + const switched = await this.rotateSessionCredential(provider, sessionId, { error, modelId, signal, apiKey: previousKey, }); - if (!rotated) return undefined; + if (!switched) { + const status = AIError.status(error); + const message = error instanceof Error ? error.message : typeof error === "string" ? error : undefined; + // Preserve no-sibling quota backoff instead of re-resolving an + // already-blocked fallback. Hard-auth declines still re-resolve + // because a peer may have refreshed the failed bearer. + if (AIError.isUsageLimit(error) || isUsageLimitOutcome(status, message)) return undefined; + } return this.getApiKey(provider, sessionId, { baseUrl, modelId, signal }); } return this.getApiKey(provider, sessionId, { baseUrl, modelId, forceRefresh: true, signal }); diff --git a/packages/ai/test/auth-gateway-thinking-loop.test.ts b/packages/ai/test/auth-gateway-thinking-loop.test.ts index 5f069da5edb..e2e623b643d 100644 --- a/packages/ai/test/auth-gateway-thinking-loop.test.ts +++ b/packages/ai/test/auth-gateway-thinking-loop.test.ts @@ -5,7 +5,8 @@ import * as path from "node:path"; import { scheduler } from "node:timers/promises"; import { clearCustomApis } from "@oh-my-pi/pi-ai/api-registry"; import { startAuthGateway } from "@oh-my-pi/pi-ai/auth-gateway"; -import { AuthStorage } from "@oh-my-pi/pi-ai/auth-storage"; +import { AuthStorage, SqliteAuthCredentialStore } from "@oh-my-pi/pi-ai/auth-storage"; +import { ProviderHttpError } from "@oh-my-pi/pi-ai/error"; import { createMockModel, registerMockApi } from "@oh-my-pi/pi-ai/providers/mock"; import { THINKING_LOOP_ERROR_MARKER } from "@oh-my-pi/pi-ai/utils/thinking-loop"; @@ -111,3 +112,80 @@ describe("auth-gateway non-streaming thinking-loop cook", () => { } }); }); + +describe("auth-gateway auth retry", () => { + it("treats structured generic quota errors as usage-limit blocks before invalidating credentials", async () => { + registerMockApi(); + const dir = await fs.mkdtemp(path.join(os.tmpdir(), "gw-quota-rotation-")); + const store = await SqliteAuthCredentialStore.open(path.join(dir, "auth.db")); + const storage = new AuthStorage(store); + await storage.set("mock", [ + { type: "api_key", key: "quota-key" }, + { type: "api_key", key: "healthy-key" }, + ]); + const markUsageLimitSpy = spyOn(storage, "markUsageLimitReached"); + const invalidateSpy = spyOn(storage, "invalidateCredentialMatching"); + let attempt = 0; + const mock = createMockModel({ + provider: "mock", + id: "gateway-quota-model", + handler: (_context, options) => { + attempt += 1; + if (attempt === 1) { + throw new ProviderHttpError("Generic provider failure", 429, { code: "insufficient_quota" }); + } + return { content: [`ok:${options?.apiKey ?? "missing"}`] }; + }, + }); + const handle = startAuthGateway({ + bind: "127.0.0.1:0", + bearerTokens: ["t"], + storage, + resolveModel: () => mock.model, + version: "test", + }); + try { + const res = await fetch(`${handle.url}/v1/chat/completions`, { + method: "POST", + headers: { "Content-Type": "application/json", Authorization: "Bearer t" }, + body: JSON.stringify({ + model: "gateway-quota-model", + messages: [{ role: "user", content: "hi" }], + prompt_cache_key: "gw-quota-rotation", + stream: false, + }), + }); + const body = (await res.json()) as { + choices?: Array<{ message?: { content?: string | null } }>; + }; + const attemptedKeys = mock.calls.map(call => call.options?.apiKey); + + expect(res.status).toBe(200); + expect(attemptedKeys).toHaveLength(2); + const [failedKey, retriedKey] = attemptedKeys; + if (typeof failedKey !== "string" || typeof retriedKey !== "string") { + throw new Error("expected gateway retries to use static API keys"); + } + expect(body.choices?.[0]?.message?.content).toBe(`ok:${retriedKey}`); + expect(new Set([failedKey, retriedKey]).size).toBe(2); + expect(markUsageLimitSpy.mock.calls).toHaveLength(1); + const usageLimitCall = markUsageLimitSpy.mock.calls[0]; + if (!usageLimitCall) { + throw new Error("expected usage-limit mark call"); + } + const [usageLimitProvider, usageLimitSessionId, usageLimitOptions] = usageLimitCall; + expect(usageLimitProvider).toBe("mock"); + expect(usageLimitSessionId).toBe("gw-quota-rotation"); + expect(usageLimitOptions?.apiKey).toBe(failedKey); + expect(invalidateSpy.mock.calls).toHaveLength(0); + expect(store.listAuthCredentials("mock")).toHaveLength(2); + expect(await storage.getApiKey("mock", "gw-quota-rotation")).toBe(retriedKey); + } finally { + markUsageLimitSpy.mockRestore(); + invalidateSpy.mockRestore(); + await handle.close(); + storage.close(); + await fs.rm(dir, { recursive: true, force: true }); + } + }); +}); diff --git a/packages/ai/test/auth-storage-force-refresh-rotate.test.ts b/packages/ai/test/auth-storage-force-refresh-rotate.test.ts index 3a5218a260b..f8c9f00b8d5 100644 --- a/packages/ai/test/auth-storage-force-refresh-rotate.test.ts +++ b/packages/ai/test/auth-storage-force-refresh-rotate.test.ts @@ -2,9 +2,11 @@ import { afterEach, beforeEach, describe, expect, test, vi } from "bun:test"; import * as fs from "node:fs/promises"; import * as os from "node:os"; import * as path from "node:path"; +import { withAuth } from "@oh-my-pi/pi-ai"; import { type AuthCredentialStore, AuthStorage, SqliteAuthCredentialStore } from "@oh-my-pi/pi-ai/auth-storage"; import { ProviderHttpError } from "@oh-my-pi/pi-ai/error"; import { registerOAuthProvider, unregisterOAuthProviders } from "@oh-my-pi/pi-ai/registry/oauth"; +import type { CredentialRankingStrategy, UsageProvider } from "@oh-my-pi/pi-ai/usage"; import { removeWithRetries } from "../../utils/src/temp"; const PROVIDER = "unit-rotate-oauth"; @@ -55,7 +57,7 @@ describe("AuthStorage forceRefresh + rotateSessionCredential", () => { } }); - function registerProvider(onRefresh?: () => void): void { + function registerProvider(onRefresh?: () => void, nextAccess?: () => string): void { registerOAuthProvider({ id: PROVIDER, name: "Rotate Unit", @@ -67,7 +69,7 @@ describe("AuthStorage forceRefresh + rotateSessionCredential", () => { onRefresh?.(); return { ...credentials, - access: "minted-access", + access: nextAccess?.() ?? "minted-access", refresh: "minted-refresh", expires: farExpiry(), }; @@ -178,6 +180,226 @@ describe("AuthStorage forceRefresh + rotateSessionCredential", () => { expect(laterSelections.has(sticky)).toBe(true); }); + test("resolver rotates away from the account matching a stale OAuth bearer", async () => { + if (!authStorage || !store) throw new Error("test setup failed"); + registerProvider(); + await authStorage.set(PROVIDER, [ + { type: "oauth", access: "stale-A", refresh: "ref-A", expires: farExpiry() }, + { type: "oauth", access: "stale-B", refresh: "ref-B", expires: farExpiry() }, + ]); + + const sessionId = "resolver-concurrent-refresh"; + const previousKey = await authStorage.getApiKey(PROVIDER, sessionId); + if (!previousKey) throw new Error("expected initial OAuth bearer"); + const rows = store.listAuthCredentials(PROVIDER); + const target = rows.find(row => row.credential.type === "oauth" && row.credential.access === previousKey); + const sibling = rows.find(row => row.id !== target?.id); + if (target?.credential.type !== "oauth" || sibling?.credential.type !== "oauth") { + throw new Error("expected target and sibling OAuth rows"); + } + store.updateAuthCredential(target.id, { + ...target.credential, + access: `${previousKey}-refreshed`, + }); + await authStorage.reload(); + + const retry = await authStorage.resolver(PROVIDER, { sessionId })({ + lastChance: true, + error: usageLimitError(), + previousKey, + }); + + expect(retry).toBe(sibling.credential.access); + expect(await authStorage.getApiKey(PROVIDER, sessionId)).toBe(sibling.credential.access); + }); + + test("resolver re-resolves a peer-refreshed OAuth bearer after a stale 401", async () => { + if (!authStorage || !store) throw new Error("test setup failed"); + registerProvider(); + await authStorage.set(PROVIDER, [ + { type: "oauth", access: "auth-A", refresh: "ref-A", expires: farExpiry() }, + { type: "oauth", access: "auth-B", refresh: "ref-B", expires: farExpiry() }, + ]); + + const sessionId = "resolver-stale-auth"; + const previousKey = await authStorage.getApiKey(PROVIDER, sessionId); + if (!previousKey) throw new Error("expected initial OAuth bearer"); + const target = store + .listAuthCredentials(PROVIDER) + .find(row => row.credential.type === "oauth" && row.credential.access === previousKey); + if (target?.credential.type !== "oauth") throw new Error("expected failed OAuth credential row"); + const refreshedKey = `${previousKey}-refreshed`; + store.updateAuthCredential(target.id, { ...target.credential, access: refreshedKey }); + await authStorage.reload(); + + const retry = await authStorage.resolver(PROVIDER, { sessionId })({ + lastChance: true, + error: authError(), + previousKey, + }); + + expect(retry).toBe(refreshedKey); + expect(await authStorage.getApiKey(PROVIDER, sessionId)).toBe(refreshedKey); + expect( + store + .listAuthCredentials(PROVIDER) + .some( + row => row.id === target.id && row.credential.type === "oauth" && row.credential.access === refreshedKey, + ), + ).toBe(true); + }); + + test("resolver stops when a usage-limit rotation has no unblocked sibling", async () => { + if (!authStorage) throw new Error("test setup failed"); + const getApiKey = vi + .spyOn(authStorage, "getApiKey") + .mockResolvedValueOnce("quota-blocked-B") + .mockResolvedValueOnce("quota-blocked-A"); + const rotate = vi.spyOn(authStorage, "rotateSessionCredential").mockResolvedValue(false); + const attemptedKeys: string[] = []; + + await expect( + withAuth(authStorage.resolver(PROVIDER, { sessionId: "all-quota-blocked" }), async key => { + attemptedKeys.push(key); + throw usageLimitError(); + }), + ).rejects.toThrow("usage limit"); + + expect(attemptedKeys).toEqual(["quota-blocked-B"]); + expect(getApiKey).toHaveBeenCalledTimes(1); + expect(rotate).toHaveBeenCalledTimes(1); + }); + + test("usage marking keeps a stale OAuth bearer bound to its original row", async () => { + if (!authStorage || !store) throw new Error("test setup failed"); + registerProvider(); + await authStorage.set(PROVIDER, [ + { type: "oauth", access: "quota-A", refresh: "ref-A", expires: farExpiry() }, + { type: "oauth", access: "quota-B", refresh: "ref-B", expires: farExpiry() }, + ]); + + const sessionId = "usage-concurrent-refresh"; + const previousKey = await authStorage.getApiKey(PROVIDER, sessionId); + if (!previousKey) throw new Error("expected initial OAuth bearer"); + const rows = store.listAuthCredentials(PROVIDER); + const target = rows.find(row => row.credential.type === "oauth" && row.credential.access === previousKey); + const sibling = rows.find(row => row.id !== target?.id); + if (target?.credential.type !== "oauth" || sibling?.credential.type !== "oauth") { + throw new Error("expected target and sibling OAuth rows"); + } + store.updateAuthCredential(target.id, { + ...target.credential, + access: `${previousKey}-refreshed`, + }); + await authStorage.reload(); + + const firstMark = await authStorage.markUsageLimitReached(PROVIDER, sessionId, { + credentialId: target.id, + }); + expect(firstMark.switched).toBe(true); + expect(await authStorage.getApiKey(PROVIDER, sessionId)).toBe(sibling.credential.access); + + const delayedMark = await authStorage.markUsageLimitReached(PROVIDER, sessionId, { + apiKey: previousKey, + }); + + expect(delayedMark.switched).toBe(true); + expect(await authStorage.getApiKey(PROVIDER, sessionId)).toBe(sibling.credential.access); + }); + + test("OAuth bearer identity history evicts old entries but retains recent delayed requests", async () => { + if (!authStorage || !store) throw new Error("test setup failed"); + let refreshCount = 0; + registerProvider(undefined, () => `bounded-${++refreshCount}`); + await authStorage.set(PROVIDER, [ + { type: "oauth", access: "bounded-A", refresh: "ref-A", expires: farExpiry() }, + { type: "oauth", access: "bounded-B", refresh: "ref-B", expires: farExpiry() }, + ]); + + const sessionId = "bounded-bearer-history"; + const initialKey = await authStorage.getApiKey(PROVIDER, sessionId); + if (!initialKey) throw new Error("expected initial OAuth bearer"); + const initialRows = store.listAuthCredentials(PROVIDER); + const target = initialRows.find(row => row.credential.type === "oauth" && row.credential.access === initialKey); + const sibling = initialRows.find(row => row.id !== target?.id); + if (target?.credential.type !== "oauth" || sibling?.credential.type !== "oauth") { + throw new Error("expected target and sibling OAuth rows"); + } + + const resolvedKeys = [initialKey]; + for (let index = 0; index < 9; index += 1) { + const refreshed = await authStorage.getApiKey(PROVIDER, sessionId, { forceRefresh: true }); + if (!refreshed) throw new Error("expected refreshed OAuth bearer"); + resolvedKeys.push(refreshed); + } + expect(new Set(resolvedKeys).size).toBe(10); + + const evictedMark = await authStorage.markUsageLimitReached(PROVIDER, sessionId, { + apiKey: resolvedKeys[0], + }); + expect(evictedMark.switched).toBe(false); + + const recentDelayedKey = resolvedKeys.at(-6); + if (!recentDelayedKey) throw new Error("expected retained delayed bearer"); + const retainedMark = await authStorage.markUsageLimitReached(PROVIDER, sessionId, { + apiKey: recentDelayedKey, + }); + expect(retainedMark.switched).toBe(true); + expect(await authStorage.getApiKey(PROVIDER, sessionId)).toBe(sibling.credential.access); + }); + + test("usage marking does not block a sibling when its target disappears during usage lookup", async () => { + if (!authStorage || !store) throw new Error("test setup failed"); + authStorage.close(); + const concurrentStore = await SqliteAuthCredentialStore.open(path.join(tempDir, "agent.db")); + store = concurrentStore; + let targetCredentialId: number | undefined; + let targetRemoved = false; + let concurrentStorage: AuthStorage; + const usageProvider: UsageProvider = { + id: PROVIDER, + async fetchUsage() { + if (targetCredentialId === undefined) throw new Error("expected target credential id"); + if (!targetRemoved) { + targetRemoved = true; + concurrentStore.deleteAuthCredential(targetCredentialId, "concurrent test removal"); + await concurrentStorage.reload(); + } + return { provider: PROVIDER, fetchedAt: Date.now(), limits: [] }; + }, + }; + const rankingStrategy: CredentialRankingStrategy = { + findWindowLimits: () => ({}), + windowDefaults: { primaryMs: 60_000, secondaryMs: 60_000 }, + }; + concurrentStorage = new AuthStorage(concurrentStore, { + usageProviderResolver: provider => (provider === PROVIDER ? usageProvider : undefined), + rankingStrategyResolver: provider => (provider === PROVIDER ? rankingStrategy : undefined), + }); + authStorage = concurrentStorage; + registerProvider(); + await concurrentStorage.set(PROVIDER, [ + { type: "oauth", access: "removed-A", refresh: "ref-A", expires: farExpiry() }, + { type: "oauth", access: "survivor-B", refresh: "ref-B", expires: farExpiry() }, + { type: "oauth", access: "survivor-C", refresh: "ref-C", expires: farExpiry() }, + ]); + + const rows = concurrentStore.listAuthCredentials(PROVIDER); + const target = rows[0]; + if (!target) throw new Error("expected target credential"); + targetCredentialId = target.id; + const siblings = rows.slice(1); + + const marked = await concurrentStorage.markUsageLimitReached(PROVIDER, undefined, { + credentialId: target.id, + }); + + expect(marked.switched).toBe(true); + for (const sibling of siblings) { + expect(concurrentStore.getCredentialBlock?.(sibling.id, `${PROVIDER}:oauth`, "")).toBeUndefined(); + } + }); + test("explicit missing rotation targets do not fall back to stale stickiness", async () => { if (!authStorage || !store) throw new Error("test setup failed"); await authStorage.set(PROVIDER, [ diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index fbbc95b9d7e..01d427956b3 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Fixed ModelRegistry/AuthStorage resolvers to continue replay-safe usage/account-quota turns across every distinct eligible credential and return exhaustion when no sibling switches, instead of re-resolving the same failed credential. + ## [16.4.6] - 2026-07-12 ### Added @@ -30,7 +34,6 @@ - Fixed the Model Hub role-assignment strip hiding the selected chip once the row overflowed; the strip now scrolls horizontally, truncating passed chips behind a leading ellipsis so the selection (plus one chip of lookahead) stays visible. - Fixed mouse hover and clicks in the /models Roles view landing one row above the pointer (the row mapping subtracted the status row twice). - Fixed model search keeping the most-recently-used model on top of the results: match quality now ranks first (an exact `gpt-5.5` beats the active `gpt-5.6-sol`), with MRU order only breaking ties between equally good matches. -- Fixed ModelRegistry/AuthStorage resolvers to continue replay-safe usage/account-quota turns across every distinct eligible credential and return exhaustion when no sibling switches, instead of re-resolving the same failed credential. ## [16.4.5] - 2026-07-11 diff --git a/packages/coding-agent/src/config/api-key-resolver.ts b/packages/coding-agent/src/config/api-key-resolver.ts index 8440b14c227..4bd5ca8234d 100644 --- a/packages/coding-agent/src/config/api-key-resolver.ts +++ b/packages/coding-agent/src/config/api-key-resolver.ts @@ -1,4 +1,5 @@ -import type { Api, ApiKeyResolver, AuthStorage, Model } from "@oh-my-pi/pi-ai"; +import { type Api, type ApiKeyResolver, type AuthStorage, isUsageLimitOutcome, type Model } from "@oh-my-pi/pi-ai"; +import * as AIError from "@oh-my-pi/pi-ai/error"; /** Model slice accepted by the model-form `resolver(model, sessionId)` overload. */ export type ApiKeyResolverModel = Pick, "provider" | "baseUrl" | "id">; @@ -27,7 +28,7 @@ export interface ApiKeyResolverRegistry { /** * Build an {@link ApiKeyResolver} implementing the central a/b/c auth-retry * policy: initial → resolve; step (b) → force-refresh same account; step (c) - * → rotate to a sibling credential, then re-resolve. + * → rotate to a sibling and re-resolve, unless quota exhaustion has no sibling. * * Two call forms: `resolver(provider, options?)` for provider-scoped keys, * and `resolver(model, sessionId?)` which derives `baseUrl`/`modelId` from @@ -59,13 +60,20 @@ export function createApiKeyResolver( // sibling exists we switch immediately; the precise no-sibling backoff // is owned by `markUsageLimitReached` (default + server usage-report // reset) and the outer whole-turn retry layer. - const rotated = await registry.authStorage.rotateSessionCredential(provider, sessionId, { + const switched = await registry.authStorage.rotateSessionCredential(provider, sessionId, { error, modelId, signal, apiKey: previousKey, }); - if (!rotated) return undefined; + if (!switched) { + const status = AIError.status(error); + const message = error instanceof Error ? error.message : typeof error === "string" ? error : undefined; + // No sibling for an account-quota failure: stop so the outer + // whole-turn retry layer can honor the recorded backoff. A hard + // auth decline can instead mean a peer refreshed the bearer. + if (AIError.isUsageLimit(error) || isUsageLimitOutcome(status, message)) return undefined; + } return registry.getApiKeyForProvider(provider, sessionId, { baseUrl, modelId }); } return registry.getApiKeyForProvider(provider, sessionId, { baseUrl, modelId, forceRefresh: true, signal }); diff --git a/packages/coding-agent/test/auth-storage-rotation.test.ts b/packages/coding-agent/test/auth-storage-rotation.test.ts index 162d3485133..3711f48c66f 100644 --- a/packages/coding-agent/test/auth-storage-rotation.test.ts +++ b/packages/coding-agent/test/auth-storage-rotation.test.ts @@ -9,6 +9,7 @@ import { getBundledModel } from "@oh-my-pi/pi-catalog/models"; import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry"; import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage"; import { removeSyncWithRetries, Snowflake } from "@oh-my-pi/pi-utils"; +import { createApiKeyResolver } from "../src/config/api-key-resolver"; describe("AuthStorage account rotation", () => { let tempDir: string; @@ -148,6 +149,63 @@ describe("AuthStorage account rotation", () => { expect(await authStorage.getApiKey("openai-codex", sessionId)).toBe(stickyKey); }); + test("API key resolver re-resolves after a concurrent OAuth refresh makes a 401 bearer stale", async () => { + const resolvedKeys = ["stale-access", "refreshed-access"]; + const rotationTargets: Array = []; + const registry: Parameters[0] = { + async getApiKeyForProvider() { + return resolvedKeys.shift(); + }, + authStorage: { + async rotateSessionCredential(_provider, _sessionId, options) { + rotationTargets.push(options?.apiKey); + return false; + }, + }, + }; + const resolver = createApiKeyResolver(registry, "openai-codex", { + sessionId: "concurrent-oauth-refresh", + }); + + const initial = await resolver({ lastChance: false, error: undefined }); + const refreshed = await resolver({ + lastChance: true, + error: Object.assign(new Error("401 authentication_error"), { status: 401 }), + previousKey: initial, + }); + + expect(initial).toBe("stale-access"); + expect(refreshed).toBe("refreshed-access"); + expect(rotationTargets).toEqual(["stale-access"]); + }); + + test("API key resolver stops when a usage-limit rotation has no unblocked sibling", async () => { + const resolvedKeys = ["quota-blocked-B", "quota-blocked-A"]; + const registry: Parameters[0] = { + async getApiKeyForProvider() { + return resolvedKeys.shift(); + }, + authStorage: { + async rotateSessionCredential() { + return false; + }, + }, + }; + const attemptedKeys: string[] = []; + + await expect( + withAuth(createApiKeyResolver(registry, "openai-codex"), async key => { + attemptedKeys.push(key); + throw Object.assign(new Error("You have hit your ChatGPT usage limit (pro plan). Try again later."), { + status: 429, + }); + }), + ).rejects.toThrow("usage limit"); + + expect(attemptedKeys).toEqual(["quota-blocked-B"]); + expect(resolvedKeys).toEqual(["quota-blocked-A"]); + }); + test("withAuth reaches a fourth healthy Codex OAuth sibling through ModelRegistry", async () => { await authStorage.set("openai-codex", [ {