mcp: allow any value for schemas

Experiment
WIP
DO NOT REVIEW
DO NOT SUBMIT

Author: findleyr

mcp/client.go

@@ -295,7 +295,8 @@ func (c *Client) elicit(ctx context.Context, req *ElicitRequest) (*ElicitResult,
 	}
 
 	// Validate that the requested schema only contains top-level properties without nesting
-	if err := validateElicitSchema(req.Params.RequestedSchema); err != nil {
+	schema, err := validateElicitSchema(req.Params.RequestedSchema)
+	if err != nil {
 		return nil, jsonrpc2.NewError(CodeInvalidParams, err.Error())
 	}
 
@@ -305,11 +306,11 @@ func (c *Client) elicit(ctx context.Context, req *ElicitRequest) (*ElicitResult,
 	}
 
 	// Validate elicitation result content against requested schema
-	if req.Params.RequestedSchema != nil && res.Content != nil {
+	if schema != nil && res.Content != nil {
 		// TODO: is this the correct behavior if validation fails?
 		// It isn't the *server's* params that are invalid, so why would we return
 		// this code to the server?
-		resolved, err := req.Params.RequestedSchema.Resolve(nil)
+		resolved, err := schema.Resolve(nil)
 		if err != nil {
 			return nil, jsonrpc2.NewError(CodeInvalidParams, fmt.Sprintf("failed to resolve requested schema: %v", err))
 		}
@@ -324,14 +325,23 @@ func (c *Client) elicit(ctx context.Context, req *ElicitRequest) (*ElicitResult,
 
 // validateElicitSchema validates that the schema conforms to MCP elicitation schema requirements.
 // Per the MCP specification, elicitation schemas are limited to flat objects with primitive properties only.
-func validateElicitSchema(schema *jsonschema.Schema) error {
-	if schema == nil {
-		return nil // nil schema is allowed
+func validateElicitSchema(wireSchema any) (*jsonschema.Schema, error) {
+	if wireSchema == nil {
+		return nil, nil // nil schema is allowed
+	}
+
+	data, err := json.Marshal(wireSchema)
+	if err != nil {
+		return nil, err
+	}
+	var schema *jsonschema.Schema
+	if err := json.Unmarshal(data, &schema); err != nil {
+		return nil, err
 	}
 
 	// The root schema must be of type "object" if specified
 	if schema.Type != "" && schema.Type != "object" {
-		return fmt.Errorf("elicit schema must be of type 'object', got %q", schema.Type)
+		return nil, fmt.Errorf("elicit schema must be of type 'object', got %q", schema.Type)
 	}
 
 	// Check if the schema has properties
@@ -342,12 +352,12 @@ func validateElicitSchema(schema *jsonschema.Schema) error {
 			}
 
 			if err := validateElicitProperty(propName, propSchema); err != nil {
-				return err
+				return nil, err
 			}
 		}
 	}
 
-	return nil
+	return schema, nil
 }
 
 // validateElicitProperty validates a single property in an elicitation schema.

mcp/client_list_test.go

@@ -6,6 +6,7 @@ package mcp_test
 
 import (
 	"context"
+	"encoding/json"
 	"iter"
 	"log"
 	"testing"
@@ -41,7 +42,13 @@ func TestList(t *testing.T) {
 			if err != nil {
 				t.Fatal(err)
 			}
-			tt.InputSchema = is
+			data, err := json.Marshal(is)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if err := json.Unmarshal(data, &tt.InputSchema); err != nil {
+				t.Fatal(err)
+			}
 			wantTools = append(wantTools, tt)
 		}
 		t.Run("list", func(t *testing.T) {

mcp/protocol.go

@@ -13,8 +13,6 @@ package mcp
 import (
 	"encoding/json"
 	"fmt"
-
-	"github.com/google/jsonschema-go/jsonschema"
 )
 
 // Optional annotations for the client. The client can use annotations to inform
@@ -914,13 +912,23 @@ type Tool struct {
 	// tools. It can be thought of like a "hint" to the model.
 	Description string `json:"description,omitempty"`
 	// A JSON Schema object defining the expected parameters for the tool.
-	InputSchema *jsonschema.Schema `json:"inputSchema"`
+	//
+	// From the server, this field may be set to any value that can marshal to
+	// JSON.
+	// From the client, this field will use the default JSON marshaling (a
+	// map[string]any).
+	InputSchema any `json:"inputSchema"`
 	// Intended for programmatic or logical use, but used as a display name in past
 	// specs or fallback (if title isn't present).
 	Name string `json:"name"`
 	// An optional JSON Schema object defining the structure of the tool's output
 	// returned in the structuredContent field of a CallToolResult.
-	OutputSchema *jsonschema.Schema `json:"outputSchema,omitempty"`
+	//
+	// From the server, this field may be set to any value that can marshal to
+	// JSON.
+	// From the client, this field will use the default JSON marshaling (a
+	// map[string]any).
+	OutputSchema any `json:"outputSchema,omitempty"`
 	// Intended for UI and end-user contexts — optimized to be human-readable and
 	// easily understood, even by those unfamiliar with domain-specific terminology.
 	// If not provided, Annotations.Title should be used for display if present,
@@ -1022,9 +1030,15 @@ type ElicitParams struct {
 	Meta `json:"_meta,omitempty"`
 	// The message to present to the user.
 	Message string `json:"message"`
-	// A restricted subset of JSON Schema.
+	// A JSON schema object defining the requested elicitation schema.
+	//
+	// From the server, this field may be set to any value that can marshal to
+	// JSON.
+	// From the client, this field will use the default JSON marshaling (a
+	// map[string]any).
+	//
 	// Only top-level properties are allowed, without nesting.
-	RequestedSchema *jsonschema.Schema `json:"requestedSchema"`
+	RequestedSchema any `json:"requestedSchema"`
 }
 
 func (x *ElicitParams) isParams() {}

mcp/server.go

@@ -189,10 +189,10 @@ func (s *Server) AddTool(t *Tool, h ToolHandler) {
 		// discovered until runtime, when the LLM sent bad data.
 		panic(fmt.Errorf("AddTool %q: missing input schema", t.Name))
 	}
-	if t.InputSchema.Type != "object" {
+	if s, ok := t.InputSchema.(*jsonschema.Schema); ok && s.Type != "object" {
 		panic(fmt.Errorf(`AddTool %q: input schema must have type "object"`, t.Name))
 	}
-	if t.OutputSchema != nil && t.OutputSchema.Type != "object" {
+	if s, ok := t.OutputSchema.(*jsonschema.Schema); ok && s.Type != "object" {
 		panic(fmt.Errorf(`AddTool %q: output schema must have type "object"`, t.Name))
 	}
 	st := &serverTool{tool: t, handler: h}
@@ -331,20 +331,32 @@ func toolForErr[In, Out any](t *Tool, h ToolHandlerFor[In, Out]) (*Tool, ToolHan
 //
 // TODO(rfindley): we really shouldn't ever return 'null' results. Maybe we
 // should have a jsonschema.Zero(schema) helper?
-func setSchema[T any](sfield **jsonschema.Schema, rfield **jsonschema.Resolved) (zero any, err error) {
+func setSchema[T any](sfield *any, rfield **jsonschema.Resolved) (zero any, err error) {
+	var internalSchema *jsonschema.Schema
 	if *sfield == nil {
 		rt := reflect.TypeFor[T]()
 		if rt.Kind() == reflect.Pointer {
 			rt = rt.Elem()
 			zero = reflect.Zero(rt).Interface()
 		}
 		// TODO: we should be able to pass nil opts here.
-		*sfield, err = jsonschema.ForType(rt, &jsonschema.ForOptions{})
+		internalSchema, err = jsonschema.ForType(rt, &jsonschema.ForOptions{})
+		if err == nil {
+			*sfield = internalSchema
+		}
+	} else {
+		data, err := json.Marshal(*sfield)
+		if err != nil {
+			return zero, err
+		}
+		if err := json.Unmarshal(data, &internalSchema); err != nil {
+			return zero, err
+		}
 	}
 	if err != nil {
 		return zero, err
 	}
-	*rfield, err = (*sfield).Resolve(&jsonschema.ResolveOptions{ValidateDefaults: true})
+	*rfield, err = internalSchema.Resolve(&jsonschema.ResolveOptions{ValidateDefaults: true})
 	return zero, err
 }
 

mcp/server_test.go

@@ -492,7 +492,7 @@ func TestAddTool(t *testing.T) {
 
 type schema = jsonschema.Schema
 
-func testToolForSchema[In, Out any](t *testing.T, tool *Tool, in string, out Out, wantIn, wantOut *schema, wantErrContaining string) {
+func testToolForSchema[In, Out any](t *testing.T, tool *Tool, in string, out Out, wantIn, wantOut any, wantErrContaining string) {
 	t.Helper()
 	th := func(context.Context, *CallToolRequest, In) (*CallToolResult, Out, error) {
 		return nil, out, nil

mcp/tool_example_test.go

@@ -68,7 +68,7 @@ func ExampleAddTool_customMarshalling() {
 	}
 	// Output:
 	// my_tool {
-	// 	"type": "object",
+	// 	"additionalProperties": false,
 	// 	"properties": {
 	// 		"end": {
 	// 			"type": "string"
@@ -80,7 +80,7 @@ func ExampleAddTool_customMarshalling() {
 	// 			"type": "string"
 	// 		}
 	// 	},
-	// 	"additionalProperties": false
+	// 	"type": "object"
 	// }
 }
 
@@ -200,9 +200,9 @@ func ExampleAddTool_complexSchema() {
 		}
 		// Formatting the entire schemas would be too much output.
 		// Just check that our customizations were effective.
-		fmt.Println("max days:", *t.InputSchema.Properties["days"].Maximum)
-		fmt.Println("max confidence:", *t.OutputSchema.Properties["confidence"].Maximum)
-		fmt.Println("weather types:", t.OutputSchema.Properties["dailyForecast"].Items.Properties["type"].Enum)
+		fmt.Println("max days:", jsonPath(t.InputSchema, "properties", "days", "maximum"))
+		fmt.Println("max confidence:", jsonPath(t.OutputSchema, "properties", "confidence", "maximum"))
+		fmt.Println("weather types:", jsonPath(t.OutputSchema, "properties", "dailyForecast", "items", "properties", "type", "enum"))
 	}
 	// Output:
 	// max days: 10
@@ -218,3 +218,10 @@ func connect(ctx context.Context, server *mcp.Server) (*mcp.ClientSession, error
 	client := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "v0.0.1"}, nil)
 	return client.Connect(ctx, t2, nil)
 }
+
+func jsonPath(s any, path ...string) any {
+	if len(path) == 0 {
+		return s
+	}
+	return jsonPath(s.(map[string]any)[path[0]], path[1:]...)
+}