mirror of
https://github.com/badlogic/pi-mono.git
synced 2026-11-06 23:44:36 +00:00
feat(coding-agent): add OAuth sign-in for remote MCP servers
HTTP servers without an Authorization header use OAuth. Connections send the stored token and refresh it after a 401; /mcp login runs the browser authorization code flow (PKCE, dynamic client registration) with a loopback callback or a pasted redirect URL, and /mcp logout removes the credentials. Tokens are stored per server URL in mcp-auth.json. Pre-registered clients are configured with the oauth block.
This commit is contained in:
Armin Ronacher
parent
e5a252c102
commit
8cd0ed48de
5 changed files with 885 additions and 29 deletions
|
|
@ -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<string, string>;
|
||||
oauth?: McpOAuthConfig;
|
||||
}
|
||||
|
||||
export type McpServerConfig = McpStdioServerConfig | McpHttpServerConfig;
|
||||
|
|
@ -76,6 +90,20 @@ function isStringRecord(value: unknown): value is Record<string, string> {
|
|||
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")) {
|
||||
|
|
|
|||
|
|
@ -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 <server>`.
|
||||
*/
|
||||
|
||||
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<string, string> = {};
|
||||
|
|
@ -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<McpClient> | 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<McpClient> {
|
||||
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<string, unknown>, options: McpRequestOptions): Promise<CallToolResult> {
|
||||
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<void> {
|
||||
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<void> {
|
||||
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<void> {
|
||||
if (this.client === client) this.client = undefined;
|
||||
await client.close().catch(() => undefined);
|
||||
}
|
||||
|
||||
private async open(): Promise<McpClient> {
|
||||
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<unknown> | 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<McpServerConnection | undefined> => {
|
||||
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);
|
||||
},
|
||||
});
|
||||
};
|
||||
|
|
|
|||
263
packages/coding-agent/src/extensions/mcp/oauth.ts
Normal file
263
packages/coding-agent/src/extensions/mcp/oauth.ts
Normal file
|
|
@ -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 <server>`, which runs
|
||||
* the authorization code flow (PKCE, dynamic client registration) against a loopback callback.
|
||||
*
|
||||
* Credentials live in `<agent-dir>/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<string, McpOAuthState>;
|
||||
|
||||
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<string | undefined>;
|
||||
}
|
||||
|
||||
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<string> {
|
||||
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<OAuthCallbackServer> {
|
||||
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<void> {
|
||||
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();
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue