feat: Transport setHeaders (#44)

Signed-off-by: John McBride <john@zuplo.com>

Author: jpmcb

src/client/index.ts

@@ -18,7 +18,7 @@ import type {
   Tool,
 } from "../mcp/20250618/types.js";
 import { LATEST_PROTOCOL_VERSION } from "../mcp/versions.js";
-import type { Transport } from "../transport/types.js";
+import type { Transport, TransportOptions } from "../transport/types.js";
 import type { MCPClientOptions } from "./types.js";
 
 export const DEFAULT_MCP_CLIENT_NAME = "MCP Client";
@@ -34,11 +34,13 @@ export class MCPClient {
   private protocolVersion?: string;
   private logger: Logger;
   private requestId = 1;
+  private transportOptions: TransportOptions;
 
   constructor(options: MCPClientOptions = {}) {
     this.name = options.name || DEFAULT_MCP_CLIENT_NAME;
     this.version = options.version || DEFAULT_MCP_CLIENT_VERSION;
     this.logger = options.logger || createDefaultLogger();
+    this.transportOptions = options.transportOptions || {};
     this.capabilities = {
       experimental: {},
       sampling: {},
@@ -50,6 +52,9 @@ export class MCPClient {
    * Connect to an MCP server using the provided transport
    */
   public async connect(transport: Transport): Promise<void> {
+    if (this.transportOptions.headers) {
+      transport.setHeaders(this.transportOptions.headers);
+    }
     this.transport = transport;
     await transport.connect();
   }

src/client/types.ts

@@ -1,9 +1,11 @@
 import type { Logger } from "../logger/types.js";
 import type { ClientCapabilities } from "../mcp/20250618/types.js";
+import type { TransportOptions } from "../transport/types.js";
 
 export interface MCPClientOptions {
   name?: string;
   version?: string;
   logger?: Logger;
   capabilities?: ClientCapabilities;
+  transportOptions?: TransportOptions;
 }

src/transport/httpclient.ts

@@ -45,6 +45,10 @@ export class HTTPClientTransport implements Transport {
     }
   }
 
+  setHeaders(headers: Record<string, string>): void {
+    this.headers = { ...this.headers, ...headers };
+  }
+
   async connect(): Promise<void> {
     try {
       new URL(this.url);
@@ -79,7 +83,10 @@ export class HTTPClientTransport implements Transport {
 
       const response = await this.fetch(this.url, {
         method: "POST",
-        headers: requestHeaders,
+        headers: {
+          ...this.headers,
+          ...(this.sessionId && { "Mcp-Session-Id": this.sessionId }),
+        },
         body: JSON.stringify(message),
         signal: controller.signal,
       });

src/transport/httpstreamable.ts

@@ -48,20 +48,31 @@ export class HTTPStreamableTransport implements Transport {
   messageHandler: MessageHandler | null = null;
   closeCallback: (() => void) | null = null;
 
+  private headers: Record<string, string>;
   private options: TransportOptions;
   private connected = false;
   private enableStreaming = false;
   private sessions: Map<string, Session> = new Map();
   private streams: Map<string, StreamInfo> = new Map();
   private logger: Logger;
 
+  setHeaders(headers: Record<string, string>): void {
+    this.headers = { ...this.headers, ...headers };
+  }
+
   constructor(options: TransportOptions = {}, streamable = false) {
     // Set defaults
     this.options = {
       timeout: 30 * 60 * 1000, // 30 minutes
       enableSessions: false,
       ...options,
     };
+
+    this.headers = {
+      "Content-Type": "application/json",
+      ...options.headers,
+    };
+
     this.logger = options.logger || createDefaultLogger();
 
     // Start session cleanup timer if sessions are enabled
@@ -189,7 +200,7 @@ export class HTTPStreamableTransport implements Transport {
         ),
         {
           status: 503,
-          headers: { "Content-Type": "application/json" },
+          headers: this.headers,
         }
       );
     }
@@ -205,7 +216,7 @@ export class HTTPStreamableTransport implements Transport {
         ),
         {
           status: 500,
-          headers: { "Content-Type": "application/json" },
+          headers: this.headers,
         }
       );
     }
@@ -260,7 +271,7 @@ export class HTTPStreamableTransport implements Transport {
         ),
         {
           status: 400,
-          headers: { "Content-Type": "application/json" },
+          headers: this.headers,
         }
       );
     }
@@ -295,7 +306,7 @@ export class HTTPStreamableTransport implements Transport {
         ),
         {
           status: 406,
-          headers: { "Content-Type": "application/json" },
+          headers: this.headers,
         }
       );
     }
@@ -313,7 +324,7 @@ export class HTTPStreamableTransport implements Transport {
         ),
         {
           status: 400,
-          headers: { "Content-Type": "application/json" },
+          headers: this.headers,
         }
       );
     }
@@ -347,9 +358,10 @@ export class HTTPStreamableTransport implements Transport {
         }
         return new Response(null, {
           status: 202,
-          headers: currentSession
-            ? { "Mcp-Session-Id": currentSession.id }
-            : undefined,
+          headers: {
+            ...this.headers,
+            ...(currentSession && { "Mcp-Session-Id": currentSession.id }),
+          },
         });
       }
 
@@ -373,17 +385,12 @@ export class HTTPStreamableTransport implements Transport {
 
         // Return a direct JSON response
         const responseBody = responses.length === 1 ? responses[0] : responses;
-        const headers: Record<string, string> = {
-          "Content-Type": "application/json",
-        };
-
-        if (currentSession) {
-          headers["Mcp-Session-Id"] = currentSession.id;
-        }
-
         return new Response(JSON.stringify(responseBody), {
           status: 200,
-          headers,
+          headers: {
+            ...this.headers,
+            ...(currentSession && { "Mcp-Session-Id": currentSession.id }),
+          },
         });
       }
 
@@ -407,18 +414,15 @@ export class HTTPStreamableTransport implements Transport {
         }
       }
 
-      const headers: Record<string, string> = {
-        "Content-Type": "text/event-stream",
-        "Cache-Control": "no-cache",
-        Connection: "keep-alive",
-      };
-
-      if (currentSession) {
-        headers["Mcp-Session-Id"] = currentSession.id;
-      }
-
       // Let responses be handled by the message handler and sent to the stream
-      return new Response(stream.readable, { headers });
+      return new Response(stream.readable, {
+        headers: {
+          "Cache-Control": "no-cache",
+          Connection: "keep-alive",
+          ...this.headers,
+          ...(currentSession && { "Mcp-Session-Id": currentSession.id }),
+        },
+      });
     } catch (error) {
       return new Response(
         JSON.stringify(
@@ -430,7 +434,7 @@ export class HTTPStreamableTransport implements Transport {
         ),
         {
           status: 500,
-          headers: { "Content-Type": "application/json" },
+          headers: this.headers,
         }
       );
     }
@@ -441,19 +445,19 @@ export class HTTPStreamableTransport implements Transport {
    */
   private async handleGetRequest(
     request: Request,
-    session?: Session
+    currentSession?: Session
   ): Promise<Response> {
     // Check Accept header as required by the spec
     const acceptHeader = request.headers.get("Accept") || "";
     if (!acceptHeader.includes("text/event-stream")) {
       return new Response(null, {
         status: 406, // Not Acceptable
-        headers: { "Content-Type": "application/json" },
+        headers: this.headers,
       });
     }
 
     // If no session and sessions are required, reject
-    if (this.options.enableSessions && !session) {
+    if (this.options.enableSessions && !currentSession) {
       return new Response(
         JSON.stringify(
           newJSONRPCError({
@@ -464,33 +468,29 @@ export class HTTPStreamableTransport implements Transport {
         ),
         {
           status: 400,
-          headers: { "Content-Type": "application/json" },
+          headers: this.headers,
         }
       );
     }
 
     // Create a new stream for server-to-client communication
-    const { stream, streamId } = this.createStream(session);
+    const { stream, streamId } = this.createStream(currentSession);
 
     // Check for Last-Event-ID for resumability
     const lastEventId = request.headers.get("Last-Event-ID");
-    if (lastEventId && session) {
+    if (lastEventId && currentSession) {
       // Replay messages if needed
-      await this.replayMessages(session, streamId, lastEventId);
+      await this.replayMessages(currentSession, streamId, lastEventId);
     }
 
-    // Return the stream
-    const headers: Record<string, string> = {
-      "Content-Type": "text/event-stream",
-      "Cache-Control": "no-cache",
-      Connection: "keep-alive",
-    };
-
-    if (session) {
-      headers["Mcp-Session-Id"] = session.id;
-    }
-
-    return new Response(stream.readable, { headers });
+    return new Response(stream.readable, {
+      headers: {
+        "Cache-Control": "no-cache",
+        Connection: "keep-alive",
+        ...this.headers,
+        ...(currentSession && { "Mcp-Session-Id": currentSession.id }),
+      },
+    });
   }
 
   /**

src/transport/types.ts

@@ -62,4 +62,9 @@ export interface Transport {
    * Set a session ID for subsequent requests
    */
   setSessionId(sessionId: string | undefined): void;
+
+  /**
+   * Set custom headers for requests
+   */
+  setHeaders(headers: Record<string, string>): void;
 }