diff --git a/shared/glean/mcp/src/auth-provider.ts b/shared/glean/mcp/src/auth-provider.ts index 375d507..d64de3e 100644 --- a/shared/glean/mcp/src/auth-provider.ts +++ b/shared/glean/mcp/src/auth-provider.ts @@ -6,12 +6,21 @@ 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, loadCredentials, saveCredentials } from "./token-store.js"; +import { + clearCredentials, + 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 = 2000; +const ROTATION_POLL_MS = 100; + /** * 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 @@ -48,7 +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; - authorizationUrl: string | undefined; /** @@ -67,6 +75,33 @@ export class GleanOAuthClientProvider implements OAuthClientProvider { } } + // 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 stored = loadCredentials(); + if (!stored) return; + if (stored.tokens) { + this._tokens = stored.tokens as OAuthTokens; + } + if (stored.clientInfo) { + this._clientInfo = stored.clientInfo as OAuthClientInformationMixed; + } + } + + // 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() + ROTATION_GRACE_MS; + 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 { return getCallbackUrl(); } @@ -88,6 +123,7 @@ export class GleanOAuthClientProvider implements OAuthClientProvider { } tokens(): OAuthTokens | undefined { + this.syncTokensFromDisk(); return this._tokens; } @@ -113,10 +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. + 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/remote-client.ts b/shared/glean/mcp/src/remote-client.ts index e9a87f3..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"; @@ -167,6 +168,7 @@ export async function createRemoteClient( serverUrl: string, opts: RemoteClientOptions, chatSessionId?: string, + authRetry = false, ): Promise { const authProvider = opts.authProvider; @@ -214,14 +216,45 @@ 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 withConnectLock(() => 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 are reported with structured OAuth errors + // (typically invalid_request); retry once if a sibling's grant lands in the + // grace window. + if ( + authProvider && + !authRetry && + isRefreshOAuthError(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; } @@ -229,6 +262,16 @@ export async function createRemoteClient( return client; } +// 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( client: Client, name: string, diff --git a/shared/glean/mcp/src/token-store.ts b/shared/glean/mcp/src/token-store.ts index e41c2fa..a81420b 100644 --- a/shared/glean/mcp/src/token-store.ts +++ b/shared/glean/mcp/src/token-store.ts @@ -35,11 +35,14 @@ 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.chmodSync(tmpPath, 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/shared/glean/mcp/tests/auth-provider.test.ts b/shared/glean/mcp/tests/auth-provider.test.ts index ba181f0..a27cb72 100644 --- a/shared/glean/mcp/tests/auth-provider.test.ts +++ b/shared/glean/mcp/tests/auth-provider.test.ts @@ -33,6 +33,7 @@ describe("GleanOAuthClientProvider", () => { }); afterEach(() => { + vi.useRealTimers(); fs.rmSync(gleanDir, { recursive: true, force: true }); }); @@ -74,6 +75,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 writeCredFile(tokens: unknown, clientInfo?: unknown): void { + fs.mkdirSync(gleanDir, { recursive: true }); + fs.writeFileSync(credFile, JSON.stringify({ tokens, clientInfo })); + } + + it("tokens() adopts a 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"); + const originalMtime = fs.statSync(credFile).mtime; + + // Sibling refreshes: new access + rotated refresh token on disk. + 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"); + }); + + 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. + writeCredFile(undefined, { client_id: "cid" }); + expect(provider.tokens()?.access_token).toBe("T0"); + }); + + it("invalidateCredentials('tokens') adopts a sibling's 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 is now on disk. + writeCredFile( + { 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. + 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")); + 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. + 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(() => { + writeCredFile( + { 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 () => { + 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/shared/glean/mcp/tests/remote-client-auth-retry.test.ts b/shared/glean/mcp/tests/remote-client-auth-retry.test.ts new file mode 100644 index 0000000..5914289 --- /dev/null +++ b/shared/glean/mcp/tests/remote-client-auth-retry.test.ts @@ -0,0 +1,161 @@ +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() })); + +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(); + }); + + // 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", refresh_token: "R0" }), + 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 untyped refresh-like errors as refresh failures", async () => { + connectMock.mockRejectedValue(new Error("Failed to refresh token")); + const provider = makeCollisionProvider(true /*siblingRefreshed*/); + + await expect( + createRemoteClient( + "https://acme-be.glean.com/mcp/gateway/proxy", + { authProvider: provider }, + "sess-7", + ), + ).rejects.toThrow("Failed to refresh token"); + + expect(provider.waitForSiblingRefresh).not.toHaveBeenCalled(); + }); +}); diff --git a/shared/glean/mcp/tests/token-store.test.ts b/shared/glean/mcp/tests/token-store.test.ts index 14fd6cc..9822106 100644 --- a/shared/glean/mcp/tests/token-store.test.ts +++ b/shared/glean/mcp/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, loadCredentials, saveCredentials } = + await import("../src/token-store.js"); describe("token-store", () => { const gleanDir = path.join(tmpDir, ".glean"); @@ -88,4 +87,5 @@ describe("token-store", () => { expect(fs.existsSync(credFile)).toBe(false); expect(() => clearCredentials()).not.toThrow(); }); + });