diff --git a/open-sse/services/tokenRefresh.ts b/open-sse/services/tokenRefresh.ts index 0b9edc7bfc..09a7861bc9 100755 --- a/open-sse/services/tokenRefresh.ts +++ b/open-sse/services/tokenRefresh.ts @@ -13,6 +13,11 @@ const CACHE_SECRET = "omniroute-token-cache"; // Key: "provider:sha256(refreshToken)" → Value: Promise const refreshPromiseCache = new Map(); +// Per-connection mutex: prevents parallel OAuth refresh for rotating tokens. +// Key: connectionId → Value: { promise, waiters } +// Primary dedup when credentials.connectionId is present; refreshPromiseCache is fallback. +const connectionRefreshMutex = new Map(); + type RefreshLogger = { info?: (tag: string, message: string, data?: Record) => void; warn?: (tag: string, message: string, data?: Record) => void; @@ -415,8 +420,8 @@ export async function refreshQwenToken(refreshToken, log, proxyConfig: unknown = /** * Specialized refresh for Codex (OpenAI) OAuth tokens. * OpenAI uses rotating (one-time-use) refresh tokens. - * Returns { error: 'refresh_token_reused' } when the token has already been consumed, - * so callers can stop retrying and request re-authentication. + * Returns { error: 'unrecoverable_refresh_error', code } when the token has already been + * consumed or is invalid, so callers can stop retrying and request re-authentication. */ export async function refreshCodexToken(refreshToken, log, proxyConfig: unknown = null) { try { @@ -831,9 +836,18 @@ export function isUnrecoverableRefreshError(result) { /** * Get access token for a specific provider (with deduplication). - * If a refresh is already in-flight for the same provider+token, - * subsequent calls share the existing promise instead of making - * parallel OAuth requests. + * + * Deduplication strategy (two layers): + * 1. Per-connection mutex (primary): if credentials.connectionId is present, all concurrent + * callers for that connection share one in-flight promise regardless of which token they + * loaded. This prevents refresh_token_reused errors with rotating (one-time-use) tokens, + * e.g. Codex/OpenAI, where callers that loaded credentials at different times may hold + * different token strings but refer to the same connection. + * 2. Token-hash fallback: if no connectionId, dedup by provider+sha256(refreshToken) as before. + * + * Additionally, when connectionId is present, the stale-token check reads the DB to detect + * whether another process already refreshed the token. If the DB token is still valid it is + * returned immediately without a new upstream call. */ export async function getAccessToken(provider, credentials, log, proxyConfig: unknown = null) { if (!credentials || !credentials.refreshToken || typeof credentials.refreshToken !== "string") { @@ -841,6 +855,52 @@ export async function getAccessToken(provider, credentials, log, proxyConfig: un return null; } + const connectionId = credentials.connectionId; + + // ── Layer 1: per-connection mutex ────────────────────────────────────────── + if (connectionId && typeof connectionId === "string") { + const existing = connectionRefreshMutex.get(connectionId); + if (existing) { + existing.waiters++; + log?.info?.("TOKEN_REFRESH", "Concurrent refresh detected — sharing in-flight refresh", { + provider, + connectionId, + waiters: existing.waiters, + }); + return existing.promise; + } + + const entry = { promise: null, waiters: 0 }; + entry.promise = _getAccessTokenWithStalenessCheck(provider, credentials, log, proxyConfig).finally(() => { + connectionRefreshMutex.delete(connectionId); + }); + connectionRefreshMutex.set(connectionId, entry); + return entry.promise; + } + + // ── Layer 2: token-hash fallback (no connectionId) ───────────────────────── + const cacheKey = getRefreshCacheKey(provider, credentials.refreshToken); + + if (refreshPromiseCache.has(cacheKey)) { + log?.info?.("TOKEN_REFRESH", `Reusing in-flight refresh for ${provider}`); + return refreshPromiseCache.get(cacheKey); + } + + const refreshPromise = _getAccessTokenInternal(provider, credentials, log, proxyConfig).finally( + () => { + refreshPromiseCache.delete(cacheKey); + } + ); + + refreshPromiseCache.set(cacheKey, refreshPromise); + return refreshPromise; +} + +/** + * Internal helper: performs the DB staleness check then calls the actual refresh. + * Only called from the per-connection mutex path (Layer 1 above). + */ +async function _getAccessTokenWithStalenessCheck(provider, credentials, log, proxyConfig) { // RACE CONDITION PREVENTION: // If the credentials object in memory is stale (e.g. it waited in a semaphore while another // request refreshed the token), using its OLD refreshToken will cause the provider (e.g. OpenAI) @@ -886,23 +946,7 @@ export async function getAccessToken(provider, credentials, log, proxyConfig: un } } - const cacheKey = getRefreshCacheKey(provider, credentials.refreshToken); - - // If a refresh is already in-flight, reuse it - if (refreshPromiseCache.has(cacheKey)) { - log?.info?.("TOKEN_REFRESH", `Reusing in-flight refresh for ${provider}`); - return refreshPromiseCache.get(cacheKey); - } - - // Start a new refresh and cache the promise - const refreshPromise = _getAccessTokenInternal(provider, credentials, log, proxyConfig).finally( - () => { - refreshPromiseCache.delete(cacheKey); - } - ); - - refreshPromiseCache.set(cacheKey, refreshPromise); - return refreshPromise; + return _getAccessTokenInternal(provider, credentials, log, proxyConfig); } /** @@ -1036,6 +1080,18 @@ export function isProviderBlocked(provider: string): boolean { return false; } +/** + * Get active per-connection mutex entries (for diagnostics/metrics). + * Returns a snapshot of connections that have an in-flight refresh and their waiter count. + */ +export function getConnectionRefreshMutexStatus(): Record { + const result: Record = {}; + for (const [connectionId, entry] of connectionRefreshMutex.entries()) { + result[connectionId] = { waiters: entry.waiters }; + } + return result; +} + /** * Get circuit breaker status for all providers (for diagnostics). */ diff --git a/src/sse/services/tokenRefresh.ts b/src/sse/services/tokenRefresh.ts index 6df10e2f72..244eb70558 100755 --- a/src/sse/services/tokenRefresh.ts +++ b/src/sse/services/tokenRefresh.ts @@ -17,6 +17,32 @@ import { getAllAccessTokens as _getAllAccessTokens, } from "@omniroute/open-sse/services/tokenRefresh.ts"; +// Per-connection mutex: prevents concurrent OAuth refresh for rotating tokens. +// Key = connectionId, Value = { promise: in-flight refresh, waiters: count of callers sharing it } +const connectionRefreshMutex = new Map; waiters: number }>(); + +export async function withConnectionRefreshMutex( + connectionId: string, + fn: () => Promise +): Promise { + const existing = connectionRefreshMutex.get(connectionId); + if (existing) { + existing.waiters++; + log.info("TOKEN_REFRESH", "Concurrent refresh detected — sharing in-flight refresh", { + connectionId, + waiters: existing.waiters, + }); + return existing.promise as Promise; + } + + const entry: { promise: Promise; waiters: number } = { promise: null as any, waiters: 0 }; + entry.promise = fn().finally(() => { + connectionRefreshMutex.delete(connectionId); + }); + connectionRefreshMutex.set(connectionId, entry); + return entry.promise; +} + export const TOKEN_EXPIRY_BUFFER_MS = BUFFER_MS; export const refreshAccessToken = async ( @@ -146,7 +172,12 @@ export async function checkAndRefreshToken(provider: string, credentials: any) { expiresIn: Math.round((expiresAt - now) / 1000), }); - const newCredentials = await getAccessToken(provider, updatedCredentials); + const connectionId: string | undefined = updatedCredentials.connectionId; + const newCredentials = connectionId + ? await withConnectionRefreshMutex(connectionId, () => + getAccessToken(provider, updatedCredentials) + ) + : await getAccessToken(provider, updatedCredentials); if (newCredentials && newCredentials.accessToken) { await updateProviderCredentials(updatedCredentials.connectionId, newCredentials); diff --git a/tests/unit/token-refresh-service.test.ts b/tests/unit/token-refresh-service.test.ts index 59562e58b6..3226f40f3f 100644 --- a/tests/unit/token-refresh-service.test.ts +++ b/tests/unit/token-refresh-service.test.ts @@ -24,6 +24,7 @@ const { getAllAccessTokens, isProviderBlocked, getCircuitBreakerStatus, + getConnectionRefreshMutexStatus, refreshWithRetry, } = tokenRefresh; @@ -907,3 +908,200 @@ test("isProviderBlocked clears expired circuit-breaker entries once cooldown pas assert.equal(getCircuitBreakerStatus()[provider], undefined); }); }); + +// ─── Per-connection mutex tests ──────────────────────────────────────────────── + +test("getAccessToken per-connection mutex: 5 concurrent callers fire exactly one upstream call", async () => { + const log = createLog(); + let upstreamCallCount = 0; + + await withPatchedProperties( + PROVIDERS, + { "custom-oauth-conn-mutex": { tokenUrl: "https://auth.example.com/token" } }, + async () => { + await withMockedFetch( + async () => { + upstreamCallCount++; + // Simulate 50ms upstream latency so all 5 callers are concurrent + await new Promise((r) => setTimeout(r, 50)); + return jsonResponse({ + access_token: "new-access-token", + refresh_token: "new-refresh-token", + expires_in: 3600, + }); + }, + async () => { + const credentials = { + connectionId: "conn-abc-123", + refreshToken: "old-refresh-token", + }; + + const results = await Promise.all([ + getAccessToken("custom-oauth-conn-mutex", { ...credentials }, log), + getAccessToken("custom-oauth-conn-mutex", { ...credentials }, log), + getAccessToken("custom-oauth-conn-mutex", { ...credentials }, log), + getAccessToken("custom-oauth-conn-mutex", { ...credentials }, log), + getAccessToken("custom-oauth-conn-mutex", { ...credentials }, log), + ]); + + assert.equal(upstreamCallCount, 1, "upstream called exactly once"); + for (const result of results) { + assert.equal(result?.accessToken, "new-access-token", "all callers got same token"); + assert.equal(result?.refreshToken, "new-refresh-token"); + } + // All results are the same object reference (shared promise) + assert.strictEqual(results[0], results[1]); + assert.strictEqual(results[1], results[4]); + } + ); + } + ); +}); + +test("getAccessToken per-connection mutex: logs concurrent refresh with waiter count", async () => { + const log = createLog(); + + await withPatchedProperties( + PROVIDERS, + { "custom-oauth-conn-mutex": { tokenUrl: "https://auth.example.com/token" } }, + async () => { + await withMockedFetch( + async () => { + await new Promise((r) => setTimeout(r, 20)); + return jsonResponse({ access_token: "tok", refresh_token: "rtok", expires_in: 600 }); + }, + async () => { + const credentials = { connectionId: "conn-log-test", refreshToken: "rt" }; + await Promise.all([ + getAccessToken("custom-oauth-conn-mutex", { ...credentials }, log), + getAccessToken("custom-oauth-conn-mutex", { ...credentials }, log), + getAccessToken("custom-oauth-conn-mutex", { ...credentials }, log), + ]); + + const concurrentLogs = log.entries.filter( + (e) => + e.level === "info" && + e.message === "Concurrent refresh detected — sharing in-flight refresh" + ); + assert.ok(concurrentLogs.length >= 1, "logged at least one concurrent refresh event"); + assert.ok( + concurrentLogs.some((e) => e.meta?.connectionId === "conn-log-test"), + "log includes connectionId" + ); + assert.ok( + concurrentLogs.some((e) => typeof e.meta?.waiters === "number" && e.meta.waiters >= 1), + "log includes waiter count" + ); + } + ); + } + ); +}); + +test("getAccessToken per-connection mutex: failed refresh propagates null to all waiters (idempotent error)", async () => { + const log = createLog(); + + await withPatchedProperties( + PROVIDERS, + { "custom-oauth-conn-mutex": { tokenUrl: "https://auth.example.com/token" } }, + async () => { + await withMockedFetch( + async () => { + await new Promise((r) => setTimeout(r, 20)); + // 400 response causes refreshAccessToken to return null + return new Response("bad_request", { status: 400 }); + }, + async () => { + const credentials = { connectionId: "conn-fail-test", refreshToken: "expired-rt" }; + const results = await Promise.all([ + getAccessToken("custom-oauth-conn-mutex", { ...credentials }, log), + getAccessToken("custom-oauth-conn-mutex", { ...credentials }, log), + getAccessToken("custom-oauth-conn-mutex", { ...credentials }, log), + ]); + + for (const result of results) { + assert.equal(result, null, "failed refresh returns null to all waiters"); + } + // Mutex cleaned up after failure + assert.equal( + getConnectionRefreshMutexStatus()["conn-fail-test"], + undefined, + "mutex entry removed after failure" + ); + } + ); + } + ); +}); + +test("getAccessToken per-connection mutex: different connections run independently", async () => { + const log = createLog(); + let upstreamCallCount = 0; + + await withPatchedProperties( + PROVIDERS, + { "custom-oauth-conn-mutex": { tokenUrl: "https://auth.example.com/token" } }, + async () => { + await withMockedFetch( + async () => { + upstreamCallCount++; + await new Promise((r) => setTimeout(r, 20)); + return jsonResponse({ + access_token: `access-${upstreamCallCount}`, + refresh_token: `refresh-${upstreamCallCount}`, + expires_in: 600, + }); + }, + async () => { + const [groupA, groupB] = await Promise.all([ + Promise.all([ + getAccessToken("custom-oauth-conn-mutex", { connectionId: "conn-A", refreshToken: "rt-a" }, log), + getAccessToken("custom-oauth-conn-mutex", { connectionId: "conn-A", refreshToken: "rt-a" }, log), + ]), + Promise.all([ + getAccessToken("custom-oauth-conn-mutex", { connectionId: "conn-B", refreshToken: "rt-b" }, log), + getAccessToken("custom-oauth-conn-mutex", { connectionId: "conn-B", refreshToken: "rt-b" }, log), + ]), + ]); + + assert.equal(upstreamCallCount, 2, "one upstream call per distinct connection"); + assert.strictEqual(groupA[0], groupA[1], "conn-A callers share same result"); + assert.strictEqual(groupB[0], groupB[1], "conn-B callers share same result"); + assert.notStrictEqual(groupA[0], groupB[0], "conn-A and conn-B got different results"); + } + ); + } + ); +}); + +test("getAccessToken per-connection mutex: mutex cleared after success, next call re-fires upstream", async () => { + const log = createLog(); + let upstreamCallCount = 0; + + await withPatchedProperties( + PROVIDERS, + { "custom-oauth-conn-mutex": { tokenUrl: "https://auth.example.com/token" } }, + async () => { + await withMockedFetch( + async () => { + upstreamCallCount++; + return jsonResponse({ + access_token: `access-${upstreamCallCount}`, + refresh_token: `refresh-${upstreamCallCount}`, + expires_in: 600, + }); + }, + async () => { + const credentials = { connectionId: "conn-refire", refreshToken: "rt" }; + + const first = await getAccessToken("custom-oauth-conn-mutex", { ...credentials }, log); + const second = await getAccessToken("custom-oauth-conn-mutex", { ...credentials }, log); + + assert.equal(upstreamCallCount, 2, "each sequential call fires upstream once"); + assert.equal(first?.accessToken, "access-1"); + assert.equal(second?.accessToken, "access-2"); + } + ); + } + ); +});