Merge pull request #1081 from jeanlaurent/mcp-ignore-headers

Apply custom HTTP headers to remote MCP server requests

Author: dgageot

pkg/tools/mcp/mcp.go

@@ -56,7 +56,7 @@ func NewRemoteToolset(name, url, transport string, headers map[string]string) *T
 
 	return &Toolset{
 		name:      name,
-		mcpClient: newRemoteClient(url, transport, headers, NewInMemoryTokenStore(), false),
+		mcpClient: newRemoteClient(url, transport, headers, NewInMemoryTokenStore()),
 		logID:     url,
 	}
 }

pkg/tools/mcp/remote.go

@@ -25,8 +25,8 @@ type remoteMCPClient struct {
 	mu                  sync.RWMutex
 }
 
-func newRemoteClient(url, transportType string, headers map[string]string, tokenStore OAuthTokenStore, managed bool) *remoteMCPClient {
-	slog.Debug("Creating remote MCP client", "url", url, "transport", transportType, "headers", headers, "managed", managed)
+func newRemoteClient(url, transportType string, headers map[string]string, tokenStore OAuthTokenStore) *remoteMCPClient {
+	slog.Debug("Creating remote MCP client", "url", url, "transport", transportType, "headers", headers)
 
 	if tokenStore == nil {
 		tokenStore = NewInMemoryTokenStore()
@@ -37,7 +37,7 @@ func newRemoteClient(url, transportType string, headers map[string]string, token
 		transportType: transportType,
 		headers:       headers,
 		tokenStore:    tokenStore,
-		managed:       managed,
+		managed:       false,
 	}
 }
 
@@ -123,16 +123,47 @@ func (c *remoteMCPClient) Initialize(ctx context.Context, _ *mcp.InitializeReque
 	return session.InitializeResult(), nil
 }
 
-// createHTTPClient creates an HTTP client with OAuth support
+// headerTransport is a RoundTripper that adds custom headers to all requests
+type headerTransport struct {
+	base    http.RoundTripper
+	headers map[string]string
+}
+
+func (t *headerTransport) RoundTrip(req *http.Request) (*http.Response, error) {
+	// Clone the request to avoid modifying the original
+	req = req.Clone(req.Context())
+
+	// Add custom headers
+	for key, value := range t.headers {
+		req.Header.Set(key, value)
+	}
+
+	return t.base.RoundTrip(req)
+}
+
+// createHTTPClient creates an HTTP client with custom headers and OAuth support
 func (c *remoteMCPClient) createHTTPClient() *http.Client {
+	transport := http.DefaultTransport
+
+	// Add custom headers first
+	if len(c.headers) > 0 {
+		transport = &headerTransport{
+			base:    transport,
+			headers: c.headers,
+		}
+	}
+
+	// Then wrap with OAuth support
+	transport = &oauthTransport{
+		base:       transport,
+		client:     c,
+		tokenStore: c.tokenStore,
+		baseURL:    c.url,
+		managed:    c.managed,
+	}
+
 	return &http.Client{
-		Transport: &oauthTransport{
-			base:       http.DefaultTransport,
-			client:     c,
-			tokenStore: c.tokenStore,
-			baseURL:    c.url,
-			managed:    c.managed,
-		},
+		Transport: transport,
 	}
 }
 

pkg/tools/mcp/remote_test.go

@@ -0,0 +1,194 @@
+package mcp
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// TestRemoteClientCustomHeaders verifies that custom headers passed to the remote
+// MCP client are actually applied to HTTP requests sent to the MCP server.
+func TestRemoteClientCustomHeaders(t *testing.T) {
+	var capturedRequest *http.Request
+	requestCaptured := make(chan bool, 1)
+
+	// Create a test SSE server that captures the request
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		capturedRequest = r
+
+		// Send a minimal SSE response to satisfy the client
+		w.Header().Set("Content-Type", "text/event-stream")
+		w.WriteHeader(http.StatusOK)
+		fmt.Fprintf(w, "event: endpoint\ndata: {\"uri\":\"/message\"}\n\n")
+
+		select {
+		case requestCaptured <- true:
+		default:
+		}
+
+		// Keep the connection open briefly
+		time.Sleep(100 * time.Millisecond)
+	}))
+	defer server.Close()
+
+	// Create remote client WITH custom headers
+	expectedHeaders := map[string]string{
+		"X-Test-Header": "test-value",
+		"X-API-Key":     "secret-key-12345",
+		"Authorization": "Bearer custom-token",
+	}
+
+	client := newRemoteClient(server.URL, "sse", expectedHeaders, NewInMemoryTokenStore())
+
+	ctx, cancel := context.WithTimeout(t.Context(), 2*time.Second)
+	defer cancel()
+
+	// Try to initialize (which will make the HTTP request)
+	// We don't care if it succeeds or fails, we just need it to make the request
+	_, _ = client.Initialize(ctx, nil)
+
+	// Wait for the request to be captured
+	select {
+	case <-requestCaptured:
+		// Verify that custom headers were applied
+		for key, expectedValue := range expectedHeaders {
+			actualValue := capturedRequest.Header.Get(key)
+			assert.Equal(t, expectedValue, actualValue,
+				"Expected header %s to have value %q, but got %q",
+				key, expectedValue, actualValue)
+		}
+	case <-time.After(1 * time.Second):
+		t.Fatal("Server did not receive request within timeout")
+	}
+}
+
+// TestRemoteClientHeadersWithStreamable verifies that custom headers work with streamable transport
+func TestRemoteClientHeadersWithStreamable(t *testing.T) {
+	var capturedRequest *http.Request
+	requestCaptured := make(chan bool, 1)
+
+	// Create a test server for streamable transport
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		capturedRequest = r
+
+		// Send a minimal response
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(http.StatusOK)
+		fmt.Fprintf(w, `{"jsonrpc":"2.0","result":{"protocolVersion":"1.0.0","capabilities":{},"serverInfo":{"name":"test","version":"1.0.0"}},"id":1}`)
+
+		select {
+		case requestCaptured <- true:
+		default:
+		}
+	}))
+	defer server.Close()
+
+	// Create remote client WITH custom headers using streamable transport
+	expectedHeaders := map[string]string{
+		"X-Custom-Auth": "custom-auth-value",
+	}
+
+	client := newRemoteClient(server.URL, "streamable", expectedHeaders, NewInMemoryTokenStore())
+
+	ctx, cancel := context.WithTimeout(t.Context(), 2*time.Second)
+	defer cancel()
+
+	// Try to initialize
+	_, _ = client.Initialize(ctx, nil)
+
+	// Wait for the request to be captured
+	select {
+	case <-requestCaptured:
+		// Verify that custom headers were applied
+		actualValue := capturedRequest.Header.Get("X-Custom-Auth")
+		assert.Equal(t, "custom-auth-value", actualValue,
+			"Expected header X-Custom-Auth to have value %q, but got %q",
+			"custom-auth-value", actualValue)
+	case <-time.After(1 * time.Second):
+		t.Fatal("Server did not receive request within timeout")
+	}
+}
+
+// TestRemoteClientNoHeaders verifies that the client works correctly even with no headers
+func TestRemoteClientNoHeaders(t *testing.T) {
+	var capturedRequest *http.Request
+	requestCaptured := make(chan bool, 1)
+
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		capturedRequest = r
+
+		w.Header().Set("Content-Type", "text/event-stream")
+		w.WriteHeader(http.StatusOK)
+		fmt.Fprintf(w, "event: endpoint\ndata: {\"uri\":\"/message\"}\n\n")
+
+		select {
+		case requestCaptured <- true:
+		default:
+		}
+
+		time.Sleep(100 * time.Millisecond)
+	}))
+	defer server.Close()
+
+	// Create remote client without custom headers (nil)
+	client := newRemoteClient(server.URL, "sse", nil, NewInMemoryTokenStore())
+
+	ctx, cancel := context.WithTimeout(t.Context(), 2*time.Second)
+	defer cancel()
+
+	_, _ = client.Initialize(ctx, nil)
+
+	// Wait for request
+	select {
+	case <-requestCaptured:
+		// Just verify we got the request - no custom headers should be present
+		require.NotNil(t, capturedRequest, "Request should have been captured")
+	case <-time.After(1 * time.Second):
+		t.Fatal("Server did not receive request within timeout")
+	}
+}
+
+// TestRemoteClientEmptyHeaders verifies that the client works correctly with an empty map
+func TestRemoteClientEmptyHeaders(t *testing.T) {
+	var capturedRequest *http.Request
+	requestCaptured := make(chan bool, 1)
+
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		capturedRequest = r
+
+		w.Header().Set("Content-Type", "text/event-stream")
+		w.WriteHeader(http.StatusOK)
+		fmt.Fprintf(w, "event: endpoint\ndata: {\"uri\":\"/message\"}\n\n")
+
+		select {
+		case requestCaptured <- true:
+		default:
+		}
+
+		time.Sleep(100 * time.Millisecond)
+	}))
+	defer server.Close()
+
+	// Create remote client with empty headers map
+	client := newRemoteClient(server.URL, "sse", map[string]string{}, NewInMemoryTokenStore())
+
+	ctx, cancel := context.WithTimeout(t.Context(), 2*time.Second)
+	defer cancel()
+
+	_, _ = client.Initialize(ctx, nil)
+
+	// Wait for request
+	select {
+	case <-requestCaptured:
+		// Just verify we got the request
+		require.NotNil(t, capturedRequest, "Request should have been captured")
+	case <-time.After(1 * time.Second):
+		t.Fatal("Server did not receive request within timeout")
+	}
+}