Skip to content

Commit cbfd05f

Browse files
committed
mcp: allow any value for schemas
As discussed at length in #518, coupling our schema fields to jsonschema.Schema limits our user's ability to use a different schema package, and makes it impossible for our client to talk to some backends. This CL relaxes the type of `*jsonschema.Schema` fields to be of type `any`, which is a breaking change but is unlikely to affect most server code. (It *is* a nontrivial change for clients). Fixes #518
1 parent 33ff851 commit cbfd05f

File tree

8 files changed

+201
-47
lines changed

8 files changed

+201
-47
lines changed

examples/server/toolschemas/main.go

Lines changed: 21 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -126,15 +126,13 @@ func main() {
126126
// Add the 'greeting' tool in a few different ways.
127127

128128
// First, we can just use [mcp.AddTool], and get the out-of-the-box handling
129-
// it provides:
129+
// it provides for schema inference, validation, parsing, and packing the
130+
// result.
130131
mcp.AddTool(server, &mcp.Tool{Name: "simple greeting"}, simpleGreeting)
131132

132-
// Next, we can create our schemas entirely manually, and add them using
133-
// [mcp.Server.AddTool]. Since we're working manually, we can add some
134-
// constraints on the length of the name.
135-
//
136-
// We don't need to do all this work: below, we use jsonschema.For to start
137-
// from the default schema.
133+
// Alternatively, we can create our schemas entirely manually, and add them
134+
// using [mcp.Server.AddTool]. Since we're using the 'raw' API, we have to do
135+
// the parsing and validation ourselves
138136
manual, err := newManualGreeter()
139137
if err != nil {
140138
log.Fatal(err)
@@ -145,6 +143,22 @@ func main() {
145143
OutputSchema: outputSchema,
146144
}, manual.greet)
147145

146+
// We can even use raw schema values. In this case, note that we're not
147+
// validating the input at all.
148+
server.AddTool(&mcp.Tool{
149+
Name: "unvalidated greeting",
150+
InputSchema: json.RawMessage(`{"type":"object","properties":{"user":{"type":"string"}}}`),
151+
}, func(_ context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
152+
// Note: no validation!
153+
var args struct{ User string }
154+
if err := json.Unmarshal(req.Params.Arguments, &args); err != nil {
155+
return nil, err
156+
}
157+
return &mcp.CallToolResult{
158+
Content: []mcp.Content{&mcp.TextContent{Text: "Hi " + args.User}},
159+
}, nil
160+
})
161+
148162
// Finally, note that we can also use custom schemas with a ToolHandlerFor.
149163
// We can do this in two ways: by using one of the schema values constructed
150164
// above, or by using jsonschema.For and adjusting the resulting schema.

mcp/client.go

Lines changed: 16 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -295,7 +295,8 @@ func (c *Client) elicit(ctx context.Context, req *ElicitRequest) (*ElicitResult,
295295
}
296296

297297
// Validate that the requested schema only contains top-level properties without nesting
298-
if err := validateElicitSchema(req.Params.RequestedSchema); err != nil {
298+
schema, err := validateElicitSchema(req.Params.RequestedSchema)
299+
if err != nil {
299300
return nil, jsonrpc2.NewError(CodeInvalidParams, err.Error())
300301
}
301302

@@ -305,11 +306,11 @@ func (c *Client) elicit(ctx context.Context, req *ElicitRequest) (*ElicitResult,
305306
}
306307

307308
// Validate elicitation result content against requested schema
308-
if req.Params.RequestedSchema != nil && res.Content != nil {
309+
if schema != nil && res.Content != nil {
309310
// TODO: is this the correct behavior if validation fails?
310311
// It isn't the *server's* params that are invalid, so why would we return
311312
// this code to the server?
312-
resolved, err := req.Params.RequestedSchema.Resolve(nil)
313+
resolved, err := schema.Resolve(nil)
313314
if err != nil {
314315
return nil, jsonrpc2.NewError(CodeInvalidParams, fmt.Sprintf("failed to resolve requested schema: %v", err))
315316
}
@@ -324,14 +325,19 @@ func (c *Client) elicit(ctx context.Context, req *ElicitRequest) (*ElicitResult,
324325

325326
// validateElicitSchema validates that the schema conforms to MCP elicitation schema requirements.
326327
// Per the MCP specification, elicitation schemas are limited to flat objects with primitive properties only.
327-
func validateElicitSchema(schema *jsonschema.Schema) error {
328-
if schema == nil {
329-
return nil // nil schema is allowed
328+
func validateElicitSchema(wireSchema any) (*jsonschema.Schema, error) {
329+
if wireSchema == nil {
330+
return nil, nil // nil schema is allowed
331+
}
332+
333+
var schema *jsonschema.Schema
334+
if err := remarshal(wireSchema, &schema); err != nil {
335+
return nil, err
330336
}
331337

332338
// The root schema must be of type "object" if specified
333339
if schema.Type != "" && schema.Type != "object" {
334-
return fmt.Errorf("elicit schema must be of type 'object', got %q", schema.Type)
340+
return nil, fmt.Errorf("elicit schema must be of type 'object', got %q", schema.Type)
335341
}
336342

337343
// Check if the schema has properties
@@ -342,12 +348,12 @@ func validateElicitSchema(schema *jsonschema.Schema) error {
342348
}
343349

344350
if err := validateElicitProperty(propName, propSchema); err != nil {
345-
return err
351+
return nil, err
346352
}
347353
}
348354
}
349355

350-
return nil
356+
return schema, nil
351357
}
352358

353359
// validateElicitProperty validates a single property in an elicitation schema.
@@ -383,7 +389,7 @@ func validateElicitStringProperty(propName string, propSchema *jsonschema.Schema
383389
if propSchema.Extra != nil {
384390
if enumNamesRaw, exists := propSchema.Extra["enumNames"]; exists {
385391
// Type check enumNames - should be a slice
386-
if enumNamesSlice, ok := enumNamesRaw.([]interface{}); ok {
392+
if enumNamesSlice, ok := enumNamesRaw.([]any); ok {
387393
if len(enumNamesSlice) != len(propSchema.Enum) {
388394
return fmt.Errorf("elicit schema property %q has %d enum values but %d enumNames, they must match", propName, len(propSchema.Enum), len(enumNamesSlice))
389395
}

mcp/client_list_test.go

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ package mcp_test
66

77
import (
88
"context"
9+
"encoding/json"
910
"iter"
1011
"log"
1112
"testing"
@@ -41,7 +42,13 @@ func TestList(t *testing.T) {
4142
if err != nil {
4243
t.Fatal(err)
4344
}
44-
tt.InputSchema = is
45+
data, err := json.Marshal(is)
46+
if err != nil {
47+
t.Fatal(err)
48+
}
49+
if err := json.Unmarshal(data, &tt.InputSchema); err != nil {
50+
t.Fatal(err)
51+
}
4552
wantTools = append(wantTools, tt)
4653
}
4754
t.Run("list", func(t *testing.T) {

mcp/protocol.go

Lines changed: 40 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,6 @@ package mcp
1313
import (
1414
"encoding/json"
1515
"fmt"
16-
17-
"github.com/google/jsonschema-go/jsonschema"
1816
)
1917

2018
// Optional annotations for the client. The client can use annotations to inform
@@ -913,14 +911,38 @@ type Tool struct {
913911
// This can be used by clients to improve the LLM's understanding of available
914912
// tools. It can be thought of like a "hint" to the model.
915913
Description string `json:"description,omitempty"`
916-
// A JSON Schema object defining the expected parameters for the tool.
917-
InputSchema *jsonschema.Schema `json:"inputSchema"`
914+
// InputSchema holds a JSON Schema object defining the expected parameters
915+
// for the tool.
916+
//
917+
// From the server, this field may be set to any value that JSON-marshals to
918+
// valid JSON schema (including json.RawMessage). However, for tools added
919+
// using [AddTool], which automatically validates inputs and outputs, the
920+
// schema must be in a draft the SDK understands. Currently, the SDK uses
921+
// github.com/google/jsonschema-go for inference and validation, which only
922+
// supports the 2020-12 draft of JSON schema. To do your own validation, use
923+
// [Server.AddTool].
924+
//
925+
// From the client, this field will hold the default JSON marshaling of the
926+
// server's input schema (a map[string]any).
927+
InputSchema any `json:"inputSchema"`
918928
// Intended for programmatic or logical use, but used as a display name in past
919929
// specs or fallback (if title isn't present).
920930
Name string `json:"name"`
921-
// An optional JSON Schema object defining the structure of the tool's output
922-
// returned in the structuredContent field of a CallToolResult.
923-
OutputSchema *jsonschema.Schema `json:"outputSchema,omitempty"`
931+
// OutputSchema holds an optional JSON Schema object defining the structure
932+
// of the tool's output returned in the StructuredContent field of a
933+
// CallToolResult.
934+
//
935+
// From the server, this field may be set to any value that JSON-marshals to
936+
// valid JSON schema (including json.RawMessage). However, for tools added
937+
// using [AddTool], which automatically validates inputs and outputs, the
938+
// schema must be in a draft the SDK understands. Currently, the SDK uses
939+
// github.com/google/jsonschema-go for inference and validation, which only
940+
// supports the 2020-12 draft of JSON schema. To do your own validation, use
941+
// [Server.AddTool].
942+
//
943+
// From the client, this field will hold the default JSON marshaling of the
944+
// server's output schema (a map[string]any).
945+
OutputSchema any `json:"outputSchema,omitempty"`
924946
// Intended for UI and end-user contexts — optimized to be human-readable and
925947
// easily understood, even by those unfamiliar with domain-specific terminology.
926948
// If not provided, Annotations.Title should be used for display if present,
@@ -1022,9 +1044,18 @@ type ElicitParams struct {
10221044
Meta `json:"_meta,omitempty"`
10231045
// The message to present to the user.
10241046
Message string `json:"message"`
1025-
// A restricted subset of JSON Schema.
1047+
// A JSON schema object defining the requested elicitation schema.
1048+
//
1049+
// From the server, this field may be set to any value that can JSON-marshal
1050+
// to valid JSON schema (including json.RawMessage for raw schema values).
1051+
// Internally, the SDK uses github.com/google/jsonschema-go for validation,
1052+
// which only supports the 2020-12 draft of the JSON schema spec.
1053+
//
1054+
// From the client, this field will use the default JSON marshaling (a
1055+
// map[string]any).
1056+
//
10261057
// Only top-level properties are allowed, without nesting.
1027-
RequestedSchema *jsonschema.Schema `json:"requestedSchema"`
1058+
RequestedSchema any `json:"requestedSchema"`
10281059
}
10291060

10301061
func (x *ElicitParams) isParams() {}

mcp/server.go

Lines changed: 46 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -164,9 +164,9 @@ func (s *Server) RemovePrompts(names ...string) {
164164
//
165165
// The tool's input schema must be non-nil and have the type "object". For a tool
166166
// that takes no input, or one where any input is valid, set [Tool.InputSchema] to
167-
// &jsonschema.Schema{Type: "object"}.
167+
// `{"type": "object"}`, using your preferred library or `json.RawMessage`.
168168
//
169-
// If present, the output schema must also have type "object".
169+
// If present, [Tool.OutputSchema] must also have type "object".
170170
//
171171
// When the handler is invoked as part of a CallTool request, req.Params.Arguments
172172
// will be a json.RawMessage.
@@ -189,11 +189,29 @@ func (s *Server) AddTool(t *Tool, h ToolHandler) {
189189
// discovered until runtime, when the LLM sent bad data.
190190
panic(fmt.Errorf("AddTool %q: missing input schema", t.Name))
191191
}
192-
if t.InputSchema.Type != "object" {
192+
if s, ok := t.InputSchema.(*jsonschema.Schema); ok && s.Type != "object" {
193193
panic(fmt.Errorf(`AddTool %q: input schema must have type "object"`, t.Name))
194+
} else {
195+
var m map[string]any
196+
if err := remarshal(t.InputSchema, &m); err != nil {
197+
panic(fmt.Errorf("AddTool %q: can't marshal input schema to a JSON object: %v", t.Name, err))
198+
}
199+
if typ := m["type"]; typ != "object" {
200+
panic(fmt.Errorf(`AddTool %q: input schema must have type "object" (got %v)`, t.Name, typ))
201+
}
194202
}
195-
if t.OutputSchema != nil && t.OutputSchema.Type != "object" {
196-
panic(fmt.Errorf(`AddTool %q: output schema must have type "object"`, t.Name))
203+
if t.OutputSchema != nil {
204+
if s, ok := t.OutputSchema.(*jsonschema.Schema); ok && s.Type != "object" {
205+
panic(fmt.Errorf(`AddTool %q: output schema must have type "object"`, t.Name))
206+
} else {
207+
var m map[string]any
208+
if err := remarshal(t.OutputSchema, &m); err != nil {
209+
panic(fmt.Errorf("AddTool %q: can't marshal output schema to a JSON object: %v", t.Name, err))
210+
}
211+
if typ := m["type"]; typ != "object" {
212+
panic(fmt.Errorf(`AddTool %q: output schema must have type "object" (got %v)`, t.Name, typ))
213+
}
214+
}
197215
}
198216
st := &serverTool{tool: t, handler: h}
199217
// Assume there was a change, since add replaces existing tools.
@@ -331,36 +349,50 @@ func toolForErr[In, Out any](t *Tool, h ToolHandlerFor[In, Out]) (*Tool, ToolHan
331349
//
332350
// TODO(rfindley): we really shouldn't ever return 'null' results. Maybe we
333351
// should have a jsonschema.Zero(schema) helper?
334-
func setSchema[T any](sfield **jsonschema.Schema, rfield **jsonschema.Resolved) (zero any, err error) {
352+
func setSchema[T any](sfield *any, rfield **jsonschema.Resolved) (zero any, err error) {
353+
var internalSchema *jsonschema.Schema
335354
if *sfield == nil {
336355
rt := reflect.TypeFor[T]()
337356
if rt.Kind() == reflect.Pointer {
338357
rt = rt.Elem()
339358
zero = reflect.Zero(rt).Interface()
340359
}
341360
// TODO: we should be able to pass nil opts here.
342-
*sfield, err = jsonschema.ForType(rt, &jsonschema.ForOptions{})
361+
internalSchema, err = jsonschema.ForType(rt, &jsonschema.ForOptions{})
362+
if err == nil {
363+
*sfield = internalSchema
364+
}
365+
} else {
366+
if err := remarshal(*sfield, &internalSchema); err != nil {
367+
return zero, err
368+
}
343369
}
344370
if err != nil {
345371
return zero, err
346372
}
347-
*rfield, err = (*sfield).Resolve(&jsonschema.ResolveOptions{ValidateDefaults: true})
373+
*rfield, err = internalSchema.Resolve(&jsonschema.ResolveOptions{ValidateDefaults: true})
348374
return zero, err
349375
}
350376

351377
// AddTool adds a tool and typed tool handler to the server.
352378
//
353379
// If the tool's input schema is nil, it is set to the schema inferred from the
354-
// In type parameter, using [jsonschema.For]. The In type argument must be a
355-
// map or a struct, so that its inferred JSON Schema has type "object".
380+
// In type parameter. Types are inferred from Go types, and property
381+
// descriptions are read from the 'jsonschema' struct tag. Internally, the SDK
382+
// uses the github.com/google/jsonschema-go package for ineference and
383+
// validation. The In type argument must be a map or a struct, so that its
384+
// inferred JSON Schema has type "object", as required by the spec. As a
385+
// special case, if the In type is 'any', the tool's input schema is set to an
386+
// empty object schema value.
356387
//
357388
// If the tool's output schema is nil, and the Out type is not 'any', the
358389
// output schema is set to the schema inferred from the Out type argument,
359-
// which also must be a map or struct.
390+
// which must also be a map or struct. If the Out type is 'any', the output
391+
// schema is omitted.
360392
//
361-
// Unlike [Server.AddTool], AddTool does a lot automatically, and forces tools
362-
// to conform to the MCP spec. See [ToolHandlerFor] for a detailed description
363-
// of this automatic behavior.
393+
// Unlike [Server.AddTool], AddTool does a lot automatically, and forces
394+
// tools to conform to the MCP spec. See [ToolHandlerFor] for a detailed
395+
// description of this automatic behavior.
364396
func AddTool[In, Out any](s *Server, t *Tool, h ToolHandlerFor[In, Out]) {
365397
tt, hh, err := toolForErr(t, h)
366398
if err != nil {

mcp/server_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -492,7 +492,7 @@ func TestAddTool(t *testing.T) {
492492

493493
type schema = jsonschema.Schema
494494

495-
func testToolForSchema[In, Out any](t *testing.T, tool *Tool, in string, out Out, wantIn, wantOut *schema, wantErrContaining string) {
495+
func testToolForSchema[In, Out any](t *testing.T, tool *Tool, in string, out Out, wantIn, wantOut any, wantErrContaining string) {
496496
t.Helper()
497497
th := func(context.Context, *CallToolRequest, In) (*CallToolResult, Out, error) {
498498
return nil, out, nil

0 commit comments

Comments
 (0)