mcp: tool.InputSchema and tool.OutputSchema validation (SEP-2106) (#1009)

Fixes #990

docs/server.md

@@ -312,9 +312,13 @@ mcp.AddTool(server, &mcp.Tool{Name: "my_tool"}, handler)
 
 This does the following automatically:
 
-- If `Tool.InputSchema` or `Tool.OutputSchema` are unset, the input and output
-  schemas are inferred from the `In` type, which must be a struct or map.
-  Optional `jsonschema` struct tags provide argument descriptions.
+- If `Tool.InputSchema` is unset, the input schema is inferred from the `In`
+  type, which must be a struct or map.
+- If `Tool.OutputSchema` is unset and the `Out` type is not `any`, the output
+  schema is inferred from the `Out` type. Per SEP-2106, `Out` may be any Go
+  type whose inferred schema is a valid JSON Schema (struct, map, slice,
+  primitive, etc.).
+- Optional `jsonschema` struct tags provide argument and output descriptions.
 - Tool arguments are validated against the input schema.
 - Tool arguments are marshaled into the `In` value.
 - Tool output (the `Out` value) is marshaled into the result's

internal/docs/server.src.md

@@ -165,9 +165,13 @@ mcp.AddTool(server, &mcp.Tool{Name: "my_tool"}, handler)
 
 This does the following automatically:
 
-- If `Tool.InputSchema` or `Tool.OutputSchema` are unset, the input and output
-  schemas are inferred from the `In` type, which must be a struct or map.
-  Optional `jsonschema` struct tags provide argument descriptions.
+- If `Tool.InputSchema` is unset, the input schema is inferred from the `In`
+  type, which must be a struct or map.
+- If `Tool.OutputSchema` is unset and the `Out` type is not `any`, the output
+  schema is inferred from the `Out` type. Per SEP-2106, `Out` may be any Go
+  type whose inferred schema is a valid JSON Schema (struct, map, slice,
+  primitive, etc.).
+- Optional `jsonschema` struct tags provide argument and output descriptions.
 - Tool arguments are validated against the input schema.
 - Tool arguments are marshaled into the `In` value.
 - Tool output (the `Out` value) is marshaled into the result's

mcp/content.go

@@ -234,7 +234,9 @@ type ToolResultContent struct {
 	ToolUseID string
 	// Content holds the unstructured result of the tool call.
 	Content []Content
-	// StructuredContent holds an optional structured result as a JSON object.
+	// StructuredContent holds an optional structured result. Per SEP-2106, it
+	// may be any valid JSON value (object, array, or primitive) conforming to
+	// the tool's output schema.
 	StructuredContent any
 	// IsError indicates whether the tool call ended in an error.
 	IsError bool

mcp/protocol.go

@@ -270,7 +270,9 @@ type CallToolResult struct {
 	Content []Content `json:"content"`
 
 	// StructuredContent is an optional value that represents the structured
-	// result of the tool call. It must marshal to a JSON object.
+	// result of the tool call. Per SEP-2106, it may marshal to any valid JSON
+	// value (object, array, or primitive) conforming to the tool's
+	// [Tool.OutputSchema].
 	//
 	// When using a [ToolHandlerFor] with structured output, you should not
 	// populate this field. It will be automatically populated with the typed Out

mcp/server.go

@@ -232,7 +232,8 @@ func (s *Server) RemovePrompts(names ...string) {
 // that takes no input, or one where any input is valid, set [Tool.InputSchema] to
 // `{"type": "object"}`, using your preferred library or `json.RawMessage`.
 //
-// If present, [Tool.OutputSchema] must also have type "object".
+// If present, [Tool.OutputSchema] may be any valid JSON Schema (object, array,
+// primitive, or composition).
 //
 // When the handler is invoked as part of a CallTool request, req.Params.Arguments
 // will be a json.RawMessage.
@@ -279,16 +280,10 @@ func (s *Server) AddTool(t *Tool, h ToolHandler) {
 			if s == nil {
 				panic(fmt.Errorf("AddTool %q: output schema is nil", t.Name))
 			}
-			if s.Type != "object" {
-				panic(fmt.Errorf(`AddTool %q: output schema must have type "object"`, t.Name))
-			}
 		} else {
-			var m map[string]any
+			var m any
 			if err := remarshal(t.OutputSchema, &m); err != nil {
-				panic(fmt.Errorf("AddTool %q: can't marshal output schema to a JSON object: %v", t.Name, err))
-			}
-			if typ := m["type"]; typ != "object" {
-				panic(fmt.Errorf(`AddTool %q: output schema must have type "object" (got %v)`, t.Name, typ))
+				panic(fmt.Errorf("AddTool %q: can't marshal output schema to JSON: %v", t.Name, err))
 			}
 		}
 	}
@@ -340,7 +335,7 @@ func toolForErr[In, Out any](t *Tool, h ToolHandlerFor[In, Out], cache *SchemaCa
 		}
 		// Validate input and apply defaults.
 		var err error
-		input, err = applySchema(input, inputResolved)
+		input, err = applySchema(input, inputResolved, false)
 		if err != nil {
 			var errRes CallToolResult
 			errRes.SetError(fmt.Errorf("validating \"arguments\": %v", err))
@@ -402,7 +397,7 @@ func toolForErr[In, Out any](t *Tool, h ToolHandlerFor[In, Out], cache *SchemaCa
 			//
 			// We validate against the JSON, rather than the output value, as
 			// some types may have custom JSON marshalling (issue #447).
-			outJSON, err = applySchema(outJSON, outputResolved)
+			outJSON, err = applySchema(outJSON, outputResolved, true)
 			if err != nil {
 				return nil, fmt.Errorf("validating tool output: %w", err)
 			}
@@ -411,10 +406,18 @@ func toolForErr[In, Out any](t *Tool, h ToolHandlerFor[In, Out], cache *SchemaCa
 			// If the Content field isn't being used, return the serialized JSON in a
 			// TextContent block, as the spec suggests:
 			// https://modelcontextprotocol.io/specification/2025-06-18/server/tools#structured-content.
+			//
+			// Ensure a serialized-JSON TextContent fallback is present in case of servers using array
+			// or primitive structuredContent, so that pre-SEP-2106 clients can recover the structured
+			// payload from unstructured content.
 			if res.Content == nil {
 				res.Content = []Content{&TextContent{
 					Text: string(outJSON),
 				}}
+			} else if !isObjectJSON(outJSON) {
+				res.Content = append(res.Content, &TextContent{
+					Text: string(outJSON),
+				})
 			}
 		}
 		return res, nil
@@ -525,9 +528,10 @@ func setSchema[T any](sfield *any, rfield **jsonschema.Resolved, cache *SchemaCa
 // empty object schema value.
 //
 // If the tool's output schema is nil, and the Out type is not 'any', the
-// output schema is set to the schema inferred from the Out type argument,
-// which must also be a map or struct. If the Out type is 'any', the output
-// schema is omitted.
+// output schema is set to the schema inferred from the Out type argument.
+// Per SEP-2106, the Out type may be any Go type whose inferred schema is a
+// valid JSON Schema (struct, map, slice, primitive, etc.). If the Out type is
+// 'any', the output schema is omitted.
 //
 // Unlike [Server.AddTool], AddTool does a lot automatically, and forces
 // tools to conform to the MCP spec. See [ToolHandlerFor] for a detailed

mcp/server_test.go

@@ -521,7 +521,8 @@ func panics(f func()) (b bool) {
 }
 
 func TestAddTool(t *testing.T) {
-	// AddTool should panic if In or Out are not JSON objects.
+	// AddTool should panic if In is not a JSON object.
+	// Out may be of any type whose inferred JSON Schema is valid.
 	s := NewServer(testImpl, nil)
 	if !panics(func() {
 		AddTool(s, &Tool{Name: "T1"}, func(context.Context, *CallToolRequest, string) (*CallToolResult, any, error) { return nil, nil, nil })
@@ -535,12 +536,12 @@ func TestAddTool(t *testing.T) {
 	}) {
 		t.Error("good In: expected no panic")
 	}
-	if !panics(func() {
-		AddTool(s, &Tool{Name: "T2"}, func(context.Context, *CallToolRequest, map[string]any) (*CallToolResult, int, error) {
+	if panics(func() {
+		AddTool(s, &Tool{Name: "T3"}, func(context.Context, *CallToolRequest, map[string]any) (*CallToolResult, int, error) {
 			return nil, 0, nil
 		})
 	}) {
-		t.Error("bad Out: expected panic")
+		t.Error("primitive Out: expected no panic")
 	}
 }
 
@@ -653,6 +654,195 @@ func TestAddToolNilSchema(t *testing.T) {
 	}
 }
 
+// TestAddToolNonObjectOutputSchema verifies the low-level
+// Server.AddTool accepts an output schema whose root type is not "object"
+// (array or primitive) and structured content of the matching shape
+// round-trips end-to-end.
+func TestAddToolNonObjectOutputSchema(t *testing.T) {
+	ctx := context.Background()
+
+	for _, tc := range []struct {
+		name         string
+		outputSchema any
+		content      any
+		want         any
+	}{
+		{
+			name:         "array of objects",
+			outputSchema: &jsonschema.Schema{Type: "array", Items: &jsonschema.Schema{Type: "object"}},
+			content:      []any{map[string]any{"id": "u1"}, map[string]any{"id": "u2"}},
+			want:         []any{map[string]any{"id": "u1"}, map[string]any{"id": "u2"}},
+		},
+		{
+			name:         "primitive number (map-based schema)",
+			outputSchema: map[string]any{"type": "number"},
+			content:      42.0,
+			want:         42.0,
+		},
+		{
+			name:         "primitive string (RawMessage schema)",
+			outputSchema: json.RawMessage(`{"type":"string"}`),
+			content:      "hello",
+			want:         "hello",
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			server := NewServer(testImpl, nil)
+			handler := func(ctx context.Context, req *CallToolRequest) (*CallToolResult, error) {
+				return &CallToolResult{StructuredContent: tc.content}, nil
+			}
+			tool := &Tool{
+				Name:         "t",
+				InputSchema:  &jsonschema.Schema{Type: "object"},
+				OutputSchema: tc.outputSchema,
+			}
+			// Must not panic.
+			server.AddTool(tool, handler)
+
+			ct, st := NewInMemoryTransports()
+			if _, err := server.Connect(ctx, st, nil); err != nil {
+				t.Fatal(err)
+			}
+			client := NewClient(testImpl, nil)
+			cs, err := client.Connect(ctx, ct, nil)
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer cs.Close()
+
+			res, err := cs.CallTool(ctx, &CallToolParams{Name: "t"})
+			if err != nil {
+				t.Fatal(err)
+			}
+			if res.IsError {
+				t.Fatalf("unexpected tool error: %v", res.Content)
+			}
+			if diff := cmp.Diff(tc.want, res.StructuredContent); diff != "" {
+				t.Errorf("structured content mismatch (-want +got):\n%s", diff)
+			}
+		})
+	}
+}
+
+// TestAddToolGenericNonObjectOutput verifies SEP-2106 for the generic
+// AddTool[In, Out] helper: Out may be a slice or primitive whose inferred
+// JSON Schema has a non-object root.
+func TestAddToolGenericNonObjectOutput(t *testing.T) {
+	ctx := context.Background()
+
+	type item struct {
+		ID string `json:"id"`
+	}
+
+	t.Run("slice output", func(t *testing.T) {
+		server := NewServer(testImpl, nil)
+		AddTool(server, &Tool{Name: "list"},
+			func(ctx context.Context, req *CallToolRequest, _ struct{}) (*CallToolResult, []item, error) {
+				return nil, []item{{ID: "u1"}, {ID: "u2"}}, nil
+			})
+
+		ct, st := NewInMemoryTransports()
+		if _, err := server.Connect(ctx, st, nil); err != nil {
+			t.Fatal(err)
+		}
+		client := NewClient(testImpl, nil)
+		cs, err := client.Connect(ctx, ct, nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer cs.Close()
+
+		// Verify tools/list reports an array-rooted output schema.
+		// jsonschema-go infers ["null","array"] for slices since nil slices
+		// are valid; accept either form.
+		lt, err := cs.ListTools(ctx, nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if len(lt.Tools) != 1 {
+			t.Fatalf("got %d tools, want 1", len(lt.Tools))
+		}
+		gotSchema, _ := lt.Tools[0].OutputSchema.(map[string]any)
+		typeOK := false
+		switch typ := gotSchema["type"].(type) {
+		case string:
+			typeOK = typ == "array"
+		case []any:
+			for _, e := range typ {
+				if e == "array" {
+					typeOK = true
+				}
+			}
+		}
+		if !typeOK {
+			t.Errorf("output schema type = %v, want to include %q", gotSchema["type"], "array")
+		}
+
+		res, err := cs.CallTool(ctx, &CallToolParams{Name: "list"})
+		if err != nil {
+			t.Fatal(err)
+		}
+		if res.IsError {
+			t.Fatalf("unexpected tool error: %v", res.Content)
+		}
+		want := []any{map[string]any{"id": "u1"}, map[string]any{"id": "u2"}}
+		if diff := cmp.Diff(want, res.StructuredContent); diff != "" {
+			t.Errorf("structured content mismatch (-want +got):\n%s", diff)
+		}
+	})
+
+	t.Run("primitive output", func(t *testing.T) {
+		server := NewServer(testImpl, nil)
+		AddTool(server, &Tool{Name: "count"},
+			func(ctx context.Context, req *CallToolRequest, _ struct{}) (*CallToolResult, int, error) {
+				return nil, 42, nil
+			})
+
+		ct, st := NewInMemoryTransports()
+		if _, err := server.Connect(ctx, st, nil); err != nil {
+			t.Fatal(err)
+		}
+		client := NewClient(testImpl, nil)
+		cs, err := client.Connect(ctx, ct, nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer cs.Close()
+
+		res, err := cs.CallTool(ctx, &CallToolParams{Name: "count"})
+		if err != nil {
+			t.Fatal(err)
+		}
+		if res.IsError {
+			t.Fatalf("unexpected tool error: %v", res.Content)
+		}
+		if diff := cmp.Diff(float64(42), res.StructuredContent); diff != "" {
+			t.Errorf("structured content mismatch (-want +got):\n%s", diff)
+		}
+	})
+}
+
+// TestAddToolInputSchemaComposition verifies SEP-2106 (input side): composition
+// keywords such as oneOf are allowed on the input schema alongside
+// type:"object".
+func TestAddToolInputSchemaComposition(t *testing.T) {
+	server := NewServer(testImpl, nil)
+	tool := &Tool{
+		Name: "lookup",
+		InputSchema: &jsonschema.Schema{
+			Type: "object",
+			OneOf: []*jsonschema.Schema{
+				{Required: []string{"id"}, Properties: map[string]*jsonschema.Schema{"id": {Type: "string"}}},
+				{Required: []string{"name"}, Properties: map[string]*jsonschema.Schema{"name": {Type: "string"}}},
+			},
+		},
+	}
+	// Must not panic.
+	server.AddTool(tool, func(ctx context.Context, req *CallToolRequest) (*CallToolResult, error) {
+		return &CallToolResult{}, nil
+	})
+}
+
 type schema = jsonschema.Schema
 
 func testToolForSchema[In, Out any](t *testing.T, tool *Tool, in string, out Out, wantIn, wantOut any, wantErrContaining string) {

mcp/tool.go

@@ -65,8 +65,14 @@ type serverTool struct {
 // applySchema validates whether data is valid JSON according to the provided
 // schema, after applying schema defaults.
 //
-// Returns the JSON value augmented with defaults.
-func applySchema(data json.RawMessage, resolved *jsonschema.Resolved) (json.RawMessage, error) {
+// If forOutput is false, the data is treated as tool input: the schema's root
+// type must be "object" and the value is unmarshaled into a map.
+//
+// If forOutput is true, the data is treated as tool output: the schema's root
+// may be of any type (object, array, primitive, composition).
+//
+// Returns the JSON value, augmented with defaults where applicable.
+func applySchema(data json.RawMessage, resolved *jsonschema.Resolved, forOutput bool) (json.RawMessage, error) {
 	// TODO: use reflection to create the struct type to unmarshal into.
 	// Separate validation from assignment.
 
@@ -75,33 +81,81 @@ func applySchema(data json.RawMessage, resolved *jsonschema.Resolved) (json.RawM
 	// This avoids inconsistent representation due to custom marshallers, such as
 	// time.Time (issue #449).
 	//
-	// Additionally, unmarshalling into a map ensures that the resulting JSON is
+	// For input, unmarshalling into a map ensures that the resulting JSON is
 	// at least {}, even if data is empty. For example, arguments is technically
 	// an optional property of callToolParams, and we still want to apply the
 	// defaults in this case.
 	//
 	// TODO(rfindley): in which cases can resolved be nil?
-	if resolved != nil {
+	if resolved == nil {
+		return data, nil
+	}
+
+	var unmarshaled any
+	if !forOutput {
 		v := make(map[string]any)
 		if len(data) > 0 {
 			if err := internaljson.Unmarshal(data, &v); err != nil {
 				return nil, fmt.Errorf("unmarshaling arguments: %w", err)
 			}
 		}
-		if err := resolved.ApplyDefaults(&v); err != nil {
-			return nil, fmt.Errorf("applying schema defaults:\n%w", err)
-		}
-		if err := resolved.Validate(&v); err != nil {
-			return nil, err
-		}
-		// We must re-marshal with the default values applied.
-		var err error
-		data, err = json.Marshal(v)
-		if err != nil {
-			return nil, fmt.Errorf("marshalling with defaults: %v", err)
+		unmarshaled = v
+	} else {
+		if len(data) > 0 {
+			if err := internaljson.Unmarshal(data, &unmarshaled); err != nil {
+				return nil, fmt.Errorf("unmarshaling output: %w", err)
+			}
 		}
 	}
-	return data, nil
+
+	// Apply defaults only when the value is a map: jsonschema.Resolved.ApplyDefaults
+	// only operates on object properties. For object-rooted output schemas,
+	// coerce a nil result (from "null" or empty data) into {} so handlers that
+	// return a typed-nil map still validate.
+	appliedDefaults := false
+	if _, ok := unmarshaled.(map[string]any); ok {
+		if err := resolved.ApplyDefaults(&unmarshaled); err != nil {
+			return nil, fmt.Errorf("applying schema defaults:\n%w", err)
+		}
+		appliedDefaults = true
+	} else if forOutput && unmarshaled == nil && resolved.Schema().Type == "object" {
+		unmarshaled = make(map[string]any)
+		if err := resolved.ApplyDefaults(&unmarshaled); err != nil {
+			return nil, fmt.Errorf("applying schema defaults:\n%w", err)
+		}
+		appliedDefaults = true
+	}
+
+	if err := resolved.Validate(&unmarshaled); err != nil {
+		return nil, err
+	}
+
+	// Re-marshal only when defaults may have changed the value.
+	if !appliedDefaults {
+		return data, nil
+	}
+	out, err := json.Marshal(unmarshaled)
+	if err != nil {
+		return nil, fmt.Errorf("marshalling with defaults: %v", err)
+	}
+	return out, nil
+}
+
+// isObjectJSON reports whether data is a JSON object (i.e., starts with '{'
+// after any leading whitespace). Returns false for arrays, primitives, null,
+// or empty input.
+func isObjectJSON(data json.RawMessage) bool {
+	for _, b := range data {
+		switch b {
+		case ' ', '\t', '\n', '\r':
+			continue
+		case '{':
+			return true
+		default:
+			return false
+		}
+	}
+	return false
 }
 
 // validateToolName checks whether name is a valid tool name, reporting a

mcp/tool_test.go

@@ -46,7 +46,7 @@ func TestApplySchema(t *testing.T) {
 		{`{"x": 0}`, new(map[string]any), &map[string]any{"x": 0.0}},
 	} {
 		raw := json.RawMessage(tt.data)
-		raw, err = applySchema(raw, resolved)
+		raw, err = applySchema(raw, resolved, false)
 		if err != nil {
 			t.Fatal(err)
 		}
@@ -59,6 +59,83 @@ func TestApplySchema(t *testing.T) {
 	}
 }
 
+func TestApplySchemaOutput(t *testing.T) {
+	// SEP-2106: when forOutput is true, the schema may have a non-object root.
+	for _, tc := range []struct {
+		name    string
+		schema  *jsonschema.Schema
+		data    string
+		want    string
+		wantErr string
+	}{
+		{
+			name:   "array root accepts array",
+			schema: &jsonschema.Schema{Type: "array", Items: &jsonschema.Schema{Type: "integer"}},
+			data:   `[1,2,3]`,
+			want:   `[1,2,3]`,
+		},
+		{
+			name:    "array root rejects wrong element type",
+			schema:  &jsonschema.Schema{Type: "array", Items: &jsonschema.Schema{Type: "integer"}},
+			data:    `[1,"two",3]`,
+			wantErr: `"integer"`,
+		},
+		{
+			name:   "primitive number root",
+			schema: &jsonschema.Schema{Type: "number"},
+			data:   `42`,
+			want:   `42`,
+		},
+		{
+			name:   "primitive string root",
+			schema: &jsonschema.Schema{Type: "string"},
+			data:   `"hello"`,
+			want:   `"hello"`,
+		},
+		{
+			name:   "object root with defaults",
+			schema: &jsonschema.Schema{Type: "object", Properties: map[string]*jsonschema.Schema{"x": {Type: "integer", Default: json.RawMessage("3")}}},
+			data:   `{}`,
+			want:   `{"x":3}`,
+		},
+		{
+			name:   "empty object root with defaults",
+			schema: &jsonschema.Schema{Type: "object", Properties: map[string]*jsonschema.Schema{"x": {Type: "integer", Default: json.RawMessage("3")}}},
+			data:   ``,
+			want:   `{"x":3}`,
+		},
+		{
+			name:   "object root accepts null as empty object",
+			schema: &jsonschema.Schema{Type: "object"},
+			data:   `null`,
+			want:   `{}`,
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			resolved, err := tc.schema.Resolve(&jsonschema.ResolveOptions{ValidateDefaults: true})
+			if err != nil {
+				t.Fatal(err)
+			}
+			got, err := applySchema(json.RawMessage(tc.data), resolved, true)
+			if tc.wantErr != "" {
+				if err == nil {
+					t.Fatalf("got %s, want error containing %q", got, tc.wantErr)
+				}
+				if !strings.Contains(err.Error(), tc.wantErr) {
+					t.Errorf("got error %q, want containing %q", err, tc.wantErr)
+				}
+				return
+			}
+			if err != nil {
+				t.Fatal(err)
+			}
+			if string(got) != tc.want {
+				t.Errorf("got %s, want %s", got, tc.want)
+			}
+		})
+	}
+}
+
 func TestToolErrorHandling(t *testing.T) {
 	// Construct server and add both tools at the top level
 	server := NewServer(testImpl, nil)