From 10b5eaa37e889982c53533369e43f7ab20e7b2dd Mon Sep 17 00:00:00 2001 From: Pragati Agrawal Date: Wed, 29 Jul 2026 22:31:16 +0530 Subject: [PATCH 1/4] [glean-vnext] Fix intermittent re-auth from cross-process refresh-token rotation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Port of gleanwork/glean-plugins-vnext#44 (squashed; full history and E2E evidence there). Each host session runs its own plugin process sharing one credentials file. The Glean OAuth server rotates refresh tokens on every refresh with no grace period, so when one session refreshes, every other session's in-memory copy is revoked; their next refresh gets invalid_grant, the SDK wipes the SHARED store, and the user sees [SETUP_REQUIRED] — plus every other live session dies with them. Fixes (E2E-verified on an experimental pod against real prod /oauth — bug reproduced on demand with the old build, silent recovery in both race shapes with this change): - tokens()/syncTokensFromDisk: mtime-guarded re-read of the shared store so a sibling's rotated grant is picked up before the SDK refreshes. - invalidateCredentials('tokens'): adopt a newer on-disk token instead of wiping — with a grace-window poll (GLEAN_ROTATION_GRACE_MS, 2s) because the loser's invalid_grant usually lands milliseconds before the winner's write. - Connect-level sibling-refresh retry: concurrent refreshes of the same grant make fosite fail the loser with invalid_request (not invalid_grant — observed live), which the SDK rethrows raw; recognize refresh-shaped failures, wait out the grace window, retry once. - saveCredentials: temp-file + rename so concurrent writers can't leave a torn store that parses as wiped. Co-Authored-By: Claude Opus 4.8 --- sources/glean-vnext/src/auth-provider.ts | 113 +++++++++++- sources/glean-vnext/src/remote-client.ts | 43 ++++- sources/glean-vnext/src/token-store.ts | 18 +- .../glean-vnext/tests/auth-provider.test.ts | 141 +++++++++++++++ .../tests/remote-client-auth-retry.test.ts | 162 ++++++++++++++++++ sources/glean-vnext/tests/token-store.test.ts | 24 ++- 6 files changed, 492 insertions(+), 9 deletions(-) create mode 100644 sources/glean-vnext/tests/remote-client-auth-retry.test.ts diff --git a/sources/glean-vnext/src/auth-provider.ts b/sources/glean-vnext/src/auth-provider.ts index 4c47d0b..cebf440 100644 --- a/sources/glean-vnext/src/auth-provider.ts +++ b/sources/glean-vnext/src/auth-provider.ts @@ -7,10 +7,32 @@ import type { import { execFile, spawn } from "node:child_process"; import { platform } from "node:os"; import { getCallbackUrl, setExpectedState } from "./auth-callback-server.js"; -import { clearCredentials, loadCredentials, saveCredentials } from "./token-store.js"; +import { + clearCredentials, + credentialsMtimeMs, + loadCredentials, + saveCredentials, +} from "./token-store.js"; export type InvalidationScope = "all" | "client" | "tokens" | "verifier"; +// Grace window for a sibling's in-flight refresh to land on disk. +const ROTATION_GRACE_MS_DEFAULT = 2000; +const ROTATION_POLL_MS = 100; + +function rotationGraceMs(): number { + const raw = process.env.GLEAN_ROTATION_GRACE_MS; + if (raw !== undefined) { + const parsed = Number.parseInt(raw, 10); + if (Number.isFinite(parsed) && parsed >= 0) return parsed; + } + return ROTATION_GRACE_MS_DEFAULT; +} + +function sleep(ms: number): Promise { + return new Promise((resolve) => setTimeout(resolve, ms)); +} + /** * Open `url` in the user's default browser. Used for the self-open sign-in * path when the client does not support URL-mode elicitation (where the client @@ -47,6 +69,8 @@ export class GleanOAuthClientProvider implements OAuthClientProvider { // explicitly invalidating. Used to detect when a previous auth URL didn't // complete — likely because the server rejected the (stale) client_id. private _authUrlPending = false; + // mtime at last read; detects sibling rewrites of the shared store. + private _credentialsMtimeMs: number | undefined; authorizationUrl: string | undefined; @@ -64,6 +88,87 @@ export class GleanOAuthClientProvider implements OAuthClientProvider { this._tokens = stored.tokens as OAuthTokens | undefined; this._clientInfo = stored.clientInfo as OAuthClientInformationMixed | undefined; } + this._credentialsMtimeMs = credentialsMtimeMs(); + } + + // Re-read the shared store after a sibling process rewrites it, so we use + // the rotated grant instead of a stale in-memory copy. + private syncTokensFromDisk(): void { + const mtimeMs = credentialsMtimeMs(); + if (mtimeMs === undefined) return; + if ( + this._credentialsMtimeMs !== undefined && + mtimeMs <= this._credentialsMtimeMs + ) { + return; + } + const stored = loadCredentials(); + this._credentialsMtimeMs = mtimeMs; + if (!stored) return; + if (stored.tokens) { + this._tokens = stored.tokens as OAuthTokens; + } + if (stored.clientInfo) { + this._clientInfo = stored.clientInfo as OAuthClientInformationMixed; + } + } + + // On invalid_grant, adopt a sibling's newer on-disk token instead of + // clearing. Returns false when nothing newer exists. + private adoptNewerTokenFromDisk(): boolean { + const diskMtime = credentialsMtimeMs(); + if ( + diskMtime === undefined || + this._credentialsMtimeMs === undefined || + diskMtime <= this._credentialsMtimeMs + ) { + return false; + } + const stored = loadCredentials(); + const diskTokens = stored?.tokens as OAuthTokens | undefined; + if ( + !diskTokens?.access_token || + diskTokens.access_token === this._tokens?.access_token + ) { + return false; + } + this._tokens = diskTokens; + this._credentialsMtimeMs = diskMtime; + if (stored?.clientInfo) { + this._clientInfo = stored.clientInfo as OAuthClientInformationMixed; + } + console.error( + "[auth] invalid_grant, but a newer token is on disk " + + "(sibling refresh) — adopting it instead of clearing", + ); + return true; + } + + // Poll briefly for the race winner's write before clearing. Skipped when + // no refresh token was held (no race possible). + private async adoptNewerTokenWithGrace(): Promise { + if (this.adoptNewerTokenFromDisk()) return true; + if (!this._tokens?.refresh_token) return false; + const deadline = Date.now() + rotationGraceMs(); + while (Date.now() < deadline) { + await sleep(ROTATION_POLL_MS); + if (this.adoptNewerTokenFromDisk()) return true; + } + return false; + } + + // Grace-bounded wait for a sibling's refresh; covers failures the SDK does + // not route through invalidateCredentials (e.g. invalid_request collisions). + async waitForSiblingRefresh( + previousAccessToken: string | undefined, + ): Promise { + const deadline = Date.now() + rotationGraceMs(); + for (;;) { + const current = this.tokens()?.access_token; + if (current && current !== previousAccessToken) return true; + if (Date.now() >= deadline) return false; + await sleep(ROTATION_POLL_MS); + } } get redirectUrl(): string { @@ -84,9 +189,11 @@ export class GleanOAuthClientProvider implements OAuthClientProvider { saveClientInformation(info: OAuthClientInformationMixed): void { this._clientInfo = info; saveCredentials(this._tokens, this._clientInfo); + this._credentialsMtimeMs = credentialsMtimeMs(); } tokens(): OAuthTokens | undefined { + this.syncTokensFromDisk(); return this._tokens; } @@ -94,6 +201,8 @@ export class GleanOAuthClientProvider implements OAuthClientProvider { this._tokens = tokens; this._authUrlPending = false; saveCredentials(this._tokens, this._clientInfo); + // Own write must not look like a sibling change. + this._credentialsMtimeMs = credentialsMtimeMs(); this.onTokensChanged?.(tokens); } @@ -113,6 +222,8 @@ export class GleanOAuthClientProvider implements OAuthClientProvider { saveCredentials(this._tokens, undefined); break; case "tokens": + // Usually a sibling's rotation — try adopting before clearing. + if (await this.adoptNewerTokenWithGrace()) return; this._tokens = undefined; saveCredentials(undefined, this._clientInfo); break; diff --git a/sources/glean-vnext/src/remote-client.ts b/sources/glean-vnext/src/remote-client.ts index 954cf52..d825947 100644 --- a/sources/glean-vnext/src/remote-client.ts +++ b/sources/glean-vnext/src/remote-client.ts @@ -144,6 +144,7 @@ export async function createRemoteClient( serverUrl: string, opts: RemoteClientOptions, chatSessionId?: string, + authRetry = false, ): Promise { const authProvider = opts.authProvider; @@ -189,14 +190,44 @@ export async function createRemoteClient( { capabilities: {} }, ); + // Snapshot to detect a sibling's refresh between connect and failure. + const accessTokenAtConnect = authProvider?.tokens()?.access_token; + const transport = buildTransport(serverUrl, opts, chatSessionId); try { await client.connect(transport); } catch (error) { - if (error instanceof UnauthorizedError && authProvider?.authorizationUrl) { - pendingTransport = transport; - throw new AuthRequiredError(authProvider.authorizationUrl); + if (error instanceof UnauthorizedError && authProvider) { + const refreshedAccessToken = authProvider.tokens()?.access_token; + if ( + !authRetry && + refreshedAccessToken && + refreshedAccessToken !== accessTokenAtConnect + ) { + console.error( + "[auth] Auth failed but a newer token is on disk " + + "(sibling refresh) — retrying once", + ); + return createRemoteClient(serverUrl, opts, chatSessionId, true); + } + if (authProvider.authorizationUrl) { + pendingTransport = transport; + throw new AuthRequiredError(authProvider.authorizationUrl); + } + } + // Concurrent-refresh losers get errors the SDK rethrows raw (e.g. fosite + // invalid_request); retry once if a sibling's grant lands in the grace window. + if ( + authProvider && + !authRetry && + isLikelyRefreshFailure(error) && + (await authProvider.waitForSiblingRefresh(accessTokenAtConnect)) + ) { + console.error( + "[auth] Refresh failed but a sibling refreshed — retrying with its token", + ); + return createRemoteClient(serverUrl, opts, chatSessionId, true); } throw error; } @@ -204,6 +235,12 @@ export async function createRemoteClient( return client; } +// Match broadly; the caller's disk re-check gates the actual retry. +function isLikelyRefreshFailure(error: unknown): boolean { + const msg = error instanceof Error ? error.message : String(error); + return /refresh|invalid_grant|invalid_request|oauth/i.test(msg); +} + export async function callRemoteTool( client: Client, name: string, diff --git a/sources/glean-vnext/src/token-store.ts b/sources/glean-vnext/src/token-store.ts index e41c2fa..535a07d 100644 --- a/sources/glean-vnext/src/token-store.ts +++ b/sources/glean-vnext/src/token-store.ts @@ -28,6 +28,18 @@ export function loadCredentials(): StoredCredentials | undefined { } } +/** + * mtime of the credentials file (epoch ms), or undefined if unreadable. + * Cheap change probe: a single stat, no read + parse. + */ +export function credentialsMtimeMs(): number | undefined { + try { + return fs.statSync(credentialsFile()).mtimeMs; + } catch { + return undefined; + } +} + export function saveCredentials(tokens: unknown, clientInfo: unknown): void { try { const filePath = credentialsFile(); @@ -35,11 +47,13 @@ export function saveCredentials(tokens: unknown, clientInfo: unknown): void { fs.mkdirSync(dir, { recursive: true, mode: DIR_MODE }); fs.chmodSync(dir, DIR_MODE); const data: StoredCredentials = { tokens, clientInfo }; - fs.writeFileSync(filePath, JSON.stringify(data, null, 2), { + // Temp-file + rename: concurrent readers never see a half-written store. + const tmpPath = `${filePath}.${process.pid}.tmp`; + fs.writeFileSync(tmpPath, JSON.stringify(data, null, 2), { encoding: "utf-8", mode: FILE_MODE, }); - fs.chmodSync(filePath, FILE_MODE); + fs.renameSync(tmpPath, filePath); } catch (err) { const msg = err instanceof Error ? err.message : String(err); console.error(`[auth] Failed to persist credentials: ${msg}`); diff --git a/sources/glean-vnext/tests/auth-provider.test.ts b/sources/glean-vnext/tests/auth-provider.test.ts index ba181f0..3e4e032 100644 --- a/sources/glean-vnext/tests/auth-provider.test.ts +++ b/sources/glean-vnext/tests/auth-provider.test.ts @@ -28,11 +28,15 @@ describe("GleanOAuthClientProvider", () => { beforeEach(() => { delete process.env.PLUGIN_DATA_DIR; + // Skip the rotation grace window by default so invalidation tests don't + // wait out the real 2s poll; the grace test overrides this explicitly. + process.env.GLEAN_ROTATION_GRACE_MS = "0"; fs.rmSync(gleanDir, { recursive: true, force: true }); vi.clearAllMocks(); }); afterEach(() => { + delete process.env.GLEAN_ROTATION_GRACE_MS; fs.rmSync(gleanDir, { recursive: true, force: true }); }); @@ -74,6 +78,143 @@ describe("GleanOAuthClientProvider", () => { expect(raw.tokens.access_token).toBe("new_tok"); }); + // --- Cross-process sync: tokens() must pick up a sibling's rewrite. --- + + const credFile = path.join(gleanDir, "mcp-credentials.json"); + + function writeCredFileNewer(tokens: unknown, clientInfo?: unknown): void { + fs.mkdirSync(gleanDir, { recursive: true }); + fs.writeFileSync(credFile, JSON.stringify({ tokens, clientInfo })); + // Guarantee a strictly-newer mtime than any prior read, independent of + // filesystem timestamp resolution. + const future = new Date(Date.now() + 10_000); + fs.utimesSync(credFile, future, future); + } + + it("tokens() adopts a newer token written by another process", () => { + fs.mkdirSync(gleanDir, { recursive: true }); + fs.writeFileSync( + credFile, + JSON.stringify({ + tokens: { access_token: "T0", refresh_token: "R0" }, + clientInfo: { client_id: "cid" }, + }), + ); + const provider = new GleanOAuthClientProvider(); + expect(provider.tokens()?.access_token).toBe("T0"); + + // Sibling refreshes: new access + rotated refresh token on disk. + writeCredFileNewer( + { access_token: "T1", refresh_token: "R1" }, + { client_id: "cid" }, + ); + + expect(provider.tokens()?.access_token).toBe("T1"); + expect(provider.tokens()?.refresh_token).toBe("R1"); + }); + + it("tokens() keeps the in-memory token when the file is deleted", () => { + fs.mkdirSync(gleanDir, { recursive: true }); + fs.writeFileSync( + credFile, + JSON.stringify({ tokens: { access_token: "T0" }, clientInfo: {} }), + ); + const provider = new GleanOAuthClientProvider(); + expect(provider.tokens()?.access_token).toBe("T0"); + + // Transient disappearance / another process mid-write — don't self-evict. + fs.rmSync(credFile, { force: true }); + expect(provider.tokens()?.access_token).toBe("T0"); + }); + + it("tokens() does not adopt a rewrite that carries no tokens", () => { + fs.mkdirSync(gleanDir, { recursive: true }); + fs.writeFileSync( + credFile, + JSON.stringify({ tokens: { access_token: "T0" }, clientInfo: {} }), + ); + const provider = new GleanOAuthClientProvider(); + expect(provider.tokens()?.access_token).toBe("T0"); + + // A client-only rewrite (tokens dropped) must not log us out in-memory. + writeCredFileNewer(undefined, { client_id: "cid" }); + expect(provider.tokens()?.access_token).toBe("T0"); + }); + + it("invalidateCredentials('tokens') adopts a sibling's newer token instead of wiping the store", async () => { + fs.mkdirSync(gleanDir, { recursive: true }); + fs.writeFileSync( + credFile, + JSON.stringify({ + tokens: { access_token: "T0", refresh_token: "R0" }, + clientInfo: { client_id: "cid" }, + }), + ); + const provider = new GleanOAuthClientProvider(); + expect(provider.tokens()?.access_token).toBe("T0"); + + // A sibling refreshed + rotated: fresh grant now on disk with a newer mtime. + writeCredFileNewer( + { access_token: "T1", refresh_token: "R1" }, + { client_id: "cid" }, + ); + + // The SDK calls this on invalid_grant. It must NOT clear — the failure was + // just our stale token; adopt the sibling's fresh one and leave it on disk. + await provider.invalidateCredentials("tokens"); + + expect(provider.tokens()?.access_token).toBe("T1"); + expect(provider.tokens()?.refresh_token).toBe("R1"); + const raw = JSON.parse(fs.readFileSync(credFile, "utf-8")); + expect(raw.tokens.access_token).toBe("T1"); // not clobbered with undefined + }); + + it("invalidateCredentials('tokens') clears when there is no newer token on disk", async () => { + const provider = new GleanOAuthClientProvider(); + provider.saveTokens({ access_token: "T0", refresh_token: "R0" } as any); + expect(provider.tokens()?.access_token).toBe("T0"); + + // No sibling write since our snapshot → a genuine invalidation → clear. + await provider.invalidateCredentials("tokens"); + + expect(provider.tokens()).toBeUndefined(); + const raw = JSON.parse(fs.readFileSync(credFile, "utf-8")); + expect(raw.tokens).toBeUndefined(); + }); + + it("invalidateCredentials('tokens') adopts a token that lands during the grace window", async () => { + // The winner's write lands just after the loser's invalid_grant. + process.env.GLEAN_ROTATION_GRACE_MS = "2000"; + const provider = new GleanOAuthClientProvider(); + provider.saveTokens({ access_token: "T0", refresh_token: "R0" } as any); + + const invalidation = provider.invalidateCredentials("tokens"); + // Sibling's write lands mid-window. + setTimeout(() => { + writeCredFileNewer( + { access_token: "T1", refresh_token: "R1" }, + { client_id: "cid" }, + ); + }, 150); + await invalidation; + + expect(provider.tokens()?.access_token).toBe("T1"); + const raw = JSON.parse(fs.readFileSync(credFile, "utf-8")); + expect(raw.tokens.access_token).toBe("T1"); // not clobbered with undefined + }); + + it("skips the grace window when no refresh token was held (no race possible)", async () => { + process.env.GLEAN_ROTATION_GRACE_MS = "5000"; + const provider = new GleanOAuthClientProvider(); + provider.saveTokens({ access_token: "T0" } as any); // no refresh_token + + const start = Date.now(); + await provider.invalidateCredentials("tokens"); + + expect(Date.now() - start).toBeLessThan(1000); // no 5s poll + expect(provider.tokens()).toBeUndefined(); + }); + it("saveClientInformation persists to disk", () => { const provider = new GleanOAuthClientProvider(); const info = { client_id: "cid", client_secret: "sec" } as any; diff --git a/sources/glean-vnext/tests/remote-client-auth-retry.test.ts b/sources/glean-vnext/tests/remote-client-auth-retry.test.ts new file mode 100644 index 0000000..f730cc7 --- /dev/null +++ b/sources/glean-vnext/tests/remote-client-auth-retry.test.ts @@ -0,0 +1,162 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { UnauthorizedError } from "@modelcontextprotocol/sdk/client/auth.js"; + +// Control client.connect() across (re)tries. +const { connectMock } = vi.hoisted(() => ({ connectMock: vi.fn() })); + +vi.mock("@modelcontextprotocol/sdk/client/index.js", () => ({ + Client: class { + async connect(...args: unknown[]) { + return connectMock(...args); + } + }, +})); + +// Keep buildTransport cheap and side-effect free. +vi.mock("@modelcontextprotocol/sdk/client/streamableHttp.js", () => ({ + StreamableHTTPClientTransport: class { + constructor() {} + async close() {} + }, +})); + +const { createRemoteClient, AuthRequiredError } = await import( + "../src/remote-client.js" +); + +/** + * Minimal OAuthClientProvider stand-in. tokens() returns the next value in + * `seq` on each call, mirroring how the real provider re-reads disk: the + * pre-connect snapshot, then the value after a sibling may have rewritten it. + */ +function makeProvider(seq: Array<{ access_token?: string } | undefined>) { + let i = 0; + return { + tokens() { + const t = seq[Math.min(i, seq.length - 1)]; + i += 1; + return t; + }, + authorizationUrl: "https://example.com/oauth/authorize?state=s1", + pendingAuthCode: undefined, + needsFreshClient: () => false, + } as any; +} + +describe("createRemoteClient sibling-refresh retry", () => { + beforeEach(() => { + connectMock.mockReset(); + }); + + it("retries once and succeeds when a newer token appears on disk", async () => { + connectMock + .mockRejectedValueOnce(new UnauthorizedError("401")) + .mockResolvedValueOnce(undefined); + + // pre-connect snapshot T0, post-failure re-read T1 (rotated), retry snapshot T1. + const provider = makeProvider([ + { access_token: "T0" }, + { access_token: "T1" }, + { access_token: "T1" }, + ]); + + const client = await createRemoteClient( + "https://acme-be.glean.com/mcp/gateway/proxy", + { authProvider: provider }, + "sess-1", + ); + + expect(client).toBeTruthy(); + expect(connectMock).toHaveBeenCalledTimes(2); + }); + + it("does not retry when the on-disk token is unchanged", async () => { + connectMock.mockRejectedValue(new UnauthorizedError("401")); + + const provider = makeProvider([ + { access_token: "T0" }, + { access_token: "T0" }, + ]); + + await expect( + createRemoteClient( + "https://acme-be.glean.com/mcp/gateway/proxy", + { authProvider: provider }, + "sess-2", + ), + ).rejects.toBeInstanceOf(AuthRequiredError); + + expect(connectMock).toHaveBeenCalledTimes(1); + }); +}); + +describe("createRemoteClient refresh-collision retry", () => { + beforeEach(() => { + connectMock.mockReset(); + }); + + // Fosite fails concurrent-refresh losers with invalid_request (SDK rethrows raw). + const collisionError = new Error( + "The request is missing a required parameter, includes an invalid " + + "parameter value, includes a parameter more than once, or is " + + "otherwise malformed. Failed to refresh token", + ); + + function makeCollisionProvider(siblingRefreshed: boolean) { + return { + tokens: () => ({ access_token: "T0" }), + authorizationUrl: undefined, + pendingAuthCode: undefined, + needsFreshClient: () => false, + waitForSiblingRefresh: vi.fn(async () => siblingRefreshed), + invalidateCredentials: vi.fn(), + } as any; + } + + it("retries once when a sibling's refresh lands during the grace wait", async () => { + connectMock + .mockRejectedValueOnce(collisionError) + .mockResolvedValueOnce(undefined); + const provider = makeCollisionProvider(true /*siblingRefreshed*/); + + const client = await createRemoteClient( + "https://acme-be.glean.com/mcp/gateway/proxy", + { authProvider: provider }, + "sess-5", + ); + + expect(client).toBeTruthy(); + expect(connectMock).toHaveBeenCalledTimes(2); + expect(provider.waitForSiblingRefresh).toHaveBeenCalledWith("T0"); + }); + + it("rethrows when no sibling token appears within the grace window", async () => { + connectMock.mockRejectedValue(collisionError); + const provider = makeCollisionProvider(false /*siblingRefreshed*/); + + await expect( + createRemoteClient( + "https://acme-be.glean.com/mcp/gateway/proxy", + { authProvider: provider }, + "sess-6", + ), + ).rejects.toBe(collisionError); + + expect(connectMock).toHaveBeenCalledTimes(1); + }); + + it("does not treat unrelated connect errors as refresh failures", async () => { + connectMock.mockRejectedValue(new Error("socket hang up")); + const provider = makeCollisionProvider(true /*siblingRefreshed*/); + + await expect( + createRemoteClient( + "https://acme-be.glean.com/mcp/gateway/proxy", + { authProvider: provider }, + "sess-7", + ), + ).rejects.toThrow("socket hang up"); + + expect(provider.waitForSiblingRefresh).not.toHaveBeenCalled(); + }); +}); diff --git a/sources/glean-vnext/tests/token-store.test.ts b/sources/glean-vnext/tests/token-store.test.ts index 14fd6cc..88685ad 100644 --- a/sources/glean-vnext/tests/token-store.test.ts +++ b/sources/glean-vnext/tests/token-store.test.ts @@ -10,9 +10,8 @@ vi.mock("node:os", async () => { return { ...actual, homedir: () => tmpDir }; }); -const { clearCredentials, loadCredentials, saveCredentials } = await import( - "../src/token-store.js" -); +const { clearCredentials, credentialsMtimeMs, loadCredentials, saveCredentials } = + await import("../src/token-store.js"); describe("token-store", () => { const gleanDir = path.join(tmpDir, ".glean"); @@ -88,4 +87,23 @@ describe("token-store", () => { expect(fs.existsSync(credFile)).toBe(false); expect(() => clearCredentials()).not.toThrow(); }); + + it("credentialsMtimeMs returns undefined when file does not exist", () => { + expect(credentialsMtimeMs()).toBeUndefined(); + }); + + it("credentialsMtimeMs returns a number once credentials are saved", () => { + saveCredentials({ access_token: "x" }, undefined); + expect(typeof credentialsMtimeMs()).toBe("number"); + }); + + it("credentialsMtimeMs advances when the file is rewritten", () => { + saveCredentials({ access_token: "x" }, undefined); + const first = credentialsMtimeMs()!; + // Force a strictly-newer mtime rather than relying on filesystem timer + // resolution between two quick writes. + const future = new Date(Date.now() + 10_000); + fs.utimesSync(credFile, future, future); + expect(credentialsMtimeMs()!).toBeGreaterThan(first); + }); }); From d7a8cb35f5312f4c15b67a7b228379e0960c3afa Mon Sep 17 00:00:00 2001 From: Pragati Agrawal Date: Tue, 1 Sep 2026 14:34:31 +0530 Subject: [PATCH 2/4] Align token rotation handling with glean-vnext --- shared/glean/mcp/src/auth-provider.ts | 98 ++++---------------- shared/glean/mcp/src/token-store.ts | 12 --- shared/glean/mcp/tests/auth-provider.test.ts | 35 ++++--- shared/glean/mcp/tests/token-store.test.ts | 20 +--- 4 files changed, 33 insertions(+), 132 deletions(-) diff --git a/shared/glean/mcp/src/auth-provider.ts b/shared/glean/mcp/src/auth-provider.ts index 5b68b99..d64de3e 100644 --- a/shared/glean/mcp/src/auth-provider.ts +++ b/shared/glean/mcp/src/auth-provider.ts @@ -6,11 +6,11 @@ import type { } from "@modelcontextprotocol/sdk/shared/auth.js"; import { execFile, spawn } from "node:child_process"; import { randomUUID } from "node:crypto"; +import { setTimeout as sleep } from "node:timers/promises"; import { platform } from "node:os"; import { getCallbackUrl, setExpectedState } from "./auth-callback-server.js"; import { clearCredentials, - credentialsMtimeMs, loadCredentials, saveCredentials, } from "./token-store.js"; @@ -18,22 +18,9 @@ import { export type InvalidationScope = "all" | "client" | "tokens" | "verifier"; // Grace window for a sibling's in-flight refresh to land on disk. -const ROTATION_GRACE_MS_DEFAULT = 2000; +const ROTATION_GRACE_MS = 2000; const ROTATION_POLL_MS = 100; -function rotationGraceMs(): number { - const raw = process.env.GLEAN_ROTATION_GRACE_MS; - if (raw !== undefined) { - const parsed = Number.parseInt(raw, 10); - if (Number.isFinite(parsed) && parsed >= 0) return parsed; - } - return ROTATION_GRACE_MS_DEFAULT; -} - -function sleep(ms: number): Promise { - return new Promise((resolve) => setTimeout(resolve, ms)); -} - /** * Open `url` in the user's default browser. Used for the self-open sign-in * path when the client does not support URL-mode elicitation (where the client @@ -70,9 +57,6 @@ export class GleanOAuthClientProvider implements OAuthClientProvider { // explicitly invalidating. Used to detect when a previous auth URL didn't // complete — likely because the server rejected the (stale) client_id. private _authUrlPending = false; - // mtime at last read; detects sibling rewrites of the shared store. - private _credentialsMtimeMs: number | undefined; - authorizationUrl: string | undefined; /** @@ -89,22 +73,12 @@ export class GleanOAuthClientProvider implements OAuthClientProvider { this._tokens = stored.tokens as OAuthTokens | undefined; this._clientInfo = stored.clientInfo as OAuthClientInformationMixed | undefined; } - this._credentialsMtimeMs = credentialsMtimeMs(); } - // Re-read the shared store after a sibling process rewrites it, so we use - // the rotated grant instead of a stale in-memory copy. + // Re-read the shared store on every token access so a sibling's rotated + // grant is used instead of a stale in-memory copy. private syncTokensFromDisk(): void { - const mtimeMs = credentialsMtimeMs(); - if (mtimeMs === undefined) return; - if ( - this._credentialsMtimeMs !== undefined && - mtimeMs <= this._credentialsMtimeMs - ) { - return; - } const stored = loadCredentials(); - this._credentialsMtimeMs = mtimeMs; if (!stored) return; if (stored.tokens) { this._tokens = stored.tokens as OAuthTokens; @@ -114,56 +88,12 @@ export class GleanOAuthClientProvider implements OAuthClientProvider { } } - // On invalid_grant, adopt a sibling's newer on-disk token instead of - // clearing. Returns false when nothing newer exists. - private adoptNewerTokenFromDisk(): boolean { - const diskMtime = credentialsMtimeMs(); - if ( - diskMtime === undefined || - this._credentialsMtimeMs === undefined || - diskMtime <= this._credentialsMtimeMs - ) { - return false; - } - const stored = loadCredentials(); - const diskTokens = stored?.tokens as OAuthTokens | undefined; - if ( - !diskTokens?.access_token || - diskTokens.access_token === this._tokens?.access_token - ) { - return false; - } - this._tokens = diskTokens; - this._credentialsMtimeMs = diskMtime; - if (stored?.clientInfo) { - this._clientInfo = stored.clientInfo as OAuthClientInformationMixed; - } - console.error( - "[auth] invalid_grant, but a newer token is on disk " + - "(sibling refresh) — adopting it instead of clearing", - ); - return true; - } - - // Poll briefly for the race winner's write before clearing. Skipped when - // no refresh token was held (no race possible). - private async adoptNewerTokenWithGrace(): Promise { - if (this.adoptNewerTokenFromDisk()) return true; - if (!this._tokens?.refresh_token) return false; - const deadline = Date.now() + rotationGraceMs(); - while (Date.now() < deadline) { - await sleep(ROTATION_POLL_MS); - if (this.adoptNewerTokenFromDisk()) return true; - } - return false; - } - - // Grace-bounded wait for a sibling's refresh; covers failures the SDK does - // not route through invalidateCredentials (e.g. invalid_request collisions). + // Wait for a sibling's refresh to land on disk. Returns true once a + // different access token is available for adoption/retry. async waitForSiblingRefresh( previousAccessToken: string | undefined, ): Promise { - const deadline = Date.now() + rotationGraceMs(); + const deadline = Date.now() + ROTATION_GRACE_MS; for (;;) { const current = this.tokens()?.access_token; if (current && current !== previousAccessToken) return true; @@ -190,7 +120,6 @@ export class GleanOAuthClientProvider implements OAuthClientProvider { saveClientInformation(info: OAuthClientInformationMixed): void { this._clientInfo = info; saveCredentials(this._tokens, this._clientInfo); - this._credentialsMtimeMs = credentialsMtimeMs(); } tokens(): OAuthTokens | undefined { @@ -202,8 +131,6 @@ export class GleanOAuthClientProvider implements OAuthClientProvider { this._tokens = tokens; this._authUrlPending = false; saveCredentials(this._tokens, this._clientInfo); - // Own write must not look like a sibling change. - this._credentialsMtimeMs = credentialsMtimeMs(); this.onTokensChanged?.(tokens); } @@ -222,12 +149,19 @@ export class GleanOAuthClientProvider implements OAuthClientProvider { this._clientInfo = undefined; saveCredentials(this._tokens, undefined); break; - case "tokens": + case "tokens": { // Usually a sibling's rotation — try adopting before clearing. - if (await this.adoptNewerTokenWithGrace()) return; + const previousAccessToken = this._tokens?.access_token; + if ( + this._tokens?.refresh_token && + (await this.waitForSiblingRefresh(previousAccessToken)) + ) { + return; + } this._tokens = undefined; saveCredentials(undefined, this._clientInfo); break; + } case "verifier": this._codeVerifier = ""; break; diff --git a/shared/glean/mcp/src/token-store.ts b/shared/glean/mcp/src/token-store.ts index 535a07d..bc7f5c3 100644 --- a/shared/glean/mcp/src/token-store.ts +++ b/shared/glean/mcp/src/token-store.ts @@ -28,18 +28,6 @@ export function loadCredentials(): StoredCredentials | undefined { } } -/** - * mtime of the credentials file (epoch ms), or undefined if unreadable. - * Cheap change probe: a single stat, no read + parse. - */ -export function credentialsMtimeMs(): number | undefined { - try { - return fs.statSync(credentialsFile()).mtimeMs; - } catch { - return undefined; - } -} - export function saveCredentials(tokens: unknown, clientInfo: unknown): void { try { const filePath = credentialsFile(); diff --git a/shared/glean/mcp/tests/auth-provider.test.ts b/shared/glean/mcp/tests/auth-provider.test.ts index 3e4e032..a27cb72 100644 --- a/shared/glean/mcp/tests/auth-provider.test.ts +++ b/shared/glean/mcp/tests/auth-provider.test.ts @@ -28,15 +28,12 @@ describe("GleanOAuthClientProvider", () => { beforeEach(() => { delete process.env.PLUGIN_DATA_DIR; - // Skip the rotation grace window by default so invalidation tests don't - // wait out the real 2s poll; the grace test overrides this explicitly. - process.env.GLEAN_ROTATION_GRACE_MS = "0"; fs.rmSync(gleanDir, { recursive: true, force: true }); vi.clearAllMocks(); }); afterEach(() => { - delete process.env.GLEAN_ROTATION_GRACE_MS; + vi.useRealTimers(); fs.rmSync(gleanDir, { recursive: true, force: true }); }); @@ -82,16 +79,12 @@ describe("GleanOAuthClientProvider", () => { const credFile = path.join(gleanDir, "mcp-credentials.json"); - function writeCredFileNewer(tokens: unknown, clientInfo?: unknown): void { + function writeCredFile(tokens: unknown, clientInfo?: unknown): void { fs.mkdirSync(gleanDir, { recursive: true }); fs.writeFileSync(credFile, JSON.stringify({ tokens, clientInfo })); - // Guarantee a strictly-newer mtime than any prior read, independent of - // filesystem timestamp resolution. - const future = new Date(Date.now() + 10_000); - fs.utimesSync(credFile, future, future); } - it("tokens() adopts a newer token written by another process", () => { + it("tokens() adopts a token written by another process", () => { fs.mkdirSync(gleanDir, { recursive: true }); fs.writeFileSync( credFile, @@ -102,12 +95,15 @@ describe("GleanOAuthClientProvider", () => { ); const provider = new GleanOAuthClientProvider(); expect(provider.tokens()?.access_token).toBe("T0"); + const originalMtime = fs.statSync(credFile).mtime; // Sibling refreshes: new access + rotated refresh token on disk. - writeCredFileNewer( + writeCredFile( { access_token: "T1", refresh_token: "R1" }, { client_id: "cid" }, ); + // The provider must not rely on mtime to observe this rewrite. + fs.utimesSync(credFile, originalMtime, originalMtime); expect(provider.tokens()?.access_token).toBe("T1"); expect(provider.tokens()?.refresh_token).toBe("R1"); @@ -137,11 +133,11 @@ describe("GleanOAuthClientProvider", () => { expect(provider.tokens()?.access_token).toBe("T0"); // A client-only rewrite (tokens dropped) must not log us out in-memory. - writeCredFileNewer(undefined, { client_id: "cid" }); + writeCredFile(undefined, { client_id: "cid" }); expect(provider.tokens()?.access_token).toBe("T0"); }); - it("invalidateCredentials('tokens') adopts a sibling's newer token instead of wiping the store", async () => { + it("invalidateCredentials('tokens') adopts a sibling's token instead of wiping the store", async () => { fs.mkdirSync(gleanDir, { recursive: true }); fs.writeFileSync( credFile, @@ -153,8 +149,8 @@ describe("GleanOAuthClientProvider", () => { const provider = new GleanOAuthClientProvider(); expect(provider.tokens()?.access_token).toBe("T0"); - // A sibling refreshed + rotated: fresh grant now on disk with a newer mtime. - writeCredFileNewer( + // A sibling refreshed + rotated: fresh grant is now on disk. + writeCredFile( { access_token: "T1", refresh_token: "R1" }, { client_id: "cid" }, ); @@ -175,7 +171,10 @@ describe("GleanOAuthClientProvider", () => { expect(provider.tokens()?.access_token).toBe("T0"); // No sibling write since our snapshot → a genuine invalidation → clear. - await provider.invalidateCredentials("tokens"); + vi.useFakeTimers(); + const invalidation = provider.invalidateCredentials("tokens"); + await vi.advanceTimersByTimeAsync(2000); + await invalidation; expect(provider.tokens()).toBeUndefined(); const raw = JSON.parse(fs.readFileSync(credFile, "utf-8")); @@ -184,14 +183,13 @@ describe("GleanOAuthClientProvider", () => { it("invalidateCredentials('tokens') adopts a token that lands during the grace window", async () => { // The winner's write lands just after the loser's invalid_grant. - process.env.GLEAN_ROTATION_GRACE_MS = "2000"; const provider = new GleanOAuthClientProvider(); provider.saveTokens({ access_token: "T0", refresh_token: "R0" } as any); const invalidation = provider.invalidateCredentials("tokens"); // Sibling's write lands mid-window. setTimeout(() => { - writeCredFileNewer( + writeCredFile( { access_token: "T1", refresh_token: "R1" }, { client_id: "cid" }, ); @@ -204,7 +202,6 @@ describe("GleanOAuthClientProvider", () => { }); it("skips the grace window when no refresh token was held (no race possible)", async () => { - process.env.GLEAN_ROTATION_GRACE_MS = "5000"; const provider = new GleanOAuthClientProvider(); provider.saveTokens({ access_token: "T0" } as any); // no refresh_token diff --git a/shared/glean/mcp/tests/token-store.test.ts b/shared/glean/mcp/tests/token-store.test.ts index 88685ad..9822106 100644 --- a/shared/glean/mcp/tests/token-store.test.ts +++ b/shared/glean/mcp/tests/token-store.test.ts @@ -10,7 +10,7 @@ vi.mock("node:os", async () => { return { ...actual, homedir: () => tmpDir }; }); -const { clearCredentials, credentialsMtimeMs, loadCredentials, saveCredentials } = +const { clearCredentials, loadCredentials, saveCredentials } = await import("../src/token-store.js"); describe("token-store", () => { @@ -88,22 +88,4 @@ describe("token-store", () => { expect(() => clearCredentials()).not.toThrow(); }); - it("credentialsMtimeMs returns undefined when file does not exist", () => { - expect(credentialsMtimeMs()).toBeUndefined(); - }); - - it("credentialsMtimeMs returns a number once credentials are saved", () => { - saveCredentials({ access_token: "x" }, undefined); - expect(typeof credentialsMtimeMs()).toBe("number"); - }); - - it("credentialsMtimeMs advances when the file is rewritten", () => { - saveCredentials({ access_token: "x" }, undefined); - const first = credentialsMtimeMs()!; - // Force a strictly-newer mtime rather than relying on filesystem timer - // resolution between two quick writes. - const future = new Date(Date.now() + 10_000); - fs.utimesSync(credFile, future, future); - expect(credentialsMtimeMs()!).toBeGreaterThan(first); - }); }); From b9ff0fa6235e7310d519e1e16d3597664e702765 Mon Sep 17 00:00:00 2001 From: Pragati Agrawal Date: Tue, 1 Sep 2026 14:51:33 +0530 Subject: [PATCH 3/4] Harden credential temp-file permissions --- shared/glean/mcp/src/token-store.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/shared/glean/mcp/src/token-store.ts b/shared/glean/mcp/src/token-store.ts index bc7f5c3..a81420b 100644 --- a/shared/glean/mcp/src/token-store.ts +++ b/shared/glean/mcp/src/token-store.ts @@ -41,6 +41,7 @@ export function saveCredentials(tokens: unknown, clientInfo: unknown): void { encoding: "utf-8", mode: FILE_MODE, }); + fs.chmodSync(tmpPath, FILE_MODE); fs.renameSync(tmpPath, filePath); } catch (err) { const msg = err instanceof Error ? err.message : String(err); From f86bcf60c25527c016ad988523b7810494bd72f3 Mon Sep 17 00:00:00 2001 From: Pragati Agrawal Date: Tue, 1 Sep 2026 15:01:46 +0530 Subject: [PATCH 4/4] Use structured OAuth error codes for refresh retry --- shared/glean/mcp/src/remote-client.ts | 20 ++++++++++++------- .../tests/remote-client-auth-retry.test.ts | 17 ++++++++-------- 2 files changed, 21 insertions(+), 16 deletions(-) diff --git a/shared/glean/mcp/src/remote-client.ts b/shared/glean/mcp/src/remote-client.ts index 568acab..d4e5713 100644 --- a/shared/glean/mcp/src/remote-client.ts +++ b/shared/glean/mcp/src/remote-client.ts @@ -1,6 +1,7 @@ import { Client } from "@modelcontextprotocol/sdk/client/index.js"; import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js"; import { UnauthorizedError } from "@modelcontextprotocol/sdk/client/auth.js"; +import { OAuthError } from "@modelcontextprotocol/sdk/server/auth/errors.js"; import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; import type { GleanOAuthClientProvider } from "./auth-provider.js"; import { PLUGIN_VERSION } from "./version.js"; @@ -241,12 +242,13 @@ export async function createRemoteClient( throw new AuthRequiredError(authProvider.authorizationUrl); } } - // Concurrent-refresh losers get errors the SDK rethrows raw (e.g. fosite - // invalid_request); retry once if a sibling's grant lands in the grace window. + // Concurrent-refresh losers are reported with structured OAuth errors + // (typically invalid_request); retry once if a sibling's grant lands in the + // grace window. if ( authProvider && !authRetry && - isLikelyRefreshFailure(error) && + isRefreshOAuthError(error) && (await authProvider.waitForSiblingRefresh(accessTokenAtConnect)) ) { console.error( @@ -260,10 +262,14 @@ export async function createRemoteClient( return client; } -// Match broadly; the caller's disk re-check gates the actual retry. -function isLikelyRefreshFailure(error: unknown): boolean { - const msg = error instanceof Error ? error.message : String(error); - return /refresh|invalid_grant|invalid_request|oauth/i.test(msg); +// Restrict recovery to OAuth errors that can indicate a refresh race. The SDK +// preserves the response's machine-readable error code. +function isRefreshOAuthError(error: unknown): boolean { + return ( + error instanceof OAuthError && + (error.errorCode === "invalid_request" || + error.errorCode === "invalid_grant") + ); } export async function callRemoteTool( diff --git a/shared/glean/mcp/tests/remote-client-auth-retry.test.ts b/shared/glean/mcp/tests/remote-client-auth-retry.test.ts index f730cc7..5914289 100644 --- a/shared/glean/mcp/tests/remote-client-auth-retry.test.ts +++ b/shared/glean/mcp/tests/remote-client-auth-retry.test.ts @@ -1,5 +1,6 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; import { UnauthorizedError } from "@modelcontextprotocol/sdk/client/auth.js"; +import { InvalidRequestError } from "@modelcontextprotocol/sdk/server/auth/errors.js"; // Control client.connect() across (re)tries. const { connectMock } = vi.hoisted(() => ({ connectMock: vi.fn() })); @@ -95,16 +96,14 @@ describe("createRemoteClient refresh-collision retry", () => { connectMock.mockReset(); }); - // Fosite fails concurrent-refresh losers with invalid_request (SDK rethrows raw). - const collisionError = new Error( - "The request is missing a required parameter, includes an invalid " + - "parameter value, includes a parameter more than once, or is " + - "otherwise malformed. Failed to refresh token", + // The SDK preserves fosite's machine-readable OAuth error code. + const collisionError = new InvalidRequestError( + "The refresh request was rejected because another process rotated the grant.", ); function makeCollisionProvider(siblingRefreshed: boolean) { return { - tokens: () => ({ access_token: "T0" }), + tokens: () => ({ access_token: "T0", refresh_token: "R0" }), authorizationUrl: undefined, pendingAuthCode: undefined, needsFreshClient: () => false, @@ -145,8 +144,8 @@ describe("createRemoteClient refresh-collision retry", () => { expect(connectMock).toHaveBeenCalledTimes(1); }); - it("does not treat unrelated connect errors as refresh failures", async () => { - connectMock.mockRejectedValue(new Error("socket hang up")); + it("does not treat untyped refresh-like errors as refresh failures", async () => { + connectMock.mockRejectedValue(new Error("Failed to refresh token")); const provider = makeCollisionProvider(true /*siblingRefreshed*/); await expect( @@ -155,7 +154,7 @@ describe("createRemoteClient refresh-collision retry", () => { { authProvider: provider }, "sess-7", ), - ).rejects.toThrow("socket hang up"); + ).rejects.toThrow("Failed to refresh token"); expect(provider.waitForSiblingRefresh).not.toHaveBeenCalled(); });