Fix automatic rollover across quota-limited accounts

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

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<string | undefined> {
 	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,

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<string | undefined> {
 	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<string>;
+	/** 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<string | undefined> {
+	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<T>(
+	access: OAuthAccess,
+	attempt: (access: OAuthAccess) => Promise<T>,
+	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<T>(
 
 	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<boolean>;
 }
 
@@ -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<T>(
 	storage: OAuthAccessSource,
@@ -221,34 +317,76 @@ export async function withOAuthAccess<T>(
 		);
 	}
 
-	const resolveStep = async (lastChance: boolean, error: unknown): Promise<OAuthAccess | undefined> => {
-		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;