@phi-code-admin/phi-code 0.82.3 → 0.84.0

package/extensions/phi/mcp/index.ts

@@ -0,0 +1,376 @@
+/**
+ * pi-mcp — MCP client extension for the Pi coding agent.
+ *
+ * Entry point registered in package.json under "pi.extensions".
+ * Pi loads this file via jiti (TypeScript executed directly, no build step).
+ *
+ * Wires together: config → server manager → tool bridge → Pi API.
+ *
+ * Vendored into phi-code from the MIT-licensed `pi-mcp-extension` by irahardianto
+ * (https://github.com/irahardianto/pi-mcp-extension). See ./LICENSE. Adapted for
+ * phi: type imports resolve via "phi-code"; config lives in ~/.phi/agent/mcp.json
+ * and <cwd>/.phi/mcp.json (phi configDir), and it ships bundled (no install).
+ */
+
+import type { ExtensionAPI, ExtensionContext, ExtensionCommandContext } from "phi-code";
+import { loadConfig } from "./config.js";
+import { ServerManager } from "./server-manager.js";
+import type { TransportAuthCallbacks } from "./server-manager.js";
+import { ToolBridge } from "./tool-bridge.js";
+import { McpError } from "./errors.js";
+import { exec } from "node:child_process";
+
+// 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.
+ * Works on macOS, Linux, and Windows.
+ */
+function openBrowser(url: string): void {
+  const cmd = process.platform === "darwin"
+    ? `open "${url}"`
+    : process.platform === "win32"
+      ? `start "" "${url}"`
+      : `xdg-open "${url}"`;
+
+  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}`);
+      }
+    }
+  });
+}
+
+export default async function (pi: ExtensionAPI): Promise<void> {
+  // ── 1. Load and validate config ──────────────────────────────────────────
+  // cwd is available on the ExtensionContext passed to event handlers.
+  // We load config lazily on session_start to get the correct per-session cwd.
+  // For the initial load we use process.cwd() as a bootstrap path to detect
+  // whether any config exists at all.
+  let config;
+  try {
+    config = await loadConfig(process.cwd());
+  } catch (err) {
+    // Can't notify yet (no ctx), so log to stderr. The session_start handler
+    // will re-try with the real cwd and surface errors properly.
+    console.error(`[pi-mcp] Config error: ${err instanceof McpError ? err.message : String(err)}`);
+    return;
+  }
+
+  if (Object.keys(config.mcpServers).length === 0) {
+    // No servers configured — silently exit. Users can create mcp.json later.
+    return;
+  }
+
+  // ── 2. Initialize bridge components ──────────────────────────────────────
+  // Auth callbacks — opens browser and notifies user when OAuth is needed
+  const authCallbacks: TransportAuthCallbacks = {
+    onAuthRequired: (serverName: string, authorizationUrl: URL): void => {
+      console.error(
+        `[pi-mcp] OAuth required for "${serverName}". Opening browser for authorization...`,
+      );
+      openBrowser(authorizationUrl.toString());
+    },
+  };
+
+  const manager = new ServerManager(config, authCallbacks);
+  const bridge = new ToolBridge(config.settings, pi);
+
+  // Connect tool refresh callback: called on connect and on list_changed
+  manager.setToolRefreshCallback(async (serverName, client) => {
+    await bridge.refreshTools(serverName, client);
+  });
+
+  // ── 3. Session lifecycle ──────────────────────────────────────────────────
+  pi.on("session_start", async (_event, ctx: ExtensionContext) => {
+    // Reload config with the real session cwd (project config may differ)
+    let sessionConfig = config;
+    try {
+      sessionConfig = await loadConfig(ctx.cwd);
+    } catch (err) {
+      const msg = err instanceof McpError ? err.userMessage : String(err);
+      ctx.ui.notify(`pi-mcp: Config error — ${msg}`, "error");
+      return;
+    }
+
+    // If config changed (different cwd with project-level overrides),
+    // shut down old servers and rebuild the manager's server list
+    if (JSON.stringify(sessionConfig) !== JSON.stringify(config)) {
+      // Deactivate and remove all tools from old config
+      for (const server of manager.getAllServers()) {
+        bridge.removeServer(server.name);
+      }
+      // Shut down all running servers
+      await manager.shutdownAll();
+      // Rebuild server entries from new config
+      manager.rebuildServers(sessionConfig);
+    }
+
+    const eagerServers = Object.entries(sessionConfig.mcpServers).filter(
+      ([, cfg]) => cfg.lifecycle === "eager",
+    );
+
+    // Start all eager servers concurrently
+    await Promise.allSettled(
+      eagerServers.map(async ([name]) => {
+        try {
+          await manager.startServer(name, ctx.cwd);
+        } catch (err) {
+          const msg = err instanceof McpError ? err.userMessage : String(err);
+          ctx.ui.notify(`pi-mcp: Failed to start ${name} — ${msg}`, "error");
+        }
+      }),
+    );
+  });
+
+  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);
+    }
+    await manager.shutdownAll();
+  });
+
+  // ── 4. /mcp — show server status ─────────────────────────────────────────
+  pi.registerCommand("mcp", {
+    description:
+      "Show MCP server status. Usage: /mcp [server-name] for detail.",
+    handler: async (args: string, ctx: ExtensionCommandContext) => {
+      const serverName = args.trim();
+      if (serverName) {
+        // Detailed view: status + recent stderr
+        const server = manager.getServer(serverName);
+        if (!server) {
+          ctx.ui.notify(`pi-mcp: No server named "${serverName}"`, "error");
+          return;
+        }
+        const logs = manager.getServerLogs(serverName);
+        const detail = [
+          `Server: ${serverName}`,
+          `State:  ${server.state}`,
+          `Retries: ${server.retryCount}`,
+          server.lastError ? `Last error: ${server.lastError.message}` : null,
+          "",
+          "Recent output:",
+          logs,
+        ]
+          .filter(Boolean)
+          .join("\n");
+        ctx.ui.notify(detail, "info");
+      } else {
+        // Summary view: all servers
+        ctx.ui.notify(manager.getStatusSummary(), "info");
+      }
+    },
+  });
+
+  // ── 5. /mcp:stop — stop a server ─────────────────────────────────────────
+  pi.registerCommand("mcp:stop", {
+    description: "Stop an MCP server. Usage: /mcp:stop <server-name>",
+    handler: async (args: string, ctx: ExtensionCommandContext) => {
+      const serverName = args.trim();
+      if (!serverName) {
+        ctx.ui.notify("Usage: /mcp:stop <server-name>", "error");
+        return;
+      }
+      if (!manager.getServer(serverName)) {
+        ctx.ui.notify(`pi-mcp: No server named "${serverName}"`, "error");
+        return;
+      }
+      bridge.deactivateServer(serverName);
+      await manager.stopServer(serverName);
+      ctx.ui.notify(`pi-mcp: Stopped ${serverName}`, "info");
+    },
+  });
+
+  // ── 6. /mcp:start — manually start a lazy server ─────────────────────────
+  pi.registerCommand("mcp:start", {
+    description: "Start an MCP server. Usage: /mcp:start <server-name>",
+    handler: async (args: string, ctx: ExtensionCommandContext) => {
+      const serverName = args.trim();
+      if (!serverName) {
+        ctx.ui.notify("Usage: /mcp:start <server-name>", "error");
+        return;
+      }
+      if (!manager.getServer(serverName)) {
+        ctx.ui.notify(`pi-mcp: No server named "${serverName}"`, "error");
+        return;
+      }
+      try {
+        await manager.startServer(serverName, ctx.cwd);
+        ctx.ui.notify(`pi-mcp: Started ${serverName}`, "info");
+      } catch (err) {
+        const msg = err instanceof McpError ? err.userMessage : String(err);
+        ctx.ui.notify(`pi-mcp: Failed to start ${serverName} — ${msg}`, "error");
+      }
+    },
+  });
+
+  // ── 7. /mcp:auth — trigger OAuth authentication for a server ────────────────
+  pi.registerCommand("mcp:auth", {
+    description:
+      "Trigger OAuth authentication for a server. Resets credentials and opens browser for re-authorization. Usage: /mcp:auth <server-name>",
+    handler: async (args: string, ctx: ExtensionCommandContext) => {
+      const serverName = args.trim();
+      if (!serverName) {
+        // List servers with auth config
+        const authServers = manager.getAllServers().filter((s) => s.config.auth);
+        if (authServers.length === 0) {
+          ctx.ui.notify(
+            "pi-mcp: No servers with OAuth configured. Add `auth: { type: \"oauth\" }` to a server in mcp.json.",
+            "error",
+          );
+          return;
+        }
+        const lines = authServers.map(async (s) => {
+          const status = await manager.getServerAuthStatus(s.name);
+          const authIcon = status?.hasTokens ? "\u2705 authenticated" : "\u274C not authenticated";
+          const savedInfo = status?.savedAt ? ` (since ${status.savedAt})` : "";
+          return `  ${s.name}: ${authIcon}${savedInfo}`;
+        });
+        const statusLines = await Promise.all(lines);
+        ctx.ui.notify(
+          [
+            "Usage: /mcp:auth <server-name>",
+            "",
+            "OAuth-enabled servers:",
+            ...statusLines,
+          ].join("\n"),
+          "info",
+        );
+        return;
+      }
+      const server = manager.getServer(serverName);
+      if (!server) {
+        ctx.ui.notify(`pi-mcp: No server named "${serverName}"`, "error");
+        return;
+      }
+      if (!server.config.auth) {
+        ctx.ui.notify(
+          `pi-mcp: Server "${serverName}" does not have OAuth configured. Add \`auth: { type: "oauth" }\` to its config in mcp.json.`,
+          "error",
+        );
+        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: Starting OAuth flow for ${serverName}...`,
+          "info",
+        );
+
+        // 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 if (authResult === "REDIRECT") {
+          // Browser was opened, wait for the callback from the user
+          ctx.ui.notify(
+            `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);
+        }
+      }
+    },
+  });
+}