diff --git a/plugins/account-pool/src/claude-adapter.ts b/plugins/account-pool/src/claude-adapter.ts index c5ea029fc0..133cd0bdcf 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"; @@ -96,27 +97,21 @@ export function createClaudeAdapter(options: { const secret = context.secret; if ( secret.kind !== "oauth" || - secret.expiresAt === null || - secret.expiresAt > context.now() + REFRESH_WINDOW_MS + (!context.forceRefresh && + (secret.expiresAt === null || + secret.expiresAt > context.now() + REFRESH_WINDOW_MS)) ) { 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..6dd556c7d6 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"; @@ -291,26 +292,21 @@ export function createCodexAdapter(options: { const secret = context.secret; if ( secret.kind !== "oauth" || - secret.expiresAt === null || - secret.expiresAt > context.now() + REFRESH_WINDOW_MS + (!context.forceRefresh && + (secret.expiresAt === null || + secret.expiresAt > context.now() + REFRESH_WINDOW_MS)) ) { 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 7b9c83cb27..c5ac772076 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,10 @@ 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 MAX_FAILURE_DETAIL_BYTES = 1_024; +const FAILURE_DISPOSAL_TIMEOUT_MS = 250; const DROPPED_RESPONSE_HEADERS = new Set([ "content-encoding", "content-length", @@ -74,11 +79,35 @@ interface UpstreamResult { release: () => void; } +interface RefreshBackoff { + kind: "proactive" | "rejected"; + accessToken: string; + retryAt: number; + delayMs: number; + error: TransientOAuthRefreshError; +} + +type SecretUse = { kind: "normal" } | { kind: "rejected"; accessToken: string }; + +type SecretFlight = + | { kind: "refresh"; use: SecretUse; result: Promise } + | { kind: "rejection-check"; result: Promise }; + +interface FailureSummary { + status: number; + message: string; + headers: Record; +} + +class UpstreamConnectionError extends Error {} + export class AccountPoolHub { private accepting = false; + private stopped = new AbortController(); private readonly inFlightByAccount = new Map(); private readonly activeControllers = new Set(); - private readonly refreshes = new Map>(); + 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>(); @@ -86,6 +115,7 @@ export class AccountPoolHub { constructor(private readonly options: HubOptions) {} async start(signal: AbortSignal): Promise { + this.stopped = new AbortController(); this.accepting = true; while (!signal.aborted) { await this.refreshUsage(); @@ -171,7 +201,8 @@ export class AccountPoolHub { const refresh = adapter .refreshUsage({ account, - freshSecret: () => this.freshSecret(account, adapter), + freshSecret: () => + this.freshSecret(account, adapter, { kind: "normal" }), accounts: this.options.accounts, quotas: this.options.quotas, fetch: this.options.fetch, @@ -185,6 +216,7 @@ export class AccountPoolHub { async stop(): Promise { this.accepting = false; + this.stopped.abort(new Error("Account Pooler stopped accepting requests.")); if (this.inFlightCount() === 0) return; let timeout: ReturnType | null = null; await Promise.race([ @@ -246,142 +278,293 @@ export class AccountPoolHub { adapter: ProviderAdapter, hostId: string | null, ): Promise { + const signal = AbortSignal.any([request.signal, this.stopped.signal]); const attempted = new Set(); + let failure: FailureSummary | null = null; const accounts = (await this.options.accounts.list()).filter( (account) => account.provider === adapter.provider, ); + const candidateIds = new Set(accounts.map((account) => account.id)); 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); - attempted.add(selected.account.id); - if (hostId !== null) { - const changed = await this.options.accounts.recordUsed( - selected.account.id, - this.options.now(), - hostId, - ); - if (changed) this.options.onAccountsChanged(); - } - let secret: AccountSecret; - try { - secret = await this.freshSecret(selected.account, adapter); - } catch (error) { - this.markError(selected.account.id, errorMessage(error)); - continue; - } - let upstream: UpstreamResult; - try { - upstream = await this.fetchUpstream( - request, - adapter.prepareBody(body, selected.account), - selected.account, - secret, - adapter, - ); - } catch { - return adapter.errorResponse( - 502, - `Account Pooler could not reach ${adapter.upstreamName}.`, - ); - } - const observed = adapter.quotaFromHeaders( - selected.account.id, - upstream.response.headers, - this.options.quotas.get(selected.account.id), - family, - this.options.now(), - ); - this.options.quotas.put(observed); - if ( - upstream.response.status === 429 && - adapter.isQuotaRejection(upstream.response.headers) - ) { - await upstream.response.body?.cancel(); - upstream.release(); - continue; - } - if (upstream.response.status === 429) { - const waitMs = retryAfterMilliseconds( - upstream.response.headers.get("retry-after"), - this.options.now(), + try { + while (attempted.size < candidateIds.size) { + signal.throwIfAborted(); + const selected = await this.select( + adapter.provider, + candidateIds, + attempted, + family, ); - this.options.quotas.put({ - ...observed, - heldUntil: this.options.now() + waitMs, - }); - if (waitMs <= MAX_INLINE_HOLD_MS) { - await upstream.response.body?.cancel(); - upstream.release(); - await delay(waitMs); + if (selected === null) break; + attempted.add(selected.account.id); + if (hostId !== null) { + const changed = await this.options.accounts.recordUsed( + selected.account.id, + this.options.now(), + hostId, + ); + if (changed) this.options.onAccountsChanged(); + } + let secret: AccountSecret; + try { + signal.throwIfAborted(); + secret = await abortable( + this.freshSecret(selected.account, adapter, { kind: "normal" }), + signal, + ); + } catch (error) { + signal.throwIfAborted(); + if (error instanceof TransientOAuthRefreshError) { + failure = { status: 503, message: error.message, headers: {} }; + } else { + this.markError(selected.account.id, errorMessage(error)); + } + continue; + } + let authRetried = false; + let paced = false; + while (true) { + signal.throwIfAborted(); + let upstream: UpstreamResult; try { - const retry = await this.fetchUpstream( + upstream = await this.fetchUpstream( request, adapter.prepareBody(body, selected.account), selected.account, secret, adapter, ); - const retryQuota = adapter.quotaFromHeaders( - selected.account.id, - retry.response.headers, - this.options.quotas.get(selected.account.id), - family, + } catch (error) { + signal.throwIfAborted(); + if (!(error instanceof UpstreamConnectionError)) throw error; + failure = { + status: 502, + message: + "Account Pooler could not reach " + adapter.upstreamName + ".", + headers: {}, + }; + break; + } + if (request.signal.aborted) { + await this.discardUpstream(upstream, false); + signal.throwIfAborted(); + } + const { response } = upstream; + const observed = adapter.quotaFromHeaders( + selected.account.id, + response.headers, + this.options.quotas.get(selected.account.id), + family, + this.options.now(), + ); + this.options.quotas.put(observed); + if (response.status === 429) { + if (adapter.isQuotaRejection(response.headers)) { + await this.discardUpstream(upstream, false); + break; + } + const waitMs = retryAfterMilliseconds( + response.headers.get("retry-after"), this.options.now(), ); - this.options.quotas.put( - retry.response.status === 429 - ? { - ...retryQuota, - heldUntil: - this.options.now() + - retryAfterMilliseconds( - retry.response.headers.get("retry-after"), - this.options.now(), - ), - } - : retryQuota, - ); - await this.captureAuthError( - retry.response, - selected.account, - adapter, - ); - return this.clientResponse(retry); - } catch { - return adapter.errorResponse( - 502, - `Account Pooler could not reach ${adapter.upstreamName}.`, - ); + this.options.quotas.put({ + ...observed, + heldUntil: this.options.now() + waitMs, + }); + if (!paced && waitMs <= MAX_INLINE_HOLD_MS) { + paced = true; + await this.discardUpstream(upstream, false); + await waitForDelay(waitMs, signal); + continue; + } + } + if ( + response.status === 401 || + response.status === 403 || + response.status === 408 || + response.status === 500 || + response.status === 502 || + response.status === 503 || + response.status === 504 || + response.status === 529 + ) { + const retryAfter = response.headers.get("retry-after"); + const detail = await this.discardUpstream(upstream, true); + signal.throwIfAborted(); + failure = { + status: response.status, + message: + detail || + adapter.upstreamName + + " returned HTTP " + + response.status + + ".", + headers: retryAfter === null ? {} : { "retry-after": retryAfter }, + }; + if ( + response.status === 401 && + secret.kind === "oauth" && + !authRetried + ) { + authRetried = true; + try { + secret = await abortable( + this.freshSecret(selected.account, adapter, { + kind: "rejected", + accessToken: secret.accessToken, + }), + signal, + ); + } catch (error) { + signal.throwIfAborted(); + if (error instanceof TransientOAuthRefreshError) { + failure = { + status: 503, + message: error.message, + headers: {}, + }; + } else { + await this.markAuthError( + selected.account, + secret, + errorMessage(error), + signal, + ); + } + break; + } + continue; + } + if (response.status === 401 || response.status === 403) { + await this.markAuthError( + selected.account, + secret, + failure.message, + signal, + ); + } + break; } + return this.clientResponse(upstream); } } - await this.captureAuthError(upstream.response, selected.account, adapter); - return this.clientResponse(upstream); + signal.throwIfAborted(); + return failure === null + ? this.noEligibleResponse(accounts, family, adapter) + : adapter.errorResponse( + failure.status, + failure.message, + failure.headers, + ); + } catch (error) { + if (!signal.aborted) throw error; + return adapter.errorResponse( + request.signal.aborted ? 499 : 503, + request.signal.aborted + ? "Account Pooler request was canceled." + : "Account Pooler stopped accepting requests.", + ); + } + } + + private async discardUpstream( + upstream: UpstreamResult, + readDetail: boolean, + ): Promise { + const reader = upstream.response.body?.getReader(); + if (reader === undefined) { + upstream.controller.abort(); + upstream.release(); + return ""; + } + let timeout: ReturnType | undefined; + const deadline = new Promise((resolve) => { + timeout = setTimeout(() => resolve(""), FAILURE_DISPOSAL_TIMEOUT_MS); + }); + let detail = ""; + try { + if (!readDetail) return detail; + return await Promise.race([ + (async () => { + const decoder = new TextDecoder(); + let bytes = 0; + while (bytes < MAX_FAILURE_DETAIL_BYTES) { + const chunk = await reader.read(); + if (chunk.done) break; + const part = chunk.value.subarray( + 0, + MAX_FAILURE_DETAIL_BYTES - bytes, + ); + bytes += part.byteLength; + detail += decoder.decode(part, { stream: true }); + } + return (detail + decoder.decode()).trim(); + })().catch(() => detail.trim()), + deadline, + ]); + } finally { + upstream.controller.abort(); + await Promise.race([reader.cancel().catch(() => undefined), deadline]); + clearTimeout(timeout); + upstream.release(); } - return this.noEligibleResponse(accounts, family, adapter); } - private async captureAuthError( - response: Response, + private async markAuthError( account: Account, - adapter: ProviderAdapter, + rejected: AccountSecret, + message: string, + signal: AbortSignal, ): Promise { - if (response.status !== 401 && response.status !== 403) return; - const detail = await response - .clone() - .text() - .catch(() => ""); - this.markError( - account.id, - detail.trim() || - `${adapter.upstreamName} returned HTTP ${response.status}.`, - ); + if (rejected.kind !== "oauth") { + this.markError(account.id, message); + return; + } + while (true) { + signal.throwIfAborted(); + const existing = this.refreshes.get(account.id); + if (existing !== undefined) { + await abortable( + existing.result.then( + () => undefined, + () => undefined, + ), + signal, + ); + continue; + } + const flight: SecretFlight = { + kind: "rejection-check", + result: this.options.accounts + .readSecret(account.id) + .then((current) => { + const backoff = this.refreshBackoffs.get(account.id); + if ( + !signal.aborted && + current.kind === "oauth" && + current.accessToken === rejected.accessToken && + !( + backoff?.kind === "rejected" && + backoff.accessToken === current.accessToken + ) + ) { + this.markError(account.id, message); + } + }) + .finally(() => { + if (this.refreshes.get(account.id) === flight) + this.refreshes.delete(account.id); + }), + }; + this.refreshes.set(account.id, flight); + await abortable(flight.result, signal); + return; + } } private async select( provider: PoolProvider, + candidateIds: ReadonlySet, attempted: ReadonlySet, family: ModelFamily, ): Promise { @@ -391,6 +574,7 @@ export class AccountPoolHub { .filter( (account) => account.provider === provider && + candidateIds.has(account.id) && account.enabled && !attempted.has(account.id), ) @@ -420,29 +604,147 @@ export class AccountPoolHub { private async freshSecret( account: Account, adapter: ProviderAdapter, + use: SecretUse, ): Promise { - const existing = this.refreshes.get(account.id); - if (existing !== undefined) return existing; - const secret = await this.options.accounts.readSecret(account.id); - const refresh = 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 }); + while (true) { + const existing = this.refreshes.get(account.id); + if (existing !== undefined) { + if (existing.kind === "rejection-check") { + await existing.result; + continue; } - return result.secret; - }) - .finally(() => this.refreshes.delete(account.id)); - this.refreshes.set(account.id, refresh); - return refresh; + let secret: AccountSecret; + try { + secret = await existing.result; + } catch (error) { + const current = this.refreshes.get(account.id); + if (current !== undefined && current !== existing) continue; + throw error; + } + const current = this.refreshes.get(account.id); + if (current !== undefined && current !== existing) continue; + const backoff = this.refreshBackoffs.get(account.id); + if ( + secret.kind === "oauth" && + backoff?.accessToken === secret.accessToken && + backoff.kind === "rejected" + ) + continue; + if ( + use.kind === "normal" || + secret.kind !== "oauth" || + secret.accessToken !== use.accessToken || + (existing.use.kind === "rejected" && + existing.use.accessToken === use.accessToken) + ) { + return secret; + } + continue; + } + const flight: Extract = { + kind: "refresh", + use, + result: this.options.accounts + .readSecret(account.id) + .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; + } + const explicitlyRejected = + secret.kind === "oauth" && + use.kind === "rejected" && + secret.accessToken === use.accessToken; + const forceRefresh = + explicitlyRejected || backoff?.kind === "rejected"; + if (forceRefresh && secret.kind === "oauth") { + flight.use = { + kind: "rejected", + accessToken: secret.accessToken, + }; + } + const error = this.options.quotas.get(account.id).error; + if (error !== null) throw new Error(error); + if ( + backoff !== undefined && + this.options.now() < backoff.retryAt && + (!explicitlyRejected || backoff.kind === "rejected") + ) { + if ( + !forceRefresh && + 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, + forceRefresh, + }); + 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, { + kind: forceRefresh ? "rejected" : "proactive", + 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 ( + !forceRefresh && + secret.expiresAt !== null && + secret.expiresAt > this.options.now() + ) { + return secret; + } + throw error; + } + }) + .finally(() => { + if (this.refreshes.get(account.id) === flight) + this.refreshes.delete(account.id); + }), + }; + this.refreshes.set(account.id, flight); + return flight.result; + } } private async fetchUpstream( @@ -472,17 +774,20 @@ export class AccountPoolHub { try { const upstreamBody = new ArrayBuffer(body.byteLength); new Uint8Array(upstreamBody).set(body); - const response = await this.options.fetch( - adapter.upstreamUrl(request, this.options.getSettings()), - { + const url = adapter.upstreamUrl(request, this.options.getSettings()); + const headers = adapter.requestHeaders(request.headers, account, secret); + const response = await this.options + .fetch(url, { method: request.method, - headers: adapter.requestHeaders(request.headers, account, secret), + headers, ...(request.method === "GET" || request.method === "HEAD" ? {} : { body: upstreamBody }), signal: controller.signal, - }, - ); + }) + .catch(() => { + throw new UpstreamConnectionError("Upstream connection failed."); + }); return { response, controller, release }; } catch (error) { release(); @@ -699,6 +1004,20 @@ function waitForDelay( }); } -function delay(milliseconds: number): Promise { - return new Promise((resolve) => setTimeout(resolve, milliseconds)); +function abortable(operation: Promise, signal: AbortSignal): Promise { + if (signal.aborted) return Promise.reject(signal.reason); + return new Promise((resolve, reject) => { + const abort = () => reject(signal.reason); + signal.addEventListener("abort", abort, { once: true }); + operation.then( + (result) => { + signal.removeEventListener("abort", abort); + resolve(result); + }, + (error) => { + signal.removeEventListener("abort", abort); + reject(error); + }, + ); + }); } 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..df9ae9395e --- /dev/null +++ b/plugins/account-pool/src/provider-adapter.test.ts @@ -0,0 +1,149 @@ +import { describe, expect, it, vi } from "vitest"; +import { + fetchOAuthRefresh, + TransientOAuthRefreshError, +} from "./provider-adapter.js"; + +describe("OAuth refresh transport", () => { + it("classifies rejected responses even when cancellation never settles", async () => { + const cancel = vi.fn(() => new Promise(() => {})); + await expect( + fetchOAuthRefresh( + { + fetch: async () => + new Response(new ReadableStream({ cancel }), { status: 503 }), + now: () => 0, + }, + "https://auth.example/token", + { refresh_token: "refresh" }, + ), + ).rejects.toBeInstanceOf(TransientOAuthRefreshError); + expect(cancel).toHaveBeenCalledOnce(); + }); + + it.each(["request", "body"])( + "bounds a %s transport that ignores abort", + async (stage) => { + const controller = new AbortController(); + const timeout = vi + .spyOn(AbortSignal, "timeout") + .mockReturnValue(controller.signal); + try { + const refresh = fetchOAuthRefresh( + { + fetch: async () => + stage === "request" + ? new Promise(() => {}) + : new Response(new ReadableStream()), + now: () => 0, + }, + "https://auth.example/token", + { refresh_token: "refresh" }, + ); + await new Promise((resolve) => setImmediate(resolve)); + controller.abort( + new DOMException("OAuth request timed out", "TimeoutError"), + ); + await expect(refresh).rejects.toBeInstanceOf( + TransientOAuthRefreshError, + ); + } finally { + timeout.mockRestore(); + } + }, + ); + + 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..d5dd56a52f 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; @@ -15,6 +27,7 @@ export interface AdapterSecretContext { quotas: QuotaStore; fetch: typeof fetch; now: () => number; + forceRefresh: boolean; } export interface AdapterUsageContext { @@ -67,6 +80,74 @@ export interface ProviderAdapter { ): Response; } +export async function fetchOAuthRefresh( + context: Pick, + url: string, + body: Record, +): Promise { + const signal = AbortSignal.timeout(OAUTH_REFRESH_TIMEOUT_MS); + let onTimeout = () => {}; + const timedOut = new Promise((_resolve, reject) => { + onTimeout = () => reject(signal.reason); + if (signal.aborted) onTimeout(); + else signal.addEventListener("abort", onTimeout, { once: true }); + }); + try { + let response: Response; + try { + response = await Promise.race([ + context + .fetch(url, { + method: "POST", + headers: { + "content-type": "application/json", + accept: "application/json", + }, + body: JSON.stringify(body), + signal, + }) + .then((result) => { + if (signal.aborted) + void result.body?.cancel().catch(() => undefined); + return result; + }), + timedOut, + ]); + } 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(), + ); + void 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 Promise.race([response.text(), timedOut]); + } catch { + throw new TransientOAuthRefreshError( + "OAuth refresh response failed due to a network error or timeout.", + 0, + ); + } + } finally { + signal.removeEventListener("abort", onTimeout); + } +} + 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 7047d74b8c..966b1156e2 100644 --- a/plugins/account-pool/src/server.test.ts +++ b/plugins/account-pool/src/server.test.ts @@ -21,7 +21,7 @@ import type { ImportedClaudeCredentials, ImportedCodexCredentials, } from "./credentials.js"; -import { HubTokenStore } from "./store.js"; +import { AccountStore, HubTokenStore } from "./store.js"; import { createAccountPoolPlugin, helloResponse, @@ -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,54 @@ 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 }; +} + +async function createOAuthRequestFixture( + provider: "claude" | "codex", + upstreamFetch: typeof fetch, + now: () => number, +): Promise { + return createFixture({ + upstreamUrl: "https://upstream.example", + provider, + source: "import", + options: { + fetch: (input, init) => + String(input) === EMPTY_USAGE_URL + ? Promise.resolve(Response.json({})) + : upstreamFetch(input, init), + now, + refreshUrl: "https://upstream.example/oauth/token", + codexRefreshUrl: "https://upstream.example/oauth/token", + codexUsageUrl: EMPTY_USAGE_URL, + importCredentials: async () => + importedCredentials({ + accessToken: "oauth-old", + expiresAt: now() + 60 * 60 * 1_000, + }), + importCodexCredentials: async () => ({ + accessToken: "oauth-old", + refreshToken: "oauth-refresh", + idToken: null, + accountId: "chatgpt-account", + email: "codex@example.com", + expiresAt: now() + 60 * 60 * 1_000, + }), + }, + }); +} + +function deferred(): { promise: Promise; resolve: () => void } { + let resolve = () => {}; + const promise = new Promise((release) => { + resolve = release; + }); + return { promise, resolve }; } function authHeaders(key: string): Record { @@ -2404,13 +2453,854 @@ describe("Account Pool plugin", () => { expect((times[1] ?? 0) - (times[0] ?? 0)).toBeGreaterThanOrEqual(30); }); + it.each([401, 403, 408, 500, 502, 503, 504, 529, "disconnect"])( + "tries another account after a pre-stream %s failure", + async (failure) => { + const attempts: Array = []; + const upstream = await startUpstream(async (request, response) => { + await readRequestBody(request); + const key = request.headers["x-api-key"]; + attempts.push(typeof key === "string" ? key : undefined); + if (key === "sk-first") { + if (typeof failure === "string") { + request.socket.destroy(); + return; + } + response.writeHead(failure, { "content-type": "application/json" }); + response.end('{"error":{"message":"first account failed"}}'); + return; + } + response.writeHead(200, { "content-type": "application/json" }); + response.end('{"result":"second"}'); + }); + cleanups.push(upstream.close); + const fixture = await createFixture({ + upstreamUrl: upstream.url, + apiKey: "sk-first", + priority: 0, + }); + await addApiAccount(fixture, "sk-second", 100); + const response = await fixture.host.harness.behavior.fetchHttp( + "POST", + "/v1/messages", + { headers: authHeaders(fixture.key), body: "{}" }, + ); + const payload = await response.json(); + expect(response.status).toBe(200); + expect(payload).toEqual({ result: "second" }); + expect(attempts).toEqual(["sk-first", "sk-second"]); + const accounts = z + .array(accountSummarySchema) + .parse( + await fixture.host.harness.behavior.callRpc("account.list", null), + ); + expect( + accounts.find((account) => account.id === fixture.account.id)?.error, + ).toEqual(failure === 401 || failure === 403 ? expect.any(String) : null); + }, + ); + + describe.each<{ provider: "claude" | "codex"; route: string }>([ + { provider: "claude", route: "/v1/messages" }, + { provider: "codex", route: "/v1/responses" }, + ])("$provider rejected-token recovery", ({ provider, route }) => { + it.each([ + { + name: "401", + statuses: [401, 200], + refreshStatus: 200, + sameToken: false, + }, + { + name: "429 then 401", + statuses: [429, 401, 200], + refreshStatus: 200, + sameToken: false, + }, + { + name: "same-token refresh", + statuses: [401, 200], + refreshStatus: 200, + sameToken: true, + }, + { + name: "401 then 429", + statuses: [401, 429, 200], + refreshStatus: 200, + sameToken: false, + }, + { + name: "401 after both retry budgets", + statuses: [401, 429, 401], + refreshStatus: 200, + sameToken: false, + }, + { + name: "repeated 401", + statuses: [401, 401], + refreshStatus: 200, + sameToken: false, + }, + { + name: "refresh outage", + statuses: [401], + refreshStatus: 503, + sameToken: false, + }, + ])( + "bounds $name recovery and never reuses rejected fallback credentials", + async ({ statuses, refreshStatus, sameToken }) => { + let now = 1_800_000_000_000; + let refreshCalls = 0; + let oauthStatus = refreshStatus; + const newToken = sameToken + ? "oauth-old" + : testJwt({ exp: now / 1_000 + 3600 }); + const authorizations: Array = []; + const fixture = await createOAuthRequestFixture( + provider, + async (input, init) => { + if (String(input).endsWith("/oauth/token")) { + refreshCalls += 1; + return Response.json( + oauthStatus === 200 + ? { access_token: newToken, expires_in: 3600 } + : { error: "temporarily_unavailable" }, + { status: oauthStatus }, + ); + } + authorizations.push( + new Headers(init?.headers).get("authorization"), + ); + return Response.json( + { result: "upstream" }, + { + status: statuses[authorizations.length - 1] ?? 200, + headers: { "retry-after": "0" }, + }, + ); + }, + () => now, + ); + const response = await fixture.host.harness.behavior.fetchHttp( + "POST", + route, + { + headers: authHeaders(fixture.key), + body: "{}", + }, + ); + await response.text(); + expect(response.status).toBe( + refreshStatus === 503 ? 503 : statuses.at(-1), + ); + expect(refreshCalls).toBe(1); + expect(authorizations).toEqual( + statuses.map( + (_status, index) => + `Bearer ${index > statuses.indexOf(401) && refreshStatus === 200 ? newToken : "oauth-old"}`, + ), + ); + if (refreshStatus === 503) { + const held = await fixture.host.harness.behavior.fetchHttp( + "POST", + route, + { + headers: authHeaders(fixture.key), + body: "{}", + }, + ); + await held.text(); + expect(held.status).toBe(503); + expect(authorizations).toHaveLength(1); + expect(refreshCalls).toBe(1); + now += 1_000; + oauthStatus = 200; + const recovered = await fixture.host.harness.behavior.fetchHttp( + "POST", + route, + { + headers: authHeaders(fixture.key), + body: "{}", + }, + ); + await recovered.text(); + expect(recovered.status).toBe(200); + expect(refreshCalls).toBe(2); + expect(authorizations.at(-1)).toBe(`Bearer ${newToken}`); + } + }, + ); + + it("joins an unchanged normal flight once and reuses replacement tokens after a late 401", async () => { + const now = 1_800_000_000_000; + const newToken = testJwt({ exp: now / 1_000 + 3600 }); + const oldResponses: Array<() => void> = []; + const authorizations: Array = []; + let refreshCalls = 0; + let releaseRefresh = () => {}; + const refreshReleased = new Promise((resolve) => { + releaseRefresh = resolve; + }); + const fixture = await createOAuthRequestFixture( + provider, + async (input, init) => { + if (String(input).endsWith("/oauth/token")) { + refreshCalls += 1; + await refreshReleased; + return Response.json({ access_token: newToken, expires_in: 3600 }); + } + const authorization = new Headers(init?.headers).get("authorization"); + authorizations.push(authorization); + if (authorization === "Bearer oauth-old") { + await new Promise((resolve) => { + oldResponses.push(resolve); + }); + return Response.json( + { error: { message: "expired access token" } }, + { status: 401 }, + ); + } + return Response.json({ result: "refreshed" }); + }, + () => now, + ); + const requests = [1, 2].map(() => + fixture.host.harness.behavior.fetchHttp("POST", route, { + headers: authHeaders(fixture.key), + body: "{}", + }), + ); + let releaseRead = () => {}; + const readReleased = new Promise((resolve) => { + releaseRead = resolve; + }); + const originalRead = AccountStore.prototype.readSecret; + const read = vi.spyOn(AccountStore.prototype, "readSecret"); + try { + await vi.waitFor(() => expect(oldResponses).toHaveLength(2)); + read.mockImplementation(async function (this: AccountStore, id) { + const secret = await originalRead.call(this, id); + await readReleased; + return secret; + }); + read.mockClear(); + requests.push( + fixture.host.harness.behavior.fetchHttp("POST", route, { + headers: authHeaders(fixture.key), + body: "{}", + }), + ); + await vi.waitFor(() => expect(read).toHaveBeenCalledTimes(1)); + oldResponses[0]?.(); + oldResponses[1]?.(); + await new Promise((resolve) => setImmediate(resolve)); + releaseRead(); + await vi.waitFor(() => { + expect(oldResponses).toHaveLength(3); + expect(refreshCalls).toBe(1); + }); + releaseRefresh(); + const first = await Promise.all(requests.slice(0, 2)); + expect(first.map((response) => response.status)).toEqual([200, 200]); + await Promise.all(first.map((response) => response.text())); + oldResponses[2]?.(); + const late = await requests[2]; + expect(late?.status).toBe(200); + await late?.text(); + expect(refreshCalls).toBe(1); + expect( + authorizations.filter((value) => value === `Bearer ${newToken}`), + ).toHaveLength(3); + const accounts = z + .array(accountSummarySchema) + .parse( + await fixture.host.harness.behavior.callRpc("account.list", null), + ); + expect(accounts[0]?.error).toBeNull(); + } finally { + releaseRead(); + releaseRefresh(); + for (const release of oldResponses) release(); + await Promise.allSettled( + requests.map(async (request) => { + const response = await request; + await response.text(); + }), + ); + read.mockRestore(); + } + }); + }); + + it.each(["terminal", "same-token cooldown"])( + "keeps late 401 recovery bounded after a %s refresh", + async (outcome) => { + const oldResponse = deferred(); + const refreshed = deferred(); + let now = 1_800_000_000_000; + let attempts = 0; + let refreshCalls = 0; + const fixture = await createOAuthRequestFixture( + "claude", + async (input) => { + if (String(input).endsWith("/oauth/token")) { + refreshCalls += 1; + if (refreshCalls === 1) + return Response.json( + { + error: + outcome === "terminal" + ? "invalid_grant" + : "temporarily_unavailable", + }, + { status: outcome === "terminal" ? 400 : 503 }, + ); + await refreshed.promise; + return Response.json({ + access_token: "oauth-old", + expires_in: 3600, + }); + } + attempts += 1; + const attempt = attempts; + if (attempt === 1) await oldResponse.promise; + return Response.json({}, { status: attempt <= 2 ? 401 : 200 }); + }, + () => now, + ); + const send = () => + fixture.host.harness.behavior.fetchHttp("POST", "/v1/messages", { + headers: authHeaders(fixture.key), + body: "{}", + }); + const requests = [send()]; + try { + await vi.waitFor(() => expect(attempts).toBe(1)); + const rejected = await send(); + expect(rejected.status).toBe(outcome === "terminal" ? 401 : 503); + await rejected.text(); + if (outcome === "same-token cooldown") { + now += 1_000; + requests.push(send()); + await vi.waitFor(() => expect(refreshCalls).toBe(2)); + } + oldResponse.resolve(); + await new Promise((resolve) => setImmediate(resolve)); + refreshed.resolve(); + const responses = await Promise.all(requests); + await Promise.all(responses.map((response) => response.text())); + expect(responses.map((response) => response.status)).toEqual( + outcome === "terminal" ? [401] : [200, 200], + ); + expect(refreshCalls).toBe(outcome === "terminal" ? 1 : 2); + expect(attempts).toBe(outcome === "terminal" ? 2 : 4); + } finally { + oldResponse.resolve(); + refreshed.resolve(); + await Promise.allSettled( + requests.map(async (request) => (await request).text()), + ); + } + }, + ); + + it.each([false, true])( + "separates rejection checks from credential flights when the reporting request is canceled: %s", + async (cancelReporter) => { + let attempts = 0; + const fixture = await createOAuthRequestFixture( + "claude", + async (input) => { + if (String(input).endsWith("/oauth/token")) + return Response.json({ + access_token: "oauth-new", + expires_in: 3600, + }); + attempts += 1; + return Response.json({}, { status: attempts <= 2 ? 401 : 200 }); + }, + () => 1_800_000_000_000, + ); + const gate = deferred(); + const checking = deferred(); + const originalRead = AccountStore.prototype.readSecret; + let reads = 0; + const read = vi + .spyOn(AccountStore.prototype, "readSecret") + .mockImplementation(async function (this: AccountStore, id) { + const secret = await originalRead.call(this, id); + reads += 1; + if (reads === 3) { + checking.resolve(); + await gate.promise; + } + return secret; + }); + const recordUsed = vi.spyOn(AccountStore.prototype, "recordUsed"); + const controller = new AbortController(); + const requests = [ + fixture.host.harness.behavior.fetchHttp("POST", "/v1/messages", { + headers: authHeaders(fixture.key), + body: "{}", + signal: controller.signal, + }), + ]; + try { + await checking.promise; + requests.push( + fixture.host.harness.behavior.fetchHttp("POST", "/v1/messages", { + headers: authHeaders(fixture.key), + body: "{}", + }), + ); + await vi.waitFor(() => expect(recordUsed).toHaveBeenCalledTimes(2)); + await Promise.all( + recordUsed.mock.results.map((result) => result.value), + ); + await new Promise((resolve) => setImmediate(resolve)); + if (cancelReporter) controller.abort(); + gate.resolve(); + const responses = await Promise.all(requests); + await Promise.all(responses.map((response) => response.text())); + expect(responses.map((response) => response.status)).toEqual( + cancelReporter ? [499, 200] : [401, 429], + ); + expect(attempts).toBe(cancelReporter ? 3 : 2); + } finally { + gate.resolve(); + await Promise.allSettled( + requests.map(async (request) => (await request).text()), + ); + read.mockRestore(); + recordUsed.mockRestore(); + } + }, + ); + + it("cancels one forced-refresh waiter without canceling shared recovery", async () => { + const gate = deferred(); + let refreshCalls = 0; + const authorizations: Array = []; + const fixture = await createOAuthRequestFixture( + "claude", + async (input, init) => { + if (String(input).endsWith("/oauth/token")) { + refreshCalls += 1; + await gate.promise; + return Response.json({ access_token: "oauth-new", expires_in: 3600 }); + } + const authorization = new Headers(init?.headers).get("authorization"); + authorizations.push(authorization); + return Response.json( + {}, + { status: authorization === "Bearer oauth-old" ? 401 : 200 }, + ); + }, + () => 1_800_000_000_000, + ); + const controller = new AbortController(); + const requests = [controller.signal, undefined].map((signal) => + fixture.host.harness.behavior.fetchHttp("POST", "/v1/messages", { + headers: authHeaders(fixture.key), + body: "{}", + signal, + }), + ); + try { + await vi.waitFor(() => { + expect(authorizations).toHaveLength(2); + expect(refreshCalls).toBe(1); + }); + controller.abort(); + const canceled = await requests[0]; + expect(canceled?.status).toBe(499); + await canceled?.text(); + gate.resolve(); + const recovered = await requests[1]; + expect(recovered?.status).toBe(200); + await recovered?.text(); + expect(refreshCalls).toBe(1); + expect(authorizations).toEqual([ + "Bearer oauth-old", + "Bearer oauth-old", + "Bearer oauth-new", + ]); + } finally { + gate.resolve(); + await Promise.allSettled( + requests.map(async (request) => (await request).text()), + ); + } + }); + + it("does not let a late second 401 poison a newer credential", async () => { + const gate = deferred(); + let refreshCalls = 0; + let newAttempts = 0; + const fixture = await createOAuthRequestFixture( + "claude", + async (input, init) => { + if (String(input).endsWith("/oauth/token")) { + refreshCalls += 1; + return Response.json({ + access_token: `oauth-new-${refreshCalls}`, + expires_in: 3600, + }); + } + const authorization = new Headers(init?.headers).get("authorization"); + if (authorization === "Bearer oauth-new-2") return Response.json({}); + if (authorization === "Bearer oauth-new-1") { + newAttempts += 1; + if (newAttempts === 1) await gate.promise; + } + return Response.json({ error: "rejected token" }, { status: 401 }); + }, + () => 1_800_000_000_000, + ); + const send = () => + fixture.host.harness.behavior.fetchHttp("POST", "/v1/messages", { + headers: authHeaders(fixture.key), + body: "{}", + }); + const first = send(); + try { + await vi.waitFor(() => expect(newAttempts).toBe(1)); + const second = await send(); + expect(second.status).toBe(200); + await second.text(); + gate.resolve(); + const late = await first; + expect(late.status).toBe(401); + await late.text(); + const accounts = z + .array(accountSummarySchema) + .parse( + await fixture.host.harness.behavior.callRpc("account.list", null), + ); + expect(accounts[0]?.error).toBeNull(); + const next = await send(); + expect(next.status).toBe(200); + await next.text(); + expect(refreshCalls).toBe(2); + } finally { + gate.resolve(); + const response = await first; + if (!response.bodyUsed) await response.text(); + } + }); + + it("honors a newer token's rejected cooldown when an older 401 arrives late", async () => { + const gate = deferred(); + let now = 1_800_000_000_000; + let refreshCalls = 0; + let attempts = 0; + const fixture = await createOAuthRequestFixture( + "claude", + async (input) => { + if (String(input).endsWith("/oauth/token")) { + refreshCalls += 1; + return refreshCalls === 2 + ? Response.json({}, { status: 503 }) + : Response.json({ + access_token: `oauth-new-${refreshCalls}`, + expires_in: 3600, + }); + } + attempts += 1; + if (attempts === 1) await gate.promise; + return Response.json( + {}, + { status: attempts === 3 || attempts >= 5 ? 200 : 401 }, + ); + }, + () => now, + ); + const send = () => + fixture.host.harness.behavior.fetchHttp("POST", "/v1/messages", { + headers: authHeaders(fixture.key), + body: "{}", + }); + const first = send(); + try { + await vi.waitFor(() => expect(attempts).toBe(1)); + const second = await send(); + expect(second.status).toBe(200); + await second.text(); + const rejected = await send(); + expect(rejected.status).toBe(503); + await rejected.text(); + gate.resolve(); + const late = await first; + expect(late.status).toBe(503); + await late.text(); + const held = await send(); + expect(held.status).toBe(503); + await held.text(); + expect(attempts).toBe(4); + now += 1_000; + const recovered = await send(); + expect(recovered.status).toBe(200); + await recovered.text(); + expect(refreshCalls).toBe(3); + } finally { + gate.resolve(); + const response = await first; + if (!response.bodyUsed) await response.text(); + } + }); + + it.each(["never ends", "cancel rejects", "cancel hangs"])( + "bounds failed response disposal when the body %s", + async (behavior) => { + const cancel = vi.fn(() => + behavior === "cancel hangs" + ? new Promise(() => {}) + : behavior === "cancel rejects" + ? Promise.reject(new Error("cancel failed")) + : Promise.resolve(), + ); + let attempts = 0; + const fixture: Fixture = await createFixture({ + upstreamUrl: "https://upstream.example", + priority: 0, + options: { + fetch: async (_input, init) => { + attempts += 1; + if (attempts === 1) + return new Response( + new ReadableStream({ + start(controller) { + if (behavior !== "never ends") + controller.enqueue(new Uint8Array(2048).fill(65)); + }, + cancel, + }), + { status: 503 }, + ); + const status = statusSchema.parse( + await fixture.host.harness.behavior.callRpc("status.get", null), + ); + expect( + status.accounts.find( + (account) => account.id === fixture.account.id, + )?.inFlight, + ).toBe(0); + expect(init?.signal?.aborted).toBe(false); + return Response.json({}); + }, + }, + }); + await addApiAccount(fixture, "sk-backup", 100); + const response = await fixture.host.harness.behavior.fetchHttp( + "POST", + "/v1/messages", + { + headers: authHeaders(fixture.key), + body: "{}", + }, + ); + expect(response.status).toBe(200); + await response.text(); + expect(attempts).toBe(2); + expect(cancel).toHaveBeenCalledOnce(); + }, + ); + + it("finishes a successful in-flight fetch during graceful hub shutdown", async () => { + const gate = deferred(); + const started = deferred(); + const fixture = await createFixture({ + upstreamUrl: "https://upstream.example", + options: { + fetch: async () => { + started.resolve(); + await gate.promise; + return Response.json({}); + }, + }, + }); + const request = fixture.host.harness.behavior.fetchHttp( + "POST", + "/v1/messages", + { + headers: authHeaders(fixture.key), + body: "{}", + }, + ); + await started.promise; + fixture.service.controller.abort(); + await vi.waitFor(async () => { + const status = statusSchema.parse( + await fixture.host.harness.behavior.callRpc("status.get", null), + ); + expect(status.accepting).toBe(false); + }); + gate.resolve(); + const response = await request; + expect(response.status).toBe(200); + await response.text(); + await fixture.service.done; + }); + + it("accepts requests after stopping and restarting the hub service", async () => { + const fixture = await createFixture({ + upstreamUrl: "https://upstream.example", + options: { fetch: async () => Response.json({}) }, + }); + fixture.service.controller.abort(); + await fixture.service.done; + const restarted = fixture.host.harness.behavior.runService("hub"); + cleanups.push(async () => { + restarted.controller.abort(); + await restarted.done; + }); + await vi.waitFor(async () => { + const status = statusSchema.parse( + await fixture.host.harness.behavior.callRpc("status.get", null), + ); + expect(status.accepting).toBe(true); + }); + const response = await fixture.host.harness.behavior.fetchHttp( + "POST", + "/v1/messages", + { + headers: authHeaders(fixture.key), + body: "{}", + }, + ); + expect(response.status).toBe(200); + await response.text(); + }); + + it.each([true, false])( + "uses only the initial account snapshot with a healthy fourth account: %s", + async (healthyFourth) => { + const attempts: Array = []; + const fixture: Fixture = await createFixture({ + upstreamUrl: "https://upstream.example", + apiKey: "sk-1", + priority: 1, + options: { + fetch: async (_input, init) => { + const key = new Headers(init?.headers).get("x-api-key"); + attempts.push(key); + if (attempts.length === 1) + await addApiAccount(fixture, "sk-late", -1); + return Response.json( + { result: key }, + { + status: + key === "sk-late" || (healthyFourth && key === "sk-4") + ? 200 + : 503, + }, + ); + }, + }, + }); + for (const account of [2, 3, 4]) + await addApiAccount(fixture, `sk-${account}`, account); + const response = await fixture.host.harness.behavior.fetchHttp( + "POST", + "/v1/messages", + { + headers: authHeaders(fixture.key), + body: "{}", + }, + ); + await response.text(); + expect(response.status).toBe(healthyFourth ? 200 : 503); + expect(attempts).toEqual(["sk-1", "sk-2", "sk-3", "sk-4"]); + }, + ); + + it("does not start another inference after cancellation during pacing", async () => { + const controller = new AbortController(); + let attempts = 0; + const fixture = await createFixture({ + upstreamUrl: "https://upstream.example", + options: { + fetch: async () => { + attempts += 1; + return attempts === 1 + ? new Response( + new ReadableStream({ + cancel() { + controller.abort(); + }, + }), + { status: 429, headers: { "retry-after": "0" } }, + ) + : Response.json({}); + }, + }, + }); + const response = await fixture.host.harness.behavior.fetchHttp( + "POST", + "/v1/messages", + { + headers: authHeaders(fixture.key), + body: "{}", + signal: controller.signal, + }, + ); + await response.text(); + expect(attempts).toBe(1); + }); + + it("never replays a committed SSE stream on another account", async () => { + let failStream = () => {}; + let attempts = 0; + const fixture = await createFixture({ + upstreamUrl: "https://upstream.example", + options: { + fetch: async () => { + attempts += 1; + return new Response( + new ReadableStream({ + start(controller) { + controller.enqueue( + new TextEncoder().encode("data: started\n\n"), + ); + failStream = () => + controller.error(new Error("stream disconnected")); + }, + }), + { headers: { "content-type": "text/event-stream" } }, + ); + }, + }, + }); + await addApiAccount(fixture, "sk-backup"); + const response = await fixture.host.harness.behavior.fetchHttp( + "POST", + "/v1/messages", + { + headers: authHeaders(fixture.key), + body: "{}", + }, + ); + const reader = response.body?.getReader(); + if (reader === undefined) throw new Error("Missing SSE stream."); + expect(new TextDecoder().decode((await reader.read()).value)).toBe( + "data: started\n\n", + ); + failStream(); + expect(new TextDecoder().decode((await reader.read()).value)).toContain( + "event: error", + ); + expect((await reader.read()).done).toBe(true); + expect(attempts).toBe(1); + }); + it("serializes refresh, writes new tokens with 0600 mode, and uses them", async () => { + let now = 1_800_000_000_000; let refreshCalls = 0; const authorizations: Array = []; const upstream = await startUpstream(async (request, response) => { if (request.url === "/oauth/token") { refreshCalls += 1; - await new Promise((resolve) => setTimeout(resolve, 25)); response.writeHead(200, { "content-type": "application/json" }); response.end( JSON.stringify({ @@ -2431,19 +3321,53 @@ describe("Account Pool plugin", () => { upstreamUrl: upstream.url, source: "import", options: { + now: () => now, importCredentials: async () => - importedCredentials({ expiresAt: Date.now() + 1_000 }), + importedCredentials({ expiresAt: now + 10 * 60 * 1_000 }), refreshUrl: `${upstream.url}/oauth/token`, }, }); + expect(refreshCalls).toBe(0); + now += 6 * 60 * 1_000; + let releaseSecretReads = () => {}; + const secretReadsReleased = new Promise((resolve) => { + releaseSecretReads = resolve; + }); + const readSecret = AccountStore.prototype.readSecret; + const pendingSecretReads: ReturnType[] = []; + const readSecretSpy = vi + .spyOn(AccountStore.prototype, "readSecret") + .mockImplementation(async function (this: AccountStore, accountId) { + const reading = readSecret.call(this, accountId); + pendingSecretReads.push(reading); + const secret = await reading; + await secretReadsReleased; + return secret; + }); + const recordUsed = vi.spyOn(AccountStore.prototype, "recordUsed"); const requests = [1, 2].map(() => fixture.host.harness.behavior.fetchHttp("POST", "/v1/messages", { headers: authHeaders(fixture.key), body: "{}", }), ); - const responses = await Promise.all(requests); - await Promise.all(responses.map((response) => response.text())); + try { + await vi.waitFor(() => { + expect(recordUsed).toHaveBeenCalledTimes(2); + }); + await Promise.all(recordUsed.mock.results.map((result) => result.value)); + await new Promise((resolve) => setImmediate(resolve)); + await Promise.all(pendingSecretReads); + releaseSecretReads(); + const responses = await Promise.all(requests); + expect(responses.map((response) => response.status)).toEqual([200, 200]); + await Promise.all(responses.map((response) => response.text())); + } finally { + releaseSecretReads(); + await Promise.allSettled(requests); + readSecretSpy.mockRestore(); + recordUsed.mockRestore(); + } expect(refreshCalls).toBe(1); expect(authorizations).toEqual(["Bearer oauth-new", "Bearer oauth-new"]); const secretPath = path.join( @@ -2465,6 +3389,295 @@ describe("Account Pool plugin", () => { expect((await fs.stat(secretPath)).mode & 0o777).toBe(0o600); }); + it("refreshes unrelated accounts independently", async () => { + let now = 1_800_000_000_000; + let releaseFirstRefresh = () => {}; + const firstRefreshReleased = new Promise((resolve) => { + releaseFirstRefresh = resolve; + }); + const refreshes: string[] = []; + const authorizations: Array = []; + const upstream = await startUpstream(async (request, response) => { + if (request.url === "/oauth/token") { + const { refresh_token } = z + .object({ refresh_token: z.string() }) + .parse(JSON.parse((await readRequestBody(request)).toString())); + refreshes.push(refresh_token); + if (refresh_token === "refresh-1") await firstRefreshReleased; + response.writeHead(200, { "content-type": "application/json" }); + response.end( + JSON.stringify({ + access_token: `new-${refresh_token}`, + refresh_token: `next-${refresh_token}`, + expires_in: 3600, + }), + ); + return; + } + authorizations.push(request.headers.authorization); + await readRequestBody(request); + response.writeHead(200, { "content-type": "application/json" }); + response.end("{}"); + }); + cleanups.push(upstream.close); + let imported = 0; + const fixture = await createFixture({ + upstreamUrl: upstream.url, + source: "import", + options: { + now: () => now, + importCredentials: async () => { + imported += 1; + return importedCredentials({ + accessToken: `access-${imported}`, + refreshToken: `refresh-${imported}`, + email: `account-${imported}@example.com`, + expiresAt: now + 10 * 60 * 1_000, + }); + }, + refreshUrl: `${upstream.url}/oauth/token`, + }, + }); + const second = accountSchema.parse( + await fixture.host.harness.behavior.callRpc("account.add", { + provider: "claude", + source: { kind: "import" }, + label: "second", + priority: 200, + }), + ); + expect(refreshes).toEqual([]); + now += 6 * 60 * 1_000; + const requests: Promise[] = []; + try { + requests.push( + fixture.host.harness.behavior.fetchHttp("POST", "/v1/messages", { + headers: authHeaders(fixture.key), + body: "{}", + }), + ); + await vi.waitFor(() => { + expect(refreshes).toEqual(["refresh-1"]); + }); + await fixture.host.harness.behavior.callRpc("account.setPriority", { + accountId: second.id, + priority: 0, + }); + requests.push( + fixture.host.harness.behavior.fetchHttp("POST", "/v1/messages", { + headers: authHeaders(fixture.key), + body: "{}", + }), + ); + await vi.waitFor(() => { + expect(authorizations).toEqual(["Bearer new-refresh-2"]); + }); + } finally { + releaseFirstRefresh(); + await Promise.allSettled(requests); + } + const responses = await Promise.all(requests); + expect(responses.map((response) => response.status)).toEqual([200, 200]); + await Promise.all(responses.map((response) => response.text())); + expect(refreshes).toEqual(["refresh-1", "refresh-2"]); + expect(authorizations).toEqual([ + "Bearer new-refresh-2", + "Bearer new-refresh-1", + ]); + }); + + 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") {