diff --git a/src/oauth/index.ts b/src/oauth/index.ts index 29acaf210dd..257bb94055d 100644 --- a/src/oauth/index.ts +++ b/src/oauth/index.ts @@ -4,6 +4,7 @@ import type { OcxConfig, OcxProviderConfig, RefreshPolicy } from "../types"; import { getDefaultConfig, initializePersistedConfigIfMissing, loadConfig, mutatePersistedConfig, observeInitialConfigState, saveConfig } from "../config"; import { resolveProviderApiKey } from "../providers/key-store"; import { projectEmail } from "../lib/privacy"; +import { createAdmissionGate } from "../lib/admission"; import { KiroTokenRefreshError, environmentKiroRoutingMetadata, loginKiro, refreshKiroToken, settleKiroLoginTransaction } from "./kiro"; import { KiroUserIdentityUnavailableError, kiroReauthIdentityMatches } from "./kiro-account-identity"; import { @@ -131,6 +132,9 @@ export type OAuthActiveTokenObservation = | { readonly kind: "unsupported" }; const MAX_OAUTH_TOKEN_REFRESH_FLIGHTS = 32; +/** Independent of client lifetime; a cancelled caller can still be awaiting a shared refresh. */ +export const MAX_OAUTH_TOKEN_REFRESH_WAITERS = 256; +const tokenRefreshWaiters = createAdmissionGate("oauth_token_refresh_waiters", MAX_OAUTH_TOKEN_REFRESH_WAITERS); const OAUTH_TOKEN_REFRESH_FLIGHT_STALE_MS = 120_000; interface OAuthRefreshFlightEvidence { flightId: string; dispatched: boolean } interface OAuthTokenRefreshFlight extends OAuthRefreshFlightEvidence { promise: Promise; startedAt: number; abort: AbortController } @@ -608,45 +612,52 @@ async function resolveAccessSnapshotForAccount( const key = `${provider}\u0000${accountId}`; let existing = tokenRefreshes.get(key); let replacedStaleFlight: OAuthRefreshFlightEvidence | undefined; - if (existing && Date.now() - existing.startedAt <= OAUTH_TOKEN_REFRESH_FLIGHT_STALE_MS) { - logOAuthEvent("OAuth refresh joined existing operation", { provider, accountId }); - return existing.promise; - } - if (existing) { + if (existing && Date.now() - existing.startedAt > OAUTH_TOKEN_REFRESH_FLIGHT_STALE_MS) { replacedStaleFlight = { flightId: existing.flightId, dispatched: existing.dispatched }; existing.abort.abort(new OAuthTokenRefreshStaleError()); if (tokenRefreshes.get(key) === existing) tokenRefreshes.delete(key); existing = undefined; } - if (tokenRefreshes.size >= MAX_OAUTH_TOKEN_REFRESH_FLIGHTS) throw new OAuthTokenRefreshBusyError(); - - const abort = new AbortController(); - const flight: OAuthTokenRefreshFlight = { - promise: undefined as unknown as Promise, - startedAt: Date.now(), - abort, - flightId: randomUUID(), - dispatched: false, - }; - const refresh = (async (): Promise => { - const accessToken = await refreshAndPersistAccessToken(provider, accountId, def, cred, abort.signal, flight, replacedStaleFlight); - const persistedRow = getAccountCredentialWithStatus(provider, accountId); - if (!persistedRow) throw new OAuthLoginRequiredError(provider); - if (persistedRow.paused === true) throw new OAuthAccountPausedError(); - const persisted = persistedRow.credential; - if (persisted.access !== accessToken) { - throw new Error(`OAuth refresh persisted an unexpected access token for ${provider}`); + // Count every consumer, including same-flight joins, until its promise settles. Cancelling + // a client cannot recycle this permit. Stale-flight abort above still runs under pressure; + // no replacement starts until the old consumers have actually released enough capacity. + const waiter = tokenRefreshWaiters.tryAcquire(); + if (!waiter) throw new OAuthTokenRefreshBusyError(); + try { + if (existing) { + logOAuthEvent("OAuth refresh joined existing operation", { provider, accountId }); + return await existing.promise; } - return accessSnapshot(provider, accountId, persisted, oauthProvider); - })().catch(error => { - if (abort.signal.reason instanceof OAuthTokenRefreshStaleError) throw abort.signal.reason; - throw error; - }).finally(() => { - if (tokenRefreshes.get(key) === flight) tokenRefreshes.delete(key); - }); - flight.promise = refresh; - tokenRefreshes.set(key, flight); - return refresh; + if (tokenRefreshes.size >= MAX_OAUTH_TOKEN_REFRESH_FLIGHTS) throw new OAuthTokenRefreshBusyError(); + + const abort = new AbortController(); + const flight: OAuthTokenRefreshFlight = { + promise: undefined as unknown as Promise, + startedAt: Date.now(), + abort, + flightId: randomUUID(), + dispatched: false, + }; + const refresh = (async (): Promise => { + const accessToken = await refreshAndPersistAccessToken(provider, accountId, def, cred, abort.signal, flight, replacedStaleFlight); + const persistedRow = getAccountCredentialWithStatus(provider, accountId); + if (!persistedRow) throw new OAuthLoginRequiredError(provider); + if (persistedRow.paused === true) throw new OAuthAccountPausedError(); + const persisted = persistedRow.credential; + if (persisted.access !== accessToken) { + throw new Error(`OAuth refresh persisted an unexpected access token for ${provider}`); + } + return accessSnapshot(provider, accountId, persisted, oauthProvider); + })().catch(error => { + if (abort.signal.reason instanceof OAuthTokenRefreshStaleError) throw abort.signal.reason; + throw error; + }).finally(() => { + if (tokenRefreshes.get(key) === flight) tokenRefreshes.delete(key); + }); + flight.promise = refresh; + tokenRefreshes.set(key, flight); + return await refresh; + } finally { waiter.release(); } } export async function getValidAccessTokenSnapshot( diff --git a/src/server/responses/request-transport.ts b/src/server/responses/request-transport.ts index 1b6cb6f4c44..18a754a7e4d 100644 --- a/src/server/responses/request-transport.ts +++ b/src/server/responses/request-transport.ts @@ -31,6 +31,7 @@ import { } from "../../oauth/anthropic-routing"; import { OAuthAccountPausedError, + OAuthTokenRefreshBusyError, getValidAccessSnapshotForAccount, forceRefreshOAuthAccessSnapshot, getValidAccessTokenSnapshot, @@ -964,6 +965,10 @@ export async function prepareResponsesTransport( `${redactSecretString(err.message)}. Remove or reconfigure provider '${safeProviderName}' in the OpenCodex configuration.`, ); } + if (err instanceof OAuthTokenRefreshBusyError) { + return (options.abortSignal ?? req.signal).aborted ? clientCancelledResponse() + : formatErrorResponse(503, "server_error", err.message, { retryAfter: "1" }); + } if (err instanceof OAuthAccountPausedError) { return formatErrorResponse(403, "permission_error", publicOAuthAuthenticationErrorMessage(err)); } diff --git a/structure/runtime.md b/structure/runtime.md index 05b24c85b6c..fbf4deb1457 100644 --- a/structure/runtime.md +++ b/structure/runtime.md @@ -30,7 +30,7 @@ Responses admission and finalization are composed through the [core module owner Catalog HTTP acquisition follows the [proxy-routing contract](catalog.md#remote-catalog-http-proxy-routing). CLI model lists expose shared catalog estimates as `price`; `models price` retains the saved override in `cost` and adds `effectiveCost`. Explicit zero overrides win; unknown prices remain null, and automatic defaults are never persisted to `modelCosts`. Covered by `tests/cli/cli-models.test.ts` and `tests/cli/cli-models-price.test.ts`. Ordinary model discovery, startup sync, and explicit sync apply the [new-arrival catalog contract](catalog.md#shared-catalog) before publishing rows. The scoped persisted decision also updates the live config merge baseline, so subsequent visibility writes preserve manual re-enablement. A file-loaded inventory that already differs from disk can publish HTTP rows through a detached policy projection without adding unsaved live disables; retained catalog writes still refuse stale inventory, and its load provenance follows the [config contract](config.md). -OAuth refresh coordination follows the [refresh-lock identity contract](catalog.md#accounts-namespaces-and-pool-rotation): a fresh unreadable lock remains held, and release requires matching descriptor identity. A failed path-identity probe preserves the refresh callback outcome. Cooperating lock metadata changes serialize through the existing SQLite mutation transaction; release keeps the descriptor open through identity comparison and any unlink, then closes it. Failed metadata writes remove only a matching owned path after successful coordination; unknown identity, failed probes or unavailable coordination retain the path for stale recovery. Async refresh work holds no metadata transaction. +OAuth refresh resolution in `src/oauth/index.ts` admits at most 256 pending consumers across the process, including callers joining an existing refresh. The separate 32-flight ceiling remains. Consumers retain their admission until the shared promise settles, even if their client already cancelled; fresh credential snapshots do not take this admission. Exhaustion raises the existing retryable, account-free refresh-capacity error, and initial Responses preparation answers 503 with a short Retry-After. This bounds pending resolution count, not request bytes, refresh duration, or arbitrary cancelled work. The release-on-cancel turn contract is unchanged. `tests/oauth/oauth-refresh-generic-lock.test.ts` covers cancelled consumers, rejection cleanup and cached credentials under pressure. OAuth refresh coordination follows the [refresh-lock identity contract](catalog.md#accounts-namespaces-and-pool-rotation): a fresh unreadable lock remains held, and release requires matching descriptor identity. A failed path-identity probe preserves the refresh callback outcome. Cooperating lock metadata changes serialize through the existing SQLite mutation transaction; release keeps the descriptor open through identity comparison and any unlink, then closes it. Failed metadata writes remove only a matching owned path after successful coordination; unknown identity, failed probes or unavailable coordination retain the path for stale recovery. Async refresh work holds no metadata transaction. The configuration-only [plaintext V2 contract](subagents.md#plaintext-v2-agent-messages) is scoped to canonical ChatGPT Responses forwarding; other source-area behavior described here is unchanged. Cursor's localized native-shell names follow the [routing-commentary guard contract](providers/cursor.md#cursor-native-exec). diff --git a/tests/oauth/oauth-refresh-generic-lock.test.ts b/tests/oauth/oauth-refresh-generic-lock.test.ts index 63cd499e03a..c857b4a6cfa 100644 --- a/tests/oauth/oauth-refresh-generic-lock.test.ts +++ b/tests/oauth/oauth-refresh-generic-lock.test.ts @@ -4,6 +4,10 @@ import { tmpdir } from "node:os"; import { join } from "node:path"; import { getValidAccessTokenForAccount, + getValidAccessTokenSnapshot, + getValidAccessSnapshotForAccount, + MAX_OAUTH_TOKEN_REFRESH_WAITERS, + OAuthTokenRefreshBusyError, refreshGenericAccountWithLock, OAuthLoginRequiredError, OAuthTokenRefreshStaleError, @@ -13,6 +17,9 @@ import type { OAuthCredentials } from "../../src/oauth/types"; import { ChatGptTokenError, ChatGptTokenRequestError, refreshChatGPTToken } from "../../src/oauth/chatgpt"; import { getAccountCredential, getAccountSet, saveCredential } from "../../src/oauth/store"; import { removeTreeWithRetry } from "../helpers/remove-tree"; +import { getActiveTurnCount, tryAdmitTurn } from "../../src/server/lifecycle"; +import { handleResponses } from "../../src/server/responses"; +import type { OcxConfig } from "../../src/types"; // Gate/CAS refresh races can exceed the 5s default under windows-latest contention // (same flake class as kiro-oauth / oauth queue budgets). @@ -403,3 +410,128 @@ describe("ChatGPT refresh through the registered provider", () => { }); }); + +/** Real resolver calls retain refresh admission after their client turn has cancelled. */ +describe("OAuth refresh consumer admission", () => { + test("cancelled callers cannot grow one pending refresh beyond the consumer limit", async () => { + const accountId = await seedExpiredKimi(); + let finish!: () => void; + let started!: () => void; + const entered = new Promise(resolve => { started = resolve; }); + const blocked = new Promise(resolve => { finish = resolve; }); + const tracker = stubKimiRefresh(async () => { + started(); + await blocked; + return { access: "fixture-refreshed", refresh: "fixture-rotated", expires: Date.now() + 3_600_000 }; + }); + const waiting: Promise[] = []; + const rejected: unknown[] = []; + const activeBefore = getActiveTurnCount(); + try { + for (let i = 0; i < MAX_OAUTH_TOKEN_REFRESH_WAITERS + 44; i++) { + const controller = new AbortController(); + const lease = tryAdmitTurn(undefined, controller.signal); + expect(lease).not.toBeNull(); + const promise = i % 2 === 0 ? getValidAccessTokenSnapshot("kimi") + : getValidAccessSnapshotForAccount("kimi", accountId); + waiting.push(promise.catch(error => { rejected.push(error); })); + controller.abort(new Error("fixture disconnected")); + lease?.release(); + } + await entered; + await Promise.resolve(); + await Promise.resolve(); + expect(getActiveTurnCount()).toBe(activeBefore); + expect(rejected).toHaveLength(44); + expect(rejected.every(error => error instanceof OAuthTokenRefreshBusyError)).toBe(true); + expect(tracker.calls()).toBe(1); + expect(getAccountCredential("kimi", accountId)?.access).toBe("kimi-old"); + expect(getAccountSet("kimi")?.accounts.find(row => row.id === accountId)?.needsReauth).not.toBe(true); + } finally { + finish(); + await Promise.all(waiting); + } + await expect(getValidAccessTokenForAccount("kimi", accountId)).resolves.toBe("fixture-refreshed"); + expect(tracker.calls()).toBe(1); + }); + + test("rejected shared refresh releases all consumers and permits one later retry", async () => { + const accountId = await seedExpiredKimi(); + let fail!: () => void; + let started!: () => void; + const entered = new Promise(resolve => { started = resolve; }); + const blocked = new Promise((_resolve, reject) => { fail = () => reject(new Error("fixture transient failure")); }); + stubKimiRefresh(async () => { started(); await blocked; throw new Error("unreachable"); }); + const waiting = Array.from({ length: MAX_OAUTH_TOKEN_REFRESH_WAITERS }, () => + getValidAccessTokenForAccount("kimi", accountId).catch(error => error)); + try { + await entered; + // Do not await an over-budget control without a deadline: the unpatched resolver waits. + let refused: unknown; + const overflow = getValidAccessTokenForAccount("kimi", accountId).catch(error => { refused = error; }); + waiting.push(overflow); + await Promise.resolve(); await Promise.resolve(); await Promise.resolve(); + expect(refused).toBeInstanceOf(OAuthTokenRefreshBusyError); + } finally { fail(); await Promise.all(waiting); } + const tracker = stubKimiRefresh(async () => ({ + access: "fixture-retry", refresh: "fixture-next", expires: Date.now() + 3_600_000, + })); + await expect(getValidAccessTokenForAccount("kimi", accountId)).resolves.toBe("fixture-retry"); + expect(tracker.calls()).toBe(1); + expect(getAccountSet("kimi")?.accounts.find(row => row.id === accountId)?.needsReauth).not.toBe(true); + }); + + test("fresh snapshots remain usable while refresh consumers are full", async () => { + const expiredId = await seedExpiredKimi(); + await saveCredential("kimi", { access: "fixture-cached", refresh: "fixture-cached-refresh", + expires: Date.now() + 3_600_000, accountId: "fixture-cached-account" }); + const freshId = getAccountSet("kimi")!.activeAccountId; + let finish!: () => void; + let started!: () => void; + const entered = new Promise(resolve => { started = resolve; }); + const blocked = new Promise(resolve => { finish = resolve; }); + const tracker = stubKimiRefresh(async () => { started(); await blocked; return { + access: "fixture-finished", refresh: "fixture-finished-refresh", expires: Date.now() + 3_600_000, + }; }); + const waiting = Array.from({ length: MAX_OAUTH_TOKEN_REFRESH_WAITERS }, () => + getValidAccessTokenForAccount("kimi", expiredId).catch(error => error)); + try { + await entered; + await expect(getValidAccessTokenForAccount("kimi", freshId)).resolves.toBe("fixture-cached"); + expect(tracker.calls()).toBe(1); + } finally { finish(); await Promise.all(waiting); } + }); +}); + + +test("Responses maps exhausted OAuth refresh consumers to retryable capacity, not login failure", async () => { + const accountId = await seedExpiredKimi(); + let finish!: () => void; + let started!: () => void; + const entered = new Promise(resolve => { started = resolve; }); + const blocked = new Promise(resolve => { finish = resolve; }); + stubKimiRefresh(async () => { started(); await blocked; return { + access: "fixture-finished", refresh: "fixture-rotated", expires: Date.now() + 3_600_000, + }; }); + const waiting = Array.from({ length: MAX_OAUTH_TOKEN_REFRESH_WAITERS }, () => + getValidAccessTokenForAccount("kimi", accountId).catch(error => error)); + let inferenceSends = 0; + const config = { port: 0, defaultProvider: "kimi", providers: { kimi: { + adapter: "openai-chat", authMode: "oauth", baseUrl: "https://fixture.invalid/v1", models: ["fixture-model"], + fetch: async () => { inferenceSends++; throw new Error("unexpected fixture send"); }, + } } } as OcxConfig; + try { + await entered; + const response = await handleResponses(new Request("http://localhost/v1/responses", { + method: "POST", headers: { "content-type": "application/json" }, + body: JSON.stringify({ model: "kimi/fixture-model", input: "fixture", stream: false }), + }), config, { model: "", provider: "" }); + expect(response.status).toBe(503); + expect(response.headers.get("retry-after")).toBe("1"); + const body = await response.text(); + expect(body).toContain("OAuth token refresh capacity reached"); + expect(body).not.toContain("kimi-old"); + expect(inferenceSends).toBe(0); + expect(getAccountSet("kimi")?.accounts.find(row => row.id === accountId)?.needsReauth).not.toBe(true); + } finally { finish(); await Promise.all(waiting); } +});