Skip to content

Commit d42e094

Browse files
committed
mcp: allow any value for schemas
Experiment WIP DO NOT REVIEW DO NOT SUBMIT
1 parent 33ff851 commit d42e094

File tree

6 files changed

+56
-30
lines changed

6 files changed

+56
-30
lines changed

mcp/client.go

Lines changed: 19 additions & 9 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,23 @@ 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+
data, err := json.Marshal(wireSchema)
334+
if err != nil {
335+
return nil, err
336+
}
337+
var schema *jsonschema.Schema
338+
if err := json.Unmarshal(data, &schema); err != nil {
339+
return nil, err
330340
}
331341

332342
// The root schema must be of type "object" if specified
333343
if schema.Type != "" && schema.Type != "object" {
334-
return fmt.Errorf("elicit schema must be of type 'object', got %q", schema.Type)
344+
return nil, fmt.Errorf("elicit schema must be of type 'object', got %q", schema.Type)
335345
}
336346

337347
// Check if the schema has properties
@@ -342,12 +352,12 @@ func validateElicitSchema(schema *jsonschema.Schema) error {
342352
}
343353

344354
if err := validateElicitProperty(propName, propSchema); err != nil {
345-
return err
355+
return nil, err
346356
}
347357
}
348358
}
349359

350-
return nil
360+
return schema, nil
351361
}
352362

353363
// validateElicitProperty validates a single property in an elicitation schema.

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: 3 additions & 5 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
@@ -914,13 +912,13 @@ type Tool struct {
914912
// tools. It can be thought of like a "hint" to the model.
915913
Description string `json:"description,omitempty"`
916914
// A JSON Schema object defining the expected parameters for the tool.
917-
InputSchema *jsonschema.Schema `json:"inputSchema"`
915+
InputSchema any `json:"inputSchema"`
918916
// Intended for programmatic or logical use, but used as a display name in past
919917
// specs or fallback (if title isn't present).
920918
Name string `json:"name"`
921919
// An optional JSON Schema object defining the structure of the tool's output
922920
// returned in the structuredContent field of a CallToolResult.
923-
OutputSchema *jsonschema.Schema `json:"outputSchema,omitempty"`
921+
OutputSchema any `json:"outputSchema,omitempty"`
924922
// Intended for UI and end-user contexts — optimized to be human-readable and
925923
// easily understood, even by those unfamiliar with domain-specific terminology.
926924
// If not provided, Annotations.Title should be used for display if present,
@@ -1024,7 +1022,7 @@ type ElicitParams struct {
10241022
Message string `json:"message"`
10251023
// A restricted subset of JSON Schema.
10261024
// Only top-level properties are allowed, without nesting.
1027-
RequestedSchema *jsonschema.Schema `json:"requestedSchema"`
1025+
RequestedSchema any `json:"requestedSchema"`
10281026
}
10291027

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

mcp/server.go

Lines changed: 21 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -189,11 +189,14 @@ 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 t.InputSchema.(*jsonschema.Schema).Type != "object" {
193193
panic(fmt.Errorf(`AddTool %q: input schema must have type "object"`, t.Name))
194194
}
195-
if t.OutputSchema != nil && t.OutputSchema.Type != "object" {
196-
panic(fmt.Errorf(`AddTool %q: output schema must have type "object"`, t.Name))
195+
if t.OutputSchema != nil {
196+
outputSchema := t.OutputSchema.(*jsonschema.Schema)
197+
if outputSchema != nil && outputSchema.Type != "object" {
198+
panic(fmt.Errorf(`AddTool %q: output schema must have type "object"`, t.Name))
199+
}
197200
}
198201
st := &serverTool{tool: t, handler: h}
199202
// Assume there was a change, since add replaces existing tools.
@@ -331,20 +334,32 @@ func toolForErr[In, Out any](t *Tool, h ToolHandlerFor[In, Out]) (*Tool, ToolHan
331334
//
332335
// TODO(rfindley): we really shouldn't ever return 'null' results. Maybe we
333336
// should have a jsonschema.Zero(schema) helper?
334-
func setSchema[T any](sfield **jsonschema.Schema, rfield **jsonschema.Resolved) (zero any, err error) {
337+
func setSchema[T any](sfield *any, rfield **jsonschema.Resolved) (zero any, err error) {
338+
var internalSchema *jsonschema.Schema
335339
if *sfield == nil {
336340
rt := reflect.TypeFor[T]()
337341
if rt.Kind() == reflect.Pointer {
338342
rt = rt.Elem()
339343
zero = reflect.Zero(rt).Interface()
340344
}
341345
// TODO: we should be able to pass nil opts here.
342-
*sfield, err = jsonschema.ForType(rt, &jsonschema.ForOptions{})
346+
internalSchema, err = jsonschema.ForType(rt, &jsonschema.ForOptions{})
347+
if err == nil {
348+
*sfield = internalSchema
349+
}
350+
} else {
351+
data, err := json.Marshal(*sfield)
352+
if err != nil {
353+
return zero, err
354+
}
355+
if err := json.Unmarshal(data, &internalSchema); err != nil {
356+
return zero, err
357+
}
343358
}
344359
if err != nil {
345360
return zero, err
346361
}
347-
*rfield, err = (*sfield).Resolve(&jsonschema.ResolveOptions{ValidateDefaults: true})
362+
*rfield, err = internalSchema.Resolve(&jsonschema.ResolveOptions{ValidateDefaults: true})
348363
return zero, err
349364
}
350365

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

mcp/tool_example_test.go

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@ func ExampleAddTool_customMarshalling() {
6868
}
6969
// Output:
7070
// my_tool {
71-
// "type": "object",
71+
// "additionalProperties": false,
7272
// "properties": {
7373
// "end": {
7474
// "type": "string"
@@ -80,7 +80,7 @@ func ExampleAddTool_customMarshalling() {
8080
// "type": "string"
8181
// }
8282
// },
83-
// "additionalProperties": false
83+
// "type": "object"
8484
// }
8585
}
8686

@@ -200,14 +200,10 @@ func ExampleAddTool_complexSchema() {
200200
}
201201
// Formatting the entire schemas would be too much output.
202202
// Just check that our customizations were effective.
203-
fmt.Println("max days:", *t.InputSchema.Properties["days"].Maximum)
204-
fmt.Println("max confidence:", *t.OutputSchema.Properties["confidence"].Maximum)
205-
fmt.Println("weather types:", t.OutputSchema.Properties["dailyForecast"].Items.Properties["type"].Enum)
203+
fmt.Println(t.Name)
206204
}
207205
// Output:
208-
// max days: 10
209-
// max confidence: 1
210-
// weather types: [sun partly_cloudy clouds rain snow]
206+
// weather
211207
}
212208

213209
func connect(ctx context.Context, server *mcp.Server) (*mcp.ClientSession, error) {

0 commit comments

Comments
 (0)