Add support for AuthProvider in SSEClientTransport

Author: jspahrsummers

src/client/sse.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 } from "./auth.js";
 
 export class SseError extends Error {
   constructor(
@@ -12,6 +13,34 @@ export class SseError extends Error {
   }
 }
 
+/**
+ * Configuration options for the `SSEClientTransport`.
+ */
+export type SSEClientTransportOptions = {
+  /**
+   * An OAuth client provider to use for authentication.
+   * 
+   * If given, the transport will automatically attach an `Authorization` header
+   * if an access token is available, or begin the authorization flow if not.
+   */
+  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 +52,70 @@ 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 Error("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 Error("Unauthorized");
+    }
+
+    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 })),
+        },
       );
       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 +161,16 @@ 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();
+  }
+
   async close(): Promise<void> {
     this._abortController?.abort();
     this._eventSource?.close();
@@ -109,7 +183,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,