mirror of
https://github.com/badlogic/pi-mono.git
synced 2026-11-05 23:37:25 +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
|
|
@ -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:<port>/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__<server>__<tool>`. The `exposure` setting controls how the model reaches them:
|
||||
|
|
@ -42,5 +73,3 @@ Each server's tools are registered as `mcp__<server>__<tool>`. 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.
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
}
|
||||
}
|
||||
292
packages/coding-agent/test/suite/agent-session-mcp-oauth.test.ts
Normal file
292
packages/coding-agent/test/suite/agent-session-mcp-oauth.test.ts
Normal file
|
|
@ -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<string> {
|
||||
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<string, string> = {}): 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<string>();
|
||||
const refreshTokens = new Set<string>();
|
||||
const challenges = new Map<string, string>();
|
||||
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<string, unknown>;
|
||||
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<void>((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<void>((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<string | undefined>;
|
||||
}): 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 <T>() => 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> | 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<string> | 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<ToolResultMessage> {
|
||||
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");
|
||||
});
|
||||
});
|
||||
Loading…
Add table
Add a link
Reference in a new issue