Merge branch 'main' into feat/sendlog

Author: pottekkat

README.md

@@ -58,7 +58,7 @@ func main() {
 }
 
 func helloHandler(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
-    name, ok := request.Params.Arguments["name"].(string)
+    name, ok := request.GetArguments()["name"].(string)
     if !ok {
         return nil, errors.New("name must be a string")
     }
@@ -97,6 +97,7 @@ MCP Go handles all the complex protocol details and server management, so you ca
   - [Session Management](#session-management)
   - [Request Hooks](#request-hooks)
   - [Tool Handler Middleware](#tool-handler-middleware)
+  - [Regenerating Server Code](#regenerating-server-code)
 - [Contributing](/CONTRIBUTING.md)
 
 ## Installation
@@ -149,9 +150,21 @@ func main() {
 
     // Add the calculator handler
     s.AddTool(calculatorTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
-        op := request.Params.Arguments["operation"].(string)
-        x := request.Params.Arguments["x"].(float64)
-        y := request.Params.Arguments["y"].(float64)
+        // Using helper functions for type-safe argument access
+        op, err := request.RequireString("operation")
+        if err != nil {
+            return mcp.NewToolResultError(err.Error()), nil
+        }
+        
+        x, err := request.RequireFloat("x")
+        if err != nil {
+            return mcp.NewToolResultError(err.Error()), nil
+        }
+        
+        y, err := request.RequireFloat("y")
+        if err != nil {
+            return mcp.NewToolResultError(err.Error()), nil
+        }
 
         var result float64
         switch op {
@@ -312,9 +325,10 @@ calculatorTool := mcp.NewTool("calculate",
 )
 
 s.AddTool(calculatorTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
-    op := request.Params.Arguments["operation"].(string)
-    x := request.Params.Arguments["x"].(float64)
-    y := request.Params.Arguments["y"].(float64)
+    args := request.GetArguments()
+    op := args["operation"].(string)
+    x := args["x"].(float64)
+    y := args["y"].(float64)
 
     var result float64
     switch op {
@@ -355,10 +369,11 @@ httpTool := mcp.NewTool("http_request",
 )
 
 s.AddTool(httpTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
-    method := request.Params.Arguments["method"].(string)
-    url := request.Params.Arguments["url"].(string)
+    args := request.GetArguments()
+    method := args["method"].(string)
+    url := args["url"].(string)
     body := ""
-    if b, ok := request.Params.Arguments["body"].(string); ok {
+    if b, ok := args["body"].(string); ok {
         body = b
     }
 
@@ -517,6 +532,10 @@ For examples, see the `examples/` directory.
 
 ## Extras
 
+### Transports
+
+MCP-Go supports stdio, SSE and streamable-HTTP transport layers.
+
 ### Session Management
 
 MCP-Go provides a robust session management system that allows you to:
@@ -742,3 +761,14 @@ Add middleware to tool call handlers using the `server.WithToolHandlerMiddleware
 
 A recovery middleware option is available to recover from panics in a tool call and can be added to the server with the `server.WithRecovery` option.
 
+### Regenerating Server Code
+
+Server hooks and request handlers are generated. Regenerate them by running:
+
+```bash
+go generate ./...
+```
+
+You need `go` installed and the `goimports` tool available. The generator runs
+`goimports` automatically to format and fix imports.
+

client/inprocess_test.go

@@ -32,7 +32,7 @@ func TestInProcessMCPClient(t *testing.T) {
 			Content: []mcp.Content{
 				mcp.TextContent{
 					Type: "text",
-					Text: "Input parameter: " + request.Params.Arguments["parameter-1"].(string),
+					Text: "Input parameter: " + request.GetArguments()["parameter-1"].(string),
 				},
 				mcp.AudioContent{
 					Type:     "audio",

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(
@@ -36,14 +44,34 @@ func TestSSEMCPClient(t *testing.T) {
 			Content: []mcp.Content{
 				mcp.TextContent{
 					Type: "text",
-					Text: "Input parameter: " + request.Params.Arguments["parameter-1"].(string),
+					Text: "Input parameter: " + request.GetArguments()["parameter-1"].(string),
+				},
+			},
+		}, 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.WithSSEContextFunc(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()
@@ -310,8 +327,11 @@ func (c *SSE) SendRequest(
 	case <-ctx.Done():
 		deleteResponseChan()
 		return nil, ctx.Err()
-	case response := <-responseChan:
-		return response, nil
+	case response, ok := <-responseChan:
+		if ok {
+			return response, nil
+		}
+		return nil, fmt.Errorf("connection has been closed")
 	}
 }
 
@@ -365,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
 
@@ -127,7 +134,6 @@ func (c *StreamableHTTP) Close() error {
 }
 
 const (
-	initializeMethod   = "initialize"
 	headerKeySessionID = "Mcp-Session-Id"
 )
 
@@ -173,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)
@@ -198,7 +209,7 @@ func (c *StreamableHTTP) SendRequest(
 		return nil, fmt.Errorf("request failed with status %d: %s", resp.StatusCode, body)
 	}
 
-	if request.Method == initializeMethod {
+	if request.Method == string(mcp.MethodInitialize) {
 		// saved the received session ID in the response
 		// empty session ID is allowed
 		if sessionID := resp.Header.Get(headerKeySessionID); sessionID != "" {
@@ -363,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)