From fbac167e46a96b26767156a00fb6573994b3b903 Mon Sep 17 00:00:00 2001 From: Zheng Li <> Date: Tue, 4 Mar 2025 15:44:18 +0900 Subject: [PATCH 1/3] add UnmarshalJSON to Tool --- mcp/tools.go | 33 +++++++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/mcp/tools.go b/mcp/tools.go index a0278e263..94eb24786 100644 --- a/mcp/tools.go +++ b/mcp/tools.go @@ -103,6 +103,39 @@ func (t Tool) MarshalJSON() ([]byte, error) { return json.Marshal(m) } + +// UnmarshalJSON implements the json.Unmarshaler interface for Tool. +// It handles unmarshaling InputSchema while supporting raw schema formats. +func (t *Tool) UnmarshalJSON(b []byte) error { + // Temporary structure for decoding + var raw struct { + Name string `json:"name"` + Description string `json:"description"` + InputSchema json.RawMessage `json:"inputSchema"` + } + + // Unmarshal into the temporary structure + if err := json.Unmarshal(b, &raw); err != nil { + return err + } + + // Assign name and description + t.Name = raw.Name + t.Description = raw.Description + + // Try to unmarshal InputSchema into structured format + var schema ToolInputSchema + if err := json.Unmarshal(raw.InputSchema, &schema); err == nil { + // Successfully parsed structured schema + t.InputSchema = schema + t.RawInputSchema = nil + } else { + return err + } + + return nil +} + type ToolInputSchema struct { Type string `json:"type"` Properties map[string]interface{} `json:"properties,omitempty"` From 2b5060f4d5ffb8583db156f0a748bec3b39c968f Mon Sep 17 00:00:00 2001 From: Zheng Li <> Date: Tue, 4 Mar 2025 15:59:04 +0900 Subject: [PATCH 2/3] https://github.com/mark3labs/mcp-go/pull/35#discussion_r1978726095 --- mcp/tools.go | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/mcp/tools.go b/mcp/tools.go index 94eb24786..736dcc882 100644 --- a/mcp/tools.go +++ b/mcp/tools.go @@ -125,12 +125,14 @@ func (t *Tool) UnmarshalJSON(b []byte) error { // Try to unmarshal InputSchema into structured format var schema ToolInputSchema - if err := json.Unmarshal(raw.InputSchema, &schema); err == nil { + if err := json.Unmarshal(raw.InputSchema, &schema); err == nil && schema.Type != "" { // Successfully parsed structured schema t.InputSchema = schema t.RawInputSchema = nil } else { - return err + // Failed to parse structured schema, use raw schema + t.InputSchema = ToolInputSchema{} + t.RawInputSchema = raw.InputSchema } return nil From 045451ff047129f7994863c059c3a5680b0d5b16 Mon Sep 17 00:00:00 2001 From: Zheng Li <> Date: Tue, 4 Mar 2025 17:01:36 +0900 Subject: [PATCH 3/3] add test --- mcp/tools.go | 37 +------------------------ mcp/tools_test.go | 69 +++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 70 insertions(+), 36 deletions(-) diff --git a/mcp/tools.go b/mcp/tools.go index 736dcc882..83b6582fd 100644 --- a/mcp/tools.go +++ b/mcp/tools.go @@ -72,7 +72,7 @@ type Tool struct { // A human-readable description of the tool. Description string `json:"description,omitempty"` // A JSON Schema object defining the expected parameters for the tool. - InputSchema ToolInputSchema `json:"-"` // Hide this from JSON marshaling + InputSchema ToolInputSchema `json:"inputSchema"` // Alternative to InputSchema - allows arbitrary JSON Schema to be provided RawInputSchema json.RawMessage `json:"-"` // Hide this from JSON marshaling } @@ -103,41 +103,6 @@ func (t Tool) MarshalJSON() ([]byte, error) { return json.Marshal(m) } - -// UnmarshalJSON implements the json.Unmarshaler interface for Tool. -// It handles unmarshaling InputSchema while supporting raw schema formats. -func (t *Tool) UnmarshalJSON(b []byte) error { - // Temporary structure for decoding - var raw struct { - Name string `json:"name"` - Description string `json:"description"` - InputSchema json.RawMessage `json:"inputSchema"` - } - - // Unmarshal into the temporary structure - if err := json.Unmarshal(b, &raw); err != nil { - return err - } - - // Assign name and description - t.Name = raw.Name - t.Description = raw.Description - - // Try to unmarshal InputSchema into structured format - var schema ToolInputSchema - if err := json.Unmarshal(raw.InputSchema, &schema); err == nil && schema.Type != "" { - // Successfully parsed structured schema - t.InputSchema = schema - t.RawInputSchema = nil - } else { - // Failed to parse structured schema, use raw schema - t.InputSchema = ToolInputSchema{} - t.RawInputSchema = raw.InputSchema - } - - return nil -} - type ToolInputSchema struct { Type string `json:"type"` Properties map[string]interface{} `json:"properties,omitempty"` diff --git a/mcp/tools_test.go b/mcp/tools_test.go index 38791f3c1..25bfb7827 100644 --- a/mcp/tools_test.go +++ b/mcp/tools_test.go @@ -73,3 +73,72 @@ func TestToolWithRawSchema(t *testing.T) { assert.True(t, ok) assert.Contains(t, required, "query") } + +func TestUnmarshalToolWithRawSchema(t *testing.T) { + // Create a complex raw schema + rawSchema := json.RawMessage(`{ + "type": "object", + "properties": { + "query": {"type": "string", "description": "Search query"}, + "limit": {"type": "integer", "minimum": 1, "maximum": 50} + }, + "required": ["query"] + }`) + + // Create a tool with raw schema + tool := NewToolWithRawSchema("search-tool", "Search API", rawSchema) + + // Marshal to JSON + data, err := json.Marshal(tool) + assert.NoError(t, err) + + // Unmarshal to verify the structure + var toolUnmarshalled Tool + err = json.Unmarshal(data, &toolUnmarshalled) + assert.NoError(t, err) + + // Verify tool properties + assert.Equal(t, tool.Name, toolUnmarshalled.Name) + assert.Equal(t, tool.Description, toolUnmarshalled.Description) + + // Verify schema was properly included + assert.Equal(t, "object", toolUnmarshalled.InputSchema.Type) + assert.Contains(t, toolUnmarshalled.InputSchema.Properties, "query") + assert.Subset(t, toolUnmarshalled.InputSchema.Properties["query"], map[string]interface{}{ + "type": "string", + "description": "Search query", + }) + assert.Contains(t, toolUnmarshalled.InputSchema.Properties, "limit") + assert.Subset(t, toolUnmarshalled.InputSchema.Properties["limit"], map[string]interface{}{ + "type": "integer", + "minimum": 1.0, + "maximum": 50.0, + }) + assert.Subset(t, toolUnmarshalled.InputSchema.Required, []string{"query"}) +} + +func TestUnmarshalToolWithoutRawSchema(t *testing.T) { + // Create a tool with both schemas set + tool := NewTool("dual-schema-tool", + WithDescription("A tool with both schemas set"), + WithString("input", Description("Test input")), + ) + + data, err := json.Marshal(tool) + assert.Nil(t, err) + + // Unmarshal to verify the structure + var toolUnmarshalled Tool + err = json.Unmarshal(data, &toolUnmarshalled) + assert.NoError(t, err) + + // Verify tool properties + assert.Equal(t, tool.Name, toolUnmarshalled.Name) + assert.Equal(t, tool.Description, toolUnmarshalled.Description) + assert.Subset(t, toolUnmarshalled.InputSchema.Properties["input"], map[string]interface{}{ + "type": "string", + "description": "Test input", + }) + assert.Empty(t, toolUnmarshalled.InputSchema.Required) + assert.Empty(t, toolUnmarshalled.RawInputSchema) +} \ No newline at end of file