mcp: keep x-mcp-header handling working next to type arrays

mcp/streamable_headers.go

@@ -5,6 +5,8 @@
 package mcp
 
 import (
+	"bytes"
+	"cmp"
 	"context"
 	"encoding/base64"
 	"encoding/json"
@@ -55,12 +57,25 @@ func extractName(method string, params json.RawMessage) (string, bool) {
 }
 
 // headerSchemaProperty captures the fields needed for x-mcp-header processing.
+// Type stays raw because it can also be an array, such as the
+// ["null", "string"] that jsonschema.For infers for a pointer field.
 type headerSchemaProperty struct {
-	Type       string                          `json:"type"`
+	Type       json.RawMessage                 `json:"type,omitempty"`
 	XMCPHeader json.RawMessage                 `json:"x-mcp-header,omitempty"`
 	Properties map[string]headerSchemaProperty `json:"properties,omitempty"`
 }
 
+// UnmarshalJSON accepts the boolean schemas true and false, which carry no
+// annotations, so that they do not fail decoding of the enclosing schema.
+func (p *headerSchemaProperty) UnmarshalJSON(data []byte) error {
+	if s := string(bytes.TrimSpace(data)); s == "true" || s == "false" {
+		*p = headerSchemaProperty{}
+		return nil
+	}
+	type plain headerSchemaProperty
+	return internaljson.Unmarshal(data, (*plain)(p))
+}
+
 // 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 +316,10 @@ 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)
+			var typ string
+			_ = json.Unmarshal(prop.Type, &typ)
+			if typ != "string" && typ != "integer" && typ != "boolean" {
+				return fmt.Errorf("property %q: x-mcp-header can only be applied to primitive types (integer, string, boolean), got %s", path, cmp.Or(string(prop.Type), `""`))
 			}
 			var headerName string
 			if err := json.Unmarshal(prop.XMCPHeader, &headerName); err != nil || headerName == "" {

mcp/streamable_headers_test.go

@@ -897,6 +897,35 @@ func TestValidateToolParamHeaders(t *testing.T) {
 			wantErr:    true,
 			wantErrSub: "primitive types",
 		},
+		{
+			name: "number type next to a type array is rejected",
+			tool: &Tool{
+				Name: "test",
+				InputSchema: map[string]any{
+					"type": "object",
+					"properties": map[string]any{
+						"count": map[string]any{"type": "number", "x-mcp-header": "Count"},
+						"note":  map[string]any{"type": []any{"null", "string"}},
+					},
+				},
+			},
+			wantErr:    true,
+			wantErrSub: "primitive types",
+		},
+		{
+			name: "type array on annotated property is rejected",
+			tool: &Tool{
+				Name: "test",
+				InputSchema: map[string]any{
+					"type": "object",
+					"properties": map[string]any{
+						"region": map[string]any{"type": []any{"null", "string"}, "x-mcp-header": "Region"},
+					},
+				},
+			},
+			wantErr:    true,
+			wantErrSub: `got ["null","string"]`,
+		},
 	}
 
 	for _, tt := range tests {
@@ -1317,6 +1346,53 @@ func TestExtractToolParamHeaders(t *testing.T) {
 			},
 			want: map[string]string{"user.id": "UserId"},
 		},
+		{
+			name: "sibling property with a type array",
+			tool: &Tool{
+				Name: "test",
+				InputSchema: 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"}},
+					},
+				},
+			},
+			want: map[string]string{"region": "Region"},
+		},
+		{
+			name: "sibling boolean schema",
+			tool: &Tool{
+				Name: "test",
+				InputSchema: map[string]any{
+					"type": "object",
+					"properties": map[string]any{
+						"region": map[string]any{"type": "string", "x-mcp-header": "Region"},
+						"extra":  true,
+					},
+				},
+			},
+			want: map[string]string{"region": "Region"},
+		},
+		{
+			name: "jsonschema.For with a pointer field",
+			tool: &Tool{
+				Name: "test",
+				InputSchema: func() *jsonschema.Schema {
+					type in struct {
+						Region string  `json:"region"`
+						Note   *string `json:"note,omitempty"`
+					}
+					s, err := jsonschema.For[in](nil)
+					if err != nil {
+						panic(err)
+					}
+					s.Properties["region"].Extra = map[string]any{"x-mcp-header": "Region"}
+					return s
+				}(),
+			},
+			want: map[string]string{"region": "Region"},
+		},
 	}
 
 	for _, tt := range tests {