diff --git a/packages/coding-agent/docs/mcp.md b/packages/coding-agent/docs/mcp.md index c20452523..ce2c9791b 100644 --- a/packages/coding-agent/docs/mcp.md +++ b/packages/coding-agent/docs/mcp.md @@ -23,7 +23,7 @@ Add servers to `~/.pi/agent/mcp.json`, or to `.pi/mcp.json` in a project. The fo ``` - stdio servers take `command`, `args`, `env`, and `cwd`. Relative `cwd` resolves against the session directory. -- HTTP servers take `url` and `headers`. The legacy SSE transport is not supported. +- HTTP servers take `url`, `headers`, and `oauth` (see [Sign in with OAuth](#sign-in-with-oauth)). The legacy SSE transport is not supported. - `env` and `headers` values can reference environment variables (`${NAME}`) or commands (`!command`), like provider API keys. - `timeout` sets the per-request timeout in seconds (default 60). Progress notifications from the server reset it. - `enabled: false` keeps an entry without connecting to it. @@ -32,6 +32,37 @@ Project entries replace global entries with the same name. A project `mcp.json` Pi connects when a session starts. The first prompt waits until startup connections finish. A server that drops its connection is reconnected on the next call. Run `/mcp` to see server status, tool counts, and errors. +## Sign in with OAuth + +Remote servers that use OAuth, such as Sentry, need no credentials in `mcp.json`: + +```json +{ + "mcpServers": { + "sentry": { "url": "https://mcp.sentry.dev/mcp" } + } +} +``` + +When such a server rejects the connection, `/mcp` shows it as needing sign-in. Run `/mcp login sentry` to open the authorization page in your browser. After you approve access, the browser redirects to a temporary server on `127.0.0.1` and pi connects. If the browser runs on another machine, for example over SSH, paste the URL it was redirected to into the prompt instead. + +Pi registers itself with the authorization server (dynamic client registration), stores tokens in `~/.pi/agent/mcp-auth.json`, and refreshes expired access tokens automatically. `/mcp logout sentry` deletes the stored credentials. + +OAuth applies to HTTP servers without an `Authorization` header. For authorization servers that do not support dynamic client registration, configure a pre-registered client: + +```json +{ + "mcpServers": { + "example": { + "url": "https://mcp.example.com/mcp", + "oauth": { "clientId": "my-client", "clientSecret": "${EXAMPLE_SECRET}", "callbackPort": 8765 } + } + } +} +``` + +`callbackPort` fixes the redirect URI to `http://127.0.0.1:/oauth/callback`, which must match the redirect URI registered for the client. `clientSecret` is optional and can reference environment variables or commands. + ## Exposure Each server's tools are registered as `mcp____`. The `exposure` setting controls how the model reaches them: @@ -42,5 +73,3 @@ Each server's tools are registered as `mcp____`. The `exposure` se MCP tools that declare an `outputSchema` return their `structuredContent` to codemode scripts. Other tools return their text output. Images are passed through as image content. Every MCP call goes through pi's tool pipeline, so `tool_call` and `tool_result` extension handlers, including permission gates, apply to MCP tools. Calls made from codemode scripts carry the codemode call's id as `parentToolCallId`. - -OAuth sign-in is not supported yet. Use header-based authentication for remote servers. diff --git a/packages/coding-agent/src/extensions/mcp/config.ts b/packages/coding-agent/src/extensions/mcp/config.ts index f9ec79fb9..2246335d0 100644 --- a/packages/coding-agent/src/extensions/mcp/config.ts +++ b/packages/coding-agent/src/extensions/mcp/config.ts @@ -10,10 +10,13 @@ * { * "mcpServers": { * "filesystem": { "command": "npx", "args": ["-y", "@modelcontextprotocol/server-filesystem", "."] }, - * "docs": { "url": "https://example.com/mcp", "headers": { "Authorization": "Bearer ${DOCS_TOKEN}" } } + * "docs": { "url": "https://example.com/mcp", "headers": { "Authorization": "Bearer ${DOCS_TOKEN}" } }, + * "sentry": { "url": "https://mcp.sentry.dev/mcp" } * } * } * ``` + * + * HTTP servers without an `Authorization` header use OAuth when they answer 401 (`/mcp login`). */ import { existsSync, readFileSync } from "node:fs"; @@ -45,11 +48,22 @@ export interface McpStdioServerConfig extends McpServerConfigBase { cwd?: string; } +/** OAuth client settings for servers that do not support dynamic client registration. */ +export interface McpOAuthConfig { + /** Pre-registered client id. Without it, pi registers a client with the authorization server. */ + clientId?: string; + /** May reference environment variables (`${NAME}`) or commands (`!cmd`). */ + clientSecret?: string; + /** Fixed loopback callback port, for clients registered with an exact redirect URI. */ + callbackPort?: number; +} + export interface McpHttpServerConfig extends McpServerConfigBase { type?: "http"; url: string; /** Values may reference environment variables (`${NAME}`) or commands (`!cmd`). */ headers?: Record; + oauth?: McpOAuthConfig; } export type McpServerConfig = McpStdioServerConfig | McpHttpServerConfig; @@ -76,6 +90,20 @@ function isStringRecord(value: unknown): value is Record { return isRecord(value) && Object.values(value).every((entry) => typeof entry === "string"); } +function validateOAuth(value: unknown): string | undefined { + if (value === undefined) return undefined; + if (!isRecord(value)) return "oauth must be an object"; + if (value.clientId !== undefined && typeof value.clientId !== "string") return "oauth.clientId must be a string"; + if (value.clientSecret !== undefined && typeof value.clientSecret !== "string") { + return "oauth.clientSecret must be a string"; + } + const port = value.callbackPort; + if (port !== undefined && (typeof port !== "number" || !Number.isInteger(port) || port < 1 || port > 65535)) { + return "oauth.callbackPort must be a port number"; + } + return undefined; +} + function validateServer(name: string, value: unknown): McpServerConfig | string { if (!SERVER_NAME.test(name)) return `invalid server name "${name}" (use letters, digits, "_" and "-")`; if (!isRecord(value)) return `server "${name}" must be an object`; @@ -93,6 +121,8 @@ function validateServer(name: string, value: unknown): McpServerConfig | string if (value.headers !== undefined && !isStringRecord(value.headers)) { return `server "${name}": headers must map names to strings`; } + const oauthError = validateOAuth(value.oauth); + if (oauthError) return `server "${name}": ${oauthError}`; return value as unknown as McpHttpServerConfig; } if (typeof value.command === "string" && (type === undefined || type === "stdio")) { diff --git a/packages/coding-agent/src/extensions/mcp/index.ts b/packages/coding-agent/src/extensions/mcp/index.ts index d02b7ca40..8171f595b 100644 --- a/packages/coding-agent/src/extensions/mcp/index.ts +++ b/packages/coding-agent/src/extensions/mcp/index.ts @@ -8,46 +8,93 @@ * * Every call runs through pi's tool pipeline, so `tool_call`/`tool_result` hooks and permission * extensions apply to MCP tools the same way they do to built-in tools. + * + * HTTP servers without an `Authorization` header authenticate with OAuth. A server that needs a + * sign-in shows as "needs sign-in" until the user runs `/mcp login `. */ import { basename, resolve } from "node:path"; import { pathToFileURL } from "node:url"; import { + type AuthProvider, + type CallToolResult, McpClient, + type McpRequestOptions, type Tool as McpTool, type McpTransport, StdioTransport, StreamableHttpTransport, } from "@earendil-works/pi-mcp"; +import { McpOAuthAuthorizationRequiredError, type OAuthChallenge } from "@earendil-works/pi-mcp/oauth"; +import type { AutocompleteItem } from "@earendil-works/pi-tui"; import { getAgentDir, VERSION } from "../../config.ts"; -import type { ExtensionAPI, ExtensionContext, ExtensionFactory } from "../../core/extensions/types.ts"; +import type { + ExtensionAPI, + ExtensionCommandContext, + ExtensionContext, + ExtensionFactory, +} from "../../core/extensions/types.ts"; import { resolveConfigValueOrThrow, resolveHeadersOrThrow } from "../../core/resolve-config-value.ts"; import { CODEMODE_TOOL_NAME } from "../../core/tools/codemode.ts"; +import { openBrowser } from "../../utils/open-browser.ts"; import { type LoadedMcpConfig, loadMcpConfig, type McpServerEntry } from "./config.ts"; -import { createMcpToolDefinition, createMcpToolName } from "./tools.ts"; +import { + createMcpAuthProvider, + McpOAuthCredentialStore, + type McpOAuthSettings, + McpSignInCancelledError, + signInMcpServer, +} from "./oauth.ts"; +import { createMcpToolDefinition, createMcpToolName, type McpToolCaller } from "./tools.ts"; const DEFAULT_TIMEOUT_SECONDS = 60; const STDERR_TAIL_CHARS = 2_000; -type ServerState = "connecting" | "connected" | "failed" | "closed"; +type ServerState = "connecting" | "connected" | "needs-auth" | "failed" | "closed"; + +export type McpTransportFactory = ( + entry: McpServerEntry, + cwd: string, + authProvider: AuthProvider | undefined, +) => McpTransport; export interface McpExtensionOptions { /** Defaults to reading `mcp.json` from the agent directory and the trusted project. */ loadConfig?: (ctx: ExtensionContext) => LoadedMcpConfig; /** Defaults to stdio and streamable HTTP transports built from the server config. */ - createTransport?: (entry: McpServerEntry, cwd: string) => McpTransport; + createTransport?: McpTransportFactory; + /** Defaults to `mcp-auth.json` in the agent directory. */ + credentials?: McpOAuthCredentialStore; + /** Opens the OAuth authorization URL. Defaults to the platform browser. */ + openUrl?: (url: string) => void; } function errorMessage(error: unknown): string { return error instanceof Error ? error.message : String(error); } -function createDefaultTransport(entry: McpServerEntry, cwd: string): McpTransport { +function signInRequiredMessage(name: string): string { + return `MCP server "${name}" requires sign-in. Run /mcp login ${name}.`; +} + +/** HTTP servers authenticate with OAuth unless the config supplies an `Authorization` header. */ +function usesOAuth(entry: McpServerEntry): boolean { + const { config } = entry; + if (!("url" in config)) return false; + return !Object.keys(config.headers ?? {}).some((header) => header.toLowerCase() === "authorization"); +} + +function createDefaultTransport( + entry: McpServerEntry, + cwd: string, + authProvider: AuthProvider | undefined, +): McpTransport { const { config, name } = entry; if ("url" in config) { return new StreamableHttpTransport({ url: config.url, headers: resolveHeadersOrThrow(config.headers, `MCP server "${name}"`), + authProvider, }); } const env: Record = {}; @@ -64,34 +111,67 @@ function createDefaultTransport(entry: McpServerEntry, cwd: string): McpTranspor } /** One configured server. Reconnects lazily when a call finds the connection gone. */ -class McpServerConnection { +class McpServerConnection implements McpToolCaller { readonly entry: McpServerEntry; state: ServerState = "connecting"; error: string | undefined; tools: McpTool[] = []; + /** Last OAuth challenge from the server; sign-in uses its resource metadata URL and scope. */ + challenge: OAuthChallenge | undefined; private client: McpClient | undefined; private opening: Promise | undefined; private closed = false; private readonly cwd: string; - private readonly createTransport: (entry: McpServerEntry, cwd: string) => McpTransport; + private readonly createTransport: McpTransportFactory; + private readonly authProvider: AuthProvider | undefined; private readonly onTools: (connection: McpServerConnection) => void; - constructor( - entry: McpServerEntry, - cwd: string, - createTransport: (entry: McpServerEntry, cwd: string) => McpTransport, - onTools: (connection: McpServerConnection) => void, - ) { - this.entry = entry; - this.cwd = cwd; - this.createTransport = createTransport; - this.onTools = onTools; + constructor(options: { + entry: McpServerEntry; + cwd: string; + createTransport: McpTransportFactory; + credentials: McpOAuthCredentialStore; + onTools: (connection: McpServerConnection) => void; + }) { + this.entry = options.entry; + this.cwd = options.cwd; + this.createTransport = options.createTransport; + this.onTools = options.onTools; + const url = this.oauthUrl; + this.authProvider = url + ? createMcpAuthProvider({ + serverUrl: url, + store: options.credentials.forServer(url), + settings: this.oauthSettings(), + onChallenge: (challenge) => { + this.challenge = challenge; + }, + }) + : undefined; } get timeoutMs(): number { return (this.entry.config.timeout ?? DEFAULT_TIMEOUT_SECONDS) * 1000; } + /** Server URL when the server authenticates with OAuth. */ + get oauthUrl(): string | undefined { + return usesOAuth(this.entry) && "url" in this.entry.config ? this.entry.config.url : undefined; + } + + oauthSettings(): McpOAuthSettings { + const oauth = "url" in this.entry.config ? this.entry.config.oauth : undefined; + if (!oauth) return {}; + return { + clientId: oauth.clientId, + clientSecret: + oauth.clientSecret === undefined + ? undefined + : resolveConfigValueOrThrow(oauth.clientSecret, `MCP server "${this.entry.name}" oauth.clientSecret`), + callbackPort: oauth.callbackPort, + }; + } + getClient(): Promise { if (this.closed) return Promise.reject(new Error(`MCP server "${this.entry.name}" is shut down`)); if (this.client?.connectionState === "connected") return Promise.resolve(this.client); @@ -101,6 +181,42 @@ class McpServerConnection { return this.opening; } + async callTool(name: string, args: Record, options: McpRequestOptions): Promise { + const client = await this.getClient(); + try { + return await client.callTool(name, args, options); + } catch (error) { + if (!(error instanceof McpOAuthAuthorizationRequiredError)) throw error; + await this.dropClient(client); + this.markNeedsAuth(); + throw new Error(signInRequiredMessage(this.entry.name)); + } + } + + /** Connect again with fresh credentials, for example after signing in. */ + async reconnect(): Promise { + await this.opening?.catch(() => undefined); + if (this.client) await this.dropClient(this.client); + await this.getClient(); + } + + /** Disconnect after the stored credentials were removed. */ + async signOut(): Promise { + await this.opening?.catch(() => undefined); + if (this.client) await this.dropClient(this.client); + if (!this.closed) this.markNeedsAuth(); + } + + private markNeedsAuth(): void { + this.state = "needs-auth"; + this.error = undefined; + } + + private async dropClient(client: McpClient): Promise { + if (this.client === client) this.client = undefined; + await client.close().catch(() => undefined); + } + private async open(): Promise { this.state = "connecting"; const client = new McpClient({ @@ -111,7 +227,7 @@ class McpServerConnection { }); let transport: McpTransport | undefined; try { - transport = this.createTransport(this.entry, this.cwd); + transport = this.createTransport(this.entry, this.cwd, this.authProvider); await client.connect(transport); client.onNotification("notifications/tools/list_changed", () => { void this.refreshTools(client); @@ -126,6 +242,10 @@ class McpServerConnection { return client; } catch (error) { await client.close().catch(() => undefined); + if (error instanceof McpOAuthAuthorizationRequiredError && !this.closed) { + this.markNeedsAuth(); + throw new Error(signInRequiredMessage(this.entry.name)); + } const stderr = transport instanceof StdioTransport ? transport.stderr.trim().slice(-STDERR_TAIL_CHARS) : ""; this.state = this.closed ? "closed" : "failed"; this.error = stderr ? `${errorMessage(error)}\n${stderr}` : errorMessage(error); @@ -158,20 +278,26 @@ function formatStatus(connections: readonly McpServerConnection[], configErrors: return `No MCP servers configured. Add them to ${resolve(getAgentDir(), "mcp.json")} or .pi/mcp.json.`; } const lines = connections.map((connection) => { - const exposure = connection.entry.config.exposure ?? "codemode"; + const { name, config } = connection.entry; + const exposure = config.exposure ?? "codemode"; + if (connection.state === "needs-auth") return `${name}: needs sign-in, run /mcp login ${name} (${exposure})`; const tools = connection.state === "connected" ? `, ${connection.tools.length} tools` : ""; const error = connection.error ? `\n ${connection.error.split("\n").join("\n ")}` : ""; - return `${connection.entry.name}: ${connection.state}${tools} (${exposure})${error}`; + return `${name}: ${connection.state}${tools} (${exposure})${error}`; }); for (const error of configErrors) lines.push(`config error: ${error}`); return lines.join("\n"); } +const MCP_USAGE = "Usage: /mcp, /mcp login [server], /mcp logout [server]"; + export function createMcpExtension(options: McpExtensionOptions = {}): ExtensionFactory { return (pi: ExtensionAPI) => { let connections: McpServerConnection[] = []; let configErrors: string[] = []; let pending: Promise | undefined; + const credentials = options.credentials ?? new McpOAuthCredentialStore(); + const openUrl = options.openUrl ?? openBrowser; const registerTools = (connection: McpServerConnection) => { const { name: server, config } = connection.entry; @@ -185,7 +311,7 @@ export function createMcpExtension(options: McpExtensionOptions = {}): Extension name, exposure, timeoutMs: connection.timeoutMs, - getClient: () => connection.getClient(), + getClient: async () => connection, }), ); } @@ -208,13 +334,98 @@ export function createMcpExtension(options: McpExtensionOptions = {}): Extension pi.setActiveTools([...pi.getActiveTools(), CODEMODE_TOOL_NAME]); }; + /** Resolve the server for `/mcp login|logout`, asking when the name is omitted and ambiguous. */ + const pickOAuthServer = async ( + name: string | undefined, + ctx: ExtensionCommandContext, + ): Promise => { + const candidates = connections.filter((connection) => connection.oauthUrl); + if (name) { + const connection = connections.find((candidate) => candidate.entry.name === name); + if (!connection) ctx.ui.notify(`No MCP server named "${name}".`, "error"); + else if (!connection.oauthUrl) { + ctx.ui.notify( + `MCP server "${name}" does not use OAuth. Only HTTP servers without an Authorization header do.`, + "error", + ); + } + return connection?.oauthUrl ? connection : undefined; + } + if (candidates.length === 0) { + ctx.ui.notify("No configured MCP server uses OAuth.", "info"); + return undefined; + } + const needsAuth = candidates.filter((connection) => connection.state === "needs-auth"); + if (candidates.length === 1) return candidates[0]; + if (needsAuth.length === 1) return needsAuth[0]; + const choice = await ctx.ui.select( + "MCP server", + candidates.map((connection) => connection.entry.name), + ); + return candidates.find((connection) => connection.entry.name === choice); + }; + + const login = async (connection: McpServerConnection, url: string, ctx: ExtensionCommandContext) => { + const { name } = connection.entry; + if (!ctx.hasUI) { + ctx.ui.notify(`Signing in to MCP server "${name}" requires interactive mode.`, "error"); + return; + } + try { + await signInMcpServer({ + serverUrl: url, + store: credentials.forServer(url), + settings: connection.oauthSettings(), + challenge: connection.challenge, + prompt: { + showAuthorizationUrl: (authorizationUrl) => { + ctx.ui.notify( + `Sign in to MCP server "${name}" in your browser:\n${authorizationUrl.href}`, + "info", + ); + openUrl(authorizationUrl.href); + }, + promptForRedirectUrl: (signal) => + ctx.ui.input( + `Waiting for sign-in to "${name}". If the browser cannot reach this machine, paste the URL it was redirected to.`, + "http://127.0.0.1:.../oauth/callback?code=...", + { signal }, + ), + }, + }); + } catch (error) { + if (error instanceof McpSignInCancelledError) ctx.ui.notify("MCP sign-in cancelled.", "info"); + else ctx.ui.notify(`Sign-in to MCP server "${name}" failed: ${errorMessage(error)}`, "error"); + return; + } + try { + await connection.reconnect(); + } catch (error) { + ctx.ui.notify(`Signed in, but ${errorMessage(error)}`, "error"); + return; + } + ensureCodemodeActive(ctx); + ctx.ui.notify(`Signed in to MCP server "${name}" (${connection.tools.length} tools).`, "info"); + }; + + const logout = async (connection: McpServerConnection, url: string, ctx: ExtensionCommandContext) => { + const { name } = connection.entry; + const removed = credentials.remove(url); + await connection.signOut(); + ctx.ui.notify( + removed ? `Signed out of MCP server "${name}".` : `No stored credentials for MCP server "${name}".`, + "info", + ); + }; + pi.on("session_start", (_event, ctx) => { const loaded = (options.loadConfig ?? defaultLoadConfig)(ctx); configErrors = loaded.errors; for (const error of configErrors) ctx.ui.notify(`MCP config: ${error}`, "warning"); const createTransport = options.createTransport ?? createDefaultTransport; connections = loaded.servers.map( - (entry) => new McpServerConnection(entry, ctx.cwd, createTransport, registerTools), + (entry) => + new McpServerConnection({ entry, cwd: ctx.cwd, createTransport, credentials, onTools: registerTools }), ); if (connections.length === 0) return; pending = Promise.allSettled( @@ -236,10 +447,41 @@ export function createMcpExtension(options: McpExtensionOptions = {}): Extension }); pi.registerCommand("mcp", { - description: "Show MCP server status", - handler: async (_args, ctx) => { + description: "Show MCP server status, or sign in and out of OAuth servers", + getArgumentCompletions: (prefix): AutocompleteItem[] | null => { + const [action, server, ...rest] = prefix.trimStart().split(/\s+/); + if (rest.length > 0) return null; + if (server === undefined) { + return ["login", "logout"] + .filter((item) => item.startsWith(action ?? "")) + .map((item) => ({ value: `${item} `, label: item })); + } + if (action !== "login" && action !== "logout") return null; + const items = connections + .filter((connection) => connection.oauthUrl && connection.entry.name.startsWith(server)) + .map((connection) => ({ + value: `${action} ${connection.entry.name}`, + label: connection.entry.name, + description: connection.state === "needs-auth" ? "needs sign-in" : connection.state, + })); + return items.length > 0 ? items : null; + }, + handler: async (args, ctx) => { await pending; - ctx.ui.notify(formatStatus(connections, configErrors), "info"); + const [action, name, ...extra] = args.trim().split(/\s+/).filter(Boolean); + if (action === undefined) { + ctx.ui.notify(formatStatus(connections, configErrors), "info"); + return; + } + if ((action !== "login" && action !== "logout") || extra.length > 0) { + ctx.ui.notify(MCP_USAGE, "warning"); + return; + } + const connection = await pickOAuthServer(name, ctx); + const url = connection?.oauthUrl; + if (!connection || !url) return; + if (action === "login") await login(connection, url, ctx); + else await logout(connection, url, ctx); }, }); }; diff --git a/packages/coding-agent/src/extensions/mcp/oauth.ts b/packages/coding-agent/src/extensions/mcp/oauth.ts new file mode 100644 index 000000000..d236ccca0 --- /dev/null +++ b/packages/coding-agent/src/extensions/mcp/oauth.ts @@ -0,0 +1,263 @@ +/** + * OAuth sign-in for remote MCP servers. + * + * Connections never start a browser flow on their own. They send the stored access token and, after + * a 401, try the stored refresh token. When that is not possible they fail with + * `McpOAuthAuthorizationRequiredError`, and the user signs in with `/mcp login `, which runs + * the authorization code flow (PKCE, dynamic client registration) against a loopback callback. + * + * Credentials live in `/mcp-auth.json`, keyed by server URL. + */ + +import { join } from "node:path"; +import type { AuthProvider } from "@earendil-works/pi-mcp"; +import { + adaptOAuthProvider, + authorizeMcp, + McpOAuthAuthorizationRequiredError, + McpOAuthProvider, + type McpOAuthState, + type McpOAuthStateStore, + OAuthCallbackServer, + type OAuthChallenge, + type OAuthClientInformationMixed, + parseWwwAuthenticate, +} from "@earendil-works/pi-mcp/oauth"; +import { APP_NAME, getAgentDir } from "../../config.ts"; +import { type AuthStorageBackend, FileAuthStorageBackend } from "../../core/auth-storage.ts"; + +const CALLBACK_HOST = "127.0.0.1"; +const CALLBACK_PATH = "/oauth/callback"; +/** Redirect URI for refreshes when none is stored. Refreshing never redirects the user. */ +const FALLBACK_REDIRECT_URL = `http://${CALLBACK_HOST}${CALLBACK_PATH}`; + +export interface McpOAuthSettings { + clientId?: string; + /** Already resolved. */ + clientSecret?: string; + callbackPort?: number; +} + +type StoredStates = Record; + +function parseStates(content: string | undefined): StoredStates { + if (!content?.trim()) return {}; + const parsed: unknown = JSON.parse(content); + return typeof parsed === "object" && parsed !== null && !Array.isArray(parsed) ? (parsed as StoredStates) : {}; +} + +/** Per-server OAuth state (client registration, tokens, pending PKCE verifier) in `mcp-auth.json`. */ +export class McpOAuthCredentialStore { + private readonly backend: AuthStorageBackend; + + constructor(backend: AuthStorageBackend = new FileAuthStorageBackend(join(getAgentDir(), "mcp-auth.json"))) { + this.backend = backend; + } + + forServer(serverUrl: string): McpOAuthStateStore { + const key = String(new URL(serverUrl)); + return { + load: () => this.read()[key], + save: (state) => + this.write((states) => { + states[key] = state; + }), + }; + } + + /** Returns whether credentials were stored for the server. */ + remove(serverUrl: string): boolean { + const key = String(new URL(serverUrl)); + if (!(key in this.read())) return false; + this.write((states) => { + delete states[key]; + }); + return true; + } + + private read(): StoredStates { + return this.backend.withLock((current) => ({ result: parseStates(current) })); + } + + private write(update: (states: StoredStates) => void): void { + this.backend.withLock((current) => { + const states = parseStates(current); + update(states); + return { result: undefined, next: `${JSON.stringify(states, null, 2)}\n` }; + }); + } +} + +function registeredRedirectUrls(client: OAuthClientInformationMixed | undefined): string[] { + return client && "redirect_uris" in client ? client.redirect_uris : []; +} + +function configuredRedirectUrl(settings: McpOAuthSettings): string | undefined { + return settings.callbackPort === undefined + ? undefined + : `http://${CALLBACK_HOST}:${settings.callbackPort}${CALLBACK_PATH}`; +} + +function createProvider( + serverUrl: string, + store: McpOAuthStateStore, + settings: McpOAuthSettings, + redirectUrl: string, + onRedirect: (url: URL) => void, +): McpOAuthProvider { + return new McpOAuthProvider({ + serverUrl, + redirectUrl, + clientMetadata: { client_name: APP_NAME }, + clientId: settings.clientId, + clientSecret: settings.clientSecret, + store, + onRedirect, + }); +} + +/** + * Auth provider for MCP connections: sends the stored access token and refreshes it after a 401. + * Throws `McpOAuthAuthorizationRequiredError` when the user has to sign in. `onChallenge` receives + * the server's `WWW-Authenticate` challenge so sign-in can use its resource metadata URL and scope. + */ +export function createMcpAuthProvider(options: { + serverUrl: string; + store: McpOAuthStateStore; + settings: McpOAuthSettings; + onChallenge: (challenge: OAuthChallenge) => void; +}): AuthProvider { + const { serverUrl, store, settings } = options; + return { + token: async () => (await store.load())?.tokens?.access_token, + onUnauthorized: async (context) => { + options.onChallenge(parseWwwAuthenticate(context.response.headers.get("www-authenticate"))); + const state = await store.load(); + if (!state?.tokens?.refresh_token) throw new McpOAuthAuthorizationRequiredError(); + const redirectUrl = + configuredRedirectUrl(settings) ?? + registeredRedirectUrls(state.clientInformation)[0] ?? + FALLBACK_REDIRECT_URL; + // Refreshes the tokens, or throws McpOAuthAuthorizationRequiredError if a new sign-in is needed. + await adaptOAuthProvider(createProvider(serverUrl, store, settings, redirectUrl, () => {})).onUnauthorized?.( + context, + ); + }, + }; +} + +export interface McpSignInPrompt { + /** Show the authorization URL to the user and open it in a browser. */ + showAuthorizationUrl(url: URL): void; + /** + * Ask for the redirect URL from the browser address bar, for when the browser cannot reach the + * loopback callback (for example over SSH). Aborted once the callback arrives. Resolves to + * `undefined` or an empty string when the user cancels. + */ + promptForRedirectUrl(signal: AbortSignal): Promise; +} + +export class McpSignInCancelledError extends Error { + constructor() { + super("Sign-in cancelled"); + this.name = "McpSignInCancelledError"; + } +} + +function codeFromRedirectUrl(input: string, state: string): string { + let url: URL; + try { + url = new URL(input.trim()); + } catch { + throw new Error("Expected the full redirect URL from the browser address bar"); + } + const error = url.searchParams.get("error"); + if (error) throw new Error(url.searchParams.get("error_description") ?? error); + if (url.searchParams.get("state") !== state) throw new Error("The redirect URL belongs to a different sign-in"); + const code = url.searchParams.get("code"); + if (!code) throw new Error("The redirect URL does not contain an authorization code"); + return code; +} + +/** Wait for the browser callback or a pasted redirect URL, whichever comes first. */ +async function waitForAuthorizationCode( + callback: OAuthCallbackServer, + state: string, + prompt: McpSignInPrompt, +): Promise { + const controller = new AbortController(); + const fromBrowser = callback.waitForCallback(state).then((result) => result.code); + const fromUser = prompt.promptForRedirectUrl(controller.signal).then((input) => { + if (!input?.trim()) throw new McpSignInCancelledError(); + return codeFromRedirectUrl(input, state); + }); + try { + return await Promise.race([fromBrowser, fromUser]); + } finally { + // The losing side rejects once the prompt is aborted or the callback server closes. + controller.abort(); + fromBrowser.catch(() => undefined); + fromUser.catch(() => undefined); + } +} + +async function listenForCallback(port: number | undefined, required: boolean): Promise { + try { + return await OAuthCallbackServer.listen({ host: CALLBACK_HOST, port: port ?? 0, path: CALLBACK_PATH }); + } catch (error) { + if (required || port === undefined) throw error; + return OAuthCallbackServer.listen({ host: CALLBACK_HOST, path: CALLBACK_PATH }); + } +} + +/** + * Sign in to an MCP server. Uses the stored refresh token when possible; otherwise runs the browser + * authorization code flow. Tokens are saved to `store`. + */ +export async function signInMcpServer(options: { + serverUrl: string; + store: McpOAuthStateStore; + settings: McpOAuthSettings; + challenge?: OAuthChallenge; + prompt: McpSignInPrompt; +}): Promise { + const { serverUrl, store, settings } = options; + const stored = await store.load(); + // Reuse the port of the registered redirect URI so the registered client stays valid. + const registered = registeredRedirectUrls(stored?.clientInformation)[0]; + const preferredPort = + settings.callbackPort ?? (registered ? Number(new URL(registered).port) || undefined : undefined); + const callback = await listenForCallback(preferredPort, settings.callbackPort !== undefined); + try { + if (stored) { + const next: McpOAuthState = { ...stored }; + // Every sign-in gets a fresh `state` parameter. + delete next.oauthState; + // A registered client cannot use another redirect URI, and its tokens belong to it. + if (!settings.clientId && !registeredRedirectUrls(stored.clientInformation).includes(callback.redirectUrl)) { + delete next.clientInformation; + delete next.tokens; + } + await store.save(next); + } + + let authorizationUrl: URL | undefined; + const provider = createProvider(serverUrl, store, settings, callback.redirectUrl, (url) => { + authorizationUrl = url; + }); + const flow = { + serverUrl, + resourceMetadataUrl: options.challenge?.resourceMetadataUrl, + scope: options.challenge?.scope, + }; + if ((await authorizeMcp(provider, flow)) === "AUTHORIZED") return; + if (!authorizationUrl) throw new Error("OAuth flow did not produce an authorization URL"); + + const state = await provider.state(); + options.prompt.showAuthorizationUrl(authorizationUrl); + const code = await waitForAuthorizationCode(callback, state, options.prompt); + await authorizeMcp(provider, { ...flow, authorizationCode: code }); + } finally { + await callback.close(); + } +} diff --git a/packages/coding-agent/test/suite/agent-session-mcp-oauth.test.ts b/packages/coding-agent/test/suite/agent-session-mcp-oauth.test.ts new file mode 100644 index 000000000..b0889ed14 --- /dev/null +++ b/packages/coding-agent/test/suite/agent-session-mcp-oauth.test.ts @@ -0,0 +1,292 @@ +import { createHash } from "node:crypto"; +import { createServer, type IncomingMessage, type Server, type ServerResponse } from "node:http"; +import { fauxAssistantMessage, fauxToolCall } from "@earendil-works/pi-ai"; +import type { ToolResultMessage } from "@earendil-works/pi-ai/compat"; +import { LATEST_PROTOCOL_VERSION } from "@earendil-works/pi-mcp"; +import { afterEach, describe, expect, it } from "vitest"; +import { InMemoryAuthStorageBackend } from "../../src/core/auth-storage.ts"; +import type { ExtensionUIContext } from "../../src/core/extensions/index.ts"; +import type { McpServerEntry } from "../../src/extensions/mcp/config.ts"; +import { createMcpExtension } from "../../src/extensions/mcp/index.ts"; +import { McpOAuthCredentialStore } from "../../src/extensions/mcp/oauth.ts"; +import { theme } from "../../src/modes/interactive/theme/theme.ts"; +import { createHarness, type Harness } from "./harness.ts"; + +async function readBody(request: IncomingMessage): Promise { + const chunks: Buffer[] = []; + for await (const chunk of request) chunks.push(Buffer.from(chunk)); + return Buffer.concat(chunks).toString("utf8"); +} + +function json(response: ServerResponse, status: number, body: unknown, headers: Record = {}): void { + response.writeHead(status, { "content-type": "application/json", ...headers }).end(JSON.stringify(body)); +} + +/** MCP server protected by OAuth, with its own authorization server (discovery, DCR, PKCE, refresh). */ +async function startOAuthMcpServer() { + const log: string[] = []; + const validTokens = new Set(); + const refreshTokens = new Set(); + const challenges = new Map(); + let issued = 0; + let origin = ""; + + const issueTokens = () => { + issued++; + const tokens = { access_token: `access-${issued}`, refresh_token: `refresh-${issued}` }; + validTokens.add(tokens.access_token); + refreshTokens.add(tokens.refresh_token); + return { ...tokens, token_type: "Bearer", expires_in: 3600 }; + }; + + const handleMcp = async (request: IncomingMessage, response: ServerResponse) => { + if (request.method !== "POST") { + response.writeHead(request.method === "GET" ? 405 : 200).end(); + return; + } + const token = request.headers.authorization?.replace(/^Bearer /, ""); + if (!token || !validTokens.has(token)) { + log.push(`401 ${token ?? "none"}`); + response + .writeHead(401, { + "www-authenticate": `Bearer resource_metadata="${origin}/.well-known/oauth-protected-resource/mcp"`, + }) + .end(); + return; + } + const message = JSON.parse(await readBody(request)) as { id?: number; method: string; params?: unknown }; + if (message.id === undefined) { + response.writeHead(202).end(); + return; + } + let result: unknown; + if (message.method === "initialize") { + result = { + protocolVersion: LATEST_PROTOCOL_VERSION, + capabilities: { tools: {} }, + serverInfo: { name: "issues", version: "1.0.0" }, + }; + } else if (message.method === "tools/list") { + result = { tools: [{ name: "whoami", inputSchema: { type: "object", properties: {} } }] }; + } else if (message.method === "tools/call") { + log.push(`call ${token}`); + result = { content: [{ type: "text", text: `token ${token}` }] }; + } else { + result = {}; + } + json(response, 200, { jsonrpc: "2.0", id: message.id, result }); + }; + + const handle = async (request: IncomingMessage, response: ServerResponse) => { + const url = new URL(request.url ?? "/", origin); + switch (url.pathname) { + case "/mcp": + return handleMcp(request, response); + case "/.well-known/oauth-protected-resource/mcp": + return json(response, 200, { resource: `${origin}/mcp`, authorization_servers: [origin] }); + case "/.well-known/oauth-authorization-server": + return json(response, 200, { + issuer: origin, + authorization_endpoint: `${origin}/authorize`, + token_endpoint: `${origin}/token`, + registration_endpoint: `${origin}/register`, + response_types_supported: ["code"], + code_challenge_methods_supported: ["S256"], + token_endpoint_auth_methods_supported: ["none"], + }); + case "/register": { + const metadata = JSON.parse(await readBody(request)) as Record; + log.push("register"); + return json(response, 201, { ...metadata, client_id: "client-1" }); + } + case "/authorize": { + const code = `code-${challenges.size + 1}`; + challenges.set(code, url.searchParams.get("code_challenge") ?? ""); + const redirect = new URL(url.searchParams.get("redirect_uri") ?? ""); + redirect.searchParams.set("code", code); + redirect.searchParams.set("state", url.searchParams.get("state") ?? ""); + response.writeHead(302, { location: redirect.href }).end(); + return; + } + case "/token": { + const params = new URLSearchParams(await readBody(request)); + if (params.get("grant_type") === "authorization_code") { + const challenge = challenges.get(params.get("code") ?? ""); + const verifier = createHash("sha256") + .update(params.get("code_verifier") ?? "") + .digest("base64url"); + if (!challenge || challenge !== verifier) return json(response, 400, { error: "invalid_grant" }); + challenges.delete(params.get("code") ?? ""); + log.push("token code"); + return json(response, 200, issueTokens()); + } + const refresh = params.get("refresh_token") ?? ""; + if (!refreshTokens.delete(refresh)) return json(response, 400, { error: "invalid_grant" }); + log.push("token refresh"); + return json(response, 200, issueTokens()); + } + default: + response.writeHead(404).end(); + } + }; + + const server: Server = createServer((request, response) => { + void handle(request, response).catch((error) => { + response.writeHead(500).end(String(error)); + }); + }); + await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)); + const address = server.address(); + if (!address || typeof address === "string") throw new Error("test server did not bind to TCP"); + origin = `http://127.0.0.1:${address.port}`; + return { + url: `${origin}/mcp`, + log, + /** Simulates access token expiry. */ + expireAccessTokens: () => validTokens.clear(), + close: () => + new Promise((resolve) => { + server.closeAllConnections(); + server.close(() => resolve()); + }), + }; +} + +function createUiContext(options: { + notifications: string[]; + /** Answers the paste-redirect-URL prompt; by default it waits until sign-in completes. */ + answerInput?: () => Promise; +}): ExtensionUIContext { + return { + select: async () => undefined, + confirm: async () => false, + input: (_title, _placeholder, opts) => + options.answerInput?.() ?? + new Promise((resolve) => opts?.signal?.addEventListener("abort", () => resolve(undefined), { once: true })), + notify: (message) => options.notifications.push(message), + onTerminalInput: () => () => {}, + setStatus: () => {}, + setWorkingMessage: () => {}, + setWorkingVisible: () => {}, + setWorkingIndicator: () => {}, + setHiddenThinkingLabel: () => {}, + setWidget: () => {}, + setFooter: () => {}, + setHeader: () => {}, + setTitle: () => {}, + custom: async () => undefined as T, + pasteToEditor: () => {}, + setEditorText: () => {}, + getEditorText: () => "", + editor: async () => undefined, + addAutocompleteProvider: () => {}, + setEditorComponent: () => {}, + getEditorComponent: () => undefined, + get theme() { + return theme; + }, + getAllThemes: () => [], + getTheme: () => undefined, + setTheme: () => ({ success: false, error: "not available in tests" }), + getToolsExpanded: () => false, + setToolsExpanded: () => {}, + }; +} + +describe("AgentSession MCP OAuth", () => { + const cleanups: (() => Promise | void)[] = []; + + afterEach(async () => { + while (cleanups.length > 0) await cleanups.pop()?.(); + }); + + async function setup(browser: "follow" | "paste") { + const server = await startOAuthMcpServer(); + cleanups.push(server.close); + const backend = new InMemoryAuthStorageBackend(); + const entry: McpServerEntry = { name: "issues", config: { url: server.url, exposure: "direct" }, source: "test" }; + const notifications: string[] = []; + let redirectLocation: Promise | undefined; + const harness: Harness = await createHarness({ + initialActiveToolNames: [], + extensionFactories: [ + createMcpExtension({ + loadConfig: () => ({ servers: [entry], errors: [] }), + credentials: new McpOAuthCredentialStore(backend), + openUrl: (url) => { + if (browser === "follow") { + // The browser follows the authorization redirect to the loopback callback. + void fetch(url); + } else { + // The browser cannot reach the callback; the user pastes the redirect URL. + redirectLocation = fetch(url, { redirect: "manual" }).then( + (response) => response.headers.get("location") ?? "", + ); + } + }, + }), + ], + }); + cleanups.push(() => harness.cleanup()); + await harness.session.bindExtensions({ + uiContext: createUiContext({ + notifications, + answerInput: browser === "paste" ? async () => redirectLocation : undefined, + }), + }); + return { harness, server, notifications, backend }; + } + + async function callWhoami(harness: Harness): Promise { + harness.setResponses([ + fauxAssistantMessage([fauxToolCall("mcp__issues__whoami", {})], { stopReason: "toolUse" }), + fauxAssistantMessage("done"), + ]); + const before = harness.session.messages.length; + await harness.session.prompt("who am i"); + const result = harness.session.messages + .slice(before) + .find((message): message is ToolResultMessage => message.role === "toolResult"); + if (!result) throw new Error("no tool result"); + return result; + } + + function text(message: ToolResultMessage): string { + return message.content.map((block) => (block.type === "text" ? block.text : "")).join("\n"); + } + + it("signs in through the browser, refreshes expired tokens, and signs out", async () => { + const { harness, server, notifications, backend } = await setup("follow"); + + await harness.session.prompt("/mcp"); + expect(notifications).toContain('MCP server "issues" requires sign-in. Run /mcp login issues.'); + expect(notifications.at(-1)).toBe("issues: needs sign-in, run /mcp login issues (direct)"); + + await harness.session.prompt("/mcp login issues"); + expect(notifications.at(-1)).toBe('Signed in to MCP server "issues" (1 tools).'); + expect(server.log).toEqual(["401 none", "register", "token code"]); + expect(backend.withLock((current) => ({ result: current }))).toContain('"access_token": "access-1"'); + + expect(text(await callWhoami(harness))).toBe("token access-1"); + + // An expired access token is refreshed without user interaction. + server.expireAccessTokens(); + expect(text(await callWhoami(harness))).toBe("token access-2"); + expect(server.log.slice(-3)).toEqual(["401 access-1", "token refresh", "call access-2"]); + + await harness.session.prompt("/mcp logout issues"); + expect(notifications.at(-1)).toBe('Signed out of MCP server "issues".'); + const result = await callWhoami(harness); + expect(result.isError).toBe(true); + expect(text(result)).toBe('MCP server "issues" requires sign-in. Run /mcp login issues.'); + }); + + it("accepts a pasted redirect URL when the browser cannot reach the callback", async () => { + const { harness, server, notifications } = await setup("paste"); + + await harness.session.prompt("/mcp login"); + expect(notifications.at(-1)).toBe('Signed in to MCP server "issues" (1 tools).'); + expect(text(await callWhoami(harness))).toBe(`token access-1`); + expect(server.log).toContain("token code"); + }); +});