diff --git a/packages/templates/src/templates/bb-guide-plugins.md b/packages/templates/src/templates/bb-guide-plugins.md index f20a0bae30..4f708c94fa 100644 --- a/packages/templates/src/templates/bb-guide-plugins.md +++ b/packages/templates/src/templates/bb-guide-plugins.md @@ -83,6 +83,12 @@ URLs. Use `bb pool config set ` to change one; the two URL values are QA-only overrides. Upgrading from a build that stored these values through plugin settings resets the threshold and QA overrides to their defaults. +Concurrent requests share one OAuth refresh per account. A temporary refresh +failure leaves a known-unexpired access token usable and delays the next refresh +attempt. If the token has expired and no other account can serve the request, +the hub returns HTTP 503. Refresh resumes automatically after the delay. +Rejected refresh credentials remain an account error. + The builtin Keep Awake plugin prevents macOS idle sleep while bb is running. Its settings page lets you target all hosts or selected hosts. The CLI equivalents are: diff --git a/plugins/account-pool/PLUGIN_OVERVIEW.md b/plugins/account-pool/PLUGIN_OVERVIEW.md index bafc78e825..24b735c8d4 100644 --- a/plugins/account-pool/PLUGIN_OVERVIEW.md +++ b/plugins/account-pool/PLUGIN_OVERVIEW.md @@ -11,6 +11,8 @@ Keep a Claude Code or Codex thread running when one account hits its limit. The The hub runs inside BB and serves an Anthropic Messages endpoint and an OpenAI Responses endpoint. With routing on, BB hands the Claude Code or Codex process a base URL that points at the hub and a token scoped to that machine, and the provider reports **Proxied** in its health row. An account is skipped for a request when it is at or above the switch threshold, on hold, or in error. The threshold defaults to 98 percent of a window. Account secrets stay in the BB data directory on the server machine, and the hub refreshes them in the background. +Concurrent requests share one OAuth refresh per account. During a temporary refresh outage, the hub can continue with an access token whose expiry is known and still in the future. It waits before another refresh attempt. If that token has expired and no other account can serve the request, the hub returns a temporary error and retries refresh on a later request. Rejected refresh credentials remain an account error. + ## Requirements Accounts you own and are permitted to use this way. diff --git a/plugins/account-pool/src/claude-adapter.ts b/plugins/account-pool/src/claude-adapter.ts index c5ea029fc0..09b0c88e2c 100644 --- a/plugins/account-pool/src/claude-adapter.ts +++ b/plugins/account-pool/src/claude-adapter.ts @@ -6,6 +6,7 @@ import { } from "./credentials.js"; import type { ProviderAdapter } from "./provider-adapter.js"; import { + fetchOAuthRefresh, filterRequestHeaders, mountedUpstreamUrl, } from "./provider-adapter.js"; @@ -101,22 +102,15 @@ export function createClaudeAdapter(options: { ) { return { secret, refreshed: false }; } - const response = await context.fetch(options.refreshUrl, { - method: "POST", - headers: { - "content-type": "application/json", - accept: "application/json", - }, - body: JSON.stringify({ - grant_type: "refresh_token", - refresh_token: secret.refreshToken, - client_id: OAUTH_CLIENT_ID, - }), - }); - if (!response.ok) { - throw new Error(`OAuth refresh failed with HTTP ${response.status}.`); - } - const parsed = refreshResponseSchema.parse(await response.json()); + const parsed = refreshResponseSchema.parse( + JSON.parse( + await fetchOAuthRefresh(context, options.refreshUrl, { + grant_type: "refresh_token", + refresh_token: secret.refreshToken, + client_id: OAUTH_CLIENT_ID, + }), + ), + ); const rawExpiresAt = parsed.expires_at ?? (parsed.expires_in === undefined diff --git a/plugins/account-pool/src/codex-adapter.ts b/plugins/account-pool/src/codex-adapter.ts index a96d08218c..83a015bf00 100644 --- a/plugins/account-pool/src/codex-adapter.ts +++ b/plugins/account-pool/src/codex-adapter.ts @@ -12,6 +12,7 @@ import { } from "./credentials.js"; import type { ProviderAdapter } from "./provider-adapter.js"; import { + fetchOAuthRefresh, filterRequestHeaders, mountedUpstreamUrl, } from "./provider-adapter.js"; @@ -296,21 +297,15 @@ export function createCodexAdapter(options: { ) { return { secret, refreshed: false }; } - const response = await context.fetch(options.refreshUrl, { - method: "POST", - headers: { - "content-type": "application/json", - accept: "application/json", - }, - body: JSON.stringify({ - client_id: CODEX_OAUTH_CLIENT_ID, - grant_type: "refresh_token", - refresh_token: secret.refreshToken, - }), - }); - if (!response.ok) - throw new Error(`OAuth refresh failed with HTTP ${response.status}.`); - const parsed = refreshResponseSchema.parse(await response.json()); + const parsed = refreshResponseSchema.parse( + JSON.parse( + await fetchOAuthRefresh(context, options.refreshUrl, { + client_id: CODEX_OAUTH_CLIENT_ID, + grant_type: "refresh_token", + refresh_token: secret.refreshToken, + }), + ), + ); const refreshed: AccountSecret = { kind: "oauth", accessToken: parsed.access_token, diff --git a/plugins/account-pool/src/hub.ts b/plugins/account-pool/src/hub.ts index 44aa8c3d04..595b5c4a1c 100644 --- a/plugins/account-pool/src/hub.ts +++ b/plugins/account-pool/src/hub.ts @@ -14,6 +14,7 @@ import { } from "./codex-adapter.js"; import type { ProviderAdapter } from "./provider-adapter.js"; import type { ImportedProviderAccount } from "./provider-adapter.js"; +import { TransientOAuthRefreshError } from "./provider-adapter.js"; import type { ImportedClaudeCredentials, ImportedCodexCredentials, @@ -32,6 +33,8 @@ const DEFAULT_USAGE_URL = "https://api.anthropic.com/api/oauth/usage"; const DEFAULT_PROFILE_URL = "https://api.anthropic.com/api/oauth/profile"; const DEFAULT_USAGE_REFRESH_INTERVAL_MS = 5 * 60 * 1_000; const MAX_INLINE_HOLD_MS = 20_000; +const MAX_REFRESH_BACKOFF_MS = 60_000; +const MAX_REFRESH_BACKOFFS = 1_024; const DROPPED_RESPONSE_HEADERS = new Set([ "content-encoding", "content-length", @@ -74,11 +77,19 @@ interface UpstreamResult { release: () => void; } +interface RefreshBackoff { + accessToken: string; + retryAt: number; + delayMs: number; + error: TransientOAuthRefreshError; +} + export class AccountPoolHub { private accepting = false; private readonly inFlightByAccount = new Map(); private readonly activeControllers = new Set(); private readonly refreshes = new Map>(); + private readonly refreshBackoffs = new Map(); private readonly usageRefreshes = new Map>(); private readonly lastUsageRefreshAt = new Map(); private readonly drainWaiters = new Set<() => void>(); @@ -247,14 +258,14 @@ export class AccountPoolHub { hostId: string | null, ): Promise { const attempted = new Set(); + let refreshFailure: TransientOAuthRefreshError | null = null; const accounts = (await this.options.accounts.list()).filter( (account) => account.provider === adapter.provider, ); const family = adapter.modelFamily(body); while (attempted.size < accounts.length) { const selected = await this.select(adapter.provider, attempted, family); - if (selected === null) - return this.noEligibleResponse(accounts, family, adapter); + if (selected === null) break; attempted.add(selected.account.id); if (hostId !== null) { const changed = await this.options.accounts.recordUsed( @@ -268,7 +279,11 @@ export class AccountPoolHub { try { secret = await this.freshSecret(selected.account, adapter); } catch (error) { - this.markError(selected.account.id, errorMessage(error)); + if (error instanceof TransientOAuthRefreshError) { + refreshFailure = error; + } else { + this.markError(selected.account.id, errorMessage(error)); + } continue; } let upstream: UpstreamResult; @@ -360,7 +375,9 @@ export class AccountPoolHub { await this.captureAuthError(upstream.response, selected.account, adapter); return this.clientResponse(upstream); } - return this.noEligibleResponse(accounts, family, adapter); + return refreshFailure === null + ? this.noEligibleResponse(accounts, family, adapter) + : adapter.errorResponse(503, refreshFailure.message); } private async captureAuthError( @@ -425,22 +442,74 @@ export class AccountPoolHub { if (existing !== undefined) return existing; const refresh = this.options.accounts .readSecret(account.id) - .then((secret) => - adapter.refreshSecret({ - account, - secret, - accounts: this.options.accounts, - quotas: this.options.quotas, - fetch: this.options.fetch, - now: this.options.now, - }), - ) - .then((result) => { - if (result.refreshed) { - const quota = this.options.quotas.get(account.id); - this.options.quotas.put({ ...quota, error: null }); + .then(async (secret) => { + let backoff = this.refreshBackoffs.get(account.id); + if ( + secret.kind !== "oauth" || + backoff?.accessToken !== secret.accessToken + ) { + this.refreshBackoffs.delete(account.id); + backoff = undefined; + } + if (backoff !== undefined && this.options.now() < backoff.retryAt) { + if ( + secret.kind === "oauth" && + secret.expiresAt !== null && + secret.expiresAt > this.options.now() + ) { + return secret; + } + throw backoff.error; + } + try { + const result = await adapter.refreshSecret({ + account, + secret, + accounts: this.options.accounts, + quotas: this.options.quotas, + fetch: this.options.fetch, + now: this.options.now, + }); + this.refreshBackoffs.delete(account.id); + if (result.refreshed) { + const quota = this.options.quotas.get(account.id); + this.options.quotas.put({ ...quota, error: null }); + } + return result.secret; + } catch (error) { + if ( + !(error instanceof TransientOAuthRefreshError) || + secret.kind !== "oauth" + ) { + this.refreshBackoffs.delete(account.id); + throw error; + } + const delayMs = Math.min( + MAX_REFRESH_BACKOFF_MS, + Math.max( + backoff === undefined ? 1_000 : backoff.delayMs * 2, + error.retryAfterMs, + ), + ); + this.refreshBackoffs.delete(account.id); + this.refreshBackoffs.set(account.id, { + accessToken: secret.accessToken, + retryAt: this.options.now() + delayMs, + delayMs, + error, + }); + while (this.refreshBackoffs.size > MAX_REFRESH_BACKOFFS) { + const oldest = this.refreshBackoffs.keys().next(); + if (!oldest.done) this.refreshBackoffs.delete(oldest.value); + } + if ( + secret.expiresAt !== null && + secret.expiresAt > this.options.now() + ) { + return secret; + } + throw error; } - return result.secret; }) .finally(() => this.refreshes.delete(account.id)); this.refreshes.set(account.id, refresh); diff --git a/plugins/account-pool/src/provider-adapter.test.ts b/plugins/account-pool/src/provider-adapter.test.ts new file mode 100644 index 0000000000..2e412b916f --- /dev/null +++ b/plugins/account-pool/src/provider-adapter.test.ts @@ -0,0 +1,101 @@ +import { describe, expect, it, vi } from "vitest"; +import { + fetchOAuthRefresh, + TransientOAuthRefreshError, +} from "./provider-adapter.js"; + +describe("OAuth refresh transport", () => { + it.each([ + { status: 400, transient: false }, + { status: 401, transient: false }, + { status: 408, transient: true }, + { status: 429, transient: true }, + { status: 500, transient: true }, + { status: 503, transient: true }, + ])( + "cancels HTTP $status responses before classifying them", + async ({ status, transient }) => { + const cancel = vi.fn(); + const response = new Response(new ReadableStream({ cancel }), { + status, + headers: { "retry-after": "120" }, + }); + const refresh = fetchOAuthRefresh( + { fetch: async () => response, now: () => 1_800_000_000_000 }, + "https://auth.example/token", + { refresh_token: "refresh" }, + ); + await expect(refresh).rejects.toThrow( + `OAuth refresh failed with HTTP ${status}.`, + ); + if (transient) { + await expect(refresh).rejects.toBeInstanceOf( + TransientOAuthRefreshError, + ); + await expect(refresh).rejects.toMatchObject({ retryAfterMs: 120_000 }); + } else { + await expect(refresh).rejects.not.toBeInstanceOf( + TransientOAuthRefreshError, + ); + } + expect(cancel).toHaveBeenCalledTimes(1); + }, + ); + + it.each(["request", "response"])( + "classifies a broken %s connection as transient", + async (stage) => { + const refresh = fetchOAuthRefresh( + { + fetch: async () => { + if (stage === "request") throw new TypeError("fetch failed"); + return new Response( + new ReadableStream({ + start(controller) { + controller.error(new TypeError("socket closed")); + }, + }), + ); + }, + now: () => 1_800_000_000_000, + }, + "https://auth.example/token", + { refresh_token: "refresh" }, + ); + await expect(refresh).rejects.toBeInstanceOf(TransientOAuthRefreshError); + }, + ); + + it("times out a pending OAuth request independently of its callers", async () => { + const controller = new AbortController(); + const timeout = vi + .spyOn(AbortSignal, "timeout") + .mockReturnValue(controller.signal); + try { + const refresh = fetchOAuthRefresh( + { + fetch: async (_input, init) => { + const signal = init?.signal; + if (signal === null || signal === undefined) + return Response.json({}); + return new Promise((_resolve, reject) => { + signal.addEventListener("abort", () => reject(signal.reason), { + once: true, + }); + }); + }, + now: () => 1_800_000_000_000, + }, + "https://auth.example/token", + { refresh_token: "refresh" }, + ); + controller.abort( + new DOMException("OAuth request timed out", "TimeoutError"), + ); + await expect(refresh).rejects.toBeInstanceOf(TransientOAuthRefreshError); + expect(timeout).toHaveBeenCalledWith(15_000); + } finally { + timeout.mockRestore(); + } + }); +}); diff --git a/plugins/account-pool/src/provider-adapter.ts b/plugins/account-pool/src/provider-adapter.ts index 713e35fb86..03d6dee92d 100644 --- a/plugins/account-pool/src/provider-adapter.ts +++ b/plugins/account-pool/src/provider-adapter.ts @@ -6,8 +6,20 @@ import type { PoolProvider, } from "./contracts.js"; import type { HubSettings } from "./hub.js"; +import { retryAfterMilliseconds } from "./quota.js"; import type { AccountStore, QuotaStore } from "./store.js"; +const OAUTH_REFRESH_TIMEOUT_MS = 15_000; + +export class TransientOAuthRefreshError extends Error { + constructor( + message: string, + readonly retryAfterMs: number, + ) { + super(message); + } +} + export interface AdapterSecretContext { account: Account; secret: AccountSecret; @@ -67,6 +79,55 @@ export interface ProviderAdapter { ): Response; } +export async function fetchOAuthRefresh( + context: Pick, + url: string, + body: Record, +): Promise { + const init: RequestInit = { + method: "POST", + headers: { + "content-type": "application/json", + accept: "application/json", + }, + body: JSON.stringify(body), + signal: AbortSignal.timeout(OAUTH_REFRESH_TIMEOUT_MS), + }; + let response: Response; + try { + response = await context.fetch(url, init); + } catch { + throw new TransientOAuthRefreshError( + "OAuth refresh failed due to a network error or timeout.", + 0, + ); + } + if (!response.ok) { + const message = `OAuth refresh failed with HTTP ${response.status}.`; + const retryAfterMs = retryAfterMilliseconds( + response.headers.get("retry-after"), + context.now(), + ); + await response.body?.cancel().catch(() => undefined); + if ( + response.status === 408 || + response.status === 429 || + response.status >= 500 + ) { + throw new TransientOAuthRefreshError(message, retryAfterMs); + } + throw new Error(message); + } + try { + return await response.text(); + } catch { + throw new TransientOAuthRefreshError( + "OAuth refresh response failed due to a network error or timeout.", + 0, + ); + } +} + export function filterRequestHeaders( inbound: Headers, allowed: ReadonlySet, diff --git a/plugins/account-pool/src/server.test.ts b/plugins/account-pool/src/server.test.ts index 04e873244e..0e2d534e29 100644 --- a/plugins/account-pool/src/server.test.ts +++ b/plugins/account-pool/src/server.test.ts @@ -179,6 +179,7 @@ function testJwt(payload: object): string { async function createFixture(args: { upstreamUrl: string; options?: AccountPoolPluginOptions; + provider?: "claude" | "codex"; source?: "api-key" | "import"; apiKey?: string; priority?: number; @@ -191,6 +192,7 @@ async function createFixture(args: { }); await host.bb.storage.kv.set("config", { anthropicUpstreamBaseUrl: args.upstreamUrl, + codexUpstreamBaseUrl: args.upstreamUrl, }); const plugin = createAccountPoolPlugin({ usageUrl: "data:application/json,{}", @@ -199,7 +201,7 @@ async function createFixture(args: { await plugin(host.bb); const accountMetadata = accountSchema.parse( await host.harness.behavior.callRpc("account.add", { - provider: "claude", + provider: args.provider ?? "claude", source: args.source === "import" ? { kind: "import" } @@ -226,7 +228,11 @@ async function createFixture(args: { await host.harness.lifecycle.dispose(); await fs.rm(dataDir, { recursive: true, force: true }); }); - return { dataDir, host, service, key: await resolveToken(host), account }; + const key = + args.provider === "codex" + ? (await resolveCodexToken(host)).token + : await resolveToken(host); + return { dataDir, host, service, key, account }; } function authHeaders(key: string): Record { @@ -2596,6 +2602,198 @@ describe("Account Pool plugin", () => { ]); }); + describe.each<{ provider: "claude" | "codex"; route: string }>([ + { provider: "claude", route: "/v1/messages" }, + { provider: "codex", route: "/v1/responses" }, + ])("$provider OAuth refresh recovery", ({ provider, route }) => { + it.each([ + { + name: "uses valid tokens during temporary failure and retries after backoff", + elapsedMinutes: 6, + failureStatus: 503, + expectedStatus: 200, + }, + { + name: "temporarily rejects expired tokens and recovers after backoff", + elapsedMinutes: 11, + failureStatus: 503, + expectedStatus: 503, + }, + { + name: "keeps invalid_grant accounts excluded despite a valid access token", + elapsedMinutes: 6, + failureStatus: 400, + expectedStatus: 429, + }, + ])("$name", async ({ elapsedMinutes, failureStatus, expectedStatus }) => { + let now = 1_800_000_000_000; + const expiresAt = now + 10 * 60 * 1_000; + const oldToken = testJwt({ exp: expiresAt / 1_000 }); + const newToken = testJwt({ exp: now / 1_000 + 3600 }); + let refreshStatus = failureStatus; + let refreshCalls = 0; + const authorizations: Array = []; + const upstream = await startUpstream(async (request, response) => { + await readRequestBody(request); + response.writeHead( + request.url === "/oauth/token" ? refreshStatus : 200, + { "content-type": "application/json" }, + ); + if (request.url === "/oauth/token") { + refreshCalls += 1; + response.end( + JSON.stringify( + refreshStatus === 200 + ? { + access_token: newToken, + refresh_token: "new-refresh", + expires_in: 3600, + } + : { + error: + refreshStatus === 400 + ? "invalid_grant" + : "temporarily_unavailable", + }, + ), + ); + return; + } + authorizations.push(request.headers.authorization); + response.end("{}"); + }); + cleanups.push(upstream.close); + const fixture = await createFixture({ + upstreamUrl: upstream.url, + provider, + source: "import", + options: { + now: () => now, + importCredentials: async () => + importedCredentials({ accessToken: oldToken, expiresAt }), + importCodexCredentials: async () => ({ + accessToken: oldToken, + refreshToken: "old-refresh", + idToken: null, + accountId: "chatgpt-account", + email: "codex@example.com", + expiresAt, + }), + refreshUrl: `${upstream.url}/oauth/token`, + codexRefreshUrl: `${upstream.url}/oauth/token`, + codexUsageUrl: EMPTY_USAGE_URL, + }, + }); + expect(refreshCalls).toBe(0); + now += elapsedMinutes * 60 * 1_000; + for (let request = 0; request < 2; request += 1) { + const response = await fixture.host.harness.behavior.fetchHttp( + "POST", + route, + { headers: authHeaders(fixture.key), body: "{}" }, + ); + await response.text(); + expect(response.status).toBe(expectedStatus); + expect(refreshCalls).toBe(1); + } + const accounts = z + .array(accountSummarySchema) + .parse( + await fixture.host.harness.behavior.callRpc("account.list", null), + ); + expect(accounts[0]?.error).toEqual( + failureStatus === 400 + ? expect.stringContaining("OAuth refresh failed") + : null, + ); + expect(authorizations).toEqual( + expectedStatus === 200 + ? [`Bearer ${oldToken}`, `Bearer ${oldToken}`] + : [], + ); + refreshStatus = 200; + now += 1_000; + const recovered = await fixture.host.harness.behavior.fetchHttp( + "POST", + route, + { headers: authHeaders(fixture.key), body: "{}" }, + ); + await recovered.text(); + expect(recovered.status).toBe(failureStatus === 400 ? 429 : 200); + expect(refreshCalls).toBe(failureStatus === 400 ? 1 : 2); + if (failureStatus !== 400) { + expect(authorizations.at(-1)).toBe(`Bearer ${newToken}`); + const recoveredAccounts = z + .array(accountSummarySchema) + .parse( + await fixture.host.harness.behavior.callRpc("account.list", null), + ); + expect(recoveredAccounts[0]?.error).toBeNull(); + } + }); + }); + + it("expires fallback tokens during refresh backoff and caps Retry-After", async () => { + let now = 1_800_000_000_000; + const expiresAt = now + 10 * 60 * 1_000; + let refreshCalls = 0; + let refreshStatus = 503; + const authorizations: Array = []; + const upstream = await startUpstream(async (request, response) => { + await readRequestBody(request); + if (request.url === "/oauth/token") { + refreshCalls += 1; + response.writeHead(refreshStatus, { + "content-type": "application/json", + "retry-after": "120", + }); + response.end( + JSON.stringify( + refreshStatus === 503 + ? { error: "temporarily_unavailable" } + : { access_token: "oauth-new", expires_in: 3600 }, + ), + ); + return; + } + authorizations.push(request.headers.authorization); + response.writeHead(200, { "content-type": "application/json" }); + response.end("{}"); + }); + cleanups.push(upstream.close); + const fixture = await createFixture({ + upstreamUrl: upstream.url, + source: "import", + options: { + now: () => now, + importCredentials: async () => importedCredentials({ expiresAt }), + refreshUrl: `${upstream.url}/oauth/token`, + }, + }); + now = expiresAt - 500; + const attempts: Array< + [advanceMs: number, expectedStatus: number, expectedRefreshes: number] + > = [ + [0, 200, 1], + [500, 503, 1], + [59_499, 503, 1], + [1, 200, 2], + ]; + for (const [advanceMs, expectedStatus, expectedRefreshes] of attempts) { + now += advanceMs; + const response = await fixture.host.harness.behavior.fetchHttp( + "POST", + "/v1/messages", + { headers: authHeaders(fixture.key), body: "{}" }, + ); + await response.text(); + expect(response.status).toBe(expectedStatus); + expect(refreshCalls).toBe(expectedRefreshes); + refreshStatus = 200; + } + expect(authorizations).toEqual(["Bearer oauth-access", "Bearer oauth-new"]); + }); + it("marks refresh and upstream authorization failures as account errors", async () => { const upstream = await startUpstream((request, response) => { if (request.url === "/oauth/token") {