diff --git a/docs/server.md b/docs/server.md index c5a7bf9b..8fe6bd37 100644 --- a/docs/server.md +++ b/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 diff --git a/internal/docs/server.src.md b/internal/docs/server.src.md index 2d822c76..26dba5a5 100644 --- a/internal/docs/server.src.md +++ b/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 diff --git a/mcp/content.go b/mcp/content.go index 95ea40d8..39c25ffd 100644 --- a/mcp/content.go +++ b/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 diff --git a/mcp/protocol.go b/mcp/protocol.go index acc7ec48..6e87f192 100644 --- a/mcp/protocol.go +++ b/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 diff --git a/mcp/server.go b/mcp/server.go index 912dea98..6cd084e7 100644 --- a/mcp/server.go +++ b/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 diff --git a/mcp/server_test.go b/mcp/server_test.go index d2814a0b..4c8b76e4 100644 --- a/mcp/server_test.go +++ b/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) { diff --git a/mcp/tool.go b/mcp/tool.go index 3ecb59d3..04cf551a 100644 --- a/mcp/tool.go +++ b/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 { + unmarshaled = v + } else { + if len(data) > 0 { + if err := internaljson.Unmarshal(data, &unmarshaled); err != nil { + return nil, fmt.Errorf("unmarshaling output: %w", err) + } + } + } + + // 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) } - if err := resolved.Validate(&v); err != nil { - return nil, 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) } - // 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) + 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 data, nil + return false } // validateToolName checks whether name is a valid tool name, reporting a diff --git a/mcp/tool_test.go b/mcp/tool_test.go index dfd859be..7a82c96d 100644 --- a/mcp/tool_test.go +++ b/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)