Merge pull request #144 from modelcontextprotocol/justin/client-auth

Client implementation of MCP auth

Author: jspahrsummers

Diff for: jest.config.js

@@ -7,6 +7,10 @@ export default {
   ...defaultEsmPreset,
   moduleNameMapper: {
     "^(\\.{1,2}/.*)\\.js$": "$1",
+    "^pkce-challenge$": "<rootDir>/src/__mocks__/pkce-challenge.ts"
   },
+  transformIgnorePatterns: [
+    "/node_modules/(?!eventsource)/"
+  ],
   testPathIgnorePatterns: ["/node_modules/", "/dist/"],
 };

Diff for: package-lock.json

@@ -48,6 +48,7 @@
   "dependencies": {
     "content-type": "^1.0.5",
     "eventsource": "^3.0.2",
+    "pkce-challenge": "^4.1.0",
     "raw-body": "^3.0.0",
     "zod": "^3.23.8",
     "zod-to-json-schema": "^3.24.1"
@@ -73,4 +74,4 @@
   "resolutions": {
     "strip-ansi": "6.0.1"
   }
-}
+}

Diff for: package.json

@@ -0,0 +1,6 @@
+export default function pkceChallenge() {
+  return {
+    code_verifier: "test_verifier",
+    code_challenge: "test_challenge",
+  };
+}

Diff for: src/__mocks__/pkce-challenge.ts

@@ -1,6 +1,7 @@
 import { EventSource, type ErrorEvent, type EventSourceInit } from "eventsource";
 import { Transport } from "../shared/transport.js";
 import { JSONRPCMessage, JSONRPCMessageSchema } from "../types.js";
+import { auth, AuthResult, OAuthClientProvider, UnauthorizedError } from "./auth.js";
 
 export class SseError extends Error {
   constructor(
@@ -12,6 +13,42 @@ export class SseError extends Error {
   }
 }
 
+/**
+ * Configuration options for the `SSEClientTransport`.
+ */
+export type SSEClientTransportOptions = {
+  /**
+   * An OAuth client provider to use for authentication.
+   * 
+   * When an `authProvider` is specified and the SSE connection is started:
+   * 1. The connection is attempted with any existing access token from the `authProvider`.
+   * 2. If the access token has expired, the `authProvider` is used to refresh the token.
+   * 3. If token refresh fails or no access token exists, and auth is required, `OAuthClientProvider.redirectToAuthorization` is called, and an `UnauthorizedError` will be thrown from `connect`/`start`.
+   * 
+   * After the user has finished authorizing via their user agent, and is redirected back to the MCP client application, call `SSEClientTransport.finishAuth` with the authorization code before retrying the connection.
+   * 
+   * If an `authProvider` is not provided, and auth is required, an `UnauthorizedError` will be thrown.
+   * 
+   * `UnauthorizedError` might also be thrown when sending any message over the SSE transport, indicating that the session has expired, and needs to be re-authed and reconnected.
+   */
+  authProvider?: OAuthClientProvider;
+
+  /**
+   * Customizes the initial SSE request to the server (the request that begins the stream).
+   * 
+   * NOTE: Setting this property will prevent an `Authorization` header from
+   * being automatically attached to the SSE request, if an `authProvider` is
+   * also given. This can be worked around by setting the `Authorization` header
+   * manually.
+   */
+  eventSourceInit?: EventSourceInit;
+
+  /**
+   * Customizes recurring POST requests to the server.
+   */
+  requestInit?: RequestInit;
+};
+
 /**
  * Client transport for SSE: this will connect to a server using Server-Sent Events for receiving
  * messages and make separate POST requests for sending messages.
@@ -23,35 +60,76 @@ export class SSEClientTransport implements Transport {
   private _url: URL;
   private _eventSourceInit?: EventSourceInit;
   private _requestInit?: RequestInit;
+  private _authProvider?: OAuthClientProvider;
 
   onclose?: () => void;
   onerror?: (error: Error) => void;
   onmessage?: (message: JSONRPCMessage) => void;
 
   constructor(
     url: URL,
-    opts?: { eventSourceInit?: EventSourceInit; requestInit?: RequestInit },
+    opts?: SSEClientTransportOptions,
   ) {
     this._url = url;
     this._eventSourceInit = opts?.eventSourceInit;
     this._requestInit = opts?.requestInit;
+    this._authProvider = opts?.authProvider;
   }
 
-  start(): Promise<void> {
-    if (this._eventSource) {
-      throw new Error(
-        "SSEClientTransport already started! If using Client class, note that connect() calls start() automatically.",
-      );
+  private async _authThenStart(): Promise<void> {
+    if (!this._authProvider) {
+      throw new UnauthorizedError("No auth provider");
+    }
+
+    let result: AuthResult;
+    try {
+      result = await auth(this._authProvider, { serverUrl: this._url });
+    } catch (error) {
+      this.onerror?.(error as Error);
+      throw error;
+    }
+
+    if (result !== "AUTHORIZED") {
+      throw new UnauthorizedError();
     }
 
+    return await this._startOrAuth();
+  }
+
+  private async _commonHeaders(): Promise<HeadersInit> {
+    const headers: HeadersInit = {};
+    if (this._authProvider) {
+      const tokens = await this._authProvider.tokens();
+      if (tokens) {
+        headers["Authorization"] = `Bearer ${tokens.access_token}`;
+      }
+    }
+
+    return headers;
+  }
+
+  private _startOrAuth(): Promise<void> {
     return new Promise((resolve, reject) => {
       this._eventSource = new EventSource(
         this._url.href,
-        this._eventSourceInit,
+        this._eventSourceInit ?? {
+          fetch: (url, init) => this._commonHeaders().then((headers) => fetch(url, {
+            ...init,
+            headers: {
+              ...headers,
+              Accept: "text/event-stream"
+            }
+          })),
+        },
       );
       this._abortController = new AbortController();
 
       this._eventSource.onerror = (event) => {
+        if (event.code === 401 && this._authProvider) {
+          this._authThenStart().then(resolve, reject);
+          return;
+        }
+
         const error = new SseError(event.code, event.message, event);
         reject(error);
         this.onerror?.(error);
@@ -97,6 +175,30 @@ export class SSEClientTransport implements Transport {
     });
   }
 
+  async start() {
+    if (this._eventSource) {
+      throw new Error(
+        "SSEClientTransport already started! If using Client class, note that connect() calls start() automatically.",
+      );
+    }
+
+    return await this._startOrAuth();
+  }
+
+  /**
+   * Call this method after the user has finished authorizing via their user agent and is redirected back to the MCP client application. This will exchange the authorization code for an access token, enabling the next connection attempt to successfully auth.
+   */
+  async finishAuth(authorizationCode: string): Promise<void> {
+    if (!this._authProvider) {
+      throw new UnauthorizedError("No auth provider");
+    }
+
+    const result = await auth(this._authProvider, { serverUrl: this._url, authorizationCode });
+    if (result !== "AUTHORIZED") {
+      throw new UnauthorizedError("Failed to authorize");
+    }
+  }
+
   async close(): Promise<void> {
     this._abortController?.abort();
     this._eventSource?.close();
@@ -109,7 +211,8 @@ export class SSEClientTransport implements Transport {
     }
 
     try {
-      const headers = new Headers(this._requestInit?.headers);
+      const commonHeaders = await this._commonHeaders();
+      const headers = new Headers({ ...commonHeaders, ...this._requestInit?.headers });
       headers.set("content-type", "application/json");
       const init = {
         ...this._requestInit,
@@ -120,8 +223,17 @@ export class SSEClientTransport implements Transport {
       };
 
       const response = await fetch(this._endpoint, init);
-
       if (!response.ok) {
+        if (response.status === 401 && this._authProvider) {
+          const result = await auth(this._authProvider, { serverUrl: this._url });
+          if (result !== "AUTHORIZED") {
+            throw new UnauthorizedError();
+          }
+
+          // Purposely _not_ awaited, so we don't call onerror twice
+          return this.send(message);
+        }
+
         const text = await response.text().catch(() => null);
         throw new Error(
           `Error POSTing to endpoint (HTTP ${response.status}): ${text}`,

Diff for: src/client/auth.test.ts

@@ -5,5 +5,5 @@
     "moduleResolution": "node",
     "outDir": "./dist/cjs"
   },
-  "exclude": ["**/*.test.ts"]
+  "exclude": ["**/*.test.ts", "src/__mocks__/**/*"]
 }

Diff for: src/client/auth.ts

@@ -3,5 +3,5 @@
   "compilerOptions": {
     "outDir": "./dist/esm"
   },
-  "exclude": ["**/*.test.ts"]
+  "exclude": ["**/*.test.ts", "src/__mocks__/**/*"]
 }