diff --git a/mcp/streamable_headers.go b/mcp/streamable_headers.go index 7133e26c9..b9c516018 100644 --- a/mcp/streamable_headers.go +++ b/mcp/streamable_headers.go @@ -63,17 +63,24 @@ type headerSchemaProperty struct { 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. +// UnmarshalJSON decodes schema properties while ignoring boolean subschemas +// and preserving non-string type values for primitive-type validation. 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)) + var obj struct { + Type json.RawMessage `json:"type,omitempty"` + XMCPHeader json.RawMessage `json:"x-mcp-header,omitempty"` + Properties map[string]headerSchemaProperty `json:"properties,omitempty"` + } + if err := json.Unmarshal(data, &obj); err != nil { + return err + } + *p = headerSchemaProperty{Type: obj.Type, XMCPHeader: obj.XMCPHeader, Properties: obj.Properties} + return nil } func (p headerSchemaProperty) primitiveType() bool { diff --git a/mcp/streamable_headers_test.go b/mcp/streamable_headers_test.go index d80c781f1..f6a31a775 100644 --- a/mcp/streamable_headers_test.go +++ b/mcp/streamable_headers_test.go @@ -442,6 +442,50 @@ func TestValidateToolParamHeaders(t *testing.T) { wantErr bool wantErrSub string }{ + { + name: "array type on an unannotated sibling does not hide an invalid annotation", + 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: "nullable annotated property is rejected as non-primitive", + tool: &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "note": map[string]any{"type": []any{"null", "string"}, "x-mcp-header": "Note"}, + }, + }, + }, + wantErr: true, + wantErrSub: "primitive types", + }, + { + name: "boolean subschema does not hide an invalid annotation", + tool: &Tool{ + Name: "test", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "count": map[string]any{"type": "number", "x-mcp-header": "Count"}, + "extra": true, + }, + }, + }, + wantErr: true, + wantErrSub: "primitive types", + }, { name: "valid tool with x-mcp-header", tool: &Tool{ @@ -1317,6 +1361,34 @@ func TestExtractToolParamHeaders(t *testing.T) { }, want: map[string]string{"user.id": "UserId"}, }, + { + name: "array type on an unannotated sibling does not hide annotations", + 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: "boolean subschema does not hide annotations", + 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"}, + }, } for _, tt := range tests {