pi-mcp-adapter 2.2.2 → 2.3.3

package/mcp-oauth-provider.ts

@@ -0,0 +1,265 @@
+/**
+ * MCP OAuth Provider
+ * 
+ * Implementation of the MCP SDK's OAuthClientProvider interface.
+ * Handles OAuth client registration, token storage, and authorization redirection.
+ */
+
+import type { OAuthClientProvider } from "@modelcontextprotocol/sdk/client/auth.js"
+import type {
+  OAuthClientMetadata,
+  OAuthTokens,
+  OAuthClientInformation,
+  OAuthClientInformationFull,
+} from "@modelcontextprotocol/sdk/shared/auth.js"
+import {
+  getAuthEntry,
+  getAuthForUrl,
+  updateTokens,
+  updateClientInfo,
+  updateCodeVerifier,
+  updateOAuthState,
+  clearAllCredentials,
+  clearClientInfo,
+  clearTokens,
+  type StoredTokens,
+  type StoredClientInfo,
+} from "./mcp-auth.js"
+
+// Callback server configuration
+const OAUTH_CALLBACK_PORT = 19876
+const OAUTH_CALLBACK_PATH = "/mcp/oauth/callback"
+
+/** Configuration options for OAuth */
+export interface McpOAuthConfig {
+  grantType?: "authorization_code" | "client_credentials"
+  clientId?: string
+  clientSecret?: string
+  scope?: string
+}
+
+/** Callbacks for OAuth flow interactions */
+export interface McpOAuthCallbacks {
+  onRedirect: (url: URL) => void | Promise<void>
+}
+
+/**
+ * OAuth provider implementation for MCP servers.
+ * Implements the OAuthClientProvider interface from the MCP SDK.
+ */
+export class McpOAuthProvider implements OAuthClientProvider {
+  constructor(
+    private serverName: string,
+    private serverUrl: string,
+    private config: McpOAuthConfig,
+    private callbacks: McpOAuthCallbacks,
+  ) {}
+
+  private get usesClientCredentials(): boolean {
+    return this.config.grantType === "client_credentials"
+  }
+
+  /**
+   * The redirect URL for OAuth callbacks.
+   * This must match the redirect_uri in client metadata.
+   */
+  get redirectUrl(): string | undefined {
+    if (this.usesClientCredentials) return undefined
+    return `http://127.0.0.1:${OAUTH_CALLBACK_PORT}${OAUTH_CALLBACK_PATH}`
+  }
+
+  /**
+   * Client metadata for dynamic registration.
+   * Describes this client to the OAuth authorization server.
+   */
+  get clientMetadata(): OAuthClientMetadata {
+    if (this.usesClientCredentials) {
+      return {
+        client_name: "Pi Coding Agent",
+        redirect_uris: [],
+        grant_types: ["client_credentials"],
+        token_endpoint_auth_method: this.config.clientSecret ? "client_secret_post" : "none",
+      }
+    }
+
+    const redirectUrl = this.redirectUrl
+    if (!redirectUrl) {
+      throw new Error("redirectUrl is required for authorization_code flow")
+    }
+
+    return {
+      redirect_uris: [redirectUrl],
+      client_name: "Pi Coding Agent",
+      client_uri: "https://github.com/nicobailon/pi-mcp-adapter",
+      grant_types: ["authorization_code", "refresh_token"],
+      response_types: ["code"],
+      token_endpoint_auth_method: this.config.clientSecret ? "client_secret_post" : "none",
+    }
+  }
+
+  /**
+   * Get client information (for pre-registered or dynamically registered clients).
+   * Returns undefined if no client info exists or if the server URL has changed.
+   */
+  async clientInformation(): Promise<OAuthClientInformation | undefined> {
+    // Check config first (pre-registered client)
+    if (this.config.clientId) {
+      return {
+        client_id: this.config.clientId,
+        client_secret: this.config.clientSecret,
+      }
+    }
+
+    // Check stored client info (from dynamic registration)
+    // Use getAuthForUrl to validate credentials are for the current server URL
+    const entry = await getAuthForUrl(this.serverName, this.serverUrl)
+    if (entry?.clientInfo) {
+      // Check if client secret has expired
+      if (entry.clientInfo.clientSecretExpiresAt && entry.clientInfo.clientSecretExpiresAt < Date.now() / 1000) {
+        return undefined
+      }
+      return {
+        client_id: entry.clientInfo.clientId,
+        client_secret: entry.clientInfo.clientSecret,
+      }
+    }
+
+    // No client info or URL changed - will trigger dynamic registration
+    return undefined
+  }
+
+  /**
+   * Save client information from dynamic registration.
+   */
+  async saveClientInformation(info: OAuthClientInformationFull): Promise<void> {
+    const clientInfo: StoredClientInfo = {
+      clientId: info.client_id,
+      clientSecret: info.client_secret,
+      clientIdIssuedAt: info.client_id_issued_at,
+      clientSecretExpiresAt: info.client_secret_expires_at,
+    }
+    updateClientInfo(this.serverName, clientInfo, this.serverUrl)
+  }
+
+  /**
+   * Get stored OAuth tokens.
+   * Returns undefined if no tokens exist or if the server URL has changed.
+   */
+  async tokens(): Promise<OAuthTokens | undefined> {
+    // Use getAuthForUrl to validate tokens are for the current server URL
+    const entry = await getAuthForUrl(this.serverName, this.serverUrl)
+    if (!entry?.tokens) return undefined
+
+    return {
+      access_token: entry.tokens.accessToken,
+      token_type: "Bearer",
+      refresh_token: entry.tokens.refreshToken,
+      expires_in: entry.tokens.expiresAt
+        ? Math.max(0, Math.floor(entry.tokens.expiresAt - Date.now() / 1000))
+        : undefined,
+      scope: entry.tokens.scope,
+    }
+  }
+
+  /**
+   * Save OAuth tokens.
+   */
+  async saveTokens(tokens: OAuthTokens): Promise<void> {
+    const storedTokens: StoredTokens = {
+      accessToken: tokens.access_token,
+      refreshToken: tokens.refresh_token,
+      expiresAt: tokens.expires_in ? Date.now() / 1000 + tokens.expires_in : undefined,
+      scope: tokens.scope,
+    }
+    updateTokens(this.serverName, storedTokens, this.serverUrl)
+  }
+
+  /**
+   * Redirect the user to the authorization URL.
+   * This opens the browser for the user to authenticate.
+   */
+  async redirectToAuthorization(authorizationUrl: URL): Promise<void> {
+    if (this.usesClientCredentials) {
+      throw new Error("redirectToAuthorization is not used for client_credentials flow")
+    }
+    // URL is passed to callback, not logged (may contain sensitive params)
+    await this.callbacks.onRedirect(authorizationUrl)
+  }
+
+  /**
+   * Save the PKCE code verifier.
+   */
+  async saveCodeVerifier(codeVerifier: string): Promise<void> {
+    updateCodeVerifier(this.serverName, codeVerifier)
+  }
+
+  /**
+   * Get the stored PKCE code verifier.
+   * @throws Error if no code verifier is stored
+   */
+  async codeVerifier(): Promise<string> {
+    if (this.usesClientCredentials) {
+      throw new Error("codeVerifier is not used for client_credentials flow")
+    }
+    const entry = await getAuthEntry(this.serverName)
+    if (!entry?.codeVerifier) {
+      throw new Error(`No code verifier saved for MCP server: ${this.serverName}`)
+    }
+    return entry.codeVerifier
+  }
+
+  /**
+   * Save the OAuth state parameter for CSRF protection.
+   */
+  async saveState(state: string): Promise<void> {
+    updateOAuthState(this.serverName, state)
+  }
+
+  /**
+   * Get the stored OAuth state parameter.
+   * @throws Error if no state is stored
+   */
+  async state(): Promise<string> {
+    if (this.usesClientCredentials) {
+      throw new Error("state is not used for client_credentials flow")
+    }
+    const entry = await getAuthEntry(this.serverName)
+    if (!entry?.oauthState) {
+      throw new Error(`No OAuth state saved for MCP server: ${this.serverName}`)
+    }
+    return entry.oauthState
+  }
+
+  /**
+   * Invalidate credentials when authentication fails.
+   * Clears tokens, client info, or all credentials based on the type.
+   */
+  async invalidateCredentials(type: "all" | "client" | "tokens"): Promise<void> {
+    switch (type) {
+      case "all":
+        clearAllCredentials(this.serverName)
+        break
+      case "client":
+        clearClientInfo(this.serverName)
+        break
+      case "tokens":
+        clearTokens(this.serverName)
+        break
+    }
+  }
+
+  prepareTokenRequest(scope?: string): URLSearchParams | undefined {
+    if (!this.usesClientCredentials) {
+      return undefined
+    }
+
+    const params = new URLSearchParams({ grant_type: "client_credentials" })
+    const requestedScope = scope ?? this.config.scope
+    if (requestedScope) {
+      params.set("scope", requestedScope)
+    }
+    return params
+  }
+}
+
+export { OAUTH_CALLBACK_PORT, OAUTH_CALLBACK_PATH }

package/package.json

@@ -1,16 +1,21 @@
 {
   "name": "pi-mcp-adapter",
-  "version": "2.2.2",
+  "version": "2.3.3",
   "description": "MCP (Model Context Protocol) adapter extension for Pi coding agent",
   "type": "module",
   "license": "MIT",
   "author": "Nico Bailon",
   "bin": {
-    "pi-mcp-adapter": "./cli.js"
+    "pi-mcp-adapter": "cli.js"
+  },
+  "scripts": {
+    "test": "vitest run",
+    "test:watch": "vitest",
+    "test:coverage": "vitest run --coverage"
   },
   "repository": {
     "type": "git",
-    "url": "https://github.com/nicobailon/pi-mcp-adapter"
+    "url": "git+https://github.com/nicobailon/pi-mcp-adapter.git"
   },
   "keywords": [
     "pi-package",
@@ -23,11 +28,6 @@
     "claude",
     "llm"
   ],
-  "scripts": {
-    "test": "vitest run",
-    "test:watch": "vitest",
-    "test:coverage": "vitest run --coverage"
-  },
   "pi": {
     "extensions": [
       "./index.ts"
@@ -60,6 +60,10 @@
     "glimpse-ui.ts",
     "npx-resolver.ts",
     "oauth-handler.ts",
+    "mcp-auth.ts",
+    "mcp-oauth-provider.ts",
+    "mcp-callback-server.ts",
+    "mcp-auth-flow.ts",
     "mcp-panel.ts",
     "logger.ts",
     "errors.ts",
@@ -72,13 +76,17 @@
     "@modelcontextprotocol/ext-apps": "^1.2.2",
     "@modelcontextprotocol/sdk": "^1.25.1",
     "@sinclair/typebox": "^0.32.0",
+    "open": "^10.2.0",
     "zod": "^3.25.0 || ^4.0.0"
   },
   "peerDependencies": {
     "zod": "^3.25.0 || ^4.0.0"
   },
   "devDependencies": {
+    "@types/bun": "^1.0.0",
     "@types/node": "^20.0.0",
+    "@types/open": "^6.2.1",
+    "tsx": "^4.21.0",
     "typescript": "^5.0.0",
     "vitest": "^3.0.0"
   }

package/proxy-modes.ts

@@ -99,6 +99,8 @@ export function executeStatus(state: McpExtensionState): ProxyToolResult {
     let status = "not connected";
     if (connection?.status === "connected") {
       status = "connected";
+    } else if (connection?.status === "needs-auth") {
+      status = "needs-auth";
     } else if (failedAgo !== null) {
       status = "failed";
     } else if (metadata !== undefined) {
@@ -117,6 +119,10 @@ export function executeStatus(state: McpExtensionState): ProxyToolResult {
       text += `✓ ${server.name} (${server.toolCount} tools)\n`;
       continue;
     }
+    if (server.status === "needs-auth") {
+      text += `⚠ ${server.name} (needs auth)\n`;
+      continue;
+    }
     if (server.status === "cached") {
       text += `○ ${server.name} (${server.toolCount} tools, cached)\n`;
       continue;
@@ -371,6 +377,12 @@ export async function executeConnect(state: McpExtensionState, serverName: strin
       state.ui.setStatus("mcp", `MCP: connecting to ${serverName}...`);
     }
     const connection = await state.manager.connect(serverName, definition);
+    if (connection.status === "needs-auth") {
+      return {
+        content: [{ type: "text" as const, text: `Server "${serverName}" requires OAuth authentication. Run /mcp-auth ${serverName} first.` }],
+        details: { mode: "connect", error: "auth_required", server: serverName },
+      };
+    }
     const prefix = state.config.settings?.toolPrefix ?? "server";
     const { metadata } = buildToolMetadata(connection.tools, connection.resources, definition, serverName, prefix);
     state.toolMetadata.set(serverName, metadata);
@@ -424,6 +436,13 @@ export async function executeCall(
     if (connected) {
       toolMeta = findToolByName(state.toolMetadata.get(serverName), toolName);
     } else {
+      const needsAuthConnection = state.manager.getConnection(serverName);
+      if (needsAuthConnection?.status === "needs-auth") {
+        return {
+          content: [{ type: "text" as const, text: `Server "${serverName}" requires OAuth authentication. Run /mcp-auth ${serverName} first.` }],
+          details: { mode: "call", error: "auth_required", server: serverName },
+        };
+      }
       const failedAgo = getFailureAgeSeconds(state, serverName);
       if (failedAgo !== null) {
         return {
@@ -472,6 +491,12 @@ export async function executeCall(
   }
 
   let connection = state.manager.getConnection(serverName);
+  if (connection?.status === "needs-auth") {
+    return {
+      content: [{ type: "text" as const, text: `Server "${serverName}" requires OAuth authentication. Run /mcp-auth ${serverName} first.` }],
+      details: { mode: "call", error: "auth_required", server: serverName },
+    };
+  }
   if (!connection || connection.status !== "connected") {
     const failedAgo = getFailureAgeSeconds(state, serverName);
     if (failedAgo !== null) {
@@ -494,6 +519,12 @@ export async function executeCall(
         state.ui.setStatus("mcp", `MCP: connecting to ${serverName}...`);
       }
       connection = await state.manager.connect(serverName, definition);
+      if (connection.status === "needs-auth") {
+        return {
+          content: [{ type: "text" as const, text: `Server "${serverName}" requires OAuth authentication. Run /mcp-auth ${serverName} first.` }],
+          details: { mode: "call", error: "auth_required", server: serverName },
+        };
+      }
       state.failureTracker.delete(serverName);
       updateServerMetadata(state, serverName);
       updateMetadataCache(state, serverName);

package/server-manager.ts

@@ -2,6 +2,7 @@ import { Client } from "@modelcontextprotocol/sdk/client/index.js";
 import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js";
 import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
 import { SSEClientTransport } from "@modelcontextprotocol/sdk/client/sse.js";
+import { UnauthorizedError } from "@modelcontextprotocol/sdk/client/auth.js";
 import type { ReadResourceResult } from "@modelcontextprotocol/sdk/types.js";
 import type {
   McpTool,
@@ -11,9 +12,10 @@ import type {
   Transport,
 } from "./types.js";
 import { serverStreamResultPatchNotificationSchema } from "./types.js";
-import { getStoredTokens } from "./oauth-handler.js";
 import { resolveNpxBinary } from "./npx-resolver.js";
 import { logger } from "./logger.js";
+import { McpOAuthProvider } from "./mcp-oauth-provider.js";
+import { supportsOAuth } from "./mcp-auth-flow.js";
 
 interface ServerConnection {
   client: Client;
@@ -23,7 +25,7 @@ interface ServerConnection {
   resources: McpResource[];
   lastUsedAt: number;
   inFlight: number;
-  status: "connected" | "closed";
+  status: "connected" | "closed" | "needs-auth";
 }
 
 type UiStreamListener = (serverName: string, notification: ServerStreamResultPatchNotification["params"]) => void;
@@ -96,7 +98,7 @@ export class McpServerManager {
     try {
       await client.connect(transport);
       this.attachAdapterNotificationHandlers(name, client);
-      
+
       // Discover tools and resources
       const [tools, resources] = await Promise.all([
         this.fetchAllTools(client),
@@ -114,6 +116,24 @@ export class McpServerManager {
         status: "connected",
       };
     } catch (error) {
+      // Check for UnauthorizedError - server requires OAuth
+      if (error instanceof UnauthorizedError && supportsOAuth(definition)) {
+        // Clean up both client and transport before reporting needs-auth.
+        await client.close().catch(() => {});
+        await transport.close().catch(() => {});
+
+        return {
+          client,
+          transport,
+          definition,
+          tools: [],
+          resources: [],
+          lastUsedAt: Date.now(),
+          inFlight: 0,
+          status: "needs-auth",
+        };
+      }
+      
       // Clean up both client and transport on any error
       await client.close().catch(() => {});
       await transport.close().catch(() => {});
@@ -121,11 +141,16 @@ export class McpServerManager {
     }
   }
   
-  private async createHttpTransport(definition: ServerDefinition, serverName?: string): Promise<Transport> {
+  private async createHttpTransport(
+    definition: ServerDefinition, 
+    serverName: string
+  ): Promise<Transport> {
     const url = new URL(definition.url!);
+    
+    // Build headers first (including any bearer token)
     const headers = resolveHeaders(definition.headers) ?? {};
     
-    // Add bearer token if configured
+    // For bearer auth, add the token to headers BEFORE creating requestInit
     if (definition.auth === "bearer") {
       const token = definition.bearerToken 
         ?? (definition.bearerTokenEnv ? process.env[definition.bearerTokenEnv] : undefined);
@@ -134,41 +159,58 @@ export class McpServerManager {
       }
     }
     
-    // Handle OAuth auth - use stored tokens
-    if (definition.auth === "oauth") {
-      if (!serverName) {
-        throw new Error("Server name required for OAuth authentication");
-      }
-      const tokens = getStoredTokens(serverName);
-      if (!tokens) {
-        throw new Error(
-          `No OAuth tokens found for "${serverName}". Run /mcp-auth ${serverName} to authenticate.`
-        );
-      }
-      headers["Authorization"] = `Bearer ${tokens.access_token}`;
-    }
-    
+    // Create request init with headers (Authorization now included for bearer auth)
     const requestInit = Object.keys(headers).length > 0 ? { headers } : undefined;
     
+    // For OAuth servers, create an auth provider
+    let authProvider: McpOAuthProvider | undefined;
+    if (supportsOAuth(definition)) {
+      // Extract OAuth config (handles both object and false cases)
+      const oauthConfig = definition.oauth === false ? {} : {
+        grantType: definition.oauth?.grantType,
+        clientId: definition.oauth?.clientId,
+        clientSecret: definition.oauth?.clientSecret,
+        scope: definition.oauth?.scope,
+      };
+      authProvider = new McpOAuthProvider(
+        serverName,
+        definition.url!,
+        oauthConfig,
+        {
+          onRedirect: async (_authUrl) => {
+            // URL is captured by startAuth, no need to log
+          },
+        }
+      );
+    }
+    
     // Try StreamableHTTP first (modern MCP servers)
-    const streamableTransport = new StreamableHTTPClientTransport(url, { requestInit });
+    const streamableTransport = new StreamableHTTPClientTransport(url, { 
+      requestInit,
+      authProvider,
+    });
     
     try {
       // Create a test client to verify the transport works
-      const testClient = new Client({ name: "pi-mcp-probe", version: "1.0.0" });
+      const testClient = new Client({ name: "pi-mcp-probe", version: "2.1.2" });
       await testClient.connect(streamableTransport);
       await testClient.close().catch(() => {});
       // Close probe transport before creating fresh one
       await streamableTransport.close().catch(() => {});
       
       // StreamableHTTP works - create fresh transport for actual use
-      return new StreamableHTTPClientTransport(url, { requestInit });
-    } catch {
+      return new StreamableHTTPClientTransport(url, { requestInit, authProvider });
+    } catch (error) {
       // StreamableHTTP failed, close and try SSE fallback
       await streamableTransport.close().catch(() => {});
       
+      // If this was an UnauthorizedError, don't try SSE - the server needs auth
+      if (error instanceof UnauthorizedError) {
+        throw error;
+      }
+      
       // SSE is the legacy transport
-      return new SSEClientTransport(url, { requestInit });
+      return new SSEClientTransport(url, { requestInit, authProvider });
     }
   }
   

package/types.ts

@@ -3,6 +3,7 @@ import type { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdi
 import type { SSEClientTransport } from "@modelcontextprotocol/sdk/client/sse.js";
 import type { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
 import type { TextContent, ImageContent } from "@mariozechner/pi-ai";
+import type { UiStreamMode } from "./ui-stream-types.js";
 
 // Transport type (stdio + HTTP)
 export type Transport = 
@@ -261,6 +262,18 @@ export interface McpContent {
 // Pi content block type
 export type ContentBlock = TextContent | ImageContent;
 
+// OAuth configuration (SDK handles auto-discovery and dynamic registration)
+export interface OAuthConfig {
+  /** OAuth grant type (defaults to authorization_code) */
+  grantType?: "authorization_code" | "client_credentials";
+  /** Pre-registered client ID (optional, dynamic registration used if not provided) */
+  clientId?: string;
+  /** Client secret for confidential clients */
+  clientSecret?: string;
+  /** Requested OAuth scopes */
+  scope?: string;
+}
+
 // Server configuration
 export interface ServerEntry {
   command?: string;
@@ -270,9 +283,22 @@ export interface ServerEntry {
   // HTTP fields
   url?: string;
   headers?: Record<string, string>;
-  auth?: "oauth" | "bearer";
+  /** 
+   * Authentication type:
+   * - 'oauth' - Use OAuth 2.1 (auto-discovers endpoints, supports dynamic client registration)
+   * - 'bearer' - Use static Bearer token
+   * - false - Disable authentication
+   * If not specified and url is present, OAuth will be auto-detected
+   */
+  auth?: "oauth" | "bearer" | false;
   bearerToken?: string;
   bearerTokenEnv?: string;
+  /** 
+   * OAuth configuration (optional).
+   * If not provided, the SDK will attempt dynamic client registration.
+   * Set to false to explicitly disable OAuth for this server.
+   */
+  oauth?: OAuthConfig | false;
   lifecycle?: "keep-alive" | "lazy" | "eager";
   idleTimeout?: number; // minutes, overrides global setting
   // Resource handling