diff --git a/changelog.d/maintenance/decomp-token-refresh.md b/changelog.d/maintenance/decomp-token-refresh.md new file mode 100644 index 0000000000..f9bc03062b --- /dev/null +++ b/changelog.d/maintenance/decomp-token-refresh.md @@ -0,0 +1 @@ +- chore(token-refresh): decompose `open-sse/services/tokenRefresh.ts` (999 → 724 lines) by extracting the rotation-map, CAS guard and circuit-breaker refresh logic into `tokenRefresh/*` leaves — behavior-preserving move; `tokenRefresh.ts` still re-exports the moved symbols so the public surface is unchanged diff --git a/open-sse/services/tokenRefresh.ts b/open-sse/services/tokenRefresh.ts index 70535d98dc..13c9e23677 100755 --- a/open-sse/services/tokenRefresh.ts +++ b/open-sse/services/tokenRefresh.ts @@ -5,17 +5,39 @@ // file keeps the orchestrator (refreshAccessToken, getAccessToken), the // in-flight/rotation dedup maps, the CAS guard, and refreshWithRetry — the // cross-provider plumbing. The provider-module split was originally proposed -// by KooshaPari in PR #7338 (base was too old to merge as-is); redone here on -// the current tip, credit preserved via co-authorship on the extraction -// commits. All previously-public exports are re-exported below so existing +// by KooshaPari in PR #7338, whose base was too old to merge as-is; this is an +// independent implementation of the same idea against the current tip, not a +// reuse of that diff. All previously-public exports are re-exported below so existing // importers (open-sse/index.ts, executors, src/sse/services/tokenRefresh.ts, // tests) are unaffected. import { AsyncLocalStorage } from "node:async_hooks"; -import { pbkdf2Sync } from "node:crypto"; import { PROVIDERS } from "../config/constants.ts"; import { runWithProxyContext } from "../utils/proxyFetch.ts"; -import { serializeRefresh, wasRefreshTokenRotated } from "./refreshSerializer.ts"; -import { extractOAuthErrorCode, type RefreshLogger } from "./tokenRefresh/shared.ts"; +import { serializeRefresh } from "./refreshSerializer.ts"; +import { + extractOAuthErrorCode, + isUnrecoverableRefreshError, + type RefreshLogger, +} from "./tokenRefresh/shared.ts"; +import { + getRefreshCacheKey, + lookupRotation, + recordRotation, + _getTokenRotationMapStats, + _clearTokenRotationMap, +} from "./tokenRefresh/rotationMap.ts"; +import { + runWithCasGuard, + getActiveCasGuard, + getCasGuardStats, + _resetCasGuardStats, + casGuardShouldSkipPersist, +} from "./tokenRefresh/casGuard.ts"; +import { + isProviderBlocked, + getCircuitBreakerStatus, + refreshWithRetry, +} from "./tokenRefresh/circuitBreaker.ts"; import { refreshWindsurfToken } from "./tokenRefresh/providers/windsurf.ts"; import { refreshCodebuddyCnToken } from "./tokenRefresh/providers/codebuddyCn.ts"; import { refreshClineToken } from "./tokenRefresh/providers/cline.ts"; @@ -43,6 +65,16 @@ export { refreshGitHubToken, refreshCopilotToken, extractOAuthErrorCode, + isUnrecoverableRefreshError, + isProviderBlocked, + getCircuitBreakerStatus, + refreshWithRetry, + runWithCasGuard, + getActiveCasGuard, + getCasGuardStats, + _resetCasGuardStats, + _getTokenRotationMapStats, + _clearTokenRotationMap, }; // Default token expiry buffer (refresh if expires within 5 minutes). @@ -105,8 +137,6 @@ export function getRefreshLeadMs( return REFRESH_LEAD_MS[provider] ?? TOKEN_EXPIRY_BUFFER_MS; } -const CACHE_SECRET = "omniroute-token-cache"; - // In-flight refresh promise cache to prevent race conditions // Key: "provider:sha256(refreshToken)" → Value: Promise const refreshPromiseCache = new Map(); @@ -116,75 +146,9 @@ const refreshPromiseCache = new Map(); // Primary dedup when credentials.connectionId is present; refreshPromiseCache is fallback. const connectionRefreshMutex = new Map(); -// ─── Token Rotation Map (codex-multi-auth pattern) ───────────────────────── -// -// When a rotating-token provider (Codex, Kimi, GitLab Duo, etc.) refreshes, -// the old refresh_token is consumed and a new one is issued. Any subsequent -// caller arriving with the OLD token would, without protection, hit upstream -// and trigger "refresh_token_reused" — which Auth0 treats as a security event -// and invalidates the entire token family. -// -// This in-memory map caches RECENT rotations so a stale caller can be redirected -// to the new tokens WITHOUT touching upstream. The DB staleness check inside -// the per-connection mutex covers the same scenario when connectionId is known, -// but not all callers pass connectionId (e.g., legacy code paths, retries that -// snapshot credentials before the rotation lands in DB). -// -// Ported from ndycode/codex-multi-auth (lib/refresh-queue.ts:218-248), the only -// publicly known tool that reliably sustains multiple Codex OAuth accounts. -// -// Key format: `provider:sha256(oldRefreshToken)` -// Value: { result: tokens, expiresAt: ms_since_epoch } -type RotationEntry = { - result: { accessToken: string; refreshToken: string; expiresIn?: number; expiresAt?: string }; - expiresAt: number; -}; -const tokenRotationMap = new Map(); -const ROTATION_MAP_TTL_MS = 60 * 1000; // 60 seconds — long enough to catch in-flight stale callers - -function cleanupRotationMap(now: number = Date.now()): void { - if (tokenRotationMap.size === 0) return; - for (const [key, entry] of tokenRotationMap.entries()) { - if (entry.expiresAt <= now) tokenRotationMap.delete(key); - } -} - -function lookupRotation(provider: string, refreshToken: string): RotationEntry | undefined { - cleanupRotationMap(); - const key = getRefreshCacheKey(provider, refreshToken); - const entry = tokenRotationMap.get(key); - if (!entry) return undefined; - if (entry.expiresAt <= Date.now()) { - tokenRotationMap.delete(key); - return undefined; - } - return entry; -} - -function recordRotation( - provider: string, - oldRefreshToken: string, - result: { accessToken: string; refreshToken: string; expiresIn?: number; expiresAt?: string } -): void { - if (!oldRefreshToken || !result.refreshToken || oldRefreshToken === result.refreshToken) { - return; - } - const key = getRefreshCacheKey(provider, oldRefreshToken); - tokenRotationMap.set(key, { - result, - expiresAt: Date.now() + ROTATION_MAP_TTL_MS, - }); -} - -// Exported for tests + diagnostics; not part of the public API surface. -export function _getTokenRotationMapStats(): { size: number; entries: number } { - cleanupRotationMap(); - return { size: tokenRotationMap.size, entries: tokenRotationMap.size }; -} - -export function _clearTokenRotationMap(): void { - tokenRotationMap.clear(); -} +// Token Rotation Map (codex-multi-auth pattern) lives in +// ./tokenRefresh/rotationMap.ts — see that leaf for the in-memory rotation +// cache + getRefreshCacheKey. Imported above and re-exported for tests. // AsyncLocalStorage for plumbing `onPersist` through executor.refreshCredentials // without modifying every executor's signature. The chatCore.ts / base.ts call @@ -208,88 +172,13 @@ export function getActiveOnPersist(): RefreshPersistFn | undefined { return onPersistStore.getStore(); } -// ── #4038: compare-and-swap (CAS) guard on the refresh persist ─────────────── -// Fix A makes [network refresh + DB write] atomic *for a single connection's -// mutex*. It does NOT protect against a THIRD writer (a sibling process, a -// concurrent HealthCheck, or a replica) landing a fresher rotation on the same -// `connection_id` between the moment the caller read the row and the moment this -// persist runs. Overwriting that fresher row reverts the sibling's rotation, the -// next caller loads the reverted (now-consumed) refresh_token, and Auth0/Anthropic -// revoke the whole token family (the 1352× claude/aa5dd5cf invalidation storm). -// -// The CAS guard carries the refresh_token the caller PRESENTED (the version token, -// since refresh_tokens rotate on every refresh) plus a `reread` of the row's -// current refresh_token. Right before persisting, `getAccessToken` re-reads and, if -// a concurrent writer already rotated the row past the presented token, SKIPS the -// persist so the DB stays at the fresher state. The caller still receives the new -// accessToken — upstream already authenticated the request; only the DB write is -// skipped. No active guard ⇒ behavior is byte-identical to before (opt-in). -type CasGuard = { - /** The refresh_token the caller presented for this refresh (CAS version token). */ - expectedRefreshToken: string | null; - /** Re-reads the CURRENT persisted refresh_token for this connection (decrypted). */ - reread: () => Promise; -}; -const casGuardStore = new AsyncLocalStorage(); -const casGuardStats = { skipped: 0, persisted: 0 }; +// #4038 compare-and-swap (CAS) guard on the refresh persist lives in +// ./tokenRefresh/casGuard.ts — imported above and re-exported for tests. +// casGuardShouldSkipPersist is imported and used by getAccessToken below. -export function runWithCasGuard( - guard: CasGuard | undefined | null, - fn: () => Promise -): Promise { - if (!guard) return fn(); - return casGuardStore.run(guard, fn); -} - -export function getActiveCasGuard(): CasGuard | undefined { - return casGuardStore.getStore(); -} - -/** Skip/persist counters for observability + tests. */ -export function getCasGuardStats(): { skipped: number; persisted: number } { - return { ...casGuardStats }; -} - -/** Test-only: reset the CAS counters between cases. */ -export function _resetCasGuardStats(): void { - casGuardStats.skipped = 0; - casGuardStats.persisted = 0; -} - -/** - * Returns true when the persist should be SKIPPED because a concurrent writer - * already rotated the row's refresh_token past the one we presented (CAS mismatch). - * Best-effort: any reread failure falls through to persist (never blocks recovery). - */ -async function casGuardShouldSkipPersist(log?: RefreshLogger): Promise { - const guard = getActiveCasGuard(); - if (!guard || !guard.expectedRefreshToken) return false; - let current: string | null | undefined; - try { - current = await guard.reread(); - } catch { - return false; // reread failed — fall through to persist (best-effort) - } - // wasRefreshTokenRotated is true iff both are non-empty AND current !== expected. - if (wasRefreshTokenRotated(guard.expectedRefreshToken, current)) { - casGuardStats.skipped++; - log?.warn?.( - "TOKEN_REFRESH", - "CAS guard: skipping persist — a concurrent writer already rotated the refresh_token (#4038)" - ); - return true; - } - casGuardStats.persisted++; - return false; -} - -function getRefreshCacheKey(provider, refreshToken) { - const tokenHash = pbkdf2Sync(refreshToken, CACHE_SECRET, 1000, 32, "sha256").toString("hex"); - return `${provider}:${tokenHash}`; -} - -// extractOAuthErrorCode lives in ./tokenRefresh/shared.ts (imported above, re-exported below) — -// used both by the generic orchestrator below and by every per-provider refresh module. +// extractOAuthErrorCode + isUnrecoverableRefreshError live in +// ./tokenRefresh/shared.ts (imported above, re-exported below) — used both by +// the generic orchestrator below and by every per-provider refresh module. /** * Refresh OAuth access token using refresh token @@ -485,21 +374,9 @@ export function supportsTokenRefresh(provider) { return !!(config?.refreshUrl || config?.tokenUrl); } -/** - * Check if a refresh result indicates an unrecoverable error - * (e.g. the refresh token was already consumed and cannot be reused). - * Callers should stop retrying and request re-authentication. - */ -export function isUnrecoverableRefreshError(result) { - return ( - result && - typeof result === "object" && - (result.error === "unrecoverable_refresh_error" || - result.error === "refresh_token_reused" || - result.error === "invalid_request" || - result.error === "invalid_grant") - ); -} +// isUnrecoverableRefreshError lives in ./tokenRefresh/shared.ts (imported above +// and re-exported) — used by refreshWithRetry (./tokenRefresh/circuitBreaker.ts) +// and by callers that need to classify a refresh result. /** * Get access token for a specific provider (with deduplication). @@ -829,51 +706,10 @@ export async function getAllAccessTokens(userInfo, log) { return results; } -/** - * Refresh token with retry and exponential backoff - * Retries on failure with increasing delay: 1s, 2s, 3s... - * - * Includes: - * - Per-provider circuit breaker (5 consecutive failures → 30min pause) - * - 30s timeout per refresh attempt to prevent hanging connections - * - * @param {function} refreshFn - Async function that returns token or null - * @param {number} maxRetries - Max retry attempts (default 3) - * @param {object} log - Logger instance (optional) - * @param {string} provider - Provider ID for circuit breaker tracking (optional) - * @returns {Promise} Token result or null if all retries fail - */ - -// ─── Circuit Breaker State ────────────────────────────────────────────────── -const _circuitBreaker: Record = {}; -const CIRCUIT_BREAKER_THRESHOLD = 5; // consecutive failures before tripping -const CIRCUIT_BREAKER_COOLDOWN = 30 * 60 * 1000; // 30 minutes -const REFRESH_TIMEOUT_MS = 30_000; // 30s max per refresh attempt - -interface CircuitBreakerStatusEntry { - failures: number; - blocked: boolean; - blockedUntil: string | null; - remainingMs: number; -} - -interface RefreshLoggerLike { - error?: (scope: string, message: string) => void; - warn?: (scope: string, message: string) => void; -} - -/** - * Check if a provider is circuit-breaker blocked. - */ -export function isProviderBlocked(provider: string): boolean { - const state = _circuitBreaker[provider]; - if (!state) return false; - if (!state.blockedUntil) return false; - if (state.blockedUntil > Date.now()) return true; - // Cooldown expired — reset - delete _circuitBreaker[provider]; - return false; -} +// Per-provider circuit breaker + refreshWithRetry + withTimeout live in +// ./tokenRefresh/circuitBreaker.ts — imported above and re-exported for tests. +// isProviderBlocked / getCircuitBreakerStatus / refreshWithRetry are +// re-exported from that leaf. /** * Get active per-connection mutex entries (for diagnostics/metrics). @@ -886,114 +722,3 @@ export function getConnectionRefreshMutexStatus(): Record { - const result: Record = {}; - for (const [provider, state] of Object.entries(_circuitBreaker)) { - result[provider] = { - failures: state.failures, - blocked: state.blockedUntil > Date.now(), - blockedUntil: - state.blockedUntil > Date.now() ? new Date(state.blockedUntil).toISOString() : null, - remainingMs: Math.max(0, state.blockedUntil - Date.now()), - }; - } - return result; -} - -/** - * Record a successful refresh — resets circuit breaker for provider. - */ -function recordSuccess(provider: string) { - if (_circuitBreaker[provider]) { - delete _circuitBreaker[provider]; - } -} - -/** - * Record a failed refresh — increments circuit breaker counter. - */ -function recordFailure(provider: string, log: RefreshLoggerLike | null = null) { - if (!_circuitBreaker[provider]) { - _circuitBreaker[provider] = { failures: 0, blockedUntil: 0 }; - } - _circuitBreaker[provider].failures++; - - if (_circuitBreaker[provider].failures >= CIRCUIT_BREAKER_THRESHOLD) { - _circuitBreaker[provider].blockedUntil = Date.now() + CIRCUIT_BREAKER_COOLDOWN; - log?.error?.( - "TOKEN_REFRESH", - `🔴 Circuit breaker tripped for ${provider}: ${CIRCUIT_BREAKER_THRESHOLD} consecutive failures. ` + - `Blocked for ${CIRCUIT_BREAKER_COOLDOWN / 60000}min. Provider needs re-authentication.` - ); - } -} - -/** - * Execute a function with a timeout. - */ -async function withTimeout(fn: () => Promise, timeoutMs: number): Promise { - return await new Promise((resolve, reject) => { - const timer = setTimeout(() => resolve(null), timeoutMs); - if (typeof timer === "object" && "unref" in timer) { - (timer as { unref?: () => void }).unref?.(); - } - - fn().then( - (result) => { - clearTimeout(timer); - resolve(result); - }, - (error) => { - clearTimeout(timer); - reject(error); - } - ); - }); -} - -export async function refreshWithRetry( - refreshFn, - maxRetries = 3, - log: RefreshLogger = null, - provider = "unknown" -) { - // Circuit breaker check - if (isProviderBlocked(provider)) { - log?.warn?.("TOKEN_REFRESH", `⚡ Circuit breaker active for ${provider}, skipping refresh`); - return null; - } - - for (let attempt = 0; attempt < maxRetries; attempt++) { - if (attempt > 0) { - const delay = attempt * 1000; - log?.debug?.("TOKEN_REFRESH", `Retry ${attempt}/${maxRetries} after ${delay}ms`); - await new Promise((r) => setTimeout(r, delay)); - } - - try { - const result = await withTimeout(refreshFn, REFRESH_TIMEOUT_MS); - if (isUnrecoverableRefreshError(result)) { - log?.warn?.( - "TOKEN_REFRESH", - `Unrecoverable refresh error for ${provider}: ${result.error} — skipping retries` - ); - return result; - } - if (result) { - recordSuccess(provider); - return result; - } - } catch (error) { - log?.warn?.("TOKEN_REFRESH", `Attempt ${attempt + 1}/${maxRetries} failed: ${error.message}`); - } - } - - // All retries exhausted — record failure for circuit breaker - recordFailure(provider, log); - log?.error?.("TOKEN_REFRESH", `All ${maxRetries} retry attempts failed for ${provider}`); - return null; -} diff --git a/open-sse/services/tokenRefresh/casGuard.ts b/open-sse/services/tokenRefresh/casGuard.ts new file mode 100644 index 0000000000..3b4a13eb45 --- /dev/null +++ b/open-sse/services/tokenRefresh/casGuard.ts @@ -0,0 +1,84 @@ +// @ts-nocheck +// +// Compare-and-swap (CAS) guard on the refresh persist — extracted from +// open-sse/services/tokenRefresh.ts. See ../shared.ts for provenance notes. +// +// #4038: Fix A makes [network refresh + DB write] atomic *for a single +// connection's mutex*. It does NOT protect against a THIRD writer (a sibling +// process, a concurrent HealthCheck, or a replica) landing a fresher rotation +// on the same `connection_id` between the moment the caller read the row and +// the moment this persist runs. Overwriting that fresher row reverts the +// sibling's rotation, the next caller loads the reverted (now-consumed) +// refresh_token, and Auth0/Anthropic revoke the whole token family (the 1352× +// claude/aa5dd5cf invalidation storm). +// +// The CAS guard carries the refresh_token the caller PRESENTED (the version +// token, since refresh_tokens rotate on every refresh) plus a `reread` of the +// row's current refresh_token. Right before persisting, `getAccessToken` +// re-reads and, if a concurrent writer already rotated the row past the +// presented token, SKIPS the persist so the DB stays at the fresher state. The +// caller still receives the new accessToken — upstream already authenticated +// the request; only the DB write is skipped. No active guard ⇒ behavior is +// byte-identical to before (opt-in). +import { AsyncLocalStorage } from "node:async_hooks"; +import { wasRefreshTokenRotated } from "../refreshSerializer.ts"; +import type { RefreshLogger } from "./shared.ts"; + +type CasGuard = { + /** The refresh_token the caller presented for this refresh (CAS version token). */ + expectedRefreshToken: string | null; + /** Re-reads the CURRENT persisted refresh_token for this connection (decrypted). */ + reread: () => Promise; +}; +const casGuardStore = new AsyncLocalStorage(); +const casGuardStats = { skipped: 0, persisted: 0 }; + +export function runWithCasGuard( + guard: CasGuard | undefined | null, + fn: () => Promise +): Promise { + if (!guard) return fn(); + return casGuardStore.run(guard, fn); +} + +export function getActiveCasGuard(): CasGuard | undefined { + return casGuardStore.getStore(); +} + +/** Skip/persist counters for observability + tests. */ +export function getCasGuardStats(): { skipped: number; persisted: number } { + return { ...casGuardStats }; +} + +/** Test-only: reset the CAS counters between cases. */ +export function _resetCasGuardStats(): void { + casGuardStats.skipped = 0; + casGuardStats.persisted = 0; +} + +/** + * Returns true when the persist should be SKIPPED because a concurrent writer + * already rotated the row's refresh_token past the one we presented (CAS mismatch). + * Best-effort: any reread failure falls through to persist (never blocks recovery). + */ +export async function casGuardShouldSkipPersist(log?: RefreshLogger): Promise { + const guard = getActiveCasGuard(); + if (!guard || !guard.expectedRefreshToken) return false; + let current: string | null | undefined; + try { + current = await guard.reread(); + } catch { + return false; // reread failed — fall through to persist (best-effort) + } + // wasRefreshTokenRotated is true iff both are non-empty AND current !== expected. + if (wasRefreshTokenRotated(guard.expectedRefreshToken, current)) { + casGuardStats.skipped++; + log?.warn?.( + "TOKEN_REFRESH", + "CAS guard: skipping persist — a concurrent writer already rotated the refresh_token (#4038)" + ); + return true; + } + casGuardStats.persisted++; + return false; +} diff --git a/open-sse/services/tokenRefresh/circuitBreaker.ts b/open-sse/services/tokenRefresh/circuitBreaker.ts new file mode 100644 index 0000000000..ae5b503ea6 --- /dev/null +++ b/open-sse/services/tokenRefresh/circuitBreaker.ts @@ -0,0 +1,168 @@ +// @ts-nocheck +// +// Per-provider circuit breaker + refreshWithRetry — extracted from +// open-sse/services/tokenRefresh.ts. See ../shared.ts for provenance notes. +// +// refreshWithRetry wraps a refresh attempt with exponential backoff, a 30s +// per-attempt timeout, and a per-provider circuit breaker (5 consecutive +// failures → 30min pause). Unrecoverable refresh errors (invalid_grant, +// refresh_token_reused, …) short-circuit retries so the HealthCheck can +// deactivate the account instead of looping every 60s. +import type { RefreshLogger } from "./shared.ts"; +import { isUnrecoverableRefreshError } from "./shared.ts"; + +// ─── Circuit Breaker State ────────────────────────────────────────────────── +const _circuitBreaker: Record = {}; +const CIRCUIT_BREAKER_THRESHOLD = 5; // consecutive failures before tripping +const CIRCUIT_BREAKER_COOLDOWN = 30 * 60 * 1000; // 30 minutes +const REFRESH_TIMEOUT_MS = 30_000; // 30s max per refresh attempt + +interface CircuitBreakerStatusEntry { + failures: number; + blocked: boolean; + blockedUntil: string | null; + remainingMs: number; +} + +interface RefreshLoggerLike { + error?: (scope: string, message: string) => void; + warn?: (scope: string, message: string) => void; +} + +/** + * Check if a provider is circuit-breaker blocked. + */ +export function isProviderBlocked(provider: string): boolean { + const state = _circuitBreaker[provider]; + if (!state) return false; + if (!state.blockedUntil) return false; + if (state.blockedUntil > Date.now()) return true; + // Cooldown expired — reset + delete _circuitBreaker[provider]; + return false; +} + +/** + * Get circuit breaker status for all providers (for diagnostics). + */ +export function getCircuitBreakerStatus(): Record { + const result: Record = {}; + for (const [provider, state] of Object.entries(_circuitBreaker)) { + result[provider] = { + failures: state.failures, + blocked: state.blockedUntil > Date.now(), + blockedUntil: + state.blockedUntil > Date.now() ? new Date(state.blockedUntil).toISOString() : null, + remainingMs: Math.max(0, state.blockedUntil - Date.now()), + }; + } + return result; +} + +/** + * Record a successful refresh — resets circuit breaker for provider. + */ +function recordSuccess(provider: string) { + if (_circuitBreaker[provider]) { + delete _circuitBreaker[provider]; + } +} + +/** + * Record a failed refresh — increments circuit breaker counter. + */ +function recordFailure(provider: string, log: RefreshLoggerLike | null = null) { + if (!_circuitBreaker[provider]) { + _circuitBreaker[provider] = { failures: 0, blockedUntil: 0 }; + } + _circuitBreaker[provider].failures++; + + if (_circuitBreaker[provider].failures >= CIRCUIT_BREAKER_THRESHOLD) { + _circuitBreaker[provider].blockedUntil = Date.now() + CIRCUIT_BREAKER_COOLDOWN; + log?.error?.( + "TOKEN_REFRESH", + `🔴 Circuit breaker tripped for ${provider}: ${CIRCUIT_BREAKER_THRESHOLD} consecutive failures. ` + + `Blocked for ${CIRCUIT_BREAKER_COOLDOWN / 60000}min. Provider needs re-authentication.` + ); + } +} + +/** + * Execute a function with a timeout. + */ +async function withTimeout(fn: () => Promise, timeoutMs: number): Promise { + return await new Promise((resolve, reject) => { + const timer = setTimeout(() => resolve(null), timeoutMs); + if (typeof timer === "object" && "unref" in timer) { + (timer as { unref?: () => void }).unref?.(); + } + + fn().then( + (result) => { + clearTimeout(timer); + resolve(result); + }, + (error) => { + clearTimeout(timer); + reject(error); + } + ); + }); +} + +/** + * Refresh token with retry and exponential backoff + * Retries on failure with increasing delay: 1s, 2s, 3s... + * + * Includes: + * - Per-provider circuit breaker (5 consecutive failures → 30min pause) + * - 30s timeout per refresh attempt to prevent hanging connections + * + * @param {function} refreshFn - Async function that returns token or null + * @param {number} maxRetries - Max retry attempts (default 3) + * @param {object} log - Logger instance (optional) + * @param {string} provider - Provider ID for circuit breaker tracking (optional) + * @returns {Promise} Token result or null if all retries fail + */ +export async function refreshWithRetry( + refreshFn, + maxRetries = 3, + log: RefreshLogger = null, + provider = "unknown" +) { + // Circuit breaker check + if (isProviderBlocked(provider)) { + log?.warn?.("TOKEN_REFRESH", `⚡ Circuit breaker active for ${provider}, skipping refresh`); + return null; + } + + for (let attempt = 0; attempt < maxRetries; attempt++) { + if (attempt > 0) { + const delay = attempt * 1000; + log?.debug?.("TOKEN_REFRESH", `Retry ${attempt}/${maxRetries} after ${delay}ms`); + await new Promise((r) => setTimeout(r, delay)); + } + + try { + const result = await withTimeout(refreshFn, REFRESH_TIMEOUT_MS); + if (isUnrecoverableRefreshError(result)) { + log?.warn?.( + "TOKEN_REFRESH", + `Unrecoverable refresh error for ${provider}: ${result.error} — skipping retries` + ); + return result; + } + if (result) { + recordSuccess(provider); + return result; + } + } catch (error) { + log?.warn?.("TOKEN_REFRESH", `Attempt ${attempt + 1}/${maxRetries} failed: ${error.message}`); + } + } + + // All retries exhausted — record failure for circuit breaker + recordFailure(provider, log); + log?.error?.("TOKEN_REFRESH", `All ${maxRetries} retry attempts failed for ${provider}`); + return null; +} diff --git a/open-sse/services/tokenRefresh/rotationMap.ts b/open-sse/services/tokenRefresh/rotationMap.ts new file mode 100644 index 0000000000..9d533a1b69 --- /dev/null +++ b/open-sse/services/tokenRefresh/rotationMap.ts @@ -0,0 +1,85 @@ +// @ts-nocheck +// +// Token Rotation Map (codex-multi-auth pattern) — extracted from +// open-sse/services/tokenRefresh.ts. See ../shared.ts for provenance notes. +// +// When a rotating-token provider (Codex, Kimi, GitLab Duo, etc.) refreshes, +// the old refresh_token is consumed and a new one is issued. Any subsequent +// caller arriving with the OLD token would, without protection, hit upstream +// and trigger "refresh_token_reused" — which Auth0 treats as a security event +// and invalidates the entire token family. +// +// This in-memory map caches RECENT rotations so a stale caller can be redirected +// to the new tokens WITHOUT touching upstream. The DB staleness check inside +// the per-connection mutex covers the same scenario when connectionId is known, +// but not all callers pass connectionId (e.g., legacy code paths, retries that +// snapshot credentials before the rotation lands in DB). +// +// Ported from ndycode/codex-multi-auth (lib/refresh-queue.ts:218-248), the only +// publicly known tool that reliably sustains multiple Codex OAuth accounts. +// +// Key format: `provider:sha256(oldRefreshToken)` +// Value: { result: tokens, expiresAt: ms_since_epoch } +import { pbkdf2Sync } from "node:crypto"; + +const CACHE_SECRET = "omniroute-token-cache"; + +/** + * Build the dedup/rotation cache key for a (provider, refreshToken) pair. + * Hashed so a raw refresh_token never sits in a Map key in plaintext. + */ +export function getRefreshCacheKey(provider, refreshToken) { + const tokenHash = pbkdf2Sync(refreshToken, CACHE_SECRET, 1000, 32, "sha256").toString("hex"); + return `${provider}:${tokenHash}`; +} + +type RotationEntry = { + result: { accessToken: string; refreshToken: string; expiresIn?: number; expiresAt?: string }; + expiresAt: number; +}; +const tokenRotationMap = new Map(); +const ROTATION_MAP_TTL_MS = 60 * 1000; // 60 seconds — long enough to catch in-flight stale callers + +function cleanupRotationMap(now: number = Date.now()): void { + if (tokenRotationMap.size === 0) return; + for (const [key, entry] of tokenRotationMap.entries()) { + if (entry.expiresAt <= now) tokenRotationMap.delete(key); + } +} + +export function lookupRotation(provider: string, refreshToken: string): RotationEntry | undefined { + cleanupRotationMap(); + const key = getRefreshCacheKey(provider, refreshToken); + const entry = tokenRotationMap.get(key); + if (!entry) return undefined; + if (entry.expiresAt <= Date.now()) { + tokenRotationMap.delete(key); + return undefined; + } + return entry; +} + +export function recordRotation( + provider: string, + oldRefreshToken: string, + result: { accessToken: string; refreshToken: string; expiresIn?: number; expiresAt?: string } +): void { + if (!oldRefreshToken || !result.refreshToken || oldRefreshToken === result.refreshToken) { + return; + } + const key = getRefreshCacheKey(provider, oldRefreshToken); + tokenRotationMap.set(key, { + result, + expiresAt: Date.now() + ROTATION_MAP_TTL_MS, + }); +} + +// Exported for tests + diagnostics; not part of the public API surface. +export function _getTokenRotationMapStats(): { size: number; entries: number } { + cleanupRotationMap(); + return { size: tokenRotationMap.size, entries: tokenRotationMap.size }; +} + +export function _clearTokenRotationMap(): void { + tokenRotationMap.clear(); +} diff --git a/open-sse/services/tokenRefresh/shared.ts b/open-sse/services/tokenRefresh/shared.ts index bde1a2db10..808a349050 100644 --- a/open-sse/services/tokenRefresh/shared.ts +++ b/open-sse/services/tokenRefresh/shared.ts @@ -110,3 +110,19 @@ export async function readRefreshErrorBody( const code = extractOAuthErrorCode(parsed) ?? extractOAuthErrorCode(rawText); return { rawText, code }; } + +/** + * Check if a refresh result indicates an unrecoverable error + * (e.g. the refresh token was already consumed and cannot be reused). + * Callers should stop retrying and request re-authentication. + */ +export function isUnrecoverableRefreshError(result) { + return ( + result && + typeof result === "object" && + (result.error === "unrecoverable_refresh_error" || + result.error === "refresh_token_reused" || + result.error === "invalid_request" || + result.error === "invalid_grant") + ); +} diff --git a/tests/unit/oauth-providers-error-handling.test.ts b/tests/unit/oauth-providers-error-handling.test.ts index fdc9644e8b..aa41173c27 100644 --- a/tests/unit/oauth-providers-error-handling.test.ts +++ b/tests/unit/oauth-providers-error-handling.test.ts @@ -240,8 +240,11 @@ test("P3: refreshWindsurfToken parses Firebase USER_DISABLED/TOKEN_EXPIRED error // ─── isUnrecoverableRefreshError consistency ────────────────────────────────── +// isUnrecoverableRefreshError moved to tokenRefresh/shared.ts in the god-file +// decomposition (tokenRefresh.ts re-exports it, so the public surface is unchanged); +// this source-text assertion has to follow it to the file that defines the body. test("isUnrecoverableRefreshError detects the normalized sentinel shape", async () => { - const src = await read("open-sse/services/tokenRefresh.ts"); + const src = await read("open-sse/services/tokenRefresh/shared.ts"); const fnMatch = src.match(/export\s+function\s+isUnrecoverableRefreshError\([\s\S]+?\n\}/); assert.ok(fnMatch, "isUnrecoverableRefreshError function body not found"); assert.match( diff --git a/tests/unit/token-refresh-cas-guard.test.ts b/tests/unit/token-refresh-cas-guard.test.ts new file mode 100644 index 0000000000..babe37b6c0 --- /dev/null +++ b/tests/unit/token-refresh-cas-guard.test.ts @@ -0,0 +1,124 @@ +import test from "node:test"; +import assert from "node:assert/strict"; + +// Unit tests for the CAS guard leaf extracted from tokenRefresh.ts (#4038). +// The CAS guard re-reads the row's refresh_token right before persisting and +// SKIPS the write when a concurrent writer already rotated it past the token +// the caller presented — preventing a revert that would invalidate the token +// family on rotating-token providers (Auth0/Anthropic). + +const { + runWithCasGuard, + getActiveCasGuard, + getCasGuardStats, + _resetCasGuardStats, + casGuardShouldSkipPersist, +} = await import("../../open-sse/services/tokenRefresh/casGuard.ts"); + +const silentLog = { info() {}, warn() {}, error() {} }; + +test.beforeEach(() => { + _resetCasGuardStats(); +}); + +test("getActiveCasGuard returns undefined outside a guard context", () => { + assert.equal(getActiveCasGuard(), undefined); +}); + +test("runWithCasGuard exposes the guard via getActiveCasGuard inside the closure", async () => { + const guard = { expectedRefreshToken: "R0", reread: async () => "R0" }; + await runWithCasGuard(guard, async () => { + assert.equal(getActiveCasGuard(), guard); + }); + assert.equal(getActiveCasGuard(), undefined, "guard is cleared after the closure resolves"); +}); + +test("runWithCasGuard with a null/undefined guard runs the function unchanged", async () => { + let ran = false; + await runWithCasGuard(null, async () => { + ran = true; + }); + assert.equal(ran, true); + assert.equal(getActiveCasGuard(), undefined); +}); + +test("casGuardShouldSkipPersist returns false when no guard is active", async () => { + assert.equal(await casGuardShouldSkipPersist(silentLog), false); + assert.equal(getCasGuardStats().skipped, 0); + assert.equal(getCasGuardStats().persisted, 0); +}); + +test("casGuardShouldSkipPersist returns false when the guard has no expectedRefreshToken", async () => { + const guard = { expectedRefreshToken: null, reread: async () => "R0" }; + await runWithCasGuard(guard, async () => { + assert.equal(await casGuardShouldSkipPersist(silentLog), false); + }); + assert.equal(getCasGuardStats().skipped, 0); + assert.equal(getCasGuardStats().persisted, 0); +}); + +test("casGuardShouldSkipPersist SKIPS when the row rotated past the presented token", async () => { + const guard = { expectedRefreshToken: "R0", reread: async () => "R_CONCURRENT" }; + await runWithCasGuard(guard, async () => { + assert.equal(await casGuardShouldSkipPersist(silentLog), true); + }); + assert.equal(getCasGuardStats().skipped, 1); + assert.equal(getCasGuardStats().persisted, 0); +}); + +test("casGuardShouldSkipPersist PERSISTS when the row still holds the presented token", async () => { + const guard = { expectedRefreshToken: "R0", reread: async () => "R0" }; + await runWithCasGuard(guard, async () => { + assert.equal(await casGuardShouldSkipPersist(silentLog), false); + }); + assert.equal(getCasGuardStats().skipped, 0); + assert.equal(getCasGuardStats().persisted, 1); +}); + +test("casGuardShouldSkipPersist falls through to persist when reread throws (best-effort)", async () => { + const guard = { + expectedRefreshToken: "R0", + reread: async () => { + throw new Error("db unavailable"); + }, + }; + await runWithCasGuard(guard, async () => { + assert.equal(await casGuardShouldSkipPersist(silentLog), false); + }); + // reread failure returns false (do not skip persist) WITHOUT touching the + // persisted counter — the counter only advances on a successful reread that + // confirms the row is unchanged. The key guarantee is skipped stays 0. + assert.equal(getCasGuardStats().skipped, 0, "reread failure must never block recovery"); + assert.equal(getCasGuardStats().persisted, 0); +}); + +test("casGuardShouldSkipPersist treats an empty reread as not-rotated", async () => { + const guard = { expectedRefreshToken: "R0", reread: async () => null }; + await runWithCasGuard(guard, async () => { + assert.equal(await casGuardShouldSkipPersist(silentLog), false); + }); + assert.equal(getCasGuardStats().persisted, 1); +}); + +test("getCasGuardStats returns a snapshot copy (not the live counters)", () => { + const guard = { expectedRefreshToken: "R0", reread: async () => "R0" }; + return runWithCasGuard(guard, async () => { + await casGuardShouldSkipPersist(silentLog); + const snap = getCasGuardStats(); + assert.equal(snap.persisted, 1); + // Mutating the snapshot must not affect future stats. + snap.persisted = 999; + assert.equal(getCasGuardStats().persisted, 1); + }); +}); + +test("_resetCasGuardStats zeroes both counters", async () => { + const guard = { expectedRefreshToken: "R0", reread: async () => "R_CONCURRENT" }; + await runWithCasGuard(guard, async () => { + await casGuardShouldSkipPersist(silentLog); + }); + assert.equal(getCasGuardStats().skipped, 1); + _resetCasGuardStats(); + assert.equal(getCasGuardStats().skipped, 0); + assert.equal(getCasGuardStats().persisted, 0); +}); diff --git a/tests/unit/token-refresh-circuit-breaker.test.ts b/tests/unit/token-refresh-circuit-breaker.test.ts new file mode 100644 index 0000000000..fb88a7d6c4 --- /dev/null +++ b/tests/unit/token-refresh-circuit-breaker.test.ts @@ -0,0 +1,176 @@ +import test from "node:test"; +import assert from "node:assert/strict"; + +// Unit tests for the circuit breaker + refreshWithRetry leaf extracted from +// tokenRefresh.ts. refreshWithRetry wraps a refresh attempt with exponential +// backoff, a 30s per-attempt timeout, and a per-provider circuit breaker +// (5 consecutive failures → 30min pause). Unrecoverable refresh errors +// short-circuit retries. + +const { isProviderBlocked, getCircuitBreakerStatus, refreshWithRetry } = + await import("../../open-sse/services/tokenRefresh/circuitBreaker.ts"); + +const silentLog = { + info() {}, + warn() {}, + error() {}, + debug() {}, +}; + +function makeLog() { + const entries = []; + const log = (level) => (scope, message) => entries.push({ level, scope, message }); + return { + entries, + debug: log("debug"), + info: log("info"), + warn: log("warn"), + error: log("error"), + }; +} + +test("isProviderBlocked returns false for an unknown provider", () => { + assert.equal(isProviderBlocked("never-seen"), false); +}); + +test("getCircuitBreakerStatus returns an empty object when no failures recorded", () => { + assert.deepEqual(getCircuitBreakerStatus(), {}); +}); + +test("refreshWithRetry returns the result on the first success and clears prior failures", async () => { + const provider = "cb-success-" + Math.random().toString(36).slice(2); + // Seed a failure so we can verify success clears it. + await refreshWithRetry(async () => null, 1, silentLog, provider); + assert.equal(getCircuitBreakerStatus()[provider].failures, 1); + + const result = await refreshWithRetry( + async () => ({ accessToken: "ok" }), + 3, + silentLog, + provider + ); + assert.equal(result.accessToken, "ok"); + assert.equal(getCircuitBreakerStatus()[provider], undefined, "success resets the breaker"); +}); + +test("refreshWithRetry retries to success within maxRetries", async () => { + const provider = "cb-retry-" + Math.random().toString(36).slice(2); + let attempts = 0; + const result = await refreshWithRetry( + async () => { + attempts++; + if (attempts < 2) return null; + return { accessToken: "ok-after-retry" }; + }, + 3, + silentLog, + provider + ); + assert.equal(result.accessToken, "ok-after-retry"); + assert.equal(attempts, 2); + assert.equal(getCircuitBreakerStatus()[provider], undefined); +}); + +test("refreshWithRetry bails immediately on an unrecoverable error without retrying", async () => { + const provider = "cb-unrecoverable-" + Math.random().toString(36).slice(2); + let attempts = 0; + const result = await refreshWithRetry( + async () => { + attempts++; + return { error: "invalid_grant" }; + }, + 3, + silentLog, + provider + ); + assert.equal(attempts, 1, "unrecoverable errors must not be retried"); + assert.equal(result.error, "invalid_grant"); + assert.equal( + getCircuitBreakerStatus()[provider], + undefined, + "no failure recorded for unrecoverable" + ); +}); + +test("refreshWithRetry bails immediately on refresh_token_reused", async () => { + const provider = "cb-reused-" + Math.random().toString(36).slice(2); + let attempts = 0; + const result = await refreshWithRetry( + async () => { + attempts++; + return { error: "refresh_token_reused" }; + }, + 3, + silentLog, + provider + ); + assert.equal(attempts, 1); + assert.equal(result.error, "refresh_token_reused"); +}); + +test("refreshWithRetry trips the circuit breaker after repeated failures", async () => { + const provider = "cb-trip-" + Math.random().toString(36).slice(2); + // 5 consecutive single-retry failures trip the breaker. + for (let i = 0; i < 5; i++) { + await refreshWithRetry(async () => null, 1, silentLog, provider); + } + assert.equal(isProviderBlocked(provider), true); + assert.equal(getCircuitBreakerStatus()[provider].blocked, true); + assert.ok(getCircuitBreakerStatus()[provider].blockedUntil); + + // A blocked provider short-circuits without calling refreshFn. + let called = false; + const blocked = await refreshWithRetry( + async () => { + called = true; + return { accessToken: "x" }; + }, + 1, + silentLog, + provider + ); + assert.equal(called, false, "refreshFn must not run while the breaker is open"); + assert.equal(blocked, null); +}); + +test("refreshWithRetry records a failure when all retries are exhausted", async () => { + const provider = "cb-exhaust-" + Math.random().toString(36).slice(2); + const log = makeLog(); + const result = await refreshWithRetry(async () => null, 2, log, provider); + assert.equal(result, null); + assert.equal(getCircuitBreakerStatus()[provider].failures, 1); + assert.ok( + log.entries.some((e) => e.level === "error" && /All 2 retry attempts failed/.test(e.message)) + ); +}); + +test("refreshWithRetry propagates thrown errors as retry failures (not crashes)", async () => { + const provider = "cb-throw-" + Math.random().toString(36).slice(2); + const log = makeLog(); + let attempts = 0; + const result = await refreshWithRetry( + async () => { + attempts++; + throw new Error("upstream boom"); + }, + 2, + log, + provider + ); + assert.equal(result, null); + assert.equal(attempts, 2, "thrown errors are retried, not fatal"); + assert.equal(getCircuitBreakerStatus()[provider].failures, 1); + assert.ok(log.entries.some((e) => e.level === "warn" && /failed: upstream boom/.test(e.message))); +}); + +test("refreshWithRetry defaults: maxRetries=3, provider='unknown'", async () => { + // With defaults, an always-null refresh exhausts 3 attempts and records a + // failure under the "unknown" provider. + let attempts = 0; + await refreshWithRetry(async () => { + attempts++; + return null; + }); + assert.equal(attempts, 3); + assert.ok(getCircuitBreakerStatus()["unknown"], "default provider is 'unknown'"); +}); diff --git a/tests/unit/token-refresh-rotation-map.test.ts b/tests/unit/token-refresh-rotation-map.test.ts new file mode 100644 index 0000000000..d38a9c64ac --- /dev/null +++ b/tests/unit/token-refresh-rotation-map.test.ts @@ -0,0 +1,100 @@ +import test from "node:test"; +import assert from "node:assert/strict"; + +// Unit tests for the token rotation map leaf extracted from tokenRefresh.ts. +// The rotation map caches RECENT refresh_token rotations so a stale caller can +// be redirected to the new tokens WITHOUT re-hitting upstream (which would +// trigger Auth0 family revocation on rotating-token providers like Codex). + +const { + getRefreshCacheKey, + lookupRotation, + recordRotation, + _getTokenRotationMapStats, + _clearTokenRotationMap, +} = await import("../../open-sse/services/tokenRefresh/rotationMap.ts"); + +test.beforeEach(() => { + _clearTokenRotationMap(); +}); + +test("getRefreshCacheKey is deterministic and provider-scoped", () => { + const a = getRefreshCacheKey("codex", "refresh-1"); + const b = getRefreshCacheKey("codex", "refresh-1"); + const c = getRefreshCacheKey("openai", "refresh-1"); + assert.equal(a, b, "same (provider, token) must hash to the same key"); + assert.notEqual(a, c, "different provider must produce a different key"); + assert.match(a, /^codex:/, "key is prefixed with the provider id"); + // The raw refresh token must NOT appear in the key (it is hashed). + assert.doesNotMatch(a, /refresh-1/); +}); + +test("recordRotation stores a rotation keyed by the OLD refresh token", () => { + recordRotation("codex", "old-rt", { + accessToken: "new-access", + refreshToken: "new-rt", + expiresIn: 3600, + }); + const stats = _getTokenRotationMapStats(); + assert.equal(stats.size, 1); + const hit = lookupRotation("codex", "old-rt"); + assert.ok(hit, "lookup by the old refresh token must find the cached rotation"); + assert.equal(hit.result.accessToken, "new-access"); + assert.equal(hit.result.refreshToken, "new-rt"); + assert.equal(hit.result.expiresIn, 3600); +}); + +test("recordRotation is a no-op when the refresh token did not rotate", () => { + recordRotation("codex", "same-rt", { + accessToken: "new-access", + refreshToken: "same-rt", + expiresIn: 3600, + }); + assert.equal(_getTokenRotationMapStats().size, 0, "no rotation recorded when token unchanged"); + assert.equal(lookupRotation("codex", "same-rt"), undefined); +}); + +test("recordRotation is a no-op when the old refresh token is empty", () => { + recordRotation("codex", "", { + accessToken: "new-access", + refreshToken: "new-rt", + }); + assert.equal(_getTokenRotationMapStats().size, 0); +}); + +test("recordRotation is a no-op when the new refresh token is empty", () => { + recordRotation("codex", "old-rt", { + accessToken: "new-access", + refreshToken: "", + }); + assert.equal(_getTokenRotationMapStats().size, 0); +}); + +test("lookupRotation returns undefined for an unknown token", () => { + assert.equal(lookupRotation("codex", "never-recorded"), undefined); +}); + +test("lookupRotation returns undefined for a different provider", () => { + recordRotation("codex", "shared-rt", { + accessToken: "a", + refreshToken: "new-rt", + }); + assert.equal(lookupRotation("openai", "shared-rt"), undefined, "rotation map is provider-scoped"); + assert.ok(lookupRotation("codex", "shared-rt"), "the original provider still hits"); +}); + +test("_clearTokenRotationMap empties the map", () => { + recordRotation("codex", "old-rt", { accessToken: "a", refreshToken: "new-rt" }); + assert.equal(_getTokenRotationMapStats().size, 1); + _clearTokenRotationMap(); + assert.equal(_getTokenRotationMapStats().size, 0); + assert.equal(lookupRotation("codex", "old-rt"), undefined); +}); + +test("_getTokenRotationMapStats reports the live entry count", () => { + assert.equal(_getTokenRotationMapStats().size, 0); + recordRotation("codex", "old-1", { accessToken: "a1", refreshToken: "new-1" }); + recordRotation("codex", "old-2", { accessToken: "a2", refreshToken: "new-2" }); + assert.equal(_getTokenRotationMapStats().size, 2); + assert.equal(_getTokenRotationMapStats().entries, 2); +});