fix(coding-agent): serialize MCP OAuth refreshes across processes

Refreshes hold a per-server lock file from reading the tokens to saving new ones, and use tokens another process already refreshed. Rotating refresh tokens (Cloudflare) were lost when two pi processes refreshed at once, leaving the server needing sign-in. Closing a connection waits for a running refresh so shutdown does not drop rotated tokens.

packages/coding-agent/src/extensions/mcp/oauth.ts

@@ -9,10 +9,11 @@
  * Credentials live in `<agent-dir>/mcp-auth.json`, keyed by server URL.
  */
 
+import { createHash } from "node:crypto";
+import { mkdirSync } from "node:fs";
 import { join } from "node:path";
-import type { AuthProvider } from "@earendil-works/pi-mcp";
+import type { AuthProvider, McpFetch } from "@earendil-works/pi-mcp";
 import {
-	adaptOAuthProvider,
 	authorizeMcp,
 	McpOAuthAuthorizationRequiredError,
 	McpOAuthProvider,
@@ -23,6 +24,7 @@ import {
 	type OAuthClientInformationMixed,
 	parseWwwAuthenticate,
 } from "@earendil-works/pi-mcp/oauth";
+import lockfile from "proper-lockfile";
 import { APP_NAME, getAgentDir } from "../../config.ts";
 import { type AuthStorageBackend, FileAuthStorageBackend } from "../../core/auth-storage.ts";
 
@@ -32,6 +34,13 @@ const CALLBACK_PATH = "/callback";
 const FALLBACK_REDIRECT_URL = `http://${CALLBACK_HOST}${CALLBACK_PATH}`;
 /** Access tokens this close to expiry are refreshed before they are sent. */
 const REFRESH_SKEW_MS = 30_000;
+/** Bounds each request of a refresh, so it cannot hold the refresh lock or delay shutdown for long. */
+const REFRESH_REQUEST_TIMEOUT_MS = 15_000;
+/** A refresh lock that its holder stopped renewing (the process was killed) is taken over after this. */
+const REFRESH_LOCK_STALE_MS = 20_000;
+/** How long to wait for another process's refresh: longer than a stale lock lives. */
+const REFRESH_LOCK_WAIT_MS = 25_000;
+const REFRESH_LOCK_RETRY_MS = 100;
 
 export interface McpOAuthSettings {
 	clientId?: string;
@@ -91,15 +100,23 @@ function parseStates(content: string | undefined): StoredStates {
 	return typeof parsed === "object" && parsed !== null && !Array.isArray(parsed) ? (parsed as StoredStates) : {};
 }
 
+export interface McpOAuthServerStore extends McpOAuthStateStore {
+	/** Run `fn` while no other process refreshes the server's tokens. */
+	withRefreshLock<T>(fn: () => Promise<T>): Promise<T>;
+}
+
 /** Per-server OAuth state (client registration, tokens, pending PKCE verifier) in `mcp-auth.json`. */
 export class McpOAuthCredentialStore {
 	private readonly backend: AuthStorageBackend;
+	/** Directory for the refresh lock files. Without one, refreshes are only serialized in this process. */
+	private readonly lockDir: string | undefined;
 
-	constructor(backend: AuthStorageBackend = new FileAuthStorageBackend(join(getAgentDir(), "mcp-auth.json"))) {
-		this.backend = backend;
+	constructor(backend?: AuthStorageBackend, lockDir?: string) {
+		this.backend = backend ?? new FileAuthStorageBackend(join(getAgentDir(), "mcp-auth.json"));
+		this.lockDir = backend ? lockDir : getAgentDir();
 	}
 
-	forServer(serverUrl: string): McpOAuthStateStore {
+	forServer(serverUrl: string): McpOAuthServerStore {
 		const key = String(new URL(serverUrl));
 		return {
 			load: () => this.read()[key],
@@ -107,9 +124,37 @@ export class McpOAuthCredentialStore {
 				this.write((states) => {
 					states[key] = state;
 				}),
+			withRefreshLock: (fn) => this.withRefreshLock(key, fn),
 		};
 	}
 
+	/**
+	 * A lock file per server. When the process exits, proper-lockfile removes the locks it holds; when
+	 * it is killed, the lock goes stale because it is no longer renewed, and the next process takes it over.
+	 */
+	private async withRefreshLock<T>(key: string, fn: () => Promise<T>): Promise<T> {
+		if (!this.lockDir) return fn();
+		mkdirSync(this.lockDir, { recursive: true, mode: 0o700 });
+		const hash = createHash("sha256").update(key).digest("hex").slice(0, 16);
+		const release = await lockfile.lock(join(this.lockDir, `mcp-auth-refresh-${hash}`), {
+			realpath: false,
+			stale: REFRESH_LOCK_STALE_MS,
+			retries: {
+				retries: REFRESH_LOCK_WAIT_MS / REFRESH_LOCK_RETRY_MS,
+				factor: 1,
+				minTimeout: REFRESH_LOCK_RETRY_MS,
+				maxTimeout: REFRESH_LOCK_RETRY_MS,
+			},
+			// The default throws from a timer. A lost lock at worst lets two refreshes overlap.
+			onCompromised: () => {},
+		});
+		try {
+			return await fn();
+		} finally {
+			await release().catch(() => undefined);
+		}
+	}
+
 	/** The stored tokens of a server, for noticing sign-ins done by another process. */
 	tokens(serverUrl: string): McpOAuthState["tokens"] {
 		return this.read()[String(new URL(serverUrl))]?.tokens;
@@ -160,6 +205,11 @@ function createProvider(
 	});
 }
 
+export interface McpAuthProvider extends AuthProvider {
+	/** Resolves when no refresh is running, so shutdown does not drop rotated tokens before they are saved. */
+	settled(): Promise<void>;
+}
+
 /**
  * Auth provider for MCP connections: sends the stored access token and refreshes it when it is about
  * to expire or after a 401. Throws `McpOAuthAuthorizationRequiredError` when the user has to sign in,
@@ -168,37 +218,46 @@ function createProvider(
  * `settings` is only called when a refresh is needed, so a secret that fails to resolve fails the
  * refresh instead of the whole connection setup.
  *
- * Concurrent requests share one refresh, and a 401 for a token that was already replaced just
- * retries: many servers rotate refresh tokens, so refreshing twice with the same one would fail
- * and discard the new grant.
+ * Many servers rotate refresh tokens, so two refreshes with the same refresh token lose the grant.
+ * Requests in this process share one refresh, and other processes are kept out by the store's
+ * refresh lock, held from reading the tokens to saving new ones. Tokens that changed meanwhile
+ * (another process refreshed them, or the user signed in) are used without refreshing.
  */
 export function createMcpAuthProvider(options: {
 	serverUrl: string;
-	store: McpOAuthStateStore;
+	store: McpOAuthServerStore;
 	settings: () => McpOAuthSettings;
 	onChallenge: (challenge: OAuthChallenge) => void;
-}): AuthProvider {
+}): McpAuthProvider {
 	const { serverUrl, store } = options;
 	let refreshing: Promise<void> | undefined;
 
-	const refresh = (context: Parameters<NonNullable<AuthProvider["onUnauthorized"]>>[0] | undefined) => {
-		refreshing ??= (async () => {
-			const state = await store.load();
-			if (!state?.tokens?.refresh_token) throw new McpOAuthAuthorizationRequiredError();
-			const settings = options.settings();
-			const redirectUrl =
-				callbackSettings(settings).fixedRedirectUrl ??
-				registeredRedirectUrls(state.clientInformation)[0] ??
-				FALLBACK_REDIRECT_URL;
-			const provider = createProvider(serverUrl, store, settings, redirectUrl, () => {});
-			// Refreshes the tokens, or reports that a new sign-in is needed.
-			const result = context
-				? await adaptOAuthProvider(provider).onUnauthorized?.(context)
-				: await authorizeMcp(provider, { serverUrl });
-			if (result === "REDIRECT") throw new McpOAuthAuthorizationRequiredError();
-		})().finally(() => {
-			refreshing = undefined;
-		});
+	/** Replace `staleToken`, the access token that expired or was rejected. */
+	const refresh = (staleToken: string | undefined, fetch: McpFetch = globalThis.fetch, challenge?: OAuthChallenge) => {
+		refreshing ??= store
+			.withRefreshLock(async () => {
+				const state = await store.load();
+				if (state?.tokens?.access_token !== staleToken) return;
+				if (!state?.tokens?.refresh_token) throw new McpOAuthAuthorizationRequiredError();
+				const settings = options.settings();
+				const redirectUrl =
+					callbackSettings(settings).fixedRedirectUrl ??
+					registeredRedirectUrls(state.clientInformation)[0] ??
+					FALLBACK_REDIRECT_URL;
+				const provider = createProvider(serverUrl, store, settings, redirectUrl, () => {});
+				// Refreshes the tokens, or reports that a new sign-in is needed.
+				const result = await authorizeMcp(provider, {
+					serverUrl,
+					resourceMetadataUrl: challenge?.resourceMetadataUrl,
+					scope: challenge?.scope,
+					fetch: (input, init) =>
+						fetch(input, { ...init, signal: AbortSignal.timeout(REFRESH_REQUEST_TIMEOUT_MS) }),
+				});
+				if (result === "REDIRECT") throw new McpOAuthAuthorizationRequiredError();
+			})
+			.finally(() => {
+				refreshing = undefined;
+			});
 		return refreshing;
 	};
 
@@ -206,10 +265,11 @@ export function createMcpAuthProvider(options: {
 		token: async () => {
 			await refreshing?.catch(() => undefined);
 			const state = await store.load();
+			const token = state?.tokens?.access_token;
 			const expired = state?.tokensExpireAt !== undefined && state.tokensExpireAt - REFRESH_SKEW_MS <= Date.now();
-			if (!expired || !state?.tokens?.refresh_token) return state?.tokens?.access_token;
+			if (!expired || !state?.tokens?.refresh_token) return token;
 			// Failures fall through: the request goes out with the old token and a 401 decides what happens.
-			await refresh(undefined).catch(() => undefined);
+			await refresh(token).catch(() => undefined);
 			return (await store.load())?.tokens?.access_token;
 		},
 		onUnauthorized: async (context) => {
@@ -217,11 +277,10 @@ export function createMcpAuthProvider(options: {
 			options.onChallenge(challenge);
 			// A refresh keeps the granted scope, so more scope needs a new sign-in.
 			if (challenge.error === "insufficient_scope") throw new McpOAuthAuthorizationRequiredError();
-			if (!refreshing && context.token !== undefined) {
-				const current = (await store.load())?.tokens?.access_token;
-				if (current !== undefined && current !== context.token) return;
-			}
-			await refresh(context);
+			await refresh(context.token, context.fetch, challenge);
+		},
+		settled: async () => {
+			await refreshing?.catch(() => undefined);
 		},
 	};
 }

packages/coding-agent/src/extensions/mcp/runtime.ts

@@ -32,7 +32,12 @@ import { VERSION } from "../../config.ts";
 import { resolveConfigValueOrThrow, resolveHeadersOrThrow } from "../../core/resolve-config-value.ts";
 import type { McpServerEntry } from "./config.ts";
 import type { McpServerLog } from "./log.ts";
-import { createMcpAuthProvider, type McpOAuthCredentialStore, type McpOAuthSettings } from "./oauth.ts";
+import {
+	createMcpAuthProvider,
+	type McpAuthProvider,
+	type McpOAuthCredentialStore,
+	type McpOAuthSettings,
+} from "./oauth.ts";
 import { isMcpAppResource, type McpResourceServer } from "./resources.ts";
 import type { McpToolCaller } from "./tools.ts";
 
@@ -169,7 +174,7 @@ export class McpServerConnection implements McpToolCaller, McpResourceServer {
 	private stderrTail: string | undefined;
 	private readonly cwd: string;
 	private readonly createTransport: McpTransportFactory;
-	private readonly authProvider: AuthProvider | undefined;
+	private readonly authProvider: McpAuthProvider | undefined;
 	private readonly onTools: (connection: McpServerConnection) => void;
 	private readonly onChange: ((connection: McpServerConnection) => void) | undefined;
 	private readonly log: McpServerLog | undefined;
@@ -452,5 +457,8 @@ export class McpServerConnection implements McpToolCaller, McpResourceServer {
 		const client = this.client;
 		this.client = undefined;
 		await client?.close().catch(() => undefined);
+		// A refresh the server already answered may have rotated the refresh token; exiting before the
+		// new tokens are saved would lose the grant.
+		await this.authProvider?.settled();
 	}
 }

packages/coding-agent/test/mcp-oauth-refresh.test.ts

@@ -0,0 +1,84 @@
+import { mkdtempSync, readdirSync, rmSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { afterEach, describe, expect, it } from "vitest";
+import { InMemoryAuthStorageBackend } from "../src/core/auth-storage.ts";
+import { createMcpAuthProvider, McpOAuthCredentialStore, signInMcpServer } from "../src/extensions/mcp/oauth.ts";
+import { startOAuthMcpServer } from "./suite/mcp-oauth-server.ts";
+
+describe("MCP OAuth refresh", () => {
+	const cleanups: (() => Promise<void> | void)[] = [];
+
+	afterEach(async () => {
+		while (cleanups.length > 0) await cleanups.pop()?.();
+	});
+
+	async function signedIn() {
+		const server = await startOAuthMcpServer();
+		cleanups.push(server.close);
+		const lockDir = mkdtempSync(join(tmpdir(), "pi-mcp-refresh-"));
+		cleanups.push(() => rmSync(lockDir, { recursive: true, force: true }));
+		// Stores sharing the credential file and lock directory stand in for separate pi processes.
+		const backend = new InMemoryAuthStorageBackend();
+		const process = () => {
+			const store = new McpOAuthCredentialStore(backend, lockDir).forServer(server.url);
+			const provider = createMcpAuthProvider({
+				serverUrl: server.url,
+				store,
+				settings: () => ({}),
+				onChallenge: () => {},
+			});
+			return { store, provider };
+		};
+		await signInMcpServer({
+			serverUrl: server.url,
+			store: process().store,
+			settings: {},
+			prompt: {
+				showAuthorizationUrl: (url) => void fetch(url),
+				promptForRedirectUrl: (signal) =>
+					new Promise((resolve) => signal.addEventListener("abort", () => resolve(undefined), { once: true })),
+			},
+		});
+		return { server, lockDir, process };
+	}
+
+	it("refreshes once when several processes find the same token rejected", async () => {
+		const { server, lockDir, process } = await signedIn();
+		server.expireAccessTokens();
+		const processes = [process(), process(), process()];
+
+		// The server rotates refresh tokens: a second refresh with refresh-1 would fail with invalid_grant.
+		await Promise.all(
+			processes.map(({ provider }) =>
+				provider.onUnauthorized?.({
+					response: new Response(null, { status: 401 }),
+					serverUrl: new URL(server.url),
+					fetch: globalThis.fetch,
+					token: "access-1",
+				}),
+			),
+		);
+
+		expect(server.log.filter((entry) => entry === "token refresh")).toHaveLength(1);
+		for (const { provider } of processes) expect(await provider.token()).toBe("access-2");
+		// The lock is released.
+		expect(readdirSync(lockDir)).toEqual([]);
+	});
+
+	it("waits for a running refresh to save the new tokens", async () => {
+		const { server, process } = await signedIn();
+		server.expireAccessTokens();
+		const { store, provider } = process();
+
+		const refresh = provider.onUnauthorized?.({
+			response: new Response(null, { status: 401 }),
+			serverUrl: new URL(server.url),
+			fetch: globalThis.fetch,
+			token: "access-1",
+		});
+		await provider.settled();
+		expect((await store.load())?.tokens?.access_token).toBe("access-2");
+		await refresh;
+	});
+});