codex-multi-auth 2.20.0 → 2.22.0

package/dist/lib/runtime/rotation-token-refresh.js

@@ -1,3 +1,4 @@
+import { sanitizeEmail } from "../accounts.js";
 import { queuedRefresh } from "../refresh-queue.js";
 import { isTokenInvalidationError, isTokenRefreshRetryable, } from "../request/rate-limit-decision.js";
 /** @internal */
@@ -24,49 +25,232 @@ function hasUsableAccessToken(account, now, skewMs) {
         typeof account.expires === "number" &&
         account.expires > now + Math.max(0, skewMs));
 }
-const runtimeRefreshCommitQueues = new WeakMap();
+let runtimeRefreshCommitQueues = new WeakMap();
+let runtimeRefreshManagerRedirects = new WeakMap();
+function resolveRuntimeRefreshCommitManager(accountManager) {
+    let current = accountManager;
+    const seen = new Set();
+    while (!seen.has(current)) {
+        seen.add(current);
+        const next = runtimeRefreshManagerRedirects.get(current);
+        if (!next)
+            return current;
+        current = next;
+    }
+    return current;
+}
+function uniqueMatch(accounts, predicate) {
+    const matches = accounts.filter(predicate);
+    return matches.length === 1 ? matches[0] ?? null : null;
+}
+function workspaceIds(account) {
+    return new Set([
+        account.accountId?.trim(),
+        ...(account.workspaces ?? []).map((workspace) => workspace.id.trim()),
+    ].filter((value) => Boolean(value)));
+}
+function hasWorkspaceConflict(source, candidate) {
+    const sourceIds = workspaceIds(source);
+    const candidateIds = workspaceIds(candidate);
+    return (sourceIds.size > 0 &&
+        candidateIds.size > 0 &&
+        ![...sourceIds].some((id) => candidateIds.has(id)));
+}
+/** @internal Resolve an account onto the latest runtime manager replacement. */
+export function resolveRuntimeRefreshTarget(accountManager, account) {
+    const targetManager = resolveRuntimeRefreshCommitManager(accountManager);
+    if (targetManager === accountManager)
+        return { accountManager, account };
+    const accounts = targetManager.getAccountsSnapshot();
+    // A persisted record id is the strongest row identity we have. If that row
+    // vanished in the replacement manager, never fall through to a same-email
+    // sibling workspace: the in-flight refresh belongs to the removed row.
+    // Derived in-memory ids hash the refresh token, so they are only strict
+    // when the row actually persisted one — otherwise rotation retires the id.
+    if (account.recordId && accountManager.hasPersistedRecordId(account)) {
+        const exact = uniqueMatch(accounts, (candidate) => candidate.recordId === account.recordId);
+        if (!exact)
+            return null;
+        const live = targetManager.getAccountByIndex(exact.index);
+        return live ? { accountManager: targetManager, account: live } : null;
+    }
+    const compatibleAccounts = accounts.filter((candidate) => !hasWorkspaceConflict(account, candidate));
+    const sourceEmail = sanitizeEmail(account.email);
+    const snapshot = uniqueMatch(compatibleAccounts, (candidate) => candidate.refreshToken === account.refreshToken) ??
+        (account.accountId && sourceEmail
+            ? uniqueMatch(compatibleAccounts, (candidate) => candidate.accountId === account.accountId && sanitizeEmail(candidate.email) === sourceEmail)
+            : null) ??
+        // accountId alone cannot merge logins: two members of one workspace share
+        // it. When both rows carry known-but-different emails they are different
+        // logins, and the in-flight refresh belongs to the removed one — let the
+        // caller restart selection instead of adopting the sibling's credentials.
+        (account.accountId
+            ? uniqueMatch(compatibleAccounts, (candidate) => candidate.accountId === account.accountId &&
+                (!sourceEmail ||
+                    !sanitizeEmail(candidate.email) ||
+                    sanitizeEmail(candidate.email) === sourceEmail))
+            : null) ??
+        (sourceEmail
+            ? uniqueMatch(compatibleAccounts, (candidate) => sanitizeEmail(candidate.email) === sourceEmail)
+            : null);
+    if (!snapshot)
+        return null;
+    const live = targetManager.getAccountByIndex(snapshot.index);
+    return live ? { accountManager: targetManager, account: live } : null;
+}
+/** @internal Test-only reset for process-global refresh commit/redirect state. */
+export function resetRuntimeRefreshStateForTests() {
+    runtimeRefreshCommitQueues = new WeakMap();
+    runtimeRefreshManagerRedirects = new WeakMap();
+}
+/** @internal Wait until every refresh-auth commit already targeting this manager settles. */
+export async function waitForRuntimeRefreshCommits(accountManager) {
+    const target = resolveRuntimeRefreshCommitManager(accountManager);
+    while (true) {
+        const queue = runtimeRefreshCommitQueues.get(target);
+        const pending = queue ? [...queue.values()] : [];
+        if (pending.length === 0)
+            return;
+        await Promise.allSettled(pending);
+    }
+}
+/** @internal Redirect late refresh state from requests that captured a superseded manager. */
+export function redirectRuntimeRefreshCommits(previousManager, replacementManager) {
+    if (previousManager === replacementManager)
+        return;
+    runtimeRefreshManagerRedirects.set(previousManager, resolveRuntimeRefreshCommitManager(replacementManager));
+}
 async function commitRefreshedAuthOnce(accountManager, account, auth) {
+    const commitManager = resolveRuntimeRefreshCommitManager(accountManager);
     const key = [
-        account.index,
+        account.recordId ?? "",
         account.accountId ?? "",
         account.email ?? "",
         account.refreshToken,
     ].join("\0");
-    let queue = runtimeRefreshCommitQueues.get(accountManager);
+    let queue = runtimeRefreshCommitQueues.get(commitManager);
     if (!queue) {
         queue = new Map();
-        runtimeRefreshCommitQueues.set(accountManager, queue);
+        runtimeRefreshCommitQueues.set(commitManager, queue);
     }
     const existing = queue.get(key);
     if (existing)
         return existing;
-    const pending = accountManager
+    const pending = commitManager
         .commitRefreshedAuth(account, auth)
+        .then((updatedAccount) => ({ accountManager: commitManager, account: updatedAccount }))
         .finally(() => queue?.delete(key));
     queue.set(key, pending);
     return pending;
 }
 /** @internal */
 export async function ensureFreshAccessToken(params) {
-    const { accountManager, account, family, model, now, tokenRefreshSkewMs, tokenInvalidationCooldownMs } = params;
-    if (hasUsableAccessToken(account, now, tokenRefreshSkewMs)) {
-        return { ok: true, accessToken: account.access ?? "", account };
+    const { accountManager, account, family, model, now, tokenRefreshSkewMs, tokenInvalidationCooldownMs, } = params;
+    const redirectedInitialManager = resolveRuntimeRefreshCommitManager(accountManager);
+    const resolvedInitialTarget = resolveRuntimeRefreshTarget(accountManager, account);
+    if (redirectedInitialManager !== accountManager && !resolvedInitialTarget) {
+        // The request captured a manager that has since been replaced, and the
+        // selected row no longer exists there. Even a still-fresh old token must not
+        // bypass that removal; restart selection on the replacement pool.
+        return {
+            ok: false,
+            retryable: true,
+            selectionRestart: true,
+            accountManager: redirectedInitialManager,
+        };
+    }
+    const initialTarget = resolvedInitialTarget ?? {
+        accountManager,
+        account,
+    };
+    if (initialTarget.account.enabled === false) {
+        return {
+            ok: false,
+            retryable: true,
+            selectionRestart: true,
+            accountManager: initialTarget.accountManager,
+        };
+    }
+    if (hasUsableAccessToken(initialTarget.account, now, tokenRefreshSkewMs)) {
+        return {
+            ok: true,
+            accessToken: initialTarget.account.access ?? "",
+            account: initialTarget.account,
+            accountManager: initialTarget.accountManager,
+        };
+    }
+    const refreshCredential = {
+        access: initialTarget.account.access,
+        refreshToken: initialTarget.account.refreshToken,
+        expires: initialTarget.account.expires,
+    };
+    const refreshResult = await queuedRefresh(refreshCredential.refreshToken);
+    const redirectedManager = resolveRuntimeRefreshCommitManager(initialTarget.accountManager);
+    const resolvedLatestTarget = resolveRuntimeRefreshTarget(initialTarget.accountManager, initialTarget.account);
+    if (redirectedManager !== initialTarget.accountManager && !resolvedLatestTarget) {
+        // The row that initiated this refresh disappeared while OAuth was in flight.
+        // Ignore the stale result rather than guessing another same-email workspace,
+        // but hand the caller the replacement manager so it can restart selection and
+        // discover a different healthy login from the current pool.
+        return {
+            ok: false,
+            retryable: true,
+            selectionRestart: true,
+            accountManager: redirectedManager,
+        };
+    }
+    const latestTarget = resolvedLatestTarget ?? initialTarget;
+    if (latestTarget.account.enabled === false) {
+        // An operator/runtime inventory change disabled this row while OAuth was in
+        // flight. Never let a successful refresh re-enable it in commitRefreshedAuth.
+        return {
+            ok: false,
+            retryable: true,
+            selectionRestart: true,
+            accountManager: latestTarget.accountManager,
+        };
+    }
+    const credentialChangedWhileRefreshing = latestTarget.account.access !== refreshCredential.access ||
+        latestTarget.account.refreshToken !== refreshCredential.refreshToken ||
+        latestTarget.account.expires !== refreshCredential.expires;
+    if (credentialChangedWhileRefreshing) {
+        // A concurrent refresh or external re-login repaired/replaced this credential.
+        // The result for the old OAuth tuple must never invalidate or overwrite it.
+        if (hasUsableAccessToken(latestTarget.account, now, tokenRefreshSkewMs)) {
+            return {
+                ok: true,
+                accessToken: latestTarget.account.access ?? "",
+                account: latestTarget.account,
+                accountManager: latestTarget.accountManager,
+            };
+        }
+        return {
+            ok: false,
+            retryable: true,
+            account: latestTarget.account,
+            accountManager: latestTarget.accountManager,
+        };
     }
-    const refreshResult = await queuedRefresh(account.refreshToken);
     if (refreshResult.type === "failed") {
-        accountManager.recordFailure(account, family, model);
-        accountManager.incrementAuthFailures(account);
+        latestTarget.accountManager.recordFailure(latestTarget.account, family, model);
+        latestTarget.accountManager.incrementAuthFailures(latestTarget.account);
         // If the refresh endpoint itself returns an explicit invalidation message
         // (e.g. Microsoft/Outlook SSO revokes the refresh token server-side), apply
         // the long cooldown and signal to the caller to stop rotating rather than
         // presenting other accounts' tokens from the same IP.
         const invalidated = isTokenInvalidationError(refreshResult.message ?? "");
         if (invalidated) {
-            accountManager.markAuthInvalidated(account);
+            latestTarget.accountManager.markAuthInvalidated(latestTarget.account);
         }
-        applyMonotonicAuthCooldown(accountManager, account, invalidated ? tokenInvalidationCooldownMs : DEFAULT_AUTH_FAILURE_COOLDOWN_MS);
-        accountManager.saveToDiskDebounced();
-        return { ok: false, retryable: isTokenRefreshRetryable(refreshResult), invalidated };
+        applyMonotonicAuthCooldown(latestTarget.accountManager, latestTarget.account, invalidated ? tokenInvalidationCooldownMs : DEFAULT_AUTH_FAILURE_COOLDOWN_MS);
+        latestTarget.accountManager.saveToDiskDebounced();
+        return {
+            ok: false,
+            retryable: isTokenRefreshRetryable(refreshResult),
+            invalidated,
+            account: latestTarget.account,
+            accountManager: latestTarget.accountManager,
+        };
     }
     const auth = {
         type: "oauth",
@@ -75,28 +259,35 @@ export async function ensureFreshAccessToken(params) {
         expires: refreshResult.expires,
     };
     try {
-        const updatedAccount = await commitRefreshedAuthOnce(accountManager, account, auth);
+        const committed = await commitRefreshedAuthOnce(latestTarget.accountManager, latestTarget.account, auth);
         // If commit succeeded, use the stored access token from the committed
-        // account — this also handles the dedup case where two concurrent callers
+        // account ? this also handles the dedup case where two concurrent callers
         // share one commit and both end up with the same committed token.
         // If commit returned null (account missing from storage after persist),
-        // fall back to the original account object but always use refreshResult.access:
-        // the original account may carry an expired token, and using that would
-        // cause a downstream 401 and incorrectly trigger invalidation-cooldown logic.
-        const resolvedAccount = updatedAccount ?? account;
-        const accessToken = updatedAccount !== null
-            ? (updatedAccount.access ?? refreshResult.access)
+        // fall back to the account that initiated the commit but always use
+        // refreshResult.access; that account may still carry an expired token.
+        const resolvedAccount = committed.account ?? latestTarget.account;
+        const resolvedManager = committed.account ? committed.accountManager : latestTarget.accountManager;
+        const accessToken = committed.account !== null
+            ? (committed.account.access ?? refreshResult.access)
             : refreshResult.access;
         return {
             ok: true,
             accessToken,
             account: resolvedAccount,
+            accountManager: resolvedManager,
         };
     }
     catch {
-        accountManager.recordFailure(account, family, model);
-        accountManager.markAccountCoolingDown(account, DEFAULT_AUTH_FAILURE_COOLDOWN_MS, "auth-failure");
-        accountManager.saveToDiskDebounced();
-        return { ok: false, retryable: true };
+        const failureTarget = resolveRuntimeRefreshTarget(latestTarget.accountManager, latestTarget.account) ?? latestTarget;
+        failureTarget.accountManager.recordFailure(failureTarget.account, family, model);
+        failureTarget.accountManager.markAccountCoolingDown(failureTarget.account, DEFAULT_AUTH_FAILURE_COOLDOWN_MS, "auth-failure");
+        failureTarget.accountManager.saveToDiskDebounced();
+        return {
+            ok: false,
+            retryable: true,
+            account: failureTarget.account,
+            accountManager: failureTarget.accountManager,
+        };
     }
 }