diff --git a/packages/coding-agent/src/extensions/mcp/oauth.ts b/packages/coding-agent/src/extensions/mcp/oauth.ts index 3fa9984f1..c1c98f45f 100644 --- a/packages/coding-agent/src/extensions/mcp/oauth.ts +++ b/packages/coding-agent/src/extensions/mcp/oauth.ts @@ -9,10 +9,11 @@ * Credentials live in `/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(fn: () => Promise): Promise; +} + /** 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(key: string, fn: () => Promise): Promise { + 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; +} + /** * 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 | undefined; - const refresh = (context: Parameters>[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); }, }; } diff --git a/packages/coding-agent/src/extensions/mcp/runtime.ts b/packages/coding-agent/src/extensions/mcp/runtime.ts index bed36f0bf..0342aaf2e 100644 --- a/packages/coding-agent/src/extensions/mcp/runtime.ts +++ b/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(); } } diff --git a/packages/coding-agent/test/mcp-oauth-refresh.test.ts b/packages/coding-agent/test/mcp-oauth-refresh.test.ts new file mode 100644 index 000000000..2f237fb4d --- /dev/null +++ b/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)[] = []; + + 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; + }); +});