fix: handle either $defs or definitions field when unmarshaling ToolArgumentsSchema (#645)

Older versions of the JSON Schema spec use the "definitions" field, while newer versions use "$defs". This can be seen in the Honeycomb MCP, which uses "definitions". This commit adds support for both fields, and prefers "$defs" over "definitions" if both are present.

Author: sd2k

mcp/tools.go

@@ -653,6 +653,31 @@ func (tis ToolArgumentsSchema) MarshalJSON() ([]byte, error) {
 return json.Marshal(m)
 }
 
+// UnmarshalJSON implements the json.Unmarshaler interface for ToolArgumentsSchema.
+// It handles both "$defs" (JSON Schema 2019-09+) and "definitions" (JSON Schema draft-07)
+// by reading either field and storing it in the Defs field.
+func (tis *ToolArgumentsSchema) UnmarshalJSON(data []byte) error {
+// Use a temporary type to avoid infinite recursion
+type Alias ToolArgumentsSchema
+aux := &struct {
+Definitions map[string]any `json:"definitions,omitempty"`
+*Alias
+}{
+Alias: (*Alias)(tis),
+}
+
+if err := json.Unmarshal(data, aux); err != nil {
+return err
+}
+
+// If $defs wasn't provided but definitions was, use definitions
+if tis.Defs == nil && aux.Definitions != nil {
+tis.Defs = aux.Definitions
+}
+
+return nil
+}
+
 type ToolAnnotation struct {
 // Human-readable title for the tool
 Title string `json:"title,omitempty"`

mcp/tools_test.go

@@ -1549,3 +1549,136 @@ func TestToolMetaMarshalingOmitsWhenNil(t *testing.T) {
 // Check that _meta field is not present
 assert.NotContains(t, result, "_meta", "Tool without Meta should not include _meta field")
 }
+
+func TestToolArgumentsSchema_UnmarshalWithDefinitions(t *testing.T) {
+// Test that "definitions" (JSON Schema draft-07) is properly unmarshaled into Defs field
+jsonData := `{
+"type": "object",
+"properties": {
+"operation": {
+"$ref": "#/definitions/operation_type"
+}
+},
+"required": ["operation"],
+"definitions": {
+"operation_type": {
+"type": "string",
+"enum": ["create", "read", "update", "delete"]
+}
+}
+}`
+
+var schema ToolArgumentsSchema
+err := json.Unmarshal([]byte(jsonData), &schema)
+assert.NoError(t, err)
+
+// Verify the schema was properly unmarshaled
+assert.Equal(t, "object", schema.Type)
+assert.Contains(t, schema.Properties, "operation")
+assert.Equal(t, []string{"operation"}, schema.Required)
+
+// Most importantly: verify that "definitions" was read into Defs field
+assert.NotNil(t, schema.Defs)
+assert.Contains(t, schema.Defs, "operation_type")
+
+operationType, ok := schema.Defs["operation_type"].(map[string]any)
+assert.True(t, ok)
+assert.Equal(t, "string", operationType["type"])
+assert.NotNil(t, operationType["enum"])
+}
+
+func TestToolArgumentsSchema_UnmarshalWithDefs(t *testing.T) {
+// Test that "$defs" (JSON Schema 2019-09+) is properly unmarshaled into Defs field
+jsonData := `{
+"type": "object",
+"properties": {
+"operation": {
+"$ref": "#/$defs/operation_type"
+}
+},
+"required": ["operation"],
+"$defs": {
+"operation_type": {
+"type": "string",
+"enum": ["create", "read", "update", "delete"]
+}
+}
+}`
+
+var schema ToolArgumentsSchema
+err := json.Unmarshal([]byte(jsonData), &schema)
+assert.NoError(t, err)
+
+// Verify the schema was properly unmarshaled
+assert.Equal(t, "object", schema.Type)
+assert.Contains(t, schema.Properties, "operation")
+assert.Equal(t, []string{"operation"}, schema.Required)
+
+// Verify that "$defs" was read into Defs field
+assert.NotNil(t, schema.Defs)
+assert.Contains(t, schema.Defs, "operation_type")
+
+operationType, ok := schema.Defs["operation_type"].(map[string]any)
+assert.True(t, ok)
+assert.Equal(t, "string", operationType["type"])
+assert.NotNil(t, operationType["enum"])
+}
+
+func TestToolArgumentsSchema_UnmarshalPrefersDefs(t *testing.T) {
+// Test that if both "$defs" and "definitions" are present, "$defs" takes precedence
+jsonData := `{
+"type": "object",
+"$defs": {
+"from_defs": {
+"type": "string"
+}
+},
+"definitions": {
+"from_definitions": {
+"type": "integer"
+}
+}
+}`
+
+var schema ToolArgumentsSchema
+err := json.Unmarshal([]byte(jsonData), &schema)
+assert.NoError(t, err)
+
+// $defs should take precedence
+assert.Contains(t, schema.Defs, "from_defs")
+assert.NotContains(t, schema.Defs, "from_definitions")
+}
+
+func TestToolArgumentsSchema_MarshalRoundTrip(t *testing.T) {
+// Test that marshaling and unmarshaling preserves definitions
+original := ToolArgumentsSchema{
+Type: "object",
+Properties: map[string]any{
+"field": map[string]any{
+"$ref": "#/$defs/my_type",
+},
+},
+Required: []string{"field"},
+Defs: map[string]any{
+"my_type": map[string]any{
+"type": "string",
+"enum": []string{"a", "b", "c"},
+},
+},
+}
+
+// Marshal
+data, err := json.Marshal(original)
+assert.NoError(t, err)
+
+// Unmarshal
+var unmarshaled ToolArgumentsSchema
+err = json.Unmarshal(data, &unmarshaled)
+assert.NoError(t, err)
+
+// Verify round-trip
+assert.Equal(t, original.Type, unmarshaled.Type)
+assert.Equal(t, original.Required, unmarshaled.Required)
+assert.NotNil(t, unmarshaled.Defs)
+assert.Contains(t, unmarshaled.Defs, "my_type")
+}