diff --git a/src/oauth/index.ts b/src/oauth/index.ts index 0162492f9a..8a18a0b007 100644 --- a/src/oauth/index.ts +++ b/src/oauth/index.ts @@ -322,6 +322,13 @@ export class OAuthReauthIdentityUnverifiedError extends Error { } } +class OAuthLoginSupersededError extends Error { + constructor() { + super("OAuth login was superseded before credential persistence"); + this.name = "OAuthLoginSupersededError"; + } +} + /** Project arbitrary OAuth failures onto the small, stable public error vocabulary. */ export function publicOAuthAuthenticationErrorMessage(error: unknown): string { if (error instanceof OAuthMutationBusyError) { @@ -1096,6 +1103,7 @@ interface RunLoginDeps { settleKiroLoginTransaction?: typeof settleKiroLoginTransaction; removeAccount?: typeof removeAccount; setActiveAccount?: typeof setActiveAccount; + assertCurrentOwner?: () => void; } /** Roll back only accounts created by this forced login, preserving concurrent refreshes of others. */ @@ -1145,6 +1153,7 @@ export async function runLogin( const cred: OAuthCredentials = rawCred.source ? rawCred : { ...rawCred, source: "oauth" }; const settleKiroTransaction = deps.settleKiroLoginTransaction ?? settleKiroLoginTransaction; try { + deps.assertCurrentOwner?.(); // Validate the provider row before credential persistence. A namespace claimed during the // credential write is handled again below before the latest row is re-upserted. if (provider !== "chatgpt") { @@ -1165,10 +1174,13 @@ export async function runLogin( if (!identityMatches) { throw new OAuthReauthIdentityMismatchError(); } - await (deps.saveAccountCredential ?? saveAccountCredential)(provider, opts.reauthAccountId, cred); + await (deps.saveAccountCredential ?? saveAccountCredential)(provider, opts.reauthAccountId, cred, { + assertBeforePersist: deps.assertCurrentOwner, + }); } else { await (deps.saveCredential ?? saveCredential)(provider, cred, { preserveIdentityless: opts?.forceLogin === true, + assertBeforePersist: deps.assertCurrentOwner, }); } if (provider !== "chatgpt") { @@ -1235,6 +1247,7 @@ export async function runLogin( */ const loginState = new Map(); const loginAbort = new Map(); +const kiroLoginSettling = new Set(); /** Pending paste for a login in progress: either a waiter or a stashed early submission. */ interface ManualCodeSlot { @@ -1403,13 +1416,14 @@ export async function startLoginFlow( const def = OAUTH_PROVIDERS[provider]; if (!def) throw new UnsupportedOAuthProviderError(provider); const existing = loginState.get(provider); - if (existing && !existing.done) { + if ((existing && !existing.done) || (provider === "kiro" && kiroLoginSettling.has(provider))) { throw new Error(`A login for ${provider} is already in progress`); } clearManualCodeSlot(provider); loginState.set(provider, { done: false }); const abort = new AbortController(); loginAbort.set(provider, abort); + if (provider === "kiro") kiroLoginSettling.add(provider); return new Promise((resolve, reject) => { let urlResolved = false; const ctrl: OAuthController = { @@ -1460,7 +1474,10 @@ export async function startLoginFlow( }; // Background: runLogin persists the credential + provider entry to disk. The lifecycle hook // lets a long-lived server config adopt that settled state before clients observe done=true. - void runLogin(provider, ctrl, opts).then( + const assertCurrentOwner = (): void => { + if (loginAbort.get(provider) !== abort) throw new OAuthLoginSupersededError(); + }; + void runLogin(provider, ctrl, opts, { assertCurrentOwner }).then( () => settle(), (e: unknown) => settle(e), ).catch((e: unknown) => { @@ -1471,6 +1488,8 @@ export async function startLoginFlow( const msg = publicOAuthAuthenticationErrorMessage(e); loginState.set(provider, { done: true, error: msg }); if (!urlResolved) reject(e); + }).finally(() => { + if (provider === "kiro") kiroLoginSettling.delete(provider); }); }); } diff --git a/src/oauth/store.ts b/src/oauth/store.ts index 31924afc33..ed2d3bf87f 100644 --- a/src/oauth/store.ts +++ b/src/oauth/store.ts @@ -465,10 +465,11 @@ function serializeMutation(work: () => Promise, retainedValues: readonly u drainOAuthMutations(); return result; } -export function mutateStore(fn:(store:AuthStore)=>T|Promise, retainedValues: readonly unknown[] = [], options?: { waitMs?: number }):Promise{return serializeMutation(async()=>{const guard=await createOAuthFileLock({path:getAuthStoreLockPath(),staleAfterMs:30000}).acquire();try{ +export function mutateStore(fn:(store:AuthStore)=>T|Promise, retainedValues: readonly unknown[] = [], options?: { waitMs?: number; assertBeforePersist?: () => void }):Promise{return serializeMutation(async()=>{const guard=await createOAuthFileLock({path:getAuthStoreLockPath(),staleAfterMs:30000}).acquire();try{ const { store, hadLegacy } = loadAuthStoreInternal(); if (hadLegacy) backupLegacyOnce(); const result = await fn(store); + options?.assertBeforePersist?.(); persist(store); return result; }finally{guard.release();}}, retainedValues, options?.waitMs); @@ -491,7 +492,7 @@ export function getCredential(provider: string): OAuthCredentials | null { export async function saveCredential( provider: string, cred: OAuthCredentials, - opts: { preserveIdentityless?: boolean } = {}, + opts: { preserveIdentityless?: boolean; assertBeforePersist?: () => void } = {}, ): Promise { const safe = normalizeCredential(cred); if (!safe) return; @@ -542,7 +543,7 @@ export async function saveCredential( set.accounts.push({ id, credential: safe, addedAt: Date.now() }); set.activeAccountId = id; } - }, [provider, safe]); + }, [provider, safe], { assertBeforePersist: opts.assertBeforePersist }); } /** @@ -632,7 +633,12 @@ export function getAccountCredential(provider: string, accountId: string): OAuth } /** Persist a refreshed credential for a SPECIFIC account without touching activeAccountId. */ -export async function saveAccountCredential(provider: string, accountId: string, cred: OAuthCredentials): Promise { +export async function saveAccountCredential( + provider: string, + accountId: string, + cred: OAuthCredentials, + opts: { assertBeforePersist?: () => void } = {}, +): Promise { const safe = normalizeCredential(cred); if (!safe) return; await mutateStore(store => { @@ -640,7 +646,7 @@ export async function saveAccountCredential(provider: string, accountId: string, if (!account) return; account.credential = safe; delete account.needsReauth; - }, [provider, accountId, safe]); + }, [provider, accountId, safe], { assertBeforePersist: opts.assertBeforePersist }); } export async function setActiveAccount(provider: string, accountId: string): Promise { diff --git a/tests/oauth-public-surface.test.ts b/tests/oauth-public-surface.test.ts index 788e5e4d56..a0ae4617f9 100644 --- a/tests/oauth-public-surface.test.ts +++ b/tests/oauth-public-surface.test.ts @@ -448,6 +448,113 @@ describe("legacy ChatGPT OAuth public-surface exclusion", () => { } }); + test("a superseded OAuth flow cannot commit after its replacement owns the provider", async () => { + saveConfig(config()); + const originalLogin = OAUTH_PROVIDERS.xai.login; + let loginCalls = 0; + OAUTH_PROVIDERS.xai.login = async (ctrl) => { + loginCalls += 1; + const call = loginCalls; + ctrl.onAuth({ url: `https://auth.example.test/${call}`, deviceCode: `flow-${call}` }); + return { + access: `access-${call}`, + refresh: `refresh-${call}`, + accountId: `account-${call}`, + email: `account-${call}@example.test`, + expires: Date.now() + 60_000, + }; + }; + + let releaseHead!: () => void; + let signalHeadStarted!: () => void; + const headStarted = new Promise(resolve => { signalHeadStarted = resolve; }); + const headGate = new Promise(resolve => { releaseHead = resolve; }); + const blockingMutation = oauthStore.mutateStore(async () => { + signalHeadStarted(); + await headGate; + }); + + const waitForMutationCount = async (minimum: number): Promise => { + for (let attempt = 0; attempt < 200; attempt += 1) { + if (oauthStore.oauthMutationTailSnapshot().active >= minimum) return; + await Bun.sleep(5); + } + throw new Error(`OAuth mutation queue did not reach ${minimum} active rows`); + }; + + try { + await headStarted; + await startLoginFlow("xai"); + await waitForMutationCount(2); + expect(cancelLoginFlow("xai")).toBe(true); + + await startLoginFlow("xai"); + await waitForMutationCount(3); + releaseHead(); + await blockingMutation; + + const status = await waitForOAuthDone("xai"); + expect(status).toMatchObject({ done: true, loggedIn: true }); + expect(getCredential("xai")).toMatchObject({ + access: "access-2", + accountId: "account-2", + }); + expect(oauthStore.getAccountSet("xai")?.accounts.map(account => account.credential.accountId)) + .toEqual(["account-2"]); + } finally { + releaseHead(); + await blockingMutation.catch(() => {}); + OAUTH_PROVIDERS.xai.login = originalLogin; + clearLoginState("xai"); + } + }); + + test("Kiro does not start a replacement until the canceled external CLI flow settles", async () => { + saveConfig(config()); + const originalLogin = OAUTH_PROVIDERS.kiro.login; + let loginCalls = 0; + OAUTH_PROVIDERS.kiro.login = async (ctrl) => { + loginCalls += 1; + const call = loginCalls; + ctrl.onAuth({ url: "", deviceCode: `kiro-flow-${call}` }); + if (call === 1) { + await new Promise((_, reject) => { + ctrl.signal.addEventListener("abort", () => reject(new Error("Kiro login cancelled")), { once: true }); + }); + } + return { + access: "kiro-replacement-access", + refresh: "kiro-replacement-refresh", + accountId: "kiro-replacement-account", + email: "kiro-replacement@example.test", + expires: Date.now() + 60_000, + }; + }; + + try { + await startLoginFlow("kiro"); + expect(cancelLoginFlow("kiro")).toBe(true); + await expect(startLoginFlow("kiro")).rejects.toThrow("A login for kiro is already in progress"); + + let replacement: Awaited> | undefined; + for (let attempt = 0; attempt < 200; attempt += 1) { + try { + replacement = await startLoginFlow("kiro"); + break; + } catch (error) { + if (!(error instanceof Error) || !error.message.includes("already in progress")) throw error; + await Bun.sleep(5); + } + } + expect(replacement).toMatchObject({ deviceCode: "kiro-flow-2" }); + expect(await waitForOAuthDone("kiro")).toMatchObject({ done: true, loggedIn: true }); + expect(loginCalls).toBe(2); + } finally { + OAUTH_PROVIDERS.kiro.login = originalLogin; + clearLoginState("kiro"); + } + }); + test("management OAuth safely reconciles live config after a late namespace claim", async () => { const liveConfig = config(); liveConfig.hostname = "0.0.0.0";