mcp: HTTP Header Standardization for x-mcp-header

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/streamable.go

@@ -491,6 +491,14 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque
 			http.Error(w, "failed connection", http.StatusInternalServerError)
 			return
 		}
+		transport.connection.toolLookup = func(name string) *Tool {
+			server.mu.Lock()
+			defer server.mu.Unlock()
+			if st, ok := server.tools.get(name); ok {
+				return st.tool
+			}
+			return nil
+		}
 		// Capture the user ID from the token info to enable session hijacking
 		// prevention on subsequent requests.
 		var userID string
@@ -669,6 +677,8 @@ type streamableServerConn struct {
 
 	logger *slog.Logger
 
+	toolLookup func(name string) *Tool
+
 	incoming chan jsonrpc.Message // messages from the client to the server
 
 	mu sync.Mutex // guards all fields below
@@ -1186,9 +1196,15 @@ 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 {
+		var tool *Tool
+		if jreq, ok := incoming[0].(*jsonrpc.Request); ok && jreq.Method == "tools/call" && c.toolLookup != nil {
+			if name, ok := extractName(jreq.Method, jreq.Params); ok {
+				tool = c.toolLookup(name)
+			}
+		}
+		if err := validateMcpHeaders(req.Header, incoming[0], tool); err != nil {
 			resp := &jsonrpc.Response{
 				Error: jsonrpc2.NewError(CodeHeaderMismatch, err.Error()),
 			}
@@ -1813,7 +1829,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.