diff --git a/mcp/streamable_headers.go b/mcp/streamable_headers.go index 72270208..d349d059 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 2042187c..6a18da69 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 {