From 2807cef5c96e55f94e6a4ce4056ad4d0891eb0db Mon Sep 17 00:00:00 2001
From: RafalOfficial <248946250+RafalOfficial@users.noreply.github.com>
Date: Fri, 16 Oct 2026 11:25:31 +0200
Subject: [PATCH] mcp: keep x-mcp-header handling working next to type arrays

headerSchemaProperty decoded "type" into a string, so one property
whose type is an array, such as the ["null", "string"] that
jsonschema.For infers for a pointer field, failed decoding of the whole
inputSchema. A boolean subschema did the same. The client then sent no
Mcp-Param-* headers for the tool, the server skipped validating them,
and validateParamHeaderAnnotations returned nil without checking any
annotation.

Keep "type" raw and accept boolean subschemas, so each property is
decoded on its own. An annotated property must still have a single
primitive type.

Fixes #1406

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 mcp/streamable_headers.go      | 23 ++++++++--
 mcp/streamable_headers_test.go | 76 ++++++++++++++++++++++++++++++++++
 2 files changed, 96 insertions(+), 3 deletions(-)

diff --git a/mcp/streamable_headers.go b/mcp/streamable_headers.go
index 72270208c..d349d0599 100644
--- a/mcp/streamable_headers.go
+++ b/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 == "" {
diff --git a/mcp/streamable_headers_test.go b/mcp/streamable_headers_test.go
index 2042187cf..6a18da69d 100644
--- a/mcp/streamable_headers_test.go
+++ b/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 {