mcp: support elicitation defaults (#644)

Fixes #617

---------

Co-authored-by: Markus Kusano <kusano@google.com>

Author: markus-kusano

mcp/client.go

@@ -321,10 +321,13 @@ func (c *Client) elicit(ctx context.Context, req *ElicitRequest) (*ElicitResult,
 		if err != nil {
 			return nil, jsonrpc2.NewError(codeInvalidParams, fmt.Sprintf("failed to resolve requested schema: %v", err))
 		}
-
 		if err := resolved.Validate(res.Content); err != nil {
 			return nil, jsonrpc2.NewError(codeInvalidParams, fmt.Sprintf("elicitation result content does not match requested schema: %v", err))
 		}
+		err = resolved.ApplyDefaults(&res.Content)
+		if err != nil {
+			return nil, jsonrpc2.NewError(codeInvalidParams, fmt.Sprintf("failed to apply schema defalts to elicitation result: %v", err))
+		}
 	}
 
 	return res, nil
@@ -341,6 +344,9 @@ func validateElicitSchema(wireSchema any) (*jsonschema.Schema, error) {
 	if err := remarshal(wireSchema, &schema); err != nil {
 		return nil, err
 	}
+	if schema == nil {
+		return nil, nil
+	}
 
 	// The root schema must be of type "object" if specified
 	if schema.Type != "" && schema.Type != "object" {
@@ -369,7 +375,6 @@ func validateElicitProperty(propName string, propSchema *jsonschema.Schema) erro
 	if len(propSchema.Properties) > 0 {
 		return fmt.Errorf("elicit schema property %q contains nested properties, only primitive properties are allowed", propName)
 	}
-
 	// Validate based on the property type - only primitives are supported
 	switch propSchema.Type {
 	case "string":
@@ -439,7 +444,7 @@ func validateElicitStringProperty(propName string, propSchema *jsonschema.Schema
 		}
 	}
 
-	return nil
+	return validateDefaultProperty[string](propName, propSchema)
 }
 
 // validateElicitNumberProperty validates number and integer-type properties.
@@ -450,19 +455,28 @@ func validateElicitNumberProperty(propName string, propSchema *jsonschema.Schema
 		}
 	}
 
+	intDefaultError := validateDefaultProperty[int](propName, propSchema)
+	floatDefaultError := validateDefaultProperty[float64](propName, propSchema)
+	if intDefaultError != nil && floatDefaultError != nil {
+		return fmt.Errorf("elicit schema property %q has default value that cannot be interpreted as an int or float", propName)
+	}
+
 	return nil
 }
 
 // validateElicitBooleanProperty validates boolean-type properties.
 func validateElicitBooleanProperty(propName string, propSchema *jsonschema.Schema) error {
-	// Validate default value if specified - must be a valid boolean
+	return validateDefaultProperty[bool](propName, propSchema)
+}
+
+func validateDefaultProperty[T any](propName string, propSchema *jsonschema.Schema) error {
+	// Validate default value if specified - must be a valid T
 	if propSchema.Default != nil {
-		var defaultValue bool
+		var defaultValue T
 		if err := json.Unmarshal(propSchema.Default, &defaultValue); err != nil {
-			return fmt.Errorf("elicit schema property %q has invalid default value, must be a boolean: %v", propName, err)
+			return fmt.Errorf("elicit schema property %q has invalid default value, must be a %T: %v", propName, defaultValue, err)
 		}
 	}
-
 	return nil
 }
 

mcp/mcp_test.go

@@ -1254,7 +1254,37 @@ func TestElicitationSchemaValidation(t *testing.T) {
 					"enabled": {Type: "boolean", Default: json.RawMessage(`"not-a-boolean"`)},
 				},
 			},
-			expectedError: "elicit schema property \"enabled\" has invalid default value, must be a boolean",
+			expectedError: "elicit schema property \"enabled\" has invalid default value, must be a bool",
+		},
+		{
+			name: "string with invalid default",
+			schema: &jsonschema.Schema{
+				Type: "object",
+				Properties: map[string]*jsonschema.Schema{
+					"enabled": {Type: "string", Default: json.RawMessage("true")},
+				},
+			},
+			expectedError: "elicit schema property \"enabled\" has invalid default value, must be a string",
+		},
+		{
+			name: "integer with invalid default",
+			schema: &jsonschema.Schema{
+				Type: "object",
+				Properties: map[string]*jsonschema.Schema{
+					"enabled": {Type: "integer", Default: json.RawMessage("true")},
+				},
+			},
+			expectedError: "elicit schema property \"enabled\" has default value that cannot be interpreted as an int or float",
+		},
+		{
+			name: "number with invalid default",
+			schema: &jsonschema.Schema{
+				Type: "object",
+				Properties: map[string]*jsonschema.Schema{
+					"enabled": {Type: "number", Default: json.RawMessage("true")},
+				},
+			},
+			expectedError: "elicit schema property \"enabled\" has default value that cannot be interpreted as an int or float",
 		},
 		{
 			name: "enum with mismatched enumNames length",
@@ -1459,6 +1489,100 @@ func TestElicitationCapabilityDeclaration(t *testing.T) {
 	})
 }
 
+func TestElicitationDefaultValues(t *testing.T) {
+	ctx := context.Background()
+	ct, st := NewInMemoryTransports()
+
+	s := NewServer(testImpl, nil)
+	ss, err := s.Connect(ctx, st, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer ss.Close()
+
+	c := NewClient(testImpl, &ClientOptions{
+		ElicitationHandler: func(context.Context, *ElicitRequest) (*ElicitResult, error) {
+			return &ElicitResult{Action: "accept", Content: map[string]any{"default": "response"}}, nil
+		},
+	})
+	cs, err := c.Connect(ctx, ct, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer cs.Close()
+
+	testcases := []struct {
+		name     string
+		schema   *jsonschema.Schema
+		expected map[string]any
+	}{
+		{
+			name: "boolean with default",
+			schema: &jsonschema.Schema{
+				Type: "object",
+				Properties: map[string]*jsonschema.Schema{
+					"key": {Type: "boolean", Default: json.RawMessage("true")},
+				},
+			},
+			expected: map[string]any{"key": true, "default": "response"},
+		},
+		{
+			name: "string with default",
+			schema: &jsonschema.Schema{
+				Type: "object",
+				Properties: map[string]*jsonschema.Schema{
+					"key": {Type: "string", Default: json.RawMessage("\"potato\"")},
+				},
+			},
+			expected: map[string]any{"key": "potato", "default": "response"},
+		},
+		{
+			name: "integer with default",
+			schema: &jsonschema.Schema{
+				Type: "object",
+				Properties: map[string]*jsonschema.Schema{
+					"key": {Type: "integer", Default: json.RawMessage("123")},
+				},
+			},
+			expected: map[string]any{"key": float64(123), "default": "response"},
+		},
+		{
+			name: "number with default",
+			schema: &jsonschema.Schema{
+				Type: "object",
+				Properties: map[string]*jsonschema.Schema{
+					"key": {Type: "number", Default: json.RawMessage("89.7")},
+				},
+			},
+			expected: map[string]any{"key": float64(89.7), "default": "response"},
+		},
+		{
+			name: "enum with default",
+			schema: &jsonschema.Schema{
+				Type: "object",
+				Properties: map[string]*jsonschema.Schema{
+					"key": {Type: "string", Enum: []any{"one", "two"}, Default: json.RawMessage("\"one\"")},
+				},
+			},
+			expected: map[string]any{"key": "one", "default": "response"},
+		},
+	}
+	for _, tc := range testcases {
+		t.Run(tc.name, func(t *testing.T) {
+			res, err := ss.Elicit(ctx, &ElicitParams{
+				Message:         "Test schema with defaults: " + tc.name,
+				RequestedSchema: tc.schema,
+			})
+			if err != nil {
+				t.Fatalf("expected no error for default schema %q, got: %v", tc.name, err)
+			}
+			if diff := cmp.Diff(tc.expected, res.Content); diff != "" {
+				t.Errorf("%s: did not get expected value, -want +got:\n%s", tc.name, diff)
+			}
+		})
+	}
+}
+
 func TestKeepAliveFailure(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()

mcp/server.go

@@ -1018,7 +1018,37 @@ func (ss *ServerSession) Elicit(ctx context.Context, params *ElicitParams) (*Eli
 	if err := ss.checkInitialized(methodElicit); err != nil {
 		return nil, err
 	}
-	return handleSend[*ElicitResult](ctx, methodElicit, newServerRequest(ss, orZero[Params](params)))
+
+	res, err := handleSend[*ElicitResult](ctx, methodElicit, newServerRequest(ss, orZero[Params](params)))
+	if err != nil {
+		return nil, err
+	}
+
+	if params.RequestedSchema == nil {
+		return res, nil
+	}
+	schema, err := validateElicitSchema(params.RequestedSchema)
+	if err != nil {
+		return nil, err
+	}
+	if schema == nil {
+		return res, nil
+	}
+
+	resolved, err := schema.Resolve(nil)
+	if err != nil {
+		fmt.Printf("  resolve err: %s", err)
+		return nil, err
+	}
+	if err := resolved.Validate(res.Content); err != nil {
+		return nil, fmt.Errorf("elicitation result content does not match requested schema: %v", err)
+	}
+	err = resolved.ApplyDefaults(&res.Content)
+	if err != nil {
+		return nil, fmt.Errorf("failed to apply schema defalts to elicitation result: %v", err)
+	}
+
+	return res, nil
 }
 
 // Log sends a log message to the client.