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 https://github.com/modelcontextprotocol/go-sdk/issues/905

---------

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

mcp/streamable_headers.go

@@ -5,10 +5,14 @@
 package mcp
 
 import (
+	"context"
+	"encoding/base64"
 	"encoding/json"
 	"errors"
 	"fmt"
+	"log/slog"
 	"net/http"
+	"strings"
 
 	internaljson "github.com/modelcontextprotocol/go-sdk/internal/json"
 	"github.com/modelcontextprotocol/go-sdk/jsonrpc"
@@ -20,7 +24,10 @@ const (
 	lastEventIDHeader            = "Last-Event-ID"
 	methodHeader                 = "Mcp-Method"
 	nameHeader                   = "Mcp-Name"
+	paramHeaderPrefix            = "Mcp-Param-"
 	minVersionForStandardHeaders = protocolVersion20260630
+	base64Prefix                 = "=?base64?"
+	base64Suffix                 = "?="
 )
 
 func extractName(method string, params json.RawMessage) (string, bool) {
@@ -45,9 +52,79 @@ func extractName(method string, params json.RawMessage) (string, bool) {
 	return "", false
 }
 
+// headerSchemaProperty captures the fields needed for x-mcp-header processing.
+type headerSchemaProperty struct {
+	Type       string                          `json:"type"`
+	XMCPHeader json.RawMessage                 `json:"x-mcp-header,omitempty"`
+	Properties map[string]headerSchemaProperty `json:"properties,omitempty"`
+}
+
+// unmarshalSchemaProperties normalizes any InputSchema type
+// (*jsonschema.Schema, map[string]any, or json.RawMessage) into a common
+// representation by marshaling to JSON and unmarshaling only the fields we need.
+func unmarshalSchemaProperties(schema any) map[string]headerSchemaProperty {
+	var s headerSchemaProperty
+	if err := remarshal(schema, &s); err != nil {
+		return nil
+	}
+	return s.Properties
+}
+
+// extractParamHeaderAnnotations returns a map of parameter name to header name
+// for all properties in the tool's InputSchema that have an x-mcp-header
+// annotation.
+func extractParamHeaderAnnotations(tool *Tool) map[string]string {
+	props := unmarshalSchemaProperties(tool.InputSchema)
+	if len(props) == 0 {
+		return nil
+	}
+	result := make(map[string]string)
+	for propName, prop := range props {
+		var headerName string
+		if err := json.Unmarshal(prop.XMCPHeader, &headerName); err != nil || headerName == "" {
+			continue
+		}
+		result[propName] = headerName
+	}
+	if len(result) == 0 {
+		return nil
+	}
+	return result
+}
+
+// primitiveToString conversion.
+// Returns false in the second return value if the argument is not a primitive value.
+func primitiveToString(value any) (string, bool) {
+	switch v := value.(type) {
+	case string:
+		return v, true
+	case float64:
+		return fmt.Sprintf("%g", v), true
+	case bool:
+		return fmt.Sprintf("%t", v), true
+	default:
+		return "", false
+	}
+}
+
+// unmarshalPrimitive unmarshals a JSON value into a Go primitive
+// (string, float64, or bool). Returns nil for non-primitive types.
+func unmarshalPrimitive(raw json.RawMessage) any {
+	var val any
+	if err := internaljson.Unmarshal(raw, &val); err != nil {
+		return nil
+	}
+	switch val.(type) {
+	case string, float64, bool:
+		return val
+	default:
+		return nil
+	}
+}
+
 // setStandardHeaders populates standard MCP headers.
 // It requires the protocol version header to be set.
-func setStandardHeaders(header http.Header, msg jsonrpc.Message) {
+func setStandardHeaders(ctx context.Context, header http.Header, msg jsonrpc.Message) {
 	if msg == nil {
 		return
 	}
@@ -61,10 +138,128 @@ func setStandardHeaders(header http.Header, msg jsonrpc.Message) {
 		if name, ok := extractName(msg.Method, msg.Params); ok {
 			header.Set(nameHeader, name)
 		}
+		if msg.Method == "tools/call" {
+			if tool, ok := ctx.Value(toolContextKey).(*Tool); ok && tool != nil {
+				for k, v := range generateParamHeaders(tool, msg.Params) {
+					header.Set(k, v)
+				}
+			}
+		}
 	}
 }
 
-func validateMcpHeaders(header http.Header, msg jsonrpc.Message) error {
+// generateParamHeaders reads x-mcp-header annotations from the tool's InputSchema
+// and returns the Mcp-Param-{Name} headers to be set on the HTTP request.
+func generateParamHeaders(tool *Tool, params json.RawMessage) map[string]string {
+	paramHeaders := extractParamHeaderAnnotations(tool)
+	if len(paramHeaders) == 0 {
+		return nil
+	}
+
+	var raw struct {
+		Arguments map[string]json.RawMessage `json:"arguments"`
+	}
+	if err := internaljson.Unmarshal(params, &raw); err != nil || raw.Arguments == nil {
+		return nil
+	}
+
+	res := make(map[string]string)
+	for paramName, headerName := range paramHeaders {
+		argRaw, ok := raw.Arguments[paramName]
+		if !ok {
+			continue
+		}
+		if string(argRaw) == "null" {
+			continue
+		}
+		val := unmarshalPrimitive(argRaw)
+		if val == nil {
+			continue
+		}
+		encoded, ok := encodeHeaderValue(val)
+		if !ok {
+			continue
+		}
+		res[paramHeaderPrefix+headerName] = encoded
+	}
+	return res
+}
+
+// filterValidTools returns only tools that have valid
+// x-mcp-header annotations. Invalid tools are logged and excluded.
+func filterValidTools(logger *slog.Logger, tools []*Tool) []*Tool {
+	logger = ensureLogger(logger)
+	result := make([]*Tool, 0, len(tools))
+	for _, tool := range tools {
+		if err := validateParamHeaderAnnotations(tool); err != nil {
+			logger.Error("excluding tool from tools/list", "tool", tool.Name, "error", err)
+			continue
+		}
+		result = append(result, tool)
+	}
+	return result
+}
+
+// validateParamHeaderAnnotations checks that a tool's x-mcp-header annotations
+// are valid.
+func validateParamHeaderAnnotations(tool *Tool) error {
+	props := unmarshalSchemaProperties(tool.InputSchema)
+	if len(props) == 0 {
+		return nil
+	}
+
+	seen := make(map[string]bool)
+	for propName, prop := range props {
+		if err := checkForNestedHeaders(prop, propName); err != nil {
+			return err
+		}
+		if prop.XMCPHeader == nil {
+			continue
+		}
+		var headerName string
+		if err := json.Unmarshal(prop.XMCPHeader, &headerName); err != nil || headerName == "" {
+			return fmt.Errorf("property %q: x-mcp-header must be a non-empty string", propName)
+		}
+		if err := validateHeaderName(headerName); err != nil {
+			return fmt.Errorf("property %q: %w", propName, err)
+		}
+		lower := strings.ToLower(headerName)
+		if seen[lower] {
+			return fmt.Errorf("property %q: duplicate x-mcp-header value %q (case-insensitive)", propName, headerName)
+		}
+		seen[lower] = true
+
+		if prop.Type != "string" && prop.Type != "number" && prop.Type != "integer" && prop.Type != "boolean" {
+			return fmt.Errorf("property %q: x-mcp-header can only be applied to primitive types, got %v", propName, prop.Type)
+		}
+	}
+	return nil
+}
+
+func checkForNestedHeaders(prop headerSchemaProperty, path string) error {
+	for propName, nested := range prop.Properties {
+		if nested.XMCPHeader != nil {
+			return fmt.Errorf("property %q: x-mcp-header cannot be applied to nested properties", path+"."+propName)
+		}
+		if err := checkForNestedHeaders(nested, path+"."+propName); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// validateHeaderName checks that a header name contains only valid
+// ASCII characters (excluding space and ':').
+func validateHeaderName(name string) error {
+	for _, c := range name {
+		if c <= 0x20 || c > 0x7E || c == ':' {
+			return fmt.Errorf("x-mcp-header value %q contains invalid character %q", name, c)
+		}
+	}
+	return nil
+}
+
+func validateMcpHeaders(header http.Header, msg jsonrpc.Message, toolLookup func(string) (*serverTool, bool)) error {
 	protocolVersion := header.Get(protocolVersionHeader)
 	if protocolVersion == "" || protocolVersion < minVersionForStandardHeaders {
 		return nil
@@ -80,12 +275,14 @@ func validateMcpHeaders(header http.Header, msg jsonrpc.Message) error {
 			return fmt.Errorf("header mismatch: Mcp-Method header value '%s' does not match body value '%s'", methodInHeader, msg.Method)
 		}
 
+		var nameInBody string
 		if msg.Method == "tools/call" || msg.Method == "resources/read" || msg.Method == "prompts/get" {
 			nameInHeader := header.Get(nameHeader)
 			if nameInHeader == "" {
 				return fmt.Errorf("missing required Mcp-Name header for method %q", msg.Method)
 			}
-			nameInBody, ok := extractName(msg.Method, msg.Params)
+			var ok bool
+			nameInBody, ok = extractName(msg.Method, msg.Params)
 			if !ok {
 				return fmt.Errorf("failed to extract name from parameters for method %q", msg.Method)
 			}
@@ -93,6 +290,126 @@ func validateMcpHeaders(header http.Header, msg jsonrpc.Message) error {
 				return fmt.Errorf("header mismatch: Mcp-Name header value '%s' does not match body value '%s'", nameInHeader, nameInBody)
 			}
 		}
+
+		if msg.Method == "tools/call" && toolLookup != nil {
+			if st, ok := toolLookup(nameInBody); ok && st != nil {
+				if err := validateParamHeaders(header, msg, st.tool); err != nil {
+					return err
+				}
+			}
+		}
 	}
 	return nil
 }
+
+func validateParamHeaders(header http.Header, msg *jsonrpc.Request, tool *Tool) error {
+	paramHeaders := extractParamHeaderAnnotations(tool)
+	if len(paramHeaders) == 0 {
+		return nil
+	}
+
+	var raw struct {
+		Arguments map[string]json.RawMessage `json:"arguments"`
+	}
+	if err := internaljson.Unmarshal(msg.Params, &raw); err != nil {
+		return nil
+	}
+
+	for paramName, headerName := range paramHeaders {
+		fullHeader := paramHeaderPrefix + headerName
+		headerVal := header.Get(fullHeader)
+		argRaw, argExists := raw.Arguments[paramName]
+
+		if !argExists || string(argRaw) == "null" {
+			if headerVal != "" {
+				return fmt.Errorf("header mismatch: unexpected %s header for absent or null parameter %q", fullHeader, paramName)
+			}
+			continue
+		}
+
+		if headerVal == "" {
+			return fmt.Errorf("header mismatch: missing %s header for parameter %q", fullHeader, paramName)
+		}
+
+		decoded, ok := decodeHeaderValue(headerVal)
+		if !ok {
+			return fmt.Errorf("header mismatch: %s header contains invalid Base64 encoding", fullHeader)
+		}
+
+		bodyVal := unmarshalPrimitive(argRaw)
+		if bodyVal == nil {
+			return fmt.Errorf("header mismatch: %s header present but body parameter %q is not a primitive type", fullHeader, paramName)
+		}
+		expected, ok := primitiveToString(bodyVal)
+		if !ok {
+			return fmt.Errorf("header mismatch: %s header present but body parameter %q is not a primitive type", fullHeader, paramName)
+		}
+
+		// TODO: String comparison may not work ideally for numbers
+		if decoded != expected {
+			return fmt.Errorf("header mismatch: %s header value '%s' does not match body value", fullHeader, headerVal)
+		}
+	}
+	return nil
+}
+
+// encodeHeaderValue converts a parameter value to an HTTP header-safe string
+// per the SEP-2243 encoding rules:
+//   - string: used as-is if safe ASCII, otherwise Base64 encoded
+//   - number (float64): decimal string representation
+//   - bool: lowercase "true" or "false"
+//
+// Values that contain non-ASCII characters, control characters, or
+// leading/trailing whitespace are Base64-encoded with the =?base64?...?= wrapper.
+//
+// The second return value is false if the value is not a supported primitive type.
+func encodeHeaderValue(value any) (string, bool) {
+	s, ok := primitiveToString(value)
+	if !ok {
+		return "", false
+	}
+	if requiresBase64Encoding(s) {
+		return encodeBase64(s), true
+	}
+	return s, true
+}
+
+// decodeHeaderValue decodes a header value that may be Base64-encoded
+// with the =?base64?...?= wrapper.
+//
+// The second return value is false if the header value is not a valid Base64 encoded value.
+func decodeHeaderValue(headerValue string) (string, bool) {
+	if len(headerValue) == 0 {
+		return headerValue, true
+	}
+
+	if strings.HasPrefix(strings.ToLower(headerValue), base64Prefix) &&
+		strings.HasSuffix(headerValue, base64Suffix) {
+		encoded := headerValue[len(base64Prefix) : len(headerValue)-len(base64Suffix)]
+		decoded, err := base64.StdEncoding.DecodeString(encoded)
+		if err != nil {
+			return "", false
+		}
+		return string(decoded), true
+	}
+	return headerValue, true
+}
+
+func requiresBase64Encoding(s string) bool {
+	if len(s) == 0 {
+		return false
+	}
+	if s[0] == ' ' || s[0] == '\t' || s[len(s)-1] == ' ' || s[len(s)-1] == '\t' {
+		return true
+	}
+	for _, c := range s {
+		if c < 0x20 || c > 0x7E {
+			return true
+		}
+	}
+	return false
+}
+
+func encodeBase64(s string) string {
+	return base64Prefix + base64.StdEncoding.EncodeToString([]byte(s)) + base64Suffix
+}