pi-mcp-extension 1.4.1 → 1.5.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-mcp-extension",
-  "version": "1.4.1",
+  "version": "1.5.0",
   "description": "MCP (Model Context Protocol) client extension for the Pi coding agent — connect Pi to any MCP server",
   "type": "module",
   "keywords": [

package/src/callback-server.ts

@@ -0,0 +1,263 @@
+/**
+ * MCP OAuth Callback Server
+ *
+ * Simple HTTP server that handles OAuth callbacks from the authorization server.
+ * Uses Node.js http module for compatibility with no external dependencies.
+ */
+
+import { createServer, type Server, type IncomingMessage, type ServerResponse } from "http";
+import { URL } from "url";
+
+interface PendingAuth {
+  resolve: (code: string) => void;
+  reject: (error: Error) => void;
+  timeout: ReturnType<typeof setTimeout>;
+}
+
+let server: Server | null = null;
+let actualServerPort: number | null = null;
+const pendingAuths = new Map<string, PendingAuth>();
+
+const DEFAULT_PORT = 19876;
+const CALLBACK_PATH = "/callback";
+const TIMEOUT_MS = 5 * 60 * 1000; // 5 minutes
+
+const HTML_SUCCESS = `<!DOCTYPE html>
+<html><head><title>Pi - Authorization Successful</title>
+<style>body{font-family:system-ui,-apple-system,sans-serif;display:flex;justify-content:center;align-items:center;height:100vh;margin:0;background:#1a1a2e;color:#eee}
+.container{text-align:center;padding:2rem}h1{color:#4ade80;margin-bottom:1rem}p{color:#aaa}</style>
+</head><body>
+<div class="container"><h1>✓ Authorization Successful</h1><p>You can close this window and return to Pi.</p></div>
+<script>setTimeout(()=>window.close(),2000)</script>
+</body></html>`;
+
+const HTML_ERROR = (error: string) => `<!DOCTYPE html>
+<html><head><title>Pi - Authorization Failed</title>
+<style>body{font-family:system-ui,-apple-system,sans-serif;display:flex;justify-content:center;align-items:center;height:100vh;margin:0;background:#1a1a2e;color:#eee}
+.container{text-align:center;padding:2rem}h1{color:#f87171;margin-bottom:1rem}p{color:#aaa}
+.error{color:#fca5a5;font-family:monospace;margin-top:1rem;padding:1rem;background:rgba(248,113,113,0.1);border-radius:.5rem;white-space:pre-wrap;word-break:break-word}</style>
+</head><body>
+<div class="container"><h1>✗ Authorization Failed</h1><p>An error occurred during authorization.</p><div class="error">${escapeHtml(error)}</div></div>
+</body></html>`;
+
+/**
+ * Escape HTML entities to prevent XSS attacks.
+ */
+function escapeHtml(text: string): string {
+  return text
+    .replace(/&/g, "&amp;")
+    .replace(/</g, "&lt;")
+    .replace(/>/g, "&gt;")
+    .replace(/"/g, "&quot;")
+    .replace(/'/g, "&#039;");
+}
+
+/**
+ * Handle incoming HTTP requests to the callback server.
+ */
+function handleRequest(req: IncomingMessage, res: ServerResponse): void {
+  const url = new URL(req.url || "/", `http://${req.headers.host}`);
+
+  // Only handle the callback path
+  if (url.pathname !== CALLBACK_PATH) {
+    res.writeHead(404, { "Content-Type": "text/plain" });
+    res.end("Not found");
+    return;
+  }
+
+  const code = url.searchParams.get("code");
+  const state = url.searchParams.get("state");
+  const error = url.searchParams.get("error");
+  const errorDescription = url.searchParams.get("error_description");
+
+  // Enforce state parameter presence for CSRF protection
+  if (!state) {
+    const errorMsg = "Missing required state parameter - potential CSRF attack";
+    res.writeHead(400, { "Content-Type": "text/html" });
+    res.end(HTML_ERROR(errorMsg));
+    return;
+  }
+
+  // Handle OAuth errors
+  if (error) {
+    const errorMsg = errorDescription || error;
+    res.writeHead(200, { "Content-Type": "text/html" });
+    res.end(HTML_ERROR(errorMsg));
+    if (pendingAuths.has(state)) {
+      const pending = pendingAuths.get(state)!;
+      clearTimeout(pending.timeout);
+      pendingAuths.delete(state);
+      setTimeout(() => pending.reject(new Error(errorMsg)), 0);
+    }
+    return;
+  }
+
+  // Require authorization code
+  if (!code) {
+    res.writeHead(400, { "Content-Type": "text/html" });
+    res.end(HTML_ERROR("No authorization code provided"));
+    return;
+  }
+
+  // Validate state parameter
+  if (!pendingAuths.has(state)) {
+    const errorMsg = "Invalid or expired state parameter - potential CSRF attack";
+    res.writeHead(400, { "Content-Type": "text/html" });
+    res.end(HTML_ERROR(errorMsg));
+    return;
+  }
+
+  const pending = pendingAuths.get(state)!;
+
+  // Clear timeout and resolve the pending promise
+  clearTimeout(pending.timeout);
+  pendingAuths.delete(state);
+  pending.resolve(code);
+
+  res.writeHead(200, { "Content-Type": "text/html" });
+  res.end(HTML_SUCCESS);
+}
+
+/**
+ * Ensure the callback server is running.
+ * Scans forward for an available local port if the preferred port is busy.
+ *
+ * @param preferredPort - The preferred port to use (default: 19876)
+ * @returns The actual port the server is listening on
+ */
+export async function ensureCallbackServer(preferredPort: number = DEFAULT_PORT): Promise<number> {
+  if (server) {
+    // Already running, return the tracked actual port
+    if (actualServerPort !== null) {
+      return actualServerPort;
+    }
+    // Fallback: try to get the port from the server
+    const address = server.address();
+    if (address && typeof address === "object" && "port" in address) {
+      actualServerPort = address.port;
+      return address.port;
+    }
+    // If we still can't determine the port, return the preferred port
+    // (this shouldn't happen in practice)
+    return preferredPort;
+  }
+
+  const maxAttempts = 25; // Try up to 25 ports
+
+  for (let offset = 0; offset < maxAttempts; offset++) {
+    const candidatePort = preferredPort + offset;
+    const candidateServer = createServer(handleRequest);
+
+    try {
+      await new Promise<void>((resolve, reject) => {
+        candidateServer.once("error", (err: any) => {
+          reject(err);
+        });
+
+        // Bind to 127.0.0.1 explicitly (IPv4) to avoid issues with IPv6
+        candidateServer.listen(candidatePort, "127.0.0.1", () => {
+          resolve();
+        });
+      });
+
+      server = candidateServer;
+      actualServerPort = candidatePort;
+      server.unref(); // Don't block process exit
+      return candidatePort;
+    } catch (error) {
+      const nodeError = error as NodeJS.ErrnoException;
+      await new Promise<void>((resolve) => {
+        candidateServer.close(() => resolve());
+      });
+
+      // If not EADDRINUSE, rethrow
+      if (nodeError.code !== "EADDRINUSE") {
+        throw error;
+      }
+    }
+  }
+
+  throw new Error(
+    `OAuth callback port ${preferredPort} is already in use and no free port was found in range ${preferredPort}-${preferredPort + maxAttempts - 1}`
+  );
+}
+
+/**
+ * Wait for a callback with the given OAuth state.
+ * Returns a promise that resolves with the authorization code.
+ *
+ * @param oauthState - The OAuth state parameter to wait for
+ * @param timeoutMs - Timeout in milliseconds (default: 5 minutes)
+ * @returns Promise that resolves with the authorization code
+ */
+export function waitForCallback(
+  oauthState: string,
+  timeoutMs: number = TIMEOUT_MS
+): Promise<string> {
+  return new Promise((resolve, reject) => {
+    const timeout = setTimeout(() => {
+      if (pendingAuths.has(oauthState)) {
+        pendingAuths.delete(oauthState);
+        reject(new Error("OAuth callback timeout - authorization took too long"));
+      }
+    }, timeoutMs);
+
+    pendingAuths.set(oauthState, { resolve, reject, timeout });
+  });
+}
+
+/**
+ * Cancel a pending authorization by state.
+ *
+ * @param oauthState - The OAuth state to cancel
+ */
+export function cancelCallback(oauthState: string): void {
+  const pending = pendingAuths.get(oauthState);
+  if (pending) {
+    clearTimeout(pending.timeout);
+    pendingAuths.delete(oauthState);
+    pending.reject(new Error("Authorization cancelled"));
+  }
+}
+
+/**
+ * Stop the callback server and reject all pending authorizations.
+ */
+export async function stopCallbackServer(): Promise<void> {
+  if (server) {
+    await new Promise<void>((resolve) => {
+      server!.close(() => {
+        resolve();
+      });
+    });
+    server = null;
+    actualServerPort = null;
+  }
+
+  // Reject all pending auths (defer to allow any pending operations to complete)
+  const pendingList = Array.from(pendingAuths.entries());
+  pendingAuths.clear();
+  setTimeout(() => {
+    for (const [, pending] of pendingList) {
+      clearTimeout(pending.timeout);
+      pending.reject(new Error("OAuth callback server stopped"));
+    }
+  }, 0);
+}
+
+/**
+ * Check if the callback server is running.
+ */
+export function isCallbackServerRunning(): boolean {
+  return server !== null;
+}
+
+/**
+ * Get the number of pending authorizations.
+ */
+export function getPendingAuthCount(): number {
+  return pendingAuths.size;
+}
+
+// Export constants for testing/config
+export { DEFAULT_PORT, CALLBACK_PATH, TIMEOUT_MS };

package/src/index.ts

@@ -14,9 +14,17 @@ import type { TransportAuthCallbacks } from "./server-manager.js";
 import { ToolBridge } from "./tool-bridge.js";
 import { McpError } from "./errors.js";
 import { exec } from "node:child_process";
-import { promisify } from "node:util";
 
-const execAsync = promisify(exec);
+// OAuth imports
+import {
+  ensureCallbackServer,
+  waitForCallback,
+  cancelCallback,
+  stopCallbackServer,
+} from "./callback-server.js";
+import { setCallbackPort, McpOAuthProvider } from "./oauth-provider.js";
+import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
+import { auth } from "@modelcontextprotocol/sdk/client/auth.js";
 
 /**
  * Open a URL in the user's default browser.
@@ -28,8 +36,15 @@ function openBrowser(url: string): void {
     : process.platform === "win32"
       ? `start "" "${url}"`
       : `xdg-open "${url}"`;
-  exec(cmd, (err) => {
-    if (err) console.error(`[pi-mcp] Failed to open browser: ${err.message}`);
+
+  exec(cmd, (err, stdout, stderr) => {
+    if (err) {
+      const errorMsg = `[pi-mcp] Failed to open browser: ${err.message}`;
+      console.error(errorMsg);
+      if (stderr) {
+        console.error(`[pi-mcp] Browser error output: ${stderr}`);
+      }
+    }
   });
 }
 
@@ -116,6 +131,9 @@ export default async function (pi: ExtensionAPI): Promise<void> {
   });
 
   pi.on("session_shutdown", async (_event, _ctx: ExtensionContext) => {
+    // Stop the callback server
+    await stopCallbackServer().catch(() => {});
+
     // Deactivate all tools before shutting down servers
     for (const server of manager.getAllServers()) {
       bridge.deactivateServer(server.name);
@@ -245,33 +263,108 @@ export default async function (pi: ExtensionAPI): Promise<void> {
         return;
       }
 
+      const config = server.config;
+      let oauthState: string | undefined;
+
       try {
         // Stop the server if running
         if (server.state !== "stopped") {
           bridge.deactivateServer(serverName);
           await manager.stopServer(serverName);
         }
+
+        // Validate that we have a server URL (required for OAuth)
+        if (!config.url) {
+          throw new McpError(
+            `Server "${serverName}" has OAuth configured but no URL. OAuth requires a URL-based server transport.`,
+            serverName,
+            "config",
+          );
+        }
+
         // Reset all OAuth credentials (tokens, client info, PKCE, discovery)
         await manager.resetServerAuth(serverName);
+
         ctx.ui.notify(
-          `pi-mcp: Reset credentials for ${serverName}. Starting server — a browser window will open for authentication...`,
+          `pi-mcp: Starting OAuth flow for ${serverName}...`,
           "info",
         );
-        // Start the server — this triggers the OAuth flow and opens the browser
-        await manager.startServer(serverName, ctx.cwd);
 
-        const newStatus = await manager.getServerAuthStatus(serverName);
-        if (newStatus?.hasTokens) {
+        // 1. Start the callback server
+        const port = await ensureCallbackServer();
+        setCallbackPort(port);
+
+        // 2. Generate a cryptographically secure state parameter for CSRF protection
+        oauthState = Array.from(crypto.getRandomValues(new Uint8Array(32)))
+          .map((b: number) => b.toString(16).padStart(2, "0"))
+          .join("");
+
+        // 3. Register the callback promise BEFORE opening the browser
+        const callbackPromise = waitForCallback(oauthState);
+
+        // 4. Create auth provider and transport
+        const authProvider = new McpOAuthProvider(
+          serverName,
+          config.auth || { type: "oauth" },
+          (url: URL) => {
+            console.error(`[pi-mcp] Opening browser for ${serverName}...`);
+            openBrowser(url.toString());
+          },
+        );
+
+        // CRITICAL FIX #1: Set the OAuth state on the provider before calling auth()
+        // This ensures the state parameter is included in the authorization URL
+        authProvider.setState(oauthState);
+
+        const transport = new StreamableHTTPClientTransport(
+          new URL(config.url),
+          { authProvider },
+        );
+
+        // 5. Start the auth flow - this will trigger redirectToAuthorization which opens the browser
+        // CRITICAL FIX #2: Check the return value of auth() instead of catching UnauthorizedError
+        // The SDK returns 'REDIRECT' when it needs browser interaction, not an error
+        const authResult = await auth(authProvider, { serverUrl: config.url });
+
+        if (authResult === "AUTHORIZED") {
+          // Auth succeeded without needing browser interaction (e.g., had valid tokens)
           ctx.ui.notify(`pi-mcp: ${serverName} authenticated successfully!`, "info");
-        } else {
+        } else if (authResult === "REDIRECT") {
+          // Browser was opened, wait for the callback from the user
           ctx.ui.notify(
-            `pi-mcp: ${serverName} started but authentication may still be in progress. Check the browser window.`,
+            `pi-mcp: Browser opened for ${serverName}. Complete authorization to continue...`,
             "info",
           );
+
+          // 6. Wait for the callback (this blocks until the user authorizes)
+          const code = await callbackPromise;
+
+          // 7. Complete the OAuth flow with the authorization code
+          await transport.finishAuth(code);
+
+          ctx.ui.notify(`pi-mcp: ${serverName} authenticated successfully!`, "info");
+        } else {
+          throw new McpError(
+            `Unexpected auth result: ${authResult}`,
+            serverName,
+            "protocol",
+          );
         }
+
+        // 8. Close the transport (we'll create a new one when starting the server)
+        await transport.close().catch(() => {});
+
+        // 9. Start the server with fresh auth credentials
+        await manager.startServer(serverName, ctx.cwd);
+
       } catch (err) {
         const msg = err instanceof McpError ? err.userMessage : String(err);
         ctx.ui.notify(`pi-mcp: Authentication failed for ${serverName} — ${msg}`, "error");
+
+        // Clean up on error
+        if (oauthState) {
+          cancelCallback(oauthState);
+        }
       }
     },
   });

package/src/oauth-provider.ts

@@ -13,7 +13,7 @@
 
 import type { OAuthClientProvider, OAuthDiscoveryState } from "@modelcontextprotocol/sdk/client/auth.js";
 import type { OAuthClientMetadata, OAuthClientInformationMixed, OAuthTokens } from "@modelcontextprotocol/sdk/shared/auth.js";
-import { readFile, writeFile, mkdir } from "node:fs/promises";
+import { readFile, writeFile, mkdir, unlink } from "node:fs/promises";
 import { join } from "node:path";
 import { homedir } from "node:os";
 import { createHash } from "node:crypto";
@@ -43,6 +43,20 @@ export interface AuthConfig {
   clientSecret?: string | undefined;
 }
 
+// ─── Callback Server Port ───────────────────────────────────────────────────────
+
+let callbackPort = 19876;
+
+/** Set the callback server port. Called by the auth flow when the server starts. */
+export function setCallbackPort(port: number): void {
+  callbackPort = port;
+}
+
+/** Get the current callback server port. */
+export function getCallbackPort(): number {
+  return callbackPort;
+}
+
 // ─── Persistent State Types ───────────────────────────────────────────────────
 
 interface StoredClientInfo {
@@ -115,6 +129,7 @@ export class McpOAuthProvider implements OAuthClientProvider {
   private authConfig: AuthConfig;
   private _redirectUrl: string | undefined;
   private _onAuthRequired: ((url: URL) => void | Promise<void>) | undefined;
+  private _oauthState: string | undefined;
 
   constructor(
     serverName: string,
@@ -129,16 +144,19 @@ export class McpOAuthProvider implements OAuthClientProvider {
 
   // --- redirectUrl ---
 
-  get redirectUrl(): string | URL | undefined {
-    return this._redirectUrl;
+  get redirectUrl(): string | URL {
+    // Use configured redirect URL if provided, otherwise use the callback server
+    // Use 127.0.0.1 explicitly (IPv4) to match the callback server binding
+    return this._redirectUrl || `http://127.0.0.1:${callbackPort}/callback`;
   }
 
   // --- clientMetadata ---
 
   get clientMetadata(): OAuthClientMetadata {
+    const redirectUrl = String(this.redirectUrl);
     return {
       client_name: `pi-mcp/${this.serverName}`,
-      redirect_uris: this._redirectUrl ? [this._redirectUrl] : [],
+      redirect_uris: [redirectUrl],
       grant_types: ["authorization_code", "refresh_token"],
       response_types: ["code"],
       token_endpoint_auth_method: this.authConfig.clientSecret ? "client_secret_basic" : "none",
@@ -182,6 +200,21 @@ export class McpOAuthProvider implements OAuthClientProvider {
   async tokens(): Promise<OAuthTokens | undefined> {
     const state = await loadState(this.serverName);
     if (!state.tokens) return undefined;
+
+    // Always return stored tokens — even if expired.
+    // The SDK's auth() function checks tokens?.refresh_token and attempts
+    // silent refresh before falling back to a new authorization flow.
+    // Returning undefined for expired tokens would prevent that silent refresh
+    // and force the user to re-authenticate via browser every time.
+    //
+    // Flow when we return expired tokens:
+    //   transport sends expired access_token → 401
+    //   → auth() sees refresh_token → silent refresh → success → retry
+    //
+    // Flow when we return undefined (WRONG):
+    //   transport has no token → auth() → no refresh possible → REDIRECT
+    //   → user must re-authenticate in browser
+
     // Build OAuthTokens — only include defined fields
     const tokens: Record<string, string> = {
       access_token: state.tokens.access_token,
@@ -249,6 +282,25 @@ export class McpOAuthProvider implements OAuthClientProvider {
     return state.discoveryState;
   }
 
+  // --- OAuth state parameter (CSRF protection) ---
+
+  /**
+   * Set the OAuth state parameter before calling auth().
+   * This should be called with a cryptographically random value.
+   */
+  setState(state: string): void {
+    this._oauthState = state;
+  }
+
+  /**
+   * Returns the OAuth state parameter for CSRF protection.
+   * This is called by the SDK's auth() function when building the authorization URL.
+   * Returns empty string if no state has been set (no CSRF protection).
+   */
+  async state(): Promise<string> {
+    return this._oauthState || "";
+  }
+
   // --- Credential invalidation ---
 
   async invalidateCredentials(scope: "all" | "client" | "tokens" | "verifier" | "discovery"): Promise<void> {
@@ -311,6 +363,5 @@ export async function getAuthStatus(serverName: string): Promise<{
  * Used to force re-authorization on next connection.
  */
 export async function resetAuth(serverName: string): Promise<void> {
-  const provider = new McpOAuthProvider(serverName, {});
-  await provider.invalidateCredentials("all");
+  await unlink(statePath(serverName)).catch(() => {});
 }