Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
96 changes: 96 additions & 0 deletions src/auth/store.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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<TestTokens>({
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<TestTokens>({
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 {
Expand Down
40 changes: 35 additions & 5 deletions src/auth/store.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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<TTokens extends BaseTokens> {
authPath: (home?: string) => string;
Expand Down Expand Up @@ -48,6 +49,15 @@ interface AuthFile<TTokens extends BaseTokens> {
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<string, Promise<unknown>>();

const AuthFileShape = type({
profiles: "Record<string, unknown>",
});
Expand Down Expand Up @@ -115,7 +125,7 @@ export function createAuthStore<TTokens extends BaseTokens>(
): Promise<void> {
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);
}
Expand Down Expand Up @@ -158,6 +168,26 @@ export function createAuthStore<TTokens extends BaseTokens>(
}
}

function enqueueAuthFileOp<TResult>(
home: string,
op: () => Promise<TResult>,
): Promise<TResult> {
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(
Expand All @@ -179,7 +209,7 @@ export function createAuthStore<TTokens extends BaseTokens>(
profile: AuthProfile<TTokens>,
home: string = homedir(),
): Promise<void> {
await withAuthFileLock(home, async () => {
await enqueueAuthFileOp(home, async () => {
const file = await readAuthFile(home);
file.profiles[profile.name] = profile;
await writeAuthFile(file, home);
Expand All @@ -192,7 +222,7 @@ export function createAuthStore<TTokens extends BaseTokens>(
tokens: TTokens,
home: string = homedir(),
): Promise<void> {
await withAuthFileLock(home, async () => {
await enqueueAuthFileOp(home, async () => {
const file = await readAuthFile(home);
const existing = file.profiles[name];
if (existing === undefined) return;
Expand All @@ -204,7 +234,7 @@ export function createAuthStore<TTokens extends BaseTokens>(
name: string | undefined,
home: string = homedir(),
): Promise<string[]> {
return withAuthFileLock(home, async () => {
return enqueueAuthFileOp(home, async () => {
const file = await readAuthFile(home);
if (name === undefined) {
const removed = Object.keys(file.profiles);
Expand Down
Loading