feat: add support for custom HTTP headers in client requests

This update introduces the ability to include custom HTTP headers in requests sent from the client. This enhancement facilitates more flexible and secure communication with servers by allowing clients to pass additional information in the header of each request, such as authentication tokens or custom metadata. This feature is crucial for integrating with APIs that require specific headers for access control, content negotiation, or tracking purposes.

Signed-off-by: Matthis Holleville <[email protected]>

Author: matthisholleville

client/client.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"net/http"
 	"slices"
 	"sync"
 	"sync/atomic"
@@ -130,6 +131,7 @@ func (c *Client) sendRequest(
 	ctx context.Context,
 	method string,
 	params any,
+	header http.Header,
 ) (*json.RawMessage, error) {
 	if !c.initialized && method != "initialize" {
 		return nil, fmt.Errorf("client not initialized")
@@ -142,6 +144,7 @@ func (c *Client) sendRequest(
 		ID:      mcp.NewRequestId(id),
 		Method:  method,
 		Params:  params,
+		Header: header,
 	}
 
 	response, err := c.transport.SendRequest(ctx, request)
@@ -179,7 +182,7 @@ func (c *Client) Initialize(
 		Capabilities:    capabilities,
 	}
 
-	response, err := c.sendRequest(ctx, "initialize", params)
+	response, err := c.sendRequest(ctx, "initialize", params, request.Header)
 	if err != nil {
 		return nil, err
 	}
@@ -224,7 +227,7 @@ func (c *Client) Initialize(
 }
 
 func (c *Client) Ping(ctx context.Context) error {
-	_, err := c.sendRequest(ctx, "ping", nil)
+	_, err := c.sendRequest(ctx, "ping", nil, nil)
 	return err
 }
 
@@ -305,7 +308,7 @@ func (c *Client) ReadResource(
 	ctx context.Context,
 	request mcp.ReadResourceRequest,
 ) (*mcp.ReadResourceResult, error) {
-	response, err := c.sendRequest(ctx, "resources/read", request.Params)
+	response, err := c.sendRequest(ctx, "resources/read", request.Params, request.Header)
 	if err != nil {
 		return nil, err
 	}
@@ -317,15 +320,15 @@ func (c *Client) Subscribe(
 	ctx context.Context,
 	request mcp.SubscribeRequest,
 ) error {
-	_, err := c.sendRequest(ctx, "resources/subscribe", request.Params)
+	_, err := c.sendRequest(ctx, "resources/subscribe", request.Params, request.Header)
 	return err
 }
 
 func (c *Client) Unsubscribe(
 	ctx context.Context,
 	request mcp.UnsubscribeRequest,
 ) error {
-	_, err := c.sendRequest(ctx, "resources/unsubscribe", request.Params)
+	_, err := c.sendRequest(ctx, "resources/unsubscribe", request.Params, request.Header)
 	return err
 }
 
@@ -369,7 +372,7 @@ func (c *Client) GetPrompt(
 	ctx context.Context,
 	request mcp.GetPromptRequest,
 ) (*mcp.GetPromptResult, error) {
-	response, err := c.sendRequest(ctx, "prompts/get", request.Params)
+	response, err := c.sendRequest(ctx, "prompts/get", request.Params, request.Header)
 	if err != nil {
 		return nil, err
 	}
@@ -417,7 +420,7 @@ func (c *Client) CallTool(
 	ctx context.Context,
 	request mcp.CallToolRequest,
 ) (*mcp.CallToolResult, error) {
-	response, err := c.sendRequest(ctx, "tools/call", request.Params)
+	response, err := c.sendRequest(ctx, "tools/call", request.Params, request.Header)
 	if err != nil {
 		return nil, err
 	}
@@ -429,15 +432,15 @@ func (c *Client) SetLevel(
 	ctx context.Context,
 	request mcp.SetLevelRequest,
 ) error {
-	_, err := c.sendRequest(ctx, "logging/setLevel", request.Params)
+	_, err := c.sendRequest(ctx, "logging/setLevel", request.Params, request.Header)
 	return err
 }
 
 func (c *Client) Complete(
 	ctx context.Context,
 	request mcp.CompleteRequest,
 ) (*mcp.CompleteResult, error) {
-	response, err := c.sendRequest(ctx, "completion/complete", request.Params)
+	response, err := c.sendRequest(ctx, "completion/complete", request.Params, request.Header)
 	if err != nil {
 		return nil, err
 	}
@@ -514,7 +517,7 @@ func listByPage[T any](
 	request mcp.PaginatedRequest,
 	method string,
 ) (*T, error) {
-	response, err := client.sendRequest(ctx, method, request.Params)
+	response, err := client.sendRequest(ctx, method, request.Params, nil)
 	if err != nil {
 		return nil, err
 	}

client/transport/interface.go

@@ -3,6 +3,7 @@ package transport
 import (
 	"context"
 	"encoding/json"
+	"net/http"
 
 	"github.com/mark3labs/mcp-go/mcp"
 )
@@ -59,6 +60,7 @@ type JSONRPCRequest struct {
 	ID      mcp.RequestId `json:"id"`
 	Method  string        `json:"method"`
 	Params  any           `json:"params,omitempty"`
+	Header  http.Header   `json:"-"`
 }
 
 type JSONRPCResponse struct {

client/transport/streamable_http.go

@@ -258,7 +258,7 @@ func (c *StreamableHTTP) SendRequest(
 	ctx, cancel := c.contextAwareOfClientClose(ctx)
 	defer cancel()
 
-	resp, err := c.sendHTTP(ctx, http.MethodPost, bytes.NewReader(requestBody), "application/json, text/event-stream")
+	resp, err := c.sendHTTP(ctx, http.MethodPost, bytes.NewReader(requestBody), "application/json, text/event-stream", request.Header)
 	if err != nil {
 		if errors.Is(err, ErrSessionTerminated) && request.Method == string(mcp.MethodInitialize) {
 			// If the request is initialize, should not return a SessionTerminated error
@@ -339,13 +339,19 @@ func (c *StreamableHTTP) sendHTTP(
 	method string,
 	body io.Reader,
 	acceptType string,
+	header http.Header,
 ) (resp *http.Response, err error) {
 	// Create HTTP request
 	req, err := http.NewRequestWithContext(ctx, method, c.serverURL.String(), body)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create request: %w", err)
 	}
 
+	// request headers
+	if header != nil {
+		req.Header = header
+	}
+
 	// Set headers
 	req.Header.Set("Content-Type", "application/json")
 	req.Header.Set("Accept", acceptType)
@@ -546,7 +552,7 @@ func (c *StreamableHTTP) SendNotification(ctx context.Context, notification mcp.
 	ctx, cancel := c.contextAwareOfClientClose(ctx)
 	defer cancel()
 
-	resp, err := c.sendHTTP(ctx, http.MethodPost, bytes.NewReader(requestBody), "application/json, text/event-stream")
+	resp, err := c.sendHTTP(ctx, http.MethodPost, bytes.NewReader(requestBody), "application/json, text/event-stream", nil)
 	if err != nil {
 		return fmt.Errorf("failed to send request: %w", err)
 	}
@@ -605,7 +611,7 @@ func (c *StreamableHTTP) listenForever(ctx context.Context) {
 		connectCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
 		err := c.createGETConnectionToServer(connectCtx)
 		cancel()
-		
+
 		if errors.Is(err, ErrGetMethodNotAllowed) {
 			// server does not support listening
 			c.logger.Errorf("server does not support listening")
@@ -621,7 +627,7 @@ func (c *StreamableHTTP) listenForever(ctx context.Context) {
 		if err != nil {
 			c.logger.Errorf("failed to listen to server. retry in 1 second: %v", err)
 		}
-		
+
 		// Use context-aware sleep
 		select {
 		case <-time.After(retryInterval):
@@ -639,7 +645,7 @@ var (
 )
 
 func (c *StreamableHTTP) createGETConnectionToServer(ctx context.Context) error {
-	resp, err := c.sendHTTP(ctx, http.MethodGet, nil, "text/event-stream")
+	resp, err := c.sendHTTP(ctx, http.MethodGet, nil, "text/event-stream", nil)
 	if err != nil {
 		return fmt.Errorf("failed to send request: %w", err)
 	}
@@ -704,15 +710,15 @@ func (c *StreamableHTTP) handleIncomingRequest(ctx context.Context, request JSON
 		// Create a new context with timeout for request handling, respecting parent context
 		requestCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
 		defer cancel()
-		
+
 		response, err := handler(requestCtx, request)
 		if err != nil {
 			c.logger.Errorf("error handling request %s: %v", request.Method, err)
-			
+
 			// Determine appropriate JSON-RPC error code based on error type
 			var errorCode int
 			var errorMessage string
-			
+
 			// Check for specific sampling-related errors
 			if errors.Is(err, context.Canceled) {
 				errorCode = -32800 // Request cancelled
@@ -731,7 +737,7 @@ func (c *StreamableHTTP) handleIncomingRequest(ctx context.Context, request JSON
 					errorMessage = err.Error()
 				}
 			}
-			
+
 			// Send error response
 			errorResponse := &JSONRPCResponse{
 				JSONRPC: "2.0",
@@ -771,7 +777,7 @@ func (c *StreamableHTTP) sendResponseToServer(ctx context.Context, response *JSO
 	ctx, cancel := c.contextAwareOfClientClose(ctx)
 	defer cancel()
 
-	resp, err := c.sendHTTP(ctx, http.MethodPost, bytes.NewReader(responseBody), "application/json")
+	resp, err := c.sendHTTP(ctx, http.MethodPost, bytes.NewReader(responseBody), "application/json", nil)
 	if err != nil {
 		c.logger.Errorf("failed to send response to server: %v", err)
 		return

client/transport/streamable_http_test.go

@@ -70,6 +70,7 @@ func startMockStreamableHTTPServer() (string, func()) {
 				"jsonrpc": "2.0",
 				"id":      request["id"],
 				"result":  request,
+				"headers": r.Header,
 			}); err != nil {
 				http.Error(w, "Failed to encode response", http.StatusInternalServerError)
 				return
@@ -122,6 +123,24 @@ func startMockStreamableHTTPServer() (string, func()) {
 				http.Error(w, "Failed to encode response", http.StatusInternalServerError)
 				return
 			}
+		case "debug/echo_header":
+			// Check session ID
+			if r.Header.Get("Mcp-Session-Id") != sessionID {
+				http.Error(w, "Invalid session ID", http.StatusNotFound)
+				return
+			}
+
+			// Echo back the request headersas the response result
+			w.Header().Set("Content-Type", "application/json")
+			w.WriteHeader(http.StatusOK)
+			if err := json.NewEncoder(w).Encode(map[string]any{
+				"jsonrpc": "2.0",
+				"id":      request["id"],
+				"result":  r.Header,
+			}); err != nil {
+				http.Error(w, "Failed to encode response", http.StatusInternalServerError)
+				return
+			}
 		}
 	})
 
@@ -215,6 +234,40 @@ func TestStreamableHTTP(t *testing.T) {
 		}
 	})
 
+	t.Run("SendRequestWithHeader", func(t *testing.T) {
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+
+		params := map[string]any{
+			"string": "hello world",
+			"array":  []any{1, 2, 3},
+		}
+
+		request := JSONRPCRequest{
+			JSONRPC: "2.0",
+			ID:      mcp.NewRequestId(int64(1)),
+			Method:  "debug/echo_header",
+			Params:  params,
+			Header:  http.Header{"X-Test-Header": {"test-header-value"}},
+		}
+
+		// Send the request
+		response, err := trans.SendRequest(ctx, request)
+		if err != nil {
+			t.Fatalf("SendRequest failed: %v", err)
+		}
+
+		// Parse the result to verify echo
+		var result map[string]any
+		if err := json.Unmarshal(response.Result, &result); err != nil {
+			t.Fatalf("Failed to unmarshal result: %v", err)
+		}
+
+		if hdr, ok := result["X-Test-Header"].([]any); !ok || len(hdr) == 0 || hdr[0] != "test-header-value" {
+			t.Errorf("Expected X-Test-Header to be ['test-header-value'], got %v", result["X-Test-Header"])
+		}
+	})
+
 	t.Run("SendRequestWithTimeout", func(t *testing.T) {
 		// Create a context that's already canceled
 		ctx, cancel := context.WithCancel(context.Background())