fix(mcp): process singleton headers correctly

mcp/streamable_headers.go

@@ -5,6 +5,7 @@
 package mcp
 
 import (
+	"bytes"
 	"context"
 	"encoding/base64"
 	"encoding/json"
@@ -56,11 +57,33 @@ func extractName(method string, params json.RawMessage) (string, bool) {
 
 // headerSchemaProperty captures the fields needed for x-mcp-header processing.
 type headerSchemaProperty struct {
-	Type       string                          `json:"type"`
+	// Type is raw JSON because it may be a list like ["null", "string"] or a primitive.
+	Type       json.RawMessage                 `json:"type,omitempty"`
 	XMCPHeader json.RawMessage                 `json:"x-mcp-header,omitempty"`
 	Properties map[string]headerSchemaProperty `json:"properties,omitempty"`
 }
 
+// UnmarshalJSON treats a boolean schema as a schema without annotations.
+// A field can be encoded as "region": {"type": "string", "x-mcp-header": "Region"},
+// but also simply as "extra": true.
+func (p *headerSchemaProperty) UnmarshalJSON(data []byte) error {
+	switch string(bytes.TrimSpace(data)) {
+	case "true", "false":
+		*p = headerSchemaProperty{}
+		return nil
+	}
+	type plain headerSchemaProperty
+	return internaljson.Unmarshal(data, (*plain)(p))
+}
+
+func (p headerSchemaProperty) primitiveType() bool {
+	var t string
+	if err := json.Unmarshal(p.Type, &t); err != nil {
+		return false
+	}
+	return t == "string" || t == "integer" || t == "boolean"
+}
+
 // 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.
@@ -301,8 +324,12 @@ func validateParamHeadersIn(props map[string]headerSchemaProperty, prefix string
 			path = prefix + "." + propName
 		}
 		if prop.XMCPHeader != nil {
-			if prop.Type != "string" && prop.Type != "integer" && prop.Type != "boolean" {
-				return fmt.Errorf("property %q: x-mcp-header can only be applied to primitive types (integer, string, boolean), got %q", path, prop.Type)
+			if !prop.primitiveType() {
+				got := string(prop.Type)
+				if got == "" {
+					got = `""`
+				}
+				return fmt.Errorf("property %q: x-mcp-header can only be applied to primitive types (integer, string, boolean), got %s", path, got)
 			}
 			var headerName string
 			if err := json.Unmarshal(prop.XMCPHeader, &headerName); err != nil || headerName == "" {
@@ -361,15 +388,35 @@ func isTChar(c rune) bool {
 	return false
 }
 
+// singletonHeader returns the value of a header that must appear at most once,
+// and reports whether it was present. Repeated field lines are rejected.
+func singletonHeader(header http.Header, name string) (string, bool, error) {
+	values := header.Values(name)
+	switch len(values) {
+	case 0:
+		return "", false, nil
+	case 1:
+		return values[0], true, nil
+	default:
+		return "", false, fmt.Errorf("header mismatch: %s header must not be repeated", name)
+	}
+}
+
 func validateMcpHeaders(header http.Header, msg jsonrpc.Message, toolLookup func(string) (*serverTool, bool)) error {
-	protocolVersion := header.Get(protocolVersionHeader)
+	protocolVersion, _, err := singletonHeader(header, protocolVersionHeader)
+	if err != nil {
+		return err
+	}
 	if protocolVersion == "" || protocolVersion < minVersionForStandardHeaders {
 		return nil
 	}
 
 	switch msg := msg.(type) {
 	case *jsonrpc.Request:
-		methodInHeader := header.Get(methodHeader)
+		methodInHeader, _, err := singletonHeader(header, methodHeader)
+		if err != nil {
+			return err
+		}
 		if methodInHeader == "" {
 			return errors.New("missing required Mcp-Method header")
 		}
@@ -380,10 +427,13 @@ func validateMcpHeaders(header http.Header, msg jsonrpc.Message, toolLookup func
 		var nameInBody string
 		if msg.Method == "tools/call" || msg.Method == "resources/read" || msg.Method == "prompts/get" {
 			// A name can be empty, which is sent as a present, empty header.
-			if len(header.Values(nameHeader)) == 0 {
+			nameInHeader, present, err := singletonHeader(header, nameHeader)
+			if err != nil {
+				return err
+			}
+			if !present {
 				return fmt.Errorf("missing required Mcp-Name header for method %q", msg.Method)
 			}
-			nameInHeader := header.Get(nameHeader)
 			// A name that is not header-safe may be Base64-wrapped by the client
 			// (=?base64?...?=), so decode before comparing, as with Mcp-Param-*.
 			decodedName, ok := decodeHeaderValue(nameInHeader)
@@ -427,8 +477,10 @@ func validateParamHeaders(header http.Header, msg *jsonrpc.Request, tool *Tool)
 		fullHeader := paramHeaderPrefix + b.Header
 		// An empty string argument is mirrored as a present, empty header, so
 		// presence is decided by the header being sent, not by its value.
-		headerPresent := len(header.Values(fullHeader)) > 0
-		headerVal := header.Get(fullHeader)
+		headerVal, headerPresent, err := singletonHeader(header, fullHeader)
+		if err != nil {
+			return err
+		}
 		argRaw, argExists := lookupArgument(raw.Arguments, b.Path)
 
 		if !argExists || string(argRaw) == "null" {

mcp/streamable_headers_test.go

@@ -1731,3 +1731,174 @@ func TestMcpNameHeaderEncodingRoundTrip(t *testing.T) {
 		})
 	}
 }
+
+// TestParamHeaderAnnotationsWithNonStringTypes verifies that a type array or a
+// boolean schema elsewhere in the input schema does not cause x-mcp-header
+// annotations to be ignored.
+func TestParamHeaderAnnotationsWithNonStringTypes(t *testing.T) {
+	type typedInput struct {
+		Region string  `json:"region"`
+		Note   *string `json:"note,omitempty"`
+	}
+	typed, err := jsonschema.For[typedInput](nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	typed.Properties["region"].Extra = map[string]any{"x-mcp-header": "Region"}
+
+	tests := []struct {
+		name   string
+		schema any
+	}{
+		{"jsonschema.For with pointer field", typed},
+		{
+			"type array on sibling property",
+			map[string]any{
+				"type": "object",
+				"properties": map[string]any{
+					"region": map[string]any{"type": "string", "x-mcp-header": "Region"},
+					"note":   map[string]any{"type": []any{"null", "string"}},
+				},
+			},
+		},
+		{
+			"type array on parent object",
+			json.RawMessage(`{"type":["object","null"],"properties":{"region":{"type":"string","x-mcp-header":"Region"}}}`),
+		},
+		{
+			"boolean schema on sibling property",
+			json.RawMessage(`{"type":"object","properties":{"region":{"type":"string","x-mcp-header":"Region"},"any":true,"none":false}}`),
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			tool := &Tool{Name: "query", InputSchema: tt.schema}
+			if err := validateParamHeaderAnnotations(tool); err != nil {
+				t.Fatalf("validateParamHeaderAnnotations() = %v, want nil", err)
+			}
+			got := extractParamHeaderAnnotations(tool)
+			if len(got) != 1 || strings.Join(got[0].Path, ".") != "region" || got[0].Header != "Region" {
+				t.Fatalf("extractParamHeaderAnnotations() = %v, want [{[region] Region}]", got)
+			}
+
+			params := mustMarshal(&CallToolParams{Name: "query", Arguments: map[string]any{"region": "us-west1"}})
+			if got := generateParamHeaders(tool, params); got[paramHeaderPrefix+"Region"] != "us-west1" {
+				t.Errorf("generateParamHeaders() = %v, want Mcp-Param-Region: us-west1", got)
+			}
+
+			msg := &jsonrpc.Request{Method: "tools/call", Params: params}
+			mismatch := http.Header{}
+			mismatch.Set(paramHeaderPrefix+"Region", "eu-west1")
+			if err := validateParamHeaders(mismatch, msg, tool); err == nil {
+				t.Error("validateParamHeaders() with mismatched header = nil, want error")
+			}
+			if err := validateParamHeaders(http.Header{}, msg, tool); err == nil {
+				t.Error("validateParamHeaders() with missing header = nil, want error")
+			}
+		})
+	}
+}
+
+func TestValidateParamHeaderAnnotationsTypeArray(t *testing.T) {
+	tests := []struct {
+		name       string
+		schema     string
+		wantErrSub string
+	}{
+		{
+			name:       "type array on annotated property",
+			schema:     `{"type":"object","properties":{"region":{"type":["null","string"],"x-mcp-header":"Region"}}}`,
+			wantErrSub: `got ["null","string"]`,
+		},
+		{
+			name:       "missing type on annotated property",
+			schema:     `{"type":"object","properties":{"region":{"x-mcp-header":"Region"}}}`,
+			wantErrSub: `got ""`,
+		},
+		{
+			name:       "invalid annotation next to type array",
+			schema:     `{"type":"object","properties":{"region":{"type":"string","x-mcp-header":"My Region"},"note":{"type":["null","string"]}}}`,
+			wantErrSub: "invalid character",
+		},
+		{
+			name:       "invalid annotation next to boolean schema",
+			schema:     `{"type":"object","properties":{"region":{"type":"string","x-mcp-header":"My Region"},"any":true}}`,
+			wantErrSub: "invalid character",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			err := validateParamHeaderAnnotations(&Tool{Name: "query", InputSchema: json.RawMessage(tt.schema)})
+			if err == nil {
+				t.Fatalf("validateParamHeaderAnnotations() = nil, want error containing %q", tt.wantErrSub)
+			}
+			if !strings.Contains(err.Error(), tt.wantErrSub) {
+				t.Errorf("validateParamHeaderAnnotations() = %q, want substring %q", err, tt.wantErrSub)
+			}
+		})
+	}
+}
+
+// TestValidateMcpHeadersRepeated verifies that the headers mirrored for
+// intermediaries are rejected when sent more than once, regardless of which
+// copy matches the body.
+func TestValidateMcpHeadersRepeated(t *testing.T) {
+	tool := &Tool{
+		Name: "query",
+		InputSchema: map[string]any{
+			"type": "object",
+			"properties": map[string]any{
+				"region": map[string]any{"type": "string", "x-mcp-header": "Region"},
+			},
+		},
+	}
+	lookup := func(name string) (*serverTool, bool) {
+		if name == tool.Name {
+			return &serverTool{tool: tool}, true
+		}
+		return nil, false
+	}
+	msg := &jsonrpc.Request{
+		Method: "tools/call",
+		Params: mustMarshal(&CallToolParams{Name: "query", Arguments: map[string]any{"region": "us-west1"}}),
+	}
+	validHeader := func() http.Header {
+		h := http.Header{}
+		h.Set(protocolVersionHeader, minVersionForStandardHeaders)
+		h.Set(methodHeader, "tools/call")
+		h.Set(nameHeader, "query")
+		h.Set(paramHeaderPrefix+"Region", "us-west1")
+		return h
+	}
+	if err := validateMcpHeaders(validHeader(), msg, lookup); err != nil {
+		t.Fatalf("validateMcpHeaders() with single headers = %v, want nil", err)
+	}
+
+	tests := []struct {
+		header string
+		extra  []string
+	}{
+		{protocolVersionHeader, []string{minVersionForStandardHeaders}},
+		{protocolVersionHeader, []string{protocolVersion20251125}},
+		{methodHeader, []string{"tools/call"}},
+		{nameHeader, []string{"query"}},
+		{nameHeader, []string{"other-tool"}},
+		{paramHeaderPrefix + "Region", []string{"us-west1"}},
+		{paramHeaderPrefix + "Region", []string{"eu-west1"}},
+	}
+	for _, tt := range tests {
+		t.Run(fmt.Sprintf("%s=%v", tt.header, tt.extra), func(t *testing.T) {
+			h := validHeader()
+			for _, v := range tt.extra {
+				h.Add(tt.header, v)
+			}
+			err := validateMcpHeaders(h, msg, lookup)
+			if err == nil {
+				t.Fatal("validateMcpHeaders() = nil, want error for repeated header")
+			}
+			if want := "must not be repeated"; !strings.Contains(err.Error(), want) {
+				t.Errorf("validateMcpHeaders() = %q, want substring %q", err, want)
+			}
+		})
+	}
+}

mcp/streamable_test.go

@@ -2262,6 +2262,78 @@ func TestStreamableMcpHeaderValidationErrorFormat(t *testing.T) {
 	}
 }
 
+// TestStreamableParamHeaderWithPointerField verifies that Mcp-Param-* headers
+// are mirrored and validated for a typed tool whose input struct has a pointer
+// field, for which jsonschema.For emits a type array.
+func TestStreamableParamHeaderWithPointerField(t *testing.T) {
+	type input struct {
+		Region string  `json:"region"`
+		Note   *string `json:"note,omitempty"`
+	}
+	schema, err := jsonschema.For[input](nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	schema.Properties["region"].Extra = map[string]any{"x-mcp-header": "Region"}
+
+	server := NewServer(&Implementation{Name: "testServer", Version: "v1.0.0"}, nil)
+	var handlerCalled atomic.Bool
+	AddTool(server, &Tool{Name: "query", InputSchema: schema},
+		func(ctx context.Context, req *CallToolRequest, in input) (*CallToolResult, any, error) {
+			handlerCalled.Store(true)
+			return &CallToolResult{}, nil, nil
+		})
+
+	handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, &StreamableHTTPOptions{
+		Stateless:    true,
+		JSONResponse: true,
+	})
+	defer handler.closeAll()
+	httpServer := httptest.NewServer(mustNotPanic(t, handler))
+	defer httpServer.Close()
+
+	var sentRegion string
+	var status int
+	customClient := &http.Client{
+		Transport: roundTripperFunc(func(req *http.Request) (*http.Response, error) {
+			if req.Header.Get(methodHeader) == "tools/call" {
+				sentRegion = req.Header.Get(paramHeaderPrefix + "Region")
+				req.Header.Set(paramHeaderPrefix+"Region", "eu-west1")
+			}
+			resp, err := http.DefaultTransport.RoundTrip(req)
+			if err == nil && req.Header.Get(methodHeader) == "tools/call" {
+				status = resp.StatusCode
+			}
+			return resp, err
+		}),
+	}
+
+	client := NewClient(&Implementation{Name: "testClient", Version: "v1.0.0"}, nil)
+	ctx := context.Background()
+	session, err := client.Connect(ctx, &StreamableClientTransport{Endpoint: httpServer.URL, HTTPClient: customClient},
+		&ClientSessionOptions{ProtocolVersion: minVersionForStandardHeaders})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer session.Close()
+	if _, err := session.ListTools(ctx, nil); err != nil {
+		t.Fatal(err)
+	}
+
+	if _, err := session.CallTool(ctx, &CallToolParams{Name: "query", Arguments: map[string]any{"region": "us-west1"}}); err == nil {
+		t.Error("CallTool succeeded with mismatched Mcp-Param-Region header")
+	}
+	if sentRegion != "us-west1" {
+		t.Errorf("client sent Mcp-Param-Region %q, want %q", sentRegion, "us-west1")
+	}
+	if status != http.StatusBadRequest {
+		t.Errorf("status code = %d, want %d", status, http.StatusBadRequest)
+	}
+	if handlerCalled.Load() {
+		t.Error("tool handler ran despite mismatched Mcp-Param-Region header")
+	}
+}
+
 // TestStreamableMcpHeaderVersionGating verifies that header validation
 // is skipped for protocol versions older than minVersionForStandardHeaders.
 func TestStreamableMcpHeaderVersionGating(t *testing.T) {