support customize request header (#315)

Author: dugenkui03

client/sse.go

@@ -12,6 +12,10 @@ func WithHeaders(headers map[string]string) transport.ClientOption {
 	return transport.WithHeaders(headers)
 }
 
+func WithHeaderFunc(headerFunc transport.HTTPHeaderFunc) transport.ClientOption {
+	return transport.WithHeaderFunc(headerFunc)
+}
+
 func WithHTTPClient(httpClient *http.Client) transport.ClientOption {
 	return transport.WithHTTPClient(httpClient)
 }

client/sse_test.go

@@ -2,6 +2,7 @@ package client
 
 import (
 	"context"
+	"net/http"
 	"testing"
 	"time"
 
@@ -11,6 +12,13 @@ import (
 	"github.com/mark3labs/mcp-go/server"
 )
 
+type contextKey string
+
+const (
+	testHeaderKey     contextKey = "X-Test-Header"
+	testHeaderFuncKey contextKey = "X-Test-Header-Func"
+)
+
 func TestSSEMCPClient(t *testing.T) {
 	// Create MCP server with capabilities
 	mcpServer := server.NewMCPServer(
@@ -41,9 +49,29 @@ func TestSSEMCPClient(t *testing.T) {
 			},
 		}, nil
 	})
+	mcpServer.AddTool(mcp.NewTool(
+		"test-tool-for-http-header",
+		mcp.WithDescription("Test tool for http header"),
+	), func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+		//  , X-Test-Header-Func
+		return &mcp.CallToolResult{
+			Content: []mcp.Content{
+				mcp.TextContent{
+					Type: "text",
+					Text: "context from header: " + ctx.Value(testHeaderKey).(string) + ", " + ctx.Value(testHeaderFuncKey).(string),
+				},
+			},
+		}, nil
+	})
 
 	// Initialize
-	testServer := server.NewTestServer(mcpServer)
+	testServer := server.NewTestServer(mcpServer,
+		server.WithHTTPContextFunc(func(ctx context.Context, r *http.Request) context.Context {
+			ctx = context.WithValue(ctx, testHeaderKey, r.Header.Get("X-Test-Header"))
+			ctx = context.WithValue(ctx, testHeaderFuncKey, r.Header.Get("X-Test-Header-Func"))
+			return ctx
+		}),
+	)
 	defer testServer.Close()
 
 	t.Run("Can create client", func(t *testing.T) {
@@ -250,4 +278,56 @@ func TestSSEMCPClient(t *testing.T) {
 			t.Errorf("Expected 1 content item, got %d", len(result.Content))
 		}
 	})
+
+	t.Run("CallTool with customized header", func(t *testing.T) {
+		client, err := NewSSEMCPClient(testServer.URL+"/sse",
+			WithHeaders(map[string]string{
+				"X-Test-Header": "test-header-value",
+			}),
+			WithHeaderFunc(func(ctx context.Context) map[string]string {
+				return map[string]string{
+					"X-Test-Header-Func": "test-header-func-value",
+				}
+			}),
+		)
+		if err != nil {
+			t.Fatalf("Failed to create client: %v", err)
+		}
+		defer client.Close()
+
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+
+		if err := client.Start(ctx); err != nil {
+			t.Fatalf("Failed to start client: %v", err)
+		}
+
+		// Initialize
+		initRequest := mcp.InitializeRequest{}
+		initRequest.Params.ProtocolVersion = mcp.LATEST_PROTOCOL_VERSION
+		initRequest.Params.ClientInfo = mcp.Implementation{
+			Name:    "test-client",
+			Version: "1.0.0",
+		}
+
+		_, err = client.Initialize(ctx, initRequest)
+		if err != nil {
+			t.Fatalf("Failed to initialize: %v", err)
+		}
+
+		request := mcp.CallToolRequest{}
+		request.Params.Name = "test-tool-for-http-header"
+
+		result, err := client.CallTool(ctx, request)
+		if err != nil {
+			t.Fatalf("CallTool failed: %v", err)
+		}
+
+		if len(result.Content) != 1 {
+			t.Errorf("Expected 1 content item, got %d", len(result.Content))
+		}
+		if result.Content[0].(mcp.TextContent).Text != "context from header: test-header-value, test-header-func-value" {
+			t.Errorf("Got %q, want %q", result.Content[0].(mcp.TextContent).Text, "context from header: test-header-value, test-header-func-value")
+		}
+	})
 }

client/transport/interface.go

@@ -7,6 +7,11 @@ import (
 	"github.com/mark3labs/mcp-go/mcp"
 )
 
+// HTTPHeaderFunc is a function that extracts header entries from the given context
+// and returns them as key-value pairs. This is typically used to add context values
+// as HTTP headers in outgoing requests.
+type HTTPHeaderFunc func(context.Context) map[string]string
+
 // Interface for the transport layer.
 type Interface interface {
 	// Start the connection. Start should only be called once.

client/transport/sse.go

@@ -31,6 +31,7 @@ type SSE struct {
 	notifyMu       sync.RWMutex
 	endpointChan   chan struct{}
 	headers        map[string]string
+	headerFunc     HTTPHeaderFunc
 
 	started         atomic.Bool
 	closed          atomic.Bool
@@ -45,6 +46,12 @@ func WithHeaders(headers map[string]string) ClientOption {
 	}
 }
 
+func WithHeaderFunc(headerFunc HTTPHeaderFunc) ClientOption {
+	return func(sc *SSE) {
+		sc.headerFunc = headerFunc
+	}
+}
+
 func WithHTTPClient(httpClient *http.Client) ClientOption {
 	return func(sc *SSE) {
 		sc.httpClient = httpClient
@@ -99,6 +106,11 @@ func (c *SSE) Start(ctx context.Context) error {
 	for k, v := range c.headers {
 		req.Header.Set(k, v)
 	}
+	if c.headerFunc != nil {
+		for k, v := range c.headerFunc(ctx) {
+			req.Header.Set(k, v)
+		}
+	}
 
 	resp, err := c.httpClient.Do(req)
 	if err != nil {
@@ -269,6 +281,11 @@ func (c *SSE) SendRequest(
 	for k, v := range c.headers {
 		req.Header.Set(k, v)
 	}
+	if c.headerFunc != nil {
+		for k, v := range c.headerFunc(ctx) {
+			req.Header.Set(k, v)
+		}
+	}
 
 	// Create string key for map lookup
 	idKey := request.ID.String()
@@ -368,6 +385,11 @@ func (c *SSE) SendNotification(ctx context.Context, notification mcp.JSONRPCNoti
 	for k, v := range c.headers {
 		req.Header.Set(k, v)
 	}
+	if c.headerFunc != nil {
+		for k, v := range c.headerFunc(ctx) {
+			req.Header.Set(k, v)
+		}
+	}
 
 	resp, err := c.httpClient.Do(req)
 	if err != nil {

client/transport/streamable_http.go

@@ -26,6 +26,12 @@ func WithHTTPHeaders(headers map[string]string) StreamableHTTPCOption {
 	}
 }
 
+func WithHTTPHeaderFunc(headerFunc HTTPHeaderFunc) StreamableHTTPCOption {
+	return func(sc *StreamableHTTP) {
+		sc.headerFunc = headerFunc
+	}
+}
+
 // WithHTTPTimeout sets the timeout for a HTTP request and stream.
 func WithHTTPTimeout(timeout time.Duration) StreamableHTTPCOption {
 	return func(sc *StreamableHTTP) {
@@ -52,6 +58,7 @@ type StreamableHTTP struct {
 	baseURL    *url.URL
 	httpClient *http.Client
 	headers    map[string]string
+	headerFunc HTTPHeaderFunc
 
 	sessionID atomic.Value // string
 
@@ -172,6 +179,11 @@ func (c *StreamableHTTP) SendRequest(
 	for k, v := range c.headers {
 		req.Header.Set(k, v)
 	}
+	if c.headerFunc != nil {
+		for k, v := range c.headerFunc(ctx) {
+			req.Header.Set(k, v)
+		}
+	}
 
 	// Send request
 	resp, err := c.httpClient.Do(req)
@@ -362,6 +374,11 @@ func (c *StreamableHTTP) SendNotification(ctx context.Context, notification mcp.
 	for k, v := range c.headers {
 		req.Header.Set(k, v)
 	}
+	if c.headerFunc != nil {
+		for k, v := range c.headerFunc(ctx) {
+			req.Header.Set(k, v)
+		}
+	}
 
 	// Send request
 	resp, err := c.httpClient.Do(req)