optimization: When using the custom SSE request, the Authorization header can still be automatically attached to the SSE request.

Author: chenxi-null

src/client/sse.test.ts

@@ -1,4 +1,4 @@
-import { createServer, type IncomingMessage, type Server } from "http";
+import { createServer, IncomingMessage, Server, ServerResponse } from "http";
 import { AddressInfo } from "net";
 import { JSONRPCMessage } from "../types.js";
 import { SSEClientTransport } from "./sse.js";
@@ -10,8 +10,21 @@ describe("SSEClientTransport", () => {
   let transport: SSEClientTransport;
   let baseUrl: URL;
   let lastServerRequest: IncomingMessage;
+  const serverRequests: Record<string, IncomingMessage[]> = {};
   let sendServerMessage: ((message: string) => void) | null = null;
 
+  const recordServerRequest = (req: IncomingMessage, res: ServerResponse) => {
+    lastServerRequest = req;
+
+    const key = `${req.method} ${req.url}`;
+    serverRequests[key] = serverRequests[key] || [];
+    serverRequests[key].push(req);
+
+    res.on('finish', () => {
+      console.log(`[server] ${req.method} ${req.url} -> ${res.statusCode} ${res.statusMessage}`);
+    });
+  };
+
   beforeEach((done) => {
     // Reset state
     lastServerRequest = null as unknown as IncomingMessage;
@@ -487,7 +500,7 @@ describe("SSEClientTransport", () => {
 
       let connectionAttempts = 0;
       server = createServer((req, res) => {
-        lastServerRequest = req;
+        recordServerRequest(req, res);
 
         if (req.url === "/token" && req.method === "POST") {
           // Handle token refresh request
@@ -496,7 +509,7 @@ describe("SSEClientTransport", () => {
           req.on("end", () => {
             const params = new URLSearchParams(body);
             if (params.get("grant_type") === "refresh_token" &&
-              params.get("refresh_token") === "refresh-token" &&
+              params.get("refresh_token")?.includes("refresh-token") &&
               params.get("client_id") === "test-client-id" &&
               params.get("client_secret") === "test-client-secret") {
               res.writeHead(200, { "Content-Type": "application/json" });
@@ -531,6 +544,7 @@ describe("SSEClientTransport", () => {
           });
           res.write("event: endpoint\n");
           res.write(`data: ${baseUrl.href}\n\n`);
+          res.end();
           connectionAttempts++;
           return;
         }
@@ -548,6 +562,14 @@ describe("SSEClientTransport", () => {
 
       transport = new SSEClientTransport(baseUrl, {
         authProvider: mockAuthProvider,
+        eventSourceInit: {
+          fetch: (url, init) => {
+            return fetch(url, { ...init, headers: {
+              ...init?.headers,
+              'X-Custom-Header': 'custom-value'
+            } });
+          }
+        },
       });
 
       await transport.start();
@@ -559,6 +581,9 @@ describe("SSEClientTransport", () => {
       });
       expect(connectionAttempts).toBe(1);
       expect(lastServerRequest.headers.authorization).toBe("Bearer new-token");
+      expect(serverRequests["GET /"]).toHaveLength(2);
+      expect(serverRequests["GET /"]
+        .every(req => req.headers["x-custom-header"] === "custom-value")).toBe(true);
     });
 
     it("refreshes expired token during POST request", async () => {

src/client/sse.ts

@@ -96,7 +96,7 @@ export class SSEClientTransport implements Transport {
     return await this._startOrAuth();
   }
 
-  private async _commonHeaders(): Promise<HeadersInit> {
+  private async _commonHeaders(): Promise<Record<string, string>> {
     const headers: HeadersInit = {};
     if (this._authProvider) {
       const tokens = await this._authProvider.tokens();
@@ -110,18 +110,7 @@ export class SSEClientTransport implements Transport {
 
   private _startOrAuth(): Promise<void> {
     return new Promise((resolve, reject) => {
-      this._eventSource = new EventSource(
-        this._url.href,
-        this._eventSourceInit ?? {
-          fetch: (url, init) => this._commonHeaders().then((headers) => fetch(url, {
-            ...init,
-            headers: {
-              ...headers,
-              Accept: "text/event-stream"
-            }
-          })),
-        },
-      );
+      this._eventSource = new EventSource(this._url.href, this._getEventSourceInit());
       this._abortController = new AbortController();
 
       this._eventSource.onerror = (event) => {
@@ -175,6 +164,44 @@ export class SSEClientTransport implements Transport {
     });
   }
 
+  private _getEventSourceInit(): EventSourceInit {
+    let eventSourceInit: EventSourceInit;
+
+    if (this._eventSourceInit) {
+      const originalFetch = this._eventSourceInit.fetch;
+
+      if (originalFetch && this._authProvider) {
+        // merge the new headers with the existing headers
+        eventSourceInit = {
+          ...this._eventSourceInit,
+          fetch: async (url, init) => {
+            const newHeaders: Record<string, string> = await this._commonHeaders();
+            return originalFetch(url, {
+              ...init,
+              headers: {
+                ...newHeaders,
+                ...init?.headers
+              }
+            });
+          }
+        };
+      } else {
+        eventSourceInit = this._eventSourceInit;
+      }
+    } else {
+      eventSourceInit = {
+        fetch: (url, init) => this._commonHeaders().then((headers) => fetch(url, {
+          ...init,
+          headers: {
+            ...headers,
+            Accept: "text/event-stream"
+          }
+        })),
+      };
+    }
+    return eventSourceInit;
+  }
+
   async start() {
     if (this._eventSource) {
       throw new Error(