fix(oauth): per-connection mutex for rotating refresh tokens (#1885)

Integrated into release/v3.7.9
This commit is contained in:
smartenok-ops
2026-05-03 15:19:43 -04:00
committed by GitHub
parent 1d38900f17
commit 0e8148d602
3 changed files with 308 additions and 23 deletions

View File

@@ -13,6 +13,11 @@ const CACHE_SECRET = "omniroute-token-cache";
// Key: "provider:sha256(refreshToken)" → Value: Promise<result>
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<string, unknown>) => void;
warn?: (tag: string, message: string, data?: Record<string, unknown>) => 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<string, { waiters: number }> {
const result: Record<string, { waiters: number }> = {};
for (const [connectionId, entry] of connectionRefreshMutex.entries()) {
result[connectionId] = { waiters: entry.waiters };
}
return result;
}
/**
* Get circuit breaker status for all providers (for diagnostics).
*/

View File

@@ -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<string, { promise: Promise<any>; waiters: number }>();
export async function withConnectionRefreshMutex<T>(
connectionId: string,
fn: () => Promise<T>
): Promise<T> {
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<T>;
}
const entry: { promise: Promise<T>; 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);

View File

@@ -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");
}
);
}
);
});