diff --git a/mcp/streamable_headers.go b/mcp/streamable_headers.go index 72270208..3ffcc3c1 100644 --- a/mcp/streamable_headers.go +++ b/mcp/streamable_headers.go @@ -61,6 +61,37 @@ type headerSchemaProperty struct { Properties map[string]headerSchemaProperty `json:"properties,omitempty"` } +// UnmarshalJSON decodes a single schema property without letting a sibling's +// unusual schema shape fail the whole schema. JSON Schema allows "type" to be +// an array (jsonschema.For emits ["null", "string"] for pointer fields) and a +// property to be a boolean subschema; neither carries an x-mcp-header that +// this code can use, so they are decoded leniently. A non-string "type" is +// kept as its raw JSON text, so it fails the primitive-type check. +func (p *headerSchemaProperty) UnmarshalJSON(data []byte) error { + 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 { + var boolSchema bool + if json.Unmarshal(data, &boolSchema) == nil { + *p = headerSchemaProperty{} + return nil + } + return err + } + *p = headerSchemaProperty{XMCPHeader: obj.XMCPHeader, Properties: obj.Properties} + if len(obj.Type) > 0 { + var typeName string + if err := json.Unmarshal(obj.Type, &typeName); err != nil { + typeName = string(obj.Type) + } + p.Type = typeName + } + return nil +} + // 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. diff --git a/mcp/streamable_headers_test.go b/mcp/streamable_headers_test.go index 2042187c..9a9a6fb5 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 {