diff --git a/src/auth/store.test.ts b/src/auth/store.test.ts index 169773d25..f2b980793 100644 --- a/src/auth/store.test.ts +++ b/src/auth/store.test.ts @@ -114,6 +114,102 @@ describe("createAuthStore", () => { } }); + test("keeps a same-process burst of profile saves without shared-deadline loss", async () => { + const home = await mkdtemp(join(tmpdir(), "oauth-store-burst-")); + try { + const store = createAuthStore({ + filename: "test-auth.json", + settingsDirName: TEST_SETTINGS_DIR, + isTokens: isTestTokens, + }); + + // Without the per-path queue, a large same-process burst shares one lock + // deadline from invoke time and some waiters time out. With the queue, + // each save gets its own window and all land. + const names = Array.from( + { length: 50 }, + (_, index) => `profile-${String(index)}`, + ); + const results = await Promise.allSettled( + names.map((name) => + store.saveProfile( + { + name, + tokens: { + access: `access-${name}`, + refresh: `refresh-${name}`, + expiresAt: 1, + }, + createdAt: 1, + }, + home, + ), + ), + ); + + const failures = results.flatMap((result, index) => + result.status === "rejected" + ? [`${names[index]}: ${String(result.reason)}`] + : [], + ); + expect(failures).toEqual([]); + expect( + (await store.listProfiles(home)).map((profile) => profile.name), + ).toEqual([...names].sort()); + } finally { + await rm(home, { recursive: true, force: true }); + } + }); + + test("gives queued same-process writes their own lock window", async () => { + const home = await mkdtemp(join(tmpdir(), "oauth-store-queue-")); + try { + const store = createAuthStore({ + filename: "test-auth.json", + settingsDirName: TEST_SETTINGS_DIR, + isTokens: isTestTokens, + }); + await store.saveProfile( + { + name: "work", + tokens: { access: "a", refresh: "r", expiresAt: 1 }, + createdAt: 1, + }, + home, + ); + + // Hold the lock until the head of the same-process queue times out; the + // queued write must still get its own lock window after we release. + // `second` may already be polling when `first` rejects — release must land + // inside LOCK_TIMEOUT_MS of that handoff. + const lockPath = `${store.authPath(home)}.lock`; + await writeFile(lockPath, "foreign", { mode: 0o600 }); + + const first = store.updateTokens( + "work", + { access: "first", refresh: "r1", expiresAt: 2 }, + home, + ); + const second = store.updateTokens( + "work", + { access: "second", refresh: "r2", expiresAt: 3 }, + home, + ); + + await expect(first).rejects.toThrow( + "Timed out waiting for OAuth credential lock", + ); + await rm(lockPath, { force: true }); + await expect(second).resolves.toBeUndefined(); + + expect((await store.loadProfile("work", home))?.tokens.access).toBe( + "second", + ); + } finally { + await rm(home, { recursive: true, force: true }); + } + }); + test("round-trips profiles under an injected home and survives corrupt files", async () => { const home = await mkdtemp(join(tmpdir(), "oauth-store-")); try { diff --git a/src/auth/store.ts b/src/auth/store.ts index 349997663..903d50a33 100644 --- a/src/auth/store.ts +++ b/src/auth/store.ts @@ -18,7 +18,8 @@ export type { AuthProfile, BaseTokens }; // for the same provider, so credentials are keyed by a user-chosen profile name // within a single file. Tokens are credentials, so the file is owner-only (0o600) // and the directory 0o700. Writes go through a temp file + rename so a concurrent -// reader never observes a torn file. +// reader never observes a torn file. Same-process writers also queue per auth path +// so each lock wait starts its own deadline. export interface AuthStore { authPath: (home?: string) => string; @@ -48,6 +49,15 @@ interface AuthFile { const LOCK_RETRY_MS = 25; const LOCK_TIMEOUT_MS = 1_000; +// Per-call unique temp (pid + counter). Matches mcp/auth-store — pid alone is not +// unique per call if writeAuthFile ever overlaps in-process. +let tmpWriteCounter = 0; + +// Same-process ops on one auth file queue here so a caller's lock deadline +// starts when it actually runs, not when it was invoked — otherwise one lock +// held past LOCK_TIMEOUT_MS fails the whole burst, not just the first waiter. +const updateChains = new Map>(); + const AuthFileShape = type({ profiles: "Record", }); @@ -115,7 +125,7 @@ export function createAuthStore( ): Promise { const path = authPath(home); await mkdir(dirname(path), { recursive: true, mode: 0o700 }); - const tmp = `${path}.${String(process.pid)}.tmp`; + const tmp = `${path}.${process.pid}.${(tmpWriteCounter += 1)}.tmp`; await writeFile(tmp, JSON.stringify(file, null, 2), { mode: 0o600 }); await rename(tmp, path); } @@ -158,6 +168,26 @@ export function createAuthStore( } } + function enqueueAuthFileOp( + home: string, + op: () => Promise, + ): Promise { + const path = authPath(home); + const previous = updateChains.get(path) ?? Promise.resolve(); + const run = previous.then( + () => withAuthFileLock(home, op), + () => withAuthFileLock(home, op), + ); + updateChains.set( + path, + run.then( + () => undefined, + () => undefined, + ), + ); + return run; + } + return { authPath, async listProfiles( @@ -179,7 +209,7 @@ export function createAuthStore( profile: AuthProfile, home: string = homedir(), ): Promise { - await withAuthFileLock(home, async () => { + await enqueueAuthFileOp(home, async () => { const file = await readAuthFile(home); file.profiles[profile.name] = profile; await writeAuthFile(file, home); @@ -192,7 +222,7 @@ export function createAuthStore( tokens: TTokens, home: string = homedir(), ): Promise { - await withAuthFileLock(home, async () => { + await enqueueAuthFileOp(home, async () => { const file = await readAuthFile(home); const existing = file.profiles[name]; if (existing === undefined) return; @@ -204,7 +234,7 @@ export function createAuthStore( name: string | undefined, home: string = homedir(), ): Promise { - return withAuthFileLock(home, async () => { + return enqueueAuthFileOp(home, async () => { const file = await readAuthFile(home); if (name === undefined) { const removed = Object.keys(file.profiles);