From 49b35edf6da698148b28f9c1d3e44e73e4707958 Mon Sep 17 00:00:00 2001 From: FRFlo <62435410+FRFlo@users.noreply.github.com> Date: Mon, 19 Oct 2026 15:23:33 +0200 Subject: [PATCH] feat: support OAuth for remote Executor MCP --- README.md | 10 ++ src/mcp-client.ts | 105 ++++++++++++----- src/oauth-auth.ts | 116 ++++++++++++++++++ src/oauth-callback.ts | 89 ++++++++++++++ src/tools.ts | 129 ++++++++++++++------ test/oauth-auth.test.ts | 124 ++++++++++++++++++++ test/oauth-callback.test.ts | 71 +++++++++++ test/oauth-mcp-client.integration.test.ts | 137 ++++++++++++++++++++++ 8 files changed, 715 insertions(+), 66 deletions(-) create mode 100644 src/oauth-auth.ts create mode 100644 src/oauth-callback.ts create mode 100644 test/oauth-auth.test.ts create mode 100644 test/oauth-callback.test.ts create mode 100644 test/oauth-mcp-client.integration.test.ts diff --git a/README.md b/README.md index f3f97ed..b33dbbf 100644 --- a/README.md +++ b/README.md @@ -93,6 +93,16 @@ Configure the extension in `~/.pi/agent/settings.json` or `.pi/settings.json`: You can also manage these interactively with `/executor-settings`. +### Remote Executor authentication + +When a remote MCP endpoint requires OAuth, an interactive Pi `execute` call opens +the authorization page in the system browser and completes the login through a +temporary localhost callback. OAuth client registration and tokens are stored +outside the project in `~/.pi/executor/oauth/`, in per-server JSON files with +user-only permissions where supported. These files contain credentials in +plaintext; protect your user account and do not share or commit them. Headless +calls reuse saved tokens but cannot start an interactive login. + ## Use in Pi Isolated project mode (loads only `pi-executor`, with no skills and no other extensions): diff --git a/src/mcp-client.ts b/src/mcp-client.ts index ce974a4..6e6253b 100644 --- a/src/mcp-client.ts +++ b/src/mcp-client.ts @@ -1,4 +1,5 @@ import { Client } from "@modelcontextprotocol/sdk/client/index.js"; +import type { OAuthClientProvider } from "@modelcontextprotocol/sdk/client/auth.js"; import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js"; import { ElicitRequestSchema, type ClientCapabilities } from "@modelcontextprotocol/sdk/types.js"; import type { JsonObject, JsonValue } from "./http.ts"; @@ -41,6 +42,8 @@ export type ExecutorMcpToolResult = { type ExecutorMcpClientOptions = { hasUI: boolean; + oauthProvider?: OAuthClientProvider; + waitForAuthorizationCode?: () => Promise; onElicitation?: (request: ExecutorElicitationRequest) => Promise; }; @@ -112,38 +115,80 @@ const connectExecutorMcpClient = async ( baseUrl: string, options: ExecutorMcpClientOptions, ): Promise => { - const client = new Client( - { name: "pi-executor", version: "0.0.1" }, - { capabilities: buildCapabilities(options.hasUI) }, - ); - const transport = new StreamableHTTPClientTransport(new URL("/mcp", baseUrl)); - - const onElicitation = options.onElicitation; - if (onElicitation) { - client.setRequestHandler(ElicitRequestSchema, async (request) => { - const params = request.params; - const response = await onElicitation( - params.mode === "url" - ? { - mode: "url", - message: params.message, - url: params.url, - elicitationId: params.elicitationId, - } - : { - mode: "form", - message: params.message, - requestedSchema: cloneJsonObject(params.requestedSchema), - }, - ); - - return response.content - ? { action: response.action, content: response.content } - : { action: response.action }; + let authorizationRedirected = false; + let authorizationRedirectError: unknown; + const authProvider = options.oauthProvider + ? { + ...options.oauthProvider, + redirectToAuthorization: async (url: URL) => { + authorizationRedirected = true; + try { + await options.oauthProvider?.redirectToAuthorization(url); + } catch (error) { + authorizationRedirectError = error; + throw error; + } + }, + } + : undefined; + const createConnection = () => { + const client = new Client( + { name: "pi-executor", version: "0.0.1" }, + { capabilities: buildCapabilities(options.hasUI) }, + ); + const transport = new StreamableHTTPClientTransport(new URL("/mcp", baseUrl), { + authProvider, }); - } + if (options.onElicitation) { + client.setRequestHandler(ElicitRequestSchema, async (request) => { + const params = request.params; + const response = await options.onElicitation!( + params.mode === "url" + ? { + mode: "url", + message: params.message, + url: params.url, + elicitationId: params.elicitationId, + } + : { + mode: "form", + message: params.message, + requestedSchema: cloneJsonObject(params.requestedSchema), + }, + ); - await client.connect(transport); + return response.content + ? { action: response.action, content: response.content } + : { action: response.action }; + }); + } + return { client, transport }; + }; + + let { client, transport } = createConnection(); + try { + await client.connect(transport); + } catch (error) { + if (authorizationRedirectError) { + await transport.terminateSession().catch(() => undefined); + await client.close().catch(() => undefined); + throw authorizationRedirectError; + } + if (!authorizationRedirected || !options.oauthProvider || !options.waitForAuthorizationCode) { + throw error; + } + try { + await transport.finishAuth(await options.waitForAuthorizationCode()); + await transport.terminateSession().catch(() => undefined); + await client.close().catch(() => undefined); + ({ client, transport } = createConnection()); + await client.connect(transport); + } catch (authorizationError) { + await transport.terminateSession().catch(() => undefined); + await client.close().catch(() => undefined); + throw authorizationError; + } + } const listTools = async (): Promise => { const tools: ExecutorMcpToolMetadata[] = []; diff --git a/src/oauth-auth.ts b/src/oauth-auth.ts new file mode 100644 index 0000000..ba2739e --- /dev/null +++ b/src/oauth-auth.ts @@ -0,0 +1,116 @@ +import { createHash } from "node:crypto"; +import { chmod, mkdir, readFile, writeFile } from "node:fs/promises"; +import { join } from "node:path"; +import type { + OAuthClientInformationMixed, + OAuthClientMetadata, + OAuthTokens, +} from "@modelcontextprotocol/sdk/shared/auth.js"; + +type PersistedOAuthState = { + clientInformation?: OAuthClientInformationMixed; + registeredRedirectUrl?: string; + tokens?: OAuthTokens; + codeVerifier?: string; +}; + +export type ExecutorOAuthProviderOptions = { + serverUrl: string; + storageRoot: string; + redirectUrl: string | URL; + state: string; + redirectToAuthorization: (url: URL) => void | Promise; +}; + +const readState = async (path: string): Promise => { + try { + return JSON.parse(await readFile(path, "utf8")) as PersistedOAuthState; + } catch (error) { + if ((error as NodeJS.ErrnoException).code === "ENOENT") { + return {}; + } + throw error; + } +}; + +export const createExecutorOAuthProvider = (options: ExecutorOAuthProviderOptions) => { + const origin = new URL(options.serverUrl).origin; + const fileName = `${createHash("sha256").update(origin).digest("hex")}.json`; + const statePath = join(options.storageRoot, fileName); + + const updateState = async (update: Partial): Promise => { + await mkdir(options.storageRoot, { recursive: true, mode: 0o700 }); + if (process.platform !== "win32") { + await chmod(options.storageRoot, 0o700); + } + const state = { ...(await readState(statePath)), ...update }; + await writeFile(statePath, `${JSON.stringify(state, null, 2)}\n`, { mode: 0o600 }); + if (process.platform !== "win32") { + await chmod(statePath, 0o600); + } + }; + + const clientMetadata: OAuthClientMetadata = { + client_name: "Pi Executor", + redirect_uris: [String(options.redirectUrl)], + grant_types: ["authorization_code", "refresh_token"], + response_types: ["code"], + token_endpoint_auth_method: "none", + }; + + return { + get redirectUrl() { + return options.redirectUrl; + }, + get clientMetadata() { + return clientMetadata; + }, + state: () => options.state, + clientInformation: async () => { + const state = await readState(statePath); + const clientInformation = state.clientInformation; + const registeredRedirectUrl = + state.registeredRedirectUrl ?? + (clientInformation && "redirect_uris" in clientInformation + ? clientInformation.redirect_uris?.[0] + : undefined); + if (registeredRedirectUrl && registeredRedirectUrl !== String(options.redirectUrl)) { + return undefined; + } + return clientInformation; + }, + saveClientInformation: async (clientInformation: OAuthClientInformationMixed) => + updateState({ + clientInformation, + registeredRedirectUrl: String(options.redirectUrl), + }), + tokens: async () => (await readState(statePath)).tokens, + saveTokens: async (tokens: OAuthTokens) => updateState({ tokens, codeVerifier: undefined }), + codeVerifier: async () => { + const verifier = (await readState(statePath)).codeVerifier; + if (!verifier) { + throw new Error("OAuth PKCE verifier is missing; restart authentication."); + } + return verifier; + }, + saveCodeVerifier: async (codeVerifier: string) => updateState({ codeVerifier }), + redirectToAuthorization: options.redirectToAuthorization, + invalidateCredentials: async ( + scope: "all" | "client" | "tokens" | "verifier" | "discovery", + ) => { + if (scope === "all") { + await updateState({ + clientInformation: undefined, + tokens: undefined, + codeVerifier: undefined, + }); + } else if (scope === "client") { + await updateState({ clientInformation: undefined }); + } else if (scope === "tokens") { + await updateState({ tokens: undefined }); + } else if (scope === "verifier") { + await updateState({ codeVerifier: undefined }); + } + }, + }; +}; diff --git a/src/oauth-callback.ts b/src/oauth-callback.ts new file mode 100644 index 0000000..62a91cb --- /dev/null +++ b/src/oauth-callback.ts @@ -0,0 +1,89 @@ +import { createServer, type Server } from "node:http"; + +export type OAuthCallbackOptions = { state: string; timeoutMs?: number }; +export type OAuthCallback = { + redirectUrl: string; + waitForCode(): Promise; + cancel(): void; + close(): Promise; +}; + +export async function createOAuthCallback({ + state, + timeoutMs = 5 * 60_000, +}: OAuthCallbackOptions): Promise { + let resolveCode!: (code: string) => void; + let rejectCode!: (error: Error) => void; + let settled = false; + let waiting = false; + let timer: ReturnType | undefined; + let server: Server; + const result = new Promise((resolve, reject) => { + resolveCode = resolve; + rejectCode = reject; + }); + // Prevent an early timeout/cancel from becoming an unhandled rejection before waiting. + result.catch(() => undefined); + const finish = (error?: Error, code?: string) => { + if (settled) return; + settled = true; + if (timer) clearTimeout(timer); + if (error) rejectCode(error); + else resolveCode(code!); + server.close(); + }; + server = createServer((req, res) => { + const url = new URL(req.url ?? "/", "http://127.0.0.1"); + const params = url.searchParams; + if (url.pathname !== "/callback") { + res.writeHead(404, { "Content-Type": "text/plain; charset=utf-8" }).end("Not found."); + finish(new Error("Unexpected OAuth callback path.")); + return; + } + const fail = (message: string) => { + res.writeHead(400, { "Content-Type": "text/plain; charset=utf-8" }).end(message); + finish(new Error(message)); + }; + if (params.get("state") !== state) return fail("Invalid OAuth state."); + if (params.has("error")) + return fail( + `Authorization failed: ${params.get("error_description") ?? params.get("error")}`, + ); + const code = params.get("code"); + if (!code) return fail("Missing authorization code."); + res + .writeHead(200, { "Content-Type": "text/plain; charset=utf-8" }) + .end("Authorization complete. You can close this window."); + finish(undefined, code); + }); + await new Promise((resolve, reject) => { + server.once("error", reject); + server.listen(0, "127.0.0.1", () => { + server.removeListener("error", reject); + resolve(); + }); + }); + const address = server.address(); + if (!address || typeof address === "string") + throw new Error("Unable to determine callback listener address."); + return { + redirectUrl: `http://127.0.0.1:${address.port}/callback`, + waitForCode() { + if (waiting) return Promise.reject(new Error("Callback can only be awaited once.")); + waiting = true; + if (!settled) + timer = setTimeout(() => finish(new Error("OAuth callback timeout.")), timeoutMs); + return result; + }, + cancel() { + finish(new Error("OAuth callback cancelled.")); + }, + close() { + if (!settled) finish(new Error("OAuth callback closed.")); + return new Promise((resolve) => { + if (!server.listening) return resolve(); + server.close(() => resolve()); + }); + }, + }; +} diff --git a/src/tools.ts b/src/tools.ts index 28b858e..d306064 100644 --- a/src/tools.ts +++ b/src/tools.ts @@ -1,4 +1,7 @@ import { StringEnum, Type, type Static } from "@mariozechner/pi-ai"; +import { randomBytes } from "node:crypto"; +import { homedir } from "node:os"; +import { join } from "node:path"; import { defineTool, type ExtensionAPI, @@ -8,6 +11,8 @@ import { import type { JsonObject, JsonValue } from "./http.ts"; import type { ResumeAction, ExecutorMcpInspection } from "./mcp-client.ts"; import { inspectExecutorMcp, withExecutorMcpClient } from "./mcp-client.ts"; +import { createExecutorOAuthProvider } from "./oauth-auth.ts"; +import { createOAuthCallback } from "./oauth-callback.ts"; import { buildExecutorSystemPrompt, parseJsonContent, @@ -107,6 +112,28 @@ const launchBrowser = async (url: string): Promise => { }); }; +const createRemoteOAuthSession = async (baseUrl: string, interactive: boolean) => { + const state = randomBytes(32).toString("base64url"); + const callback = interactive ? await createOAuthCallback({ state }) : undefined; + const provider = createExecutorOAuthProvider({ + serverUrl: baseUrl, + storageRoot: join(homedir(), ".pi", "executor", "oauth"), + redirectUrl: callback?.redirectUrl ?? "http://127.0.0.1/oauth-callback", + state, + redirectToAuthorization: (url) => { + if (!interactive) { + throw new Error("Remote Executor authorization requires an interactive Pi session."); + } + return launchBrowser(url.href); + }, + }); + return { + oauthProvider: provider, + ...(callback ? { waitForAuthorizationCode: callback.waitForCode } : {}), + close: () => callback?.close(), + }; +}; + const promptForInteraction = async ( interaction: { mode: "form" | "url"; @@ -275,31 +302,45 @@ const buildExecuteTool = (description: string) => }), async execute(_toolCallId, params, _signal, _onUpdate, ctx): Promise { const endpoint = await connectExecutor(ctx); - - const outcome = await withExecutorMcpClient( - endpoint.baseUrl, - { - hasUI: ctx.hasUI, - onElicitation: ctx.hasUI - ? (interaction) => - promptForInteraction( - interaction.mode === "url" - ? { - mode: "url", - message: interaction.message, - url: interaction.url, - } - : { - mode: "form", - message: interaction.message, - requestedSchema: interaction.requestedSchema, - }, - ctx, - ) - : undefined, - }, - async (client) => client.execute(params.code), - ); + const oauthSession = + endpoint.mode === "remote" + ? await createRemoteOAuthSession(endpoint.baseUrl, ctx.hasUI) + : undefined; + let outcome; + try { + outcome = await withExecutorMcpClient( + endpoint.baseUrl, + { + hasUI: ctx.hasUI, + ...(oauthSession + ? { + oauthProvider: oauthSession.oauthProvider, + waitForAuthorizationCode: oauthSession.waitForAuthorizationCode, + } + : {}), + onElicitation: ctx.hasUI + ? (interaction) => + promptForInteraction( + interaction.mode === "url" + ? { + mode: "url", + message: interaction.message, + url: interaction.url, + } + : { + mode: "form", + message: interaction.message, + requestedSchema: interaction.requestedSchema, + }, + ctx, + ) + : undefined, + }, + async (client) => client.execute(params.code), + ); + } finally { + await oauthSession?.close(); + } return toToolResult(outcome, { baseUrl: endpoint.baseUrl, scopeId: endpoint.scope.id }); }, @@ -320,17 +361,33 @@ const buildResumeTool = (description: string) => }), async execute(_toolCallId, params, _signal, _onUpdate, ctx): Promise { const endpoint = await connectExecutor(ctx); - - const outcome = await withExecutorMcpClient( - endpoint.baseUrl, - { hasUI: false }, - async (client) => - client.resume( - params.executionId, - params.action as ResumeAction, - parseJsonContent(params.content), - ), - ); + const oauthSession = + endpoint.mode === "remote" + ? await createRemoteOAuthSession(endpoint.baseUrl, false) + : undefined; + let outcome; + try { + outcome = await withExecutorMcpClient( + endpoint.baseUrl, + { + hasUI: false, + ...(oauthSession + ? { + oauthProvider: oauthSession.oauthProvider, + waitForAuthorizationCode: oauthSession.waitForAuthorizationCode, + } + : {}), + }, + async (client) => + client.resume( + params.executionId, + params.action as ResumeAction, + parseJsonContent(params.content), + ), + ); + } finally { + await oauthSession?.close(); + } return toToolResult(outcome, { baseUrl: endpoint.baseUrl, diff --git a/test/oauth-auth.test.ts b/test/oauth-auth.test.ts new file mode 100644 index 0000000..496f3f8 --- /dev/null +++ b/test/oauth-auth.test.ts @@ -0,0 +1,124 @@ +import { afterEach, describe, expect, test } from "bun:test"; +import { mkdtemp, readFile, rm, stat } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { createExecutorOAuthProvider } from "../src/oauth-auth.ts"; + +const temporaryDirectories: string[] = []; + +const createTemporaryDirectory = async (): Promise => { + const directory = await mkdtemp(join(tmpdir(), "pi-executor-oauth-test-")); + temporaryDirectories.push(directory); + return directory; +}; + +afterEach(async () => { + while (temporaryDirectories.length > 0) { + await rm(temporaryDirectories.pop()!, { recursive: true, force: true }); + } +}); + +describe("Executor OAuth provider", () => { + test("persists OAuth state by server origin with private file permissions", async () => { + const storageRoot = await createTemporaryDirectory(); + const provider = createExecutorOAuthProvider({ + serverUrl: "https://executor.example/mcp", + storageRoot, + redirectUrl: "http://127.0.0.1:12345/callback", + state: "state-1", + redirectToAuthorization: async () => undefined, + }); + + await provider.saveClientInformation?.({ + client_id: "client-1", + client_secret: "client-secret", + }); + await provider.saveTokens({ access_token: "access-secret", token_type: "Bearer" }); + await provider.saveCodeVerifier("pkce-verifier"); + + const reloaded = createExecutorOAuthProvider({ + serverUrl: "https://executor.example/another-path", + storageRoot, + redirectUrl: "http://127.0.0.1:12345/callback", + state: "state-1", + redirectToAuthorization: async () => undefined, + }); + expect(await reloaded.clientInformation()).toEqual({ + client_id: "client-1", + client_secret: "client-secret", + }); + expect(await reloaded.tokens()).toEqual({ + access_token: "access-secret", + token_type: "Bearer", + }); + expect(await reloaded.codeVerifier()).toBe("pkce-verifier"); + + const files = await Array.fromAsync( + new Bun.Glob("*.json").scan({ cwd: storageRoot, onlyFiles: true }), + ); + expect(files).toHaveLength(1); + expect(await readFile(join(storageRoot, files[0]!), "utf8")).toContain("access-secret"); + if (process.platform !== "win32") { + expect((await stat(join(storageRoot, files[0]!))).mode & 0o777).toBe(0o600); + } + }); + + test("keeps credentials isolated by origin", async () => { + const storageRoot = await createTemporaryDirectory(); + const createProvider = (serverUrl: string) => + createExecutorOAuthProvider({ + serverUrl, + storageRoot, + redirectUrl: "http://127.0.0.1:12345/callback", + state: "state-1", + redirectToAuthorization: async () => undefined, + }); + + await createProvider("https://one.example/mcp").saveTokens({ + access_token: "one-token", + token_type: "Bearer", + }); + expect(await createProvider("https://two.example/mcp").tokens()).toBeUndefined(); + }); + + test("does not reuse a client registration bound to another callback port", async () => { + const storageRoot = await createTemporaryDirectory(); + const first = createExecutorOAuthProvider({ + serverUrl: "https://executor.example/mcp", + storageRoot, + redirectUrl: "http://127.0.0.1:12345/callback", + state: "state-1", + redirectToAuthorization: async () => undefined, + }); + await first.saveClientInformation?.({ + client_id: "client-1", + redirect_uris: ["http://127.0.0.1:12345/callback"], + }); + const nextSession = createExecutorOAuthProvider({ + serverUrl: "https://executor.example/mcp", + storageRoot, + redirectUrl: "http://127.0.0.1:54321/callback", + state: "state-2", + redirectToAuthorization: async () => undefined, + }); + + expect(await nextSession.clientInformation()).toBeUndefined(); + }); + + test("opens the authorization URL through the supplied browser callback", async () => { + const storageRoot = await createTemporaryDirectory(); + const opened: URL[] = []; + const provider = createExecutorOAuthProvider({ + serverUrl: "https://executor.example/mcp", + storageRoot, + redirectUrl: "http://127.0.0.1:12345/callback", + state: "state-1", + redirectToAuthorization: async (url) => { + opened.push(url); + }, + }); + + await provider.redirectToAuthorization(new URL("https://login.example/authorize")); + expect(opened.map((url) => url.href)).toEqual(["https://login.example/authorize"]); + }); +}); diff --git a/test/oauth-callback.test.ts b/test/oauth-callback.test.ts new file mode 100644 index 0000000..6b836d8 --- /dev/null +++ b/test/oauth-callback.test.ts @@ -0,0 +1,71 @@ +import { describe, expect, test } from "bun:test"; +import { get } from "node:http"; +import { createOAuthCallback } from "../src/oauth-callback.ts"; + +const request = (url: string) => + new Promise<{ status: number; body: string }>((resolve, reject) => { + get(url, (response) => { + let body = ""; + response.setEncoding("utf8"); + response.on("data", (chunk: string) => (body += chunk)); + response.on("end", () => resolve({ status: response.statusCode ?? 0, body })); + }).on("error", reject); + }); + +describe("localhost OAuth callback", () => { + test("accepts exactly one callback with matching state and returns response", async () => { + const callback = await createOAuthCallback({ state: "expected", timeoutMs: 1000 }); + try { + expect(new URL(callback.redirectUrl).hostname).toBe("127.0.0.1"); + const result = callback.waitForCode(); + expect(await request(`${callback.redirectUrl}?code=abc&state=expected`)).toEqual({ + status: 200, + body: "Authorization complete. You can close this window.", + }); + await expect(result).resolves.toBe("abc"); + } finally { + await callback.close(); + } + }); + + test.each([ + ["wrong state", "?code=abc&state=wrong"], + ["missing code", "?state=expected"], + ["OAuth error", "?error=access_denied&state=expected"], + ])("rejects %s", async (_label, query) => { + const callback = await createOAuthCallback({ state: "expected", timeoutMs: 1000 }); + try { + const result = callback.waitForCode(); + const response = await request(callback.redirectUrl + query); + expect(response.status).toBe(400); + await expect(result).rejects.toThrow(); + } finally { + await callback.close(); + } + }); + + test("rejects callbacks sent to any path other than /callback", async () => { + const callback = await createOAuthCallback({ state: "expected", timeoutMs: 1000 }); + try { + const result = callback.waitForCode(); + const url = new URL(callback.redirectUrl); + url.pathname = "/other"; + url.search = "?code=abc&state=expected"; + const response = await request(url.href); + expect(response.status).toBe(404); + await expect(result).rejects.toThrow("Unexpected OAuth callback path."); + } finally { + await callback.close(); + } + }); + + test("rejects after timeout and supports cancellation", async () => { + const timed = await createOAuthCallback({ state: "s", timeoutMs: 10 }); + await expect(timed.waitForCode()).rejects.toThrow(/timeout/i); + const cancelled = await createOAuthCallback({ state: "s", timeoutMs: 1000 }); + const result = cancelled.waitForCode(); + cancelled.cancel(); + await expect(result).rejects.toThrow(/cancel/i); + await cancelled.close(); + }); +}); diff --git a/test/oauth-mcp-client.integration.test.ts b/test/oauth-mcp-client.integration.test.ts new file mode 100644 index 0000000..591e4b2 --- /dev/null +++ b/test/oauth-mcp-client.integration.test.ts @@ -0,0 +1,137 @@ +import { expect, test } from "bun:test"; +import { mkdtemp, rm } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { createExecutorOAuthProvider } from "../src/oauth-auth.ts"; +import { withExecutorMcpClient } from "../src/mcp-client.ts"; + +test("recovers from an OAuth challenge and calls MCP with the exchanged bearer token", async () => { + const storageRoot = await mkdtemp(join(tmpdir(), "executor-mcp-oauth-")); + const requests: string[] = []; + let authorized = false; + let registrationCount = 0; + let baseUrl = ""; + const server = Bun.serve({ + port: 0, + async fetch(request) { + const url = new URL(request.url); + requests.push(`${request.method} ${url.pathname}`); + if (url.pathname === "/.well-known/oauth-protected-resource") { + return Response.json({ + resource: `${baseUrl}/mcp`, + authorization_servers: [baseUrl], + }); + } + if (url.pathname === "/.well-known/oauth-authorization-server") { + return Response.json({ + issuer: baseUrl, + authorization_endpoint: `${baseUrl}/authorize`, + token_endpoint: `${baseUrl}/token`, + registration_endpoint: `${baseUrl}/register`, + response_types_supported: ["code"], + grant_types_supported: ["authorization_code", "refresh_token"], + code_challenge_methods_supported: ["S256"], + token_endpoint_auth_methods_supported: ["none"], + }); + } + if (url.pathname === "/register" && request.method === "POST") { + registrationCount += 1; + const registration = (await request.json()) as { redirect_uris: string[] }; + return Response.json({ + client_id: "test-client", + redirect_uris: registration.redirect_uris, + }); + } + if (url.pathname === "/token" && request.method === "POST") { + const form = new URLSearchParams(await request.text()); + expect(form.get("grant_type")).toBe("authorization_code"); + expect(form.get("code")).toBe("auth-code"); + expect(form.get("code_verifier")).toBeTruthy(); + authorized = true; + return Response.json({ + access_token: "integration-token", + token_type: "Bearer", + }); + } + if (url.pathname !== "/mcp") return new Response("Not found", { status: 404 }); + if (!authorized) { + return new Response("Unauthorized", { + status: 401, + headers: { + "WWW-Authenticate": `Bearer resource_metadata="${baseUrl}/.well-known/oauth-protected-resource"`, + }, + }); + } + if (request.headers.get("authorization") !== "Bearer integration-token") { + return new Response("Unauthorized", { status: 401 }); + } + if (request.method === "GET") return new Response("Method not allowed", { status: 405 }); + if (request.method === "DELETE") return new Response(null, { status: 200 }); + + const message = (await request.json()) as { + id?: string | number; + method: string; + }; + if (message.method === "notifications/initialized") + return new Response(null, { status: 202 }); + const result = + message.method === "initialize" + ? { + protocolVersion: "2025-03-26", + capabilities: { tools: {} }, + serverInfo: { name: "oauth-test", version: "1.0.0" }, + } + : message.method === "tools/list" + ? { tools: [{ name: "execute", description: "Run code" }] } + : message.method === "tools/call" + ? { content: [{ type: "text", text: "authenticated execution" }] } + : {}; + return Response.json({ jsonrpc: "2.0", id: message.id, result }); + }, + }); + baseUrl = `http://127.0.0.1:${server.port}`; + const provider = createExecutorOAuthProvider({ + serverUrl: baseUrl, + storageRoot, + redirectUrl: "http://127.0.0.1:12345/callback", + state: "integration-state", + redirectToAuthorization: async (authorizationUrl) => { + expect(authorizationUrl.searchParams.get("state")).toBe("integration-state"); + expect(authorizationUrl.searchParams.get("code_challenge_method")).toBe("S256"); + }, + }); + + try { + const result = await withExecutorMcpClient( + baseUrl, + { + hasUI: true, + oauthProvider: provider, + waitForAuthorizationCode: async () => "auth-code", + }, + (client) => client.execute("return 1"), + ); + expect(result.text).toBe("authenticated execution"); + const reconnected = await withExecutorMcpClient( + baseUrl, + { + hasUI: true, + oauthProvider: { + ...provider, + redirectToAuthorization: async () => { + throw new Error("Saved OAuth credentials should avoid a new browser login."); + }, + }, + }, + (client) => client.execute("return 2"), + ); + expect(reconnected.text).toBe("authenticated execution"); + expect(registrationCount).toBe(1); + expect(requests).toContain("POST /token"); + expect(requests.filter((request) => request === "POST /mcp").length).toBeGreaterThan(1); + expect(await provider.tokens()).toMatchObject({ access_token: "integration-token" }); + } finally { + server.stop(true); + await rm(storageRoot, { recursive: true, force: true }); + } +});