feat: add support for custom HTTP headers in client requests (#546)

* 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]>

* feat(client/transport): enhance HTTP request flexibility

Enhanced the flexibility of HTTP requests in the streamable HTTP client by allowing additional headers to be specified. This change aims to support more diverse server requirements and improve the adaptability of our client transport layer.

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

* fix: Improve OAuth error handling and test readability in HTTP transport

- Enhanced OAuth error detection by using `errors.Is` for more reliable error handling.
- Corrected a typo in a comment and improved code readability in tests by using a variable for headers.

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

* fix: improve variable naming for clarity in streamable_http_test

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

* feat: Ensure system headers are preserved in streamable HTTP tests

To maintain consistency and ensure the integrity of HTTP headers during tests, system headers like Content-Type are now verified to be preserved. This change enhances the reliability of our testing framework by ensuring essential headers are not inadvertently removed or altered during the testing process.

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

---------

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

Author: matthisholleville

client/client.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"net/http"
 	"slices"
 	"sync"
 	"sync/atomic"
@@ -140,6 +141,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")
@@ -152,6 +154,7 @@ func (c *Client) sendRequest(
 		ID:      mcp.NewRequestId(id),
 		Method:  method,
 		Params:  params,
+		Header: header,
 	}
 
 	response, err := c.transport.SendRequest(ctx, request)
@@ -193,7 +196,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
 	}
@@ -238,7 +241,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
 }
 
@@ -319,7 +322,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
 	}
@@ -331,15 +334,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
 }
 
@@ -383,7 +386,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
 	}
@@ -431,7 +434,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
 	}
@@ -443,15 +446,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
 	}
@@ -591,7 +594,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:"-"`
 }
 
 // JSONRPCResponse represents a JSON-RPC 2.0 response message.

client/transport/streamable_http.go

@@ -265,7 +265,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
@@ -346,13 +346,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)
@@ -375,7 +381,7 @@ func (c *StreamableHTTP) sendHTTP(
 		authHeader, err := c.oauthHandler.GetAuthorizationHeader(ctx)
 		if err != nil {
 			// If we get an authorization error, return a specific error that can be handled by the client
-			if err.Error() == "no valid token available, authorization required" {
+			if errors.Is(err, ErrOAuthAuthorizationRequired) {
 				return nil, &OAuthAuthorizationRequiredError{
 					Handler: c.oauthHandler,
 				}
@@ -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)
 	}
@@ -642,7 +648,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)
 	}
@@ -757,7 +763,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, text/event-stream")
+	resp, err := c.sendHTTP(ctx, http.MethodPost, bytes.NewReader(responseBody), "application/json, text/event-stream", 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 headers as 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,46 @@ 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},
+		}
+
+		hdr := http.Header{"X-Test-Header": {"test-header-value"}}
+		request := JSONRPCRequest{
+			JSONRPC: "2.0",
+			ID:      mcp.NewRequestId(int64(1)),
+			Method:  "debug/echo_header",
+			Params:  params,
+			Header:  hdr,
+		}
+
+		// 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 headerValues, ok := result["X-Test-Header"].([]any); !ok || len(headerValues) == 0 || headerValues[0] != "test-header-value" {
+			t.Errorf("Expected X-Test-Header to be ['test-header-value'], got %v", result["X-Test-Header"])
+		}
+
+		// Verify system headers are still present
+		if contentType, ok := result["Content-Type"].([]any); !ok || len(contentType) == 0 {
+			t.Errorf("Expected Content-Type header to be preserved")
+		}
+	})
+
 	t.Run("SendRequestWithTimeout", func(t *testing.T) {
 		// Create a context that's already canceled
 		ctx, cancel := context.WithCancel(context.Background())