Support OAuth authentication for remote Executor MCP

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):

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<string>;
   onElicitation?: (request: ExecutorElicitationRequest) => Promise<ExecutorElicitationResponse>;
 };
 
@@ -112,38 +115,80 @@ const connectExecutorMcpClient = async (
   baseUrl: string,
   options: ExecutorMcpClientOptions,
 ): Promise<ConnectedExecutorMcpClient> => {
-  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<ExecutorMcpToolMetadata[]> => {
     const tools: ExecutorMcpToolMetadata[] = [];

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<void>;
+};
+
+const readState = async (path: string): Promise<PersistedOAuthState> => {
+  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<PersistedOAuthState>): Promise<void> => {
+    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 });
+      }
+    },
+  };
+};

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<string>;
+  cancel(): void;
+  close(): Promise<void>;
+};
+
+export async function createOAuthCallback({
+  state,
+  timeoutMs = 5 * 60_000,
+}: OAuthCallbackOptions): Promise<OAuthCallback> {
+  let resolveCode!: (code: string) => void;
+  let rejectCode!: (error: Error) => void;
+  let settled = false;
+  let waiting = false;
+  let timer: ReturnType<typeof setTimeout> | undefined;
+  let server: Server;
+  const result = new Promise<string>((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<void>((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());
+      });
+    },
+  };
+}