pi-mcp-adapter 2.2.1 → 2.3.2

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