diff --git a/packages/ai/src/auth-retry.ts b/packages/ai/src/auth-retry.ts index 2e4bda417..972448b2e 100644 --- a/packages/ai/src/auth-retry.ts +++ b/packages/ai/src/auth-retry.ts @@ -3,7 +3,7 @@ import type { OAuthAccess } from "./auth-storage"; import * as AIError from "./error"; import { isAuthRetryableError, isInvalidatedOAuthTokenError } from "./error/auth-classify"; import { isUsageLimit } from "./error/flags"; -import { isUsageLimitOutcome } from "./error/rate-limit"; +import { isConcurrencyCapExclusion, isUsageLimitOutcome } from "./error/rate-limit"; /** * Context passed to an {@link ApiKeyResolver} on each resolution attempt. @@ -93,11 +93,14 @@ export const AUTH_RETRY_MAX_ATTEMPTS = 64; function isDirectCredentialRotationError(error: unknown): boolean { if (isUsageLimit(error) || isInvalidatedOAuthTokenError(error)) return true; const status = AIError.status(error); - // 403: the token is valid but access was denied, so refreshing the same - // credential can't help — rotate straight through the sibling pool. - if (status === 403) return true; const message = error instanceof Error ? error.message : typeof error === "string" ? error : undefined; - if (status === undefined && message !== undefined && extractHttpStatusFromError({ message }) === 403) return true; + // A 403 normally means a valid token lacks access, so rotate through + // siblings. A concurrency-cap 403 is transient instead; do not burn a + // sibling before the caller's backoff layer can retry it. + const isForbidden = + status === 403 || + (status === undefined && message !== undefined && extractHttpStatusFromError({ message }) === 403); + if (isForbidden && !isConcurrencyCapExclusion(status, message)) return true; return isUsageLimitOutcome(status, message); } diff --git a/packages/ai/src/error/rate-limit.ts b/packages/ai/src/error/rate-limit.ts index 920e389de..7be4e71d9 100644 --- a/packages/ai/src/error/rate-limit.ts +++ b/packages/ai/src/error/rate-limit.ts @@ -30,11 +30,12 @@ const OPENROUTER_DAILY_FREE_LIMIT_PATTERN = /\bfree[-_ ]models[-_ ]per[-_ ]day\b const RESOURCE_EXHAUSTED_PATTERN = /resource.?exhausted/gi; const CONCURRENT_LIMIT_PATTERN = // Require an actual cap signal (limit/quota/exceeded/reached) near "concurrent". - // Bare nouns ("concurrent request is not supported", "only one concurrent - // invocation is supported") are deterministic 4xx feature rejections, not - // transient caps — matching them here would set Flag.Transient and retry the - // rejection instead of surfacing it. - /\bconcurren\w*\b[^\n]{0,60}\b(?:limit|quota|exceed\w*|reach\w*)\b|\b(?:limit|quota|exceed\w*|reach\w*)\b[^\n]{0,60}\bconcurren\w*\b/i; + // The first two alternatives rely on `\b`, which treats `_` as a word char, so + // structured snake_case codes ("concurrent_limit_exceeded", + // "concurrent_requests_limit_reached", "concurrency_quota_exceeded") need the + // third alternative. Bare space-separated concurrency feature rejections stay + // excluded because they neither use `[-_]` nor carry a cap keyword. + /\bconcurren\w*\b[^\n]{0,60}\b(?:limit|quota|exceed\w*|reach\w*)\b|\b(?:limit|quota|exceed\w*|reach\w*)\b[^\n]{0,60}\bconcurren\w*\b|\bconcurren[a-z]*[-_](?:[a-z]+[_-])*(?:limit|quota|exceed\w*|reach\w*)/i; const ACCOUNT_SCOPED_403_PATTERN = // The bare "limit will reset" / "will reset in" phrasing also appears on // statusless per-minute transients ("Rate limit will reset in 30 seconds"), @@ -186,7 +187,7 @@ export function isUsageLimitOutcome(status: number | undefined, message: string // still an exhausted billing cap and must rotate; gate the exclusion on the // status not being that categorical billing cap. const isBillingCapStatus = status === 402; - if (message && parseRateLimitReason(message) === "CONCURRENT_LIMIT" && !isBillingCapStatus) return false; + if (isConcurrencyCapExclusion(status, message)) return false; if (message && matchesUsageLimitText(message)) return true; // A 403 is normally an auth failure, but several providers deliver an // account-scoped cap with it (Devin/Codeium Connect `permission_denied`, @@ -242,3 +243,12 @@ export function matchesUsageLimitText(errorMessage: string): boolean { export function isAccountScopedCapText(message: string): boolean { return ACCOUNT_SCOPED_403_PATTERN.test(message); } + +/** + * A concurrency cap on a non-billing status is shed-and-backoff, not + * credential-rotatable. This mirrors the exclusion in {@link isUsageLimitOutcome} + * for the 403 auth-retry entry points. A 402 remains a categorical billing cap. + */ +export function isConcurrencyCapExclusion(status: number | undefined, message: string | undefined): boolean { + return message !== undefined && parseRateLimitReason(message) === "CONCURRENT_LIMIT" && status !== 402; +} diff --git a/packages/ai/src/stream.ts b/packages/ai/src/stream.ts index aff282f07..d8c3e1a28 100644 --- a/packages/ai/src/stream.ts +++ b/packages/ai/src/stream.ts @@ -21,7 +21,7 @@ import { createAuthRetryKeyState, isApiKeyResolver, resolveNextAuthRetryKey } fr import * as AIError from "./error"; import { ProviderHttpError } from "./error"; import { isInvalidatedOAuthTokenError } from "./error/auth-classify"; -import { isUsageLimitOutcome } from "./error/rate-limit"; +import { isConcurrencyCapExclusion, isUsageLimitOutcome } from "./error/rate-limit"; import type { BedrockOptions } from "./providers/amazon-bedrock"; import type { AnthropicOptions } from "./providers/anthropic"; import { coworkFetch } from "./providers/cowork-fetch"; @@ -995,7 +995,7 @@ function isRetryableUpstreamError(error: unknown, status: number | undefined, me // instead of burning siblings. if (AIError.isUsageLimit(error)) return true; if (isInvalidatedOAuthTokenError(error)) return true; - if (status === 401 || status === 403) return true; + if (status === 401 || (status === 403 && !isConcurrencyCapExclusion(status, message))) return true; return isUsageLimitOutcome(status, message); } diff --git a/packages/ai/test/auth-retry.test.ts b/packages/ai/test/auth-retry.test.ts index e4fbdb652..ad8e9bcb7 100644 --- a/packages/ai/test/auth-retry.test.ts +++ b/packages/ai/test/auth-retry.test.ts @@ -275,6 +275,32 @@ describe("withAuth", () => { expect(contexts.map(ctx => ctx.lastChance)).toEqual([false, true, true, true]); }); + it("does not directly rotate through every sibling on a 403 concurrency cap", async () => { + const keys: string[] = []; + const contexts: ApiKeyResolveContext[] = []; + const pool = ["k0", "k1", "k2", "k3"]; + let resolveIndex = 0; + const concurrencyCap = Object.assign(new Error("concurrent requests limit reached"), { status: 403 }); + + await expect( + withAuth( + ctx => { + contexts.push(ctx); + return ctx.error === undefined ? pool[0] : pool[++resolveIndex]; + }, + async key => { + keys.push(key); + throw concurrencyCap; + }, + ), + ).rejects.toBe(concurrencyCap); + + // The concurrency classification takes precedence over plain-403 direct + // rotation: refresh once, then take only the legacy sibling switch. + expect(keys).toEqual(["k0", "k1", "k2"]); + expect(contexts.map(ctx => ctx.lastChance)).toEqual([false, false, true]); + }); + it("surfaces the last 403 when every sibling is denied", async () => { const errors = [authError(403), authError(403)]; const resolved = ["k0", "k1", "k0"]; diff --git a/packages/ai/test/rate-limit-utils.test.ts b/packages/ai/test/rate-limit-utils.test.ts index 787563b0e..8d2feb84c 100644 --- a/packages/ai/test/rate-limit-utils.test.ts +++ b/packages/ai/test/rate-limit-utils.test.ts @@ -1,8 +1,9 @@ import { describe, expect, it } from "bun:test"; import { ProviderHttpError } from "@oh-my-pi/pi-ai/error"; -import { classify, classifyMessage, Flag, is, isUsageLimit, retriable } from "@oh-my-pi/pi-ai/error/flags"; +import { classify, Flag, is, isUsageLimit, retriable } from "@oh-my-pi/pi-ai/error/flags"; import { calculateRateLimitBackoffMs, + isConcurrencyCapExclusion, isUsageLimitOutcome, isUsageLimitStatus, parseRateLimitReason, @@ -61,6 +62,9 @@ describe("parseRateLimitReason", () => { it("classifies concurrent request caps separately from rate limits and quota exhaustion", () => { expect(parseRateLimitReason("Number of concurrent requests exceeded")).toBe("CONCURRENT_LIMIT"); expect(parseRateLimitReason("Maximum concurrent invocation limit reached")).toBe("CONCURRENT_LIMIT"); + expect(parseRateLimitReason("concurrent_limit_exceeded")).toBe("CONCURRENT_LIMIT"); + expect(parseRateLimitReason("concurrent_requests_limit_reached")).toBe("CONCURRENT_LIMIT"); + expect(parseRateLimitReason("concurrency_quota_exceeded")).toBe("CONCURRENT_LIMIT"); expect(parseRateLimitReason("Rate limit reached for gpt-4o")).toBe("RATE_LIMIT_EXCEEDED"); expect(parseRateLimitReason("Your quota will reset at 07-28")).toBe("QUOTA_EXHAUSTED"); }); @@ -309,30 +313,6 @@ describe("isUsageLimitOutcome", () => { expect(isUsageLimit(message)).toBe(false); }); - // agent-session gates Copilot credential removal on AuthFailed && !UsageLimit. - // A 403 whose body is a recognized account cap carries UsageLimit, so the gate - // suppresses removal and retains the still-valid credential; a plain 401 has no - // UsageLimit, so removal proceeds. Pin the gate condition for both shapes. - it("keeps Copilot credentials for a 403 account cap, removes on a 401", () => { - const cap = new ProviderHttpError( - "Reached overall message rate limit. Your limit will reset in 13 minutes.", - 403, - ); - const capId = classifyMessage({ - errorId: classify(cap), - errorMessage: "GitHub Copilot access denied (HTTP 403).", - errorStatus: 403, - }); - expect(is(capId, Flag.AuthFailed) && !is(capId, Flag.UsageLimit)).toBe(false); - const auth = new ProviderHttpError("401 Unauthorized", 401); - const authId = classifyMessage({ - errorId: classify(auth), - errorMessage: "GitHub Copilot authentication failed (HTTP 401).", - errorStatus: 401, - }); - expect(is(authId, Flag.AuthFailed) && !is(authId, Flag.UsageLimit)).toBe(true); - }); - it("rotates on xAI Grok Build 402 usage-balance exhaustion regardless of status", () => { const message = "402 Grok Build usage balance exhausted"; expect(isUsageLimitOutcome(402, message)).toBe(true); @@ -365,6 +345,14 @@ describe("isUsageLimitOutcome", () => { expect(isUsageLimit(message)).toBe(false); }); + it("excludes non-billing concurrency caps from credential rotation", () => { + const message = "concurrent requests limit reached"; + expect(isConcurrencyCapExclusion(403, message)).toBe(true); + expect(isConcurrencyCapExclusion(undefined, message)).toBe(true); + expect(isConcurrencyCapExclusion(402, message)).toBe(false); + expect(isConcurrencyCapExclusion(403, "Forbidden")).toBe(false); + }); + // The same bare concurrency wording can reach turn recovery without a // preserved HTTP status (Vertex/Bedrock paths that bypass API-key // resolution). The body misses TRANSIENT_TRANSPORT_PATTERN, so without an diff --git a/packages/ai/test/stream-auth-retry.test.ts b/packages/ai/test/stream-auth-retry.test.ts index 83e06c66d..90e782f35 100644 --- a/packages/ai/test/stream-auth-retry.test.ts +++ b/packages/ai/test/stream-auth-retry.test.ts @@ -120,6 +120,41 @@ describe("streamSimple resolver auth retry", () => { expect((contexts[1]!.error as { status?: number }).status).toBe(401); }); + it("surfaces a 403 concurrency cap for transient backoff without rotating credentials", async () => { + const keys: unknown[] = []; + const contexts: ApiKeyResolveContext[] = []; + const concurrencyCap = Object.assign(new Error("concurrent requests limit reached"), { status: 403 }); + registerCustomApi( + API, + (_model: Model, _context: Context, options?: SimpleStreamOptions) => { + pushKey(keys, options); + const stream = new AssistantMessageEventStream(); + queueMicrotask(() => stream.fail(concurrencyCap)); + return stream; + }, + SOURCE_ID, + ); + + const stream = streamSimple(model(), context, { + apiKey: async ctx => { + contexts.push(ctx); + return ctx.error === undefined ? "old-key" : ctx.lastChance ? "sibling-key" : "refresh-key"; + }, + }); + await expect( + (async () => { + for await (const _event of stream) { + // drain + } + })(), + ).rejects.toBe(concurrencyCap); + + expect(keys).toEqual(["old-key"]); + expect(contexts.map(ctx => ({ lastChance: ctx.lastChance, hasError: ctx.error !== undefined }))).toEqual([ + { lastChance: false, hasError: false }, + ]); + }); + it("buffers the start event and retries on a 401 error event before content", async () => { const keys: unknown[] = []; const eventTypes: string[] = []; diff --git a/packages/coding-agent/test/agent-session-copilot-credential-removal.test.ts b/packages/coding-agent/test/agent-session-copilot-credential-removal.test.ts new file mode 100644 index 000000000..24547b579 --- /dev/null +++ b/packages/coding-agent/test/agent-session-copilot-credential-removal.test.ts @@ -0,0 +1,72 @@ +import { expect, it, spyOn } from "bun:test"; +import { Agent } from "@oh-my-pi/pi-agent-core"; +import type { AssistantMessage } from "@oh-my-pi/pi-ai"; +import { getBundledModel } from "@oh-my-pi/pi-catalog/models"; +import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry"; +import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; +import { AgentSession } from "@oh-my-pi/pi-coding-agent/session/agent-session"; +import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage"; +import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager"; +import { TempDir } from "@oh-my-pi/pi-utils"; + +it("removes a Copilot credential for 401 but retains it for a 403 account cap", async () => { + const tempDir = TempDir.createSync("@pi-copilot-credential-removal-"); + const authStorage = await AuthStorage.create(tempDir.join("testauth.db")); + const modelRegistry = new ModelRegistry(authStorage, tempDir.join("models.yml")); + const model = getBundledModel("anthropic", "claude-sonnet-4-5"); + if (!model) throw new Error("Expected bundled Anthropic test model to exist"); + + const agent = new Agent({ + initialState: { model, systemPrompt: ["Test"], tools: [], messages: [] }, + }); + const session = new AgentSession({ + agent, + sessionManager: SessionManager.inMemory(), + settings: Settings.isolated({ "compaction.enabled": false }), + modelRegistry, + }); + const removeSpy = spyOn(authStorage, "remove").mockResolvedValue(undefined); + + try { + const unauthorized: AssistantMessage = { + role: "assistant", + content: [], + api: "openai-responses", + provider: "github-copilot", + model: "gpt-5-mini", + stopReason: "error", + errorMessage: "GitHub Copilot authentication failed (HTTP 401).", + errorStatus: 401, + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + timestamp: Date.now(), + }; + agent.emitExternalEvent({ type: "message_end", message: unauthorized }); + agent.emitExternalEvent({ type: "agent_end", messages: [unauthorized] }); + await session.waitForIdle(); + expect(removeSpy).toHaveBeenCalledWith("github-copilot"); + + removeSpy.mockClear(); + const accountCap: AssistantMessage = { + ...unauthorized, + errorMessage: "Reached overall message rate limit. Your limit will reset in 13 minutes.", + errorStatus: 403, + timestamp: Date.now(), + }; + agent.emitExternalEvent({ type: "message_end", message: accountCap }); + agent.emitExternalEvent({ type: "agent_end", messages: [accountCap] }); + await session.waitForIdle(); + expect(removeSpy).not.toHaveBeenCalled(); + } finally { + await session.dispose(); + removeSpy.mockRestore(); + authStorage.close(); + tempDir.removeSync(); + } +});