From 20540ac506f693028dd9ae2cccd5542b9412ac6c Mon Sep 17 00:00:00 2001 From: Yaroslav Shevchuk Date: Mon, 26 Oct 2026 09:56:31 +0000 Subject: [PATCH] process singleton headers correctly --- mcp/streamable_headers.go | 70 ++++++++++++-- mcp/streamable_headers_test.go | 171 +++++++++++++++++++++++++++++++++ mcp/streamable_test.go | 72 ++++++++++++++ 3 files changed, 304 insertions(+), 9 deletions(-) diff --git a/mcp/streamable_headers.go b/mcp/streamable_headers.go index 72270208c..7133e26c9 100644 --- a/mcp/streamable_headers.go +++ b/mcp/streamable_headers.go @@ -5,6 +5,7 @@ package mcp import ( + "bytes" "context" "encoding/base64" "encoding/json" @@ -56,11 +57,33 @@ func extractName(method string, params json.RawMessage) (string, bool) { // headerSchemaProperty captures the fields needed for x-mcp-header processing. type headerSchemaProperty struct { - Type string `json:"type"` + // Type is raw JSON because it may be a list like ["null", "string"] or a primitive. + Type json.RawMessage `json:"type,omitempty"` XMCPHeader json.RawMessage `json:"x-mcp-header,omitempty"` 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. +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)) +} + +func (p headerSchemaProperty) primitiveType() bool { + var t string + if err := json.Unmarshal(p.Type, &t); err != nil { + return false + } + return t == "string" || t == "integer" || t == "boolean" +} + // 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 +324,12 @@ 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) + if !prop.primitiveType() { + got := string(prop.Type) + if got == "" { + got = `""` + } + return fmt.Errorf("property %q: x-mcp-header can only be applied to primitive types (integer, string, boolean), got %s", path, got) } var headerName string if err := json.Unmarshal(prop.XMCPHeader, &headerName); err != nil || headerName == "" { @@ -361,15 +388,35 @@ func isTChar(c rune) bool { return false } +// singletonHeader returns the value of a header that must appear at most once, +// and reports whether it was present. Repeated field lines are rejected. +func singletonHeader(header http.Header, name string) (string, bool, error) { + values := header.Values(name) + switch len(values) { + case 0: + return "", false, nil + case 1: + return values[0], true, nil + default: + return "", false, fmt.Errorf("header mismatch: %s header must not be repeated", name) + } +} + func validateMcpHeaders(header http.Header, msg jsonrpc.Message, toolLookup func(string) (*serverTool, bool)) error { - protocolVersion := header.Get(protocolVersionHeader) + protocolVersion, _, err := singletonHeader(header, protocolVersionHeader) + if err != nil { + return err + } if protocolVersion == "" || protocolVersion < minVersionForStandardHeaders { return nil } switch msg := msg.(type) { case *jsonrpc.Request: - methodInHeader := header.Get(methodHeader) + methodInHeader, _, err := singletonHeader(header, methodHeader) + if err != nil { + return err + } if methodInHeader == "" { return errors.New("missing required Mcp-Method header") } @@ -380,10 +427,13 @@ func validateMcpHeaders(header http.Header, msg jsonrpc.Message, toolLookup func var nameInBody string if msg.Method == "tools/call" || msg.Method == "resources/read" || msg.Method == "prompts/get" { // A name can be empty, which is sent as a present, empty header. - if len(header.Values(nameHeader)) == 0 { + nameInHeader, present, err := singletonHeader(header, nameHeader) + if err != nil { + return err + } + if !present { return fmt.Errorf("missing required Mcp-Name header for method %q", msg.Method) } - nameInHeader := header.Get(nameHeader) // A name that is not header-safe may be Base64-wrapped by the client // (=?base64?...?=), so decode before comparing, as with Mcp-Param-*. decodedName, ok := decodeHeaderValue(nameInHeader) @@ -427,8 +477,10 @@ func validateParamHeaders(header http.Header, msg *jsonrpc.Request, tool *Tool) fullHeader := paramHeaderPrefix + b.Header // An empty string argument is mirrored as a present, empty header, so // presence is decided by the header being sent, not by its value. - headerPresent := len(header.Values(fullHeader)) > 0 - headerVal := header.Get(fullHeader) + headerVal, headerPresent, err := singletonHeader(header, fullHeader) + if err != nil { + return err + } argRaw, argExists := lookupArgument(raw.Arguments, b.Path) if !argExists || string(argRaw) == "null" { diff --git a/mcp/streamable_headers_test.go b/mcp/streamable_headers_test.go index 2042187cf..d80c781f1 100644 --- a/mcp/streamable_headers_test.go +++ b/mcp/streamable_headers_test.go @@ -1731,3 +1731,174 @@ func TestMcpNameHeaderEncodingRoundTrip(t *testing.T) { }) } } + +// TestParamHeaderAnnotationsWithNonStringTypes verifies that a type array or a +// boolean schema elsewhere in the input schema does not cause x-mcp-header +// annotations to be ignored. +func TestParamHeaderAnnotationsWithNonStringTypes(t *testing.T) { + type typedInput struct { + Region string `json:"region"` + Note *string `json:"note,omitempty"` + } + typed, err := jsonschema.For[typedInput](nil) + if err != nil { + t.Fatal(err) + } + typed.Properties["region"].Extra = map[string]any{"x-mcp-header": "Region"} + + tests := []struct { + name string + schema any + }{ + {"jsonschema.For with pointer field", typed}, + { + "type array on sibling property", + 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"}}, + }, + }, + }, + { + "type array on parent object", + json.RawMessage(`{"type":["object","null"],"properties":{"region":{"type":"string","x-mcp-header":"Region"}}}`), + }, + { + "boolean schema on sibling property", + json.RawMessage(`{"type":"object","properties":{"region":{"type":"string","x-mcp-header":"Region"},"any":true,"none":false}}`), + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + tool := &Tool{Name: "query", InputSchema: tt.schema} + if err := validateParamHeaderAnnotations(tool); err != nil { + t.Fatalf("validateParamHeaderAnnotations() = %v, want nil", err) + } + got := extractParamHeaderAnnotations(tool) + if len(got) != 1 || strings.Join(got[0].Path, ".") != "region" || got[0].Header != "Region" { + t.Fatalf("extractParamHeaderAnnotations() = %v, want [{[region] Region}]", got) + } + + params := mustMarshal(&CallToolParams{Name: "query", Arguments: map[string]any{"region": "us-west1"}}) + if got := generateParamHeaders(tool, params); got[paramHeaderPrefix+"Region"] != "us-west1" { + t.Errorf("generateParamHeaders() = %v, want Mcp-Param-Region: us-west1", got) + } + + msg := &jsonrpc.Request{Method: "tools/call", Params: params} + mismatch := http.Header{} + mismatch.Set(paramHeaderPrefix+"Region", "eu-west1") + if err := validateParamHeaders(mismatch, msg, tool); err == nil { + t.Error("validateParamHeaders() with mismatched header = nil, want error") + } + if err := validateParamHeaders(http.Header{}, msg, tool); err == nil { + t.Error("validateParamHeaders() with missing header = nil, want error") + } + }) + } +} + +func TestValidateParamHeaderAnnotationsTypeArray(t *testing.T) { + tests := []struct { + name string + schema string + wantErrSub string + }{ + { + name: "type array on annotated property", + schema: `{"type":"object","properties":{"region":{"type":["null","string"],"x-mcp-header":"Region"}}}`, + wantErrSub: `got ["null","string"]`, + }, + { + name: "missing type on annotated property", + schema: `{"type":"object","properties":{"region":{"x-mcp-header":"Region"}}}`, + wantErrSub: `got ""`, + }, + { + name: "invalid annotation next to type array", + schema: `{"type":"object","properties":{"region":{"type":"string","x-mcp-header":"My Region"},"note":{"type":["null","string"]}}}`, + wantErrSub: "invalid character", + }, + { + name: "invalid annotation next to boolean schema", + schema: `{"type":"object","properties":{"region":{"type":"string","x-mcp-header":"My Region"},"any":true}}`, + wantErrSub: "invalid character", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validateParamHeaderAnnotations(&Tool{Name: "query", InputSchema: json.RawMessage(tt.schema)}) + if err == nil { + t.Fatalf("validateParamHeaderAnnotations() = nil, want error containing %q", tt.wantErrSub) + } + if !strings.Contains(err.Error(), tt.wantErrSub) { + t.Errorf("validateParamHeaderAnnotations() = %q, want substring %q", err, tt.wantErrSub) + } + }) + } +} + +// TestValidateMcpHeadersRepeated verifies that the headers mirrored for +// intermediaries are rejected when sent more than once, regardless of which +// copy matches the body. +func TestValidateMcpHeadersRepeated(t *testing.T) { + tool := &Tool{ + Name: "query", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "region": map[string]any{"type": "string", "x-mcp-header": "Region"}, + }, + }, + } + lookup := func(name string) (*serverTool, bool) { + if name == tool.Name { + return &serverTool{tool: tool}, true + } + return nil, false + } + msg := &jsonrpc.Request{ + Method: "tools/call", + Params: mustMarshal(&CallToolParams{Name: "query", Arguments: map[string]any{"region": "us-west1"}}), + } + validHeader := func() http.Header { + h := http.Header{} + h.Set(protocolVersionHeader, minVersionForStandardHeaders) + h.Set(methodHeader, "tools/call") + h.Set(nameHeader, "query") + h.Set(paramHeaderPrefix+"Region", "us-west1") + return h + } + if err := validateMcpHeaders(validHeader(), msg, lookup); err != nil { + t.Fatalf("validateMcpHeaders() with single headers = %v, want nil", err) + } + + tests := []struct { + header string + extra []string + }{ + {protocolVersionHeader, []string{minVersionForStandardHeaders}}, + {protocolVersionHeader, []string{protocolVersion20251125}}, + {methodHeader, []string{"tools/call"}}, + {nameHeader, []string{"query"}}, + {nameHeader, []string{"other-tool"}}, + {paramHeaderPrefix + "Region", []string{"us-west1"}}, + {paramHeaderPrefix + "Region", []string{"eu-west1"}}, + } + for _, tt := range tests { + t.Run(fmt.Sprintf("%s=%v", tt.header, tt.extra), func(t *testing.T) { + h := validHeader() + for _, v := range tt.extra { + h.Add(tt.header, v) + } + err := validateMcpHeaders(h, msg, lookup) + if err == nil { + t.Fatal("validateMcpHeaders() = nil, want error for repeated header") + } + if want := "must not be repeated"; !strings.Contains(err.Error(), want) { + t.Errorf("validateMcpHeaders() = %q, want substring %q", err, want) + } + }) + } +} diff --git a/mcp/streamable_test.go b/mcp/streamable_test.go index 19f475e15..129f91073 100644 --- a/mcp/streamable_test.go +++ b/mcp/streamable_test.go @@ -2262,6 +2262,78 @@ func TestStreamableMcpHeaderValidationErrorFormat(t *testing.T) { } } +// TestStreamableParamHeaderWithPointerField verifies that Mcp-Param-* headers +// are mirrored and validated for a typed tool whose input struct has a pointer +// field, for which jsonschema.For emits a type array. +func TestStreamableParamHeaderWithPointerField(t *testing.T) { + type input struct { + Region string `json:"region"` + Note *string `json:"note,omitempty"` + } + schema, err := jsonschema.For[input](nil) + if err != nil { + t.Fatal(err) + } + schema.Properties["region"].Extra = map[string]any{"x-mcp-header": "Region"} + + server := NewServer(&Implementation{Name: "testServer", Version: "v1.0.0"}, nil) + var handlerCalled atomic.Bool + AddTool(server, &Tool{Name: "query", InputSchema: schema}, + func(ctx context.Context, req *CallToolRequest, in input) (*CallToolResult, any, error) { + handlerCalled.Store(true) + return &CallToolResult{}, nil, nil + }) + + handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, &StreamableHTTPOptions{ + Stateless: true, + JSONResponse: true, + }) + defer handler.closeAll() + httpServer := httptest.NewServer(mustNotPanic(t, handler)) + defer httpServer.Close() + + var sentRegion string + var status int + customClient := &http.Client{ + Transport: roundTripperFunc(func(req *http.Request) (*http.Response, error) { + if req.Header.Get(methodHeader) == "tools/call" { + sentRegion = req.Header.Get(paramHeaderPrefix + "Region") + req.Header.Set(paramHeaderPrefix+"Region", "eu-west1") + } + resp, err := http.DefaultTransport.RoundTrip(req) + if err == nil && req.Header.Get(methodHeader) == "tools/call" { + status = resp.StatusCode + } + return resp, err + }), + } + + client := NewClient(&Implementation{Name: "testClient", Version: "v1.0.0"}, nil) + ctx := context.Background() + session, err := client.Connect(ctx, &StreamableClientTransport{Endpoint: httpServer.URL, HTTPClient: customClient}, + &ClientSessionOptions{ProtocolVersion: minVersionForStandardHeaders}) + if err != nil { + t.Fatal(err) + } + defer session.Close() + if _, err := session.ListTools(ctx, nil); err != nil { + t.Fatal(err) + } + + if _, err := session.CallTool(ctx, &CallToolParams{Name: "query", Arguments: map[string]any{"region": "us-west1"}}); err == nil { + t.Error("CallTool succeeded with mismatched Mcp-Param-Region header") + } + if sentRegion != "us-west1" { + t.Errorf("client sent Mcp-Param-Region %q, want %q", sentRegion, "us-west1") + } + if status != http.StatusBadRequest { + t.Errorf("status code = %d, want %d", status, http.StatusBadRequest) + } + if handlerCalled.Load() { + t.Error("tool handler ran despite mismatched Mcp-Param-Region header") + } +} + // TestStreamableMcpHeaderVersionGating verifies that header validation // is skipped for protocol versions older than minVersionForStandardHeaders. func TestStreamableMcpHeaderVersionGating(t *testing.T) {