mcp: HTTP Header Standardization for x-mcp-header (#915)

## Description

Implements
[SEP-2243](https://modelcontextprotocol.io/seps/2243-http-standardization)
(HTTP Header Standardization) for x-mcp-param custom header.

Fixes #905

---------

Co-authored-by: Maciej Kisiel <mkisiel@google.com>

Author: guglielmo-san

mcp/client.go

@@ -160,6 +160,12 @@ type ClientOptions struct {
 	KeepAlive time.Duration
 }
 
+// toolContextKeyType is the context key type for passing tool definitions
+// from CallTool to the transport layer.
+type toolContextKeyType struct{}
+
+var toolContextKey = toolContextKeyType{}
+
 // bind implements the binder[*ClientSession] interface, so that Clients can
 // be connected using [connect].
 func (c *Client) bind(mcpConn Connection, conn *jsonrpc2.Connection, state *clientSessionState, onClose func()) *ClientSession {
@@ -318,6 +324,13 @@ type ClientSession struct {
 	// Pending URL elicitations waiting for completion notifications.
 	pendingElicitationsMu sync.Mutex
 	pendingElicitations   map[string]chan struct{}
+
+	// toolCacheMu guards toolCache.
+	toolCacheMu sync.RWMutex
+	// toolCache stores tool definitions keyed by name.
+	// It is used to look up x-mcp-header annotations when
+	// constructing Mcp-Param-* headers for tools/call requests.
+	toolCache map[string]*Tool
 }
 
 type clientSessionState struct {
@@ -363,6 +376,21 @@ func (cs *ClientSession) Wait() error {
 	return cs.conn.Wait()
 }
 
+func (cs *ClientSession) cacheTools(tools []*Tool) {
+	cs.toolCacheMu.Lock()
+	defer cs.toolCacheMu.Unlock()
+	cs.toolCache = make(map[string]*Tool, len(tools))
+	for _, tool := range tools {
+		cs.toolCache[tool.Name] = tool
+	}
+}
+
+func (cs *ClientSession) getCachedTool(name string) *Tool {
+	cs.toolCacheMu.RLock()
+	defer cs.toolCacheMu.RUnlock()
+	return cs.toolCache[name]
+}
+
 // registerElicitationWaiter registers a waiter for an elicitation complete
 // notification with the given elicitation ID. It returns two functions: an await
 // function that waits for the notification or context cancellation, and a cleanup
@@ -981,7 +1009,13 @@ func (cs *ClientSession) GetPrompt(ctx context.Context, params *GetPromptParams)
 
 // ListTools lists tools that are currently available on the server.
 func (cs *ClientSession) ListTools(ctx context.Context, params *ListToolsParams) (*ListToolsResult, error) {
-	return handleSend[*ListToolsResult](ctx, methodListTools, newClientRequest(cs, orZero[Params](params)))
+	result, err := handleSend[*ListToolsResult](ctx, methodListTools, newClientRequest(cs, orZero[Params](params)))
+	if err != nil {
+		return nil, err
+	}
+	result.Tools = filterValidTools(cs.client.opts.Logger, result.Tools)
+	cs.cacheTools(result.Tools)
+	return result, nil
 }
 
 // CallTool calls the tool with the given parameters.
@@ -995,6 +1029,9 @@ func (cs *ClientSession) CallTool(ctx context.Context, params *CallToolParams) (
 		// Avoid sending nil over the wire.
 		params.Arguments = map[string]any{}
 	}
+	if tool := cs.getCachedTool(params.Name); tool != nil {
+		ctx = context.WithValue(ctx, toolContextKey, tool)
+	}
 	return handleSend[*CallToolResult](ctx, methodCallTool, newClientRequest(cs, orZero[Params](params)))
 }
 

mcp/client_test.go

@@ -440,6 +440,80 @@ func TestClientCapabilities(t *testing.T) {
 	}
 }
 
+func TestToolCache(t *testing.T) {
+	tool1 := &Tool{Name: "tool1", Description: "first"}
+	tool2 := &Tool{Name: "tool2", Description: "second"}
+	tool1Updated := &Tool{Name: "tool1", Description: "updated"}
+
+	testCases := []struct {
+		name         string
+		cacheBatches [][]*Tool
+		lookup       string
+		want         *Tool
+	}{
+		{
+			name:   "empty cache",
+			lookup: "tool1",
+			want:   nil,
+		},
+		{
+			name:         "single tool found",
+			cacheBatches: [][]*Tool{{tool1}},
+			lookup:       "tool1",
+			want:         tool1,
+		},
+		{
+			name:         "unknown tool",
+			cacheBatches: [][]*Tool{{tool1}},
+			lookup:       "nonexistent",
+			want:         nil,
+		},
+		{
+			name:         "multiple tools single batch",
+			cacheBatches: [][]*Tool{{tool1, tool2}},
+			lookup:       "tool2",
+			want:         tool2,
+		},
+		{
+			name:         "replace clears old entries",
+			cacheBatches: [][]*Tool{{tool1}, {tool2}},
+			lookup:       "tool1",
+			want:         nil,
+		},
+		{
+			name:         "replace keeps new entries",
+			cacheBatches: [][]*Tool{{tool1}, {tool2}},
+			lookup:       "tool2",
+			want:         tool2,
+		},
+		{
+			name:         "overwrite existing entry",
+			cacheBatches: [][]*Tool{{tool1}, {tool1Updated}},
+			lookup:       "tool1",
+			want:         tool1Updated,
+		},
+		{
+			name:         "empty batch no-op",
+			cacheBatches: [][]*Tool{{}},
+			lookup:       "tool1",
+			want:         nil,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			cs := &ClientSession{}
+			for _, batch := range tc.cacheBatches {
+				cs.cacheTools(batch)
+			}
+			got := cs.getCachedTool(tc.lookup)
+			if diff := cmp.Diff(tc.want, got); diff != "" {
+				t.Errorf("getCachedTool(%q) mismatch (-want +got):\n%s", tc.lookup, diff)
+			}
+		})
+	}
+}
+
 func TestClientCapabilitiesOverWire(t *testing.T) {
 	testCases := []struct {
 		name             string

mcp/server.go

@@ -280,6 +280,9 @@ func (s *Server) AddTool(t *Tool, h ToolHandler) {
 			}
 		}
 	}
+	if err := validateParamHeaderAnnotations(t); err != nil {
+		panic(fmt.Errorf("AddTool %q: invalid parameter header annotations: %v", t.Name, err))
+	}
 	st := &serverTool{tool: t, handler: h}
 	// Assume there was a change, since add replaces existing tools.
 	// (It's possible a tool was replaced with an identical one, but not worth checking.)
@@ -753,10 +756,15 @@ func (s *Server) listTools(_ context.Context, req *ListToolsRequest) (*ListTools
 	})
 }
 
-func (s *Server) callTool(ctx context.Context, req *CallToolRequest) (*CallToolResult, error) {
+// getServerTool looks up a server tool by name.
+func (s *Server) getServerTool(name string) (*serverTool, bool) {
 	s.mu.Lock()
-	st, ok := s.tools.get(req.Params.Name)
-	s.mu.Unlock()
+	defer s.mu.Unlock()
+	return s.tools.get(name)
+}
+
+func (s *Server) callTool(ctx context.Context, req *CallToolRequest) (*CallToolResult, error) {
+	st, ok := s.getServerTool(req.Params.Name)
 	if !ok {
 		return nil, &jsonrpc.Error{
 			Code:    jsonrpc.CodeInvalidParams,

mcp/streamable.go

@@ -491,6 +491,7 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque
 			http.Error(w, "failed connection", http.StatusInternalServerError)
 			return
 		}
+		transport.connection.toolLookup = server.getServerTool
 		// Capture the user ID from the token info to enable session hijacking
 		// prevention on subsequent requests.
 		var userID string
@@ -669,6 +670,8 @@ type streamableServerConn struct {
 
 	logger *slog.Logger
 
+	toolLookup func(name string) (*serverTool, bool)
+
 	incoming chan jsonrpc.Message // messages from the client to the server
 
 	mu sync.Mutex // guards all fields below
@@ -1202,9 +1205,9 @@ func (c *streamableServerConn) servePOST(w http.ResponseWriter, req *http.Reques
 		}
 	}
 
-	// Validate MCP standard headers (Mcp-Method, Mcp-Name)
+	// Validate MCP standard headers (Mcp-Method, Mcp-Name, Mcp-Param-*)
 	if !isBatch && len(incoming) == 1 {
-		if err := validateMcpHeaders(req.Header, incoming[0]); err != nil {
+		if err := validateMcpHeaders(req.Header, incoming[0], c.toolLookup); err != nil {
 			resp := &jsonrpc.Response{
 				Error: jsonrpc2.NewError(CodeHeaderMismatch, err.Error()),
 			}
@@ -1829,7 +1832,7 @@ func (c *streamableClientConn) Write(ctx context.Context, msg jsonrpc.Message) e
 		}
 		// Keep this after the setMCPHeaders call to ensure that the
 		// protocol version header is set.
-		setStandardHeaders(req.Header, msg)
+		setStandardHeaders(ctx, req.Header, msg)
 		resp, err := c.client.Do(req)
 		if err != nil {
 			// Any error from client.Do means the request didn't reach the server.