diff --git a/internal/common/config/mcp.go b/internal/common/config/mcp.go index 541cc480..0fe52f39 100644 --- a/internal/common/config/mcp.go +++ b/internal/common/config/mcp.go @@ -62,6 +62,7 @@ type ( ToolConfig struct { Name string `json:"name" yaml:"name"` + Title string `json:"title,omitempty" yaml:"title,omitempty"` Description string `json:"description,omitempty" yaml:"description,omitempty"` Method string `json:"method" yaml:"method"` Endpoint string `json:"endpoint" yaml:"endpoint"` @@ -71,6 +72,7 @@ type ( RequestBody string `json:"requestBody" yaml:"requestBody"` ResponseBody string `json:"responseBody" yaml:"responseBody"` InputSchema map[string]any `json:"inputSchema,omitempty" yaml:"inputSchema,omitempty"` + OutputSchema map[string]any `json:"outputSchema,omitempty" yaml:"outputSchema,omitempty"` } MCPServerConfig struct { @@ -186,8 +188,9 @@ func (t *ToolConfig) ToToolSchema() mcp.ToolSchema { } } - return mcp.ToolSchema{ + result := mcp.ToolSchema{ Name: t.Name, + Title: t.Title, Description: t.Description, InputSchema: mcp.ToolInputSchema{ Type: "object", @@ -195,6 +198,16 @@ func (t *ToolConfig) ToToolSchema() mcp.ToolSchema { Required: required, }, } + + // Add output schema if provided + if t.OutputSchema != nil { + result.OutputSchema = &mcp.ToolInputSchema{ + Type: "object", + Properties: t.OutputSchema, + } + } + + return result } // ToPromptSchema converts a PromptConfig to a PromptSchema diff --git a/internal/common/config/mcp_test.go b/internal/common/config/mcp_test.go new file mode 100644 index 00000000..80067e7e --- /dev/null +++ b/internal/common/config/mcp_test.go @@ -0,0 +1,139 @@ +package config + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestToolConfigToToolSchema(t *testing.T) { + // Test ToolConfig with Title and OutputSchema + toolConfig := ToolConfig{ + Name: "get_weather_data", + Title: "Weather Data Retriever", + Description: "Get current weather data for a location", + Method: "GET", + Endpoint: "https://api.weather.com/data", + Args: []ArgConfig{ + { + Name: "location", + Type: "string", + Description: "City name or zip code", + Required: true, + }, + }, + OutputSchema: map[string]any{ + "temperature": map[string]any{ + "type": "number", + "description": "Temperature in celsius", + }, + "conditions": map[string]any{ + "type": "string", + "description": "Weather conditions description", + }, + "humidity": map[string]any{ + "type": "number", + "description": "Humidity percentage", + }, + }, + } + + toolSchema := toolConfig.ToToolSchema() + + // Verify basic fields + assert.Equal(t, "get_weather_data", toolSchema.Name) + assert.Equal(t, "Weather Data Retriever", toolSchema.Title) + assert.Equal(t, "Get current weather data for a location", toolSchema.Description) + + // Verify InputSchema + assert.Equal(t, "object", toolSchema.InputSchema.Type) + assert.Len(t, toolSchema.InputSchema.Properties, 1) + assert.Contains(t, toolSchema.InputSchema.Properties, "location") + assert.Equal(t, []string{"location"}, toolSchema.InputSchema.Required) + + // Verify OutputSchema + assert.NotNil(t, toolSchema.OutputSchema) + assert.Equal(t, "object", toolSchema.OutputSchema.Type) + assert.Len(t, toolSchema.OutputSchema.Properties, 3) + assert.Contains(t, toolSchema.OutputSchema.Properties, "temperature") + assert.Contains(t, toolSchema.OutputSchema.Properties, "conditions") + assert.Contains(t, toolSchema.OutputSchema.Properties, "humidity") +} + +func TestToolConfigToToolSchemaWithoutOutputSchema(t *testing.T) { + // Test ToolConfig without Title and OutputSchema (backward compatibility) + toolConfig := ToolConfig{ + Name: "simple_tool", + Description: "A simple tool", + Method: "POST", + Endpoint: "https://api.example.com/tool", + Args: []ArgConfig{ + { + Name: "param1", + Type: "string", + Required: false, + }, + }, + } + + toolSchema := toolConfig.ToToolSchema() + + // Verify basic fields + assert.Equal(t, "simple_tool", toolSchema.Name) + assert.Equal(t, "", toolSchema.Title) // Should be empty when not provided + assert.Equal(t, "A simple tool", toolSchema.Description) + + // Verify InputSchema + assert.Equal(t, "object", toolSchema.InputSchema.Type) + assert.Len(t, toolSchema.InputSchema.Properties, 1) + assert.Contains(t, toolSchema.InputSchema.Properties, "param1") + assert.Empty(t, toolSchema.InputSchema.Required) // No required args + + // Verify OutputSchema is nil + assert.Nil(t, toolSchema.OutputSchema) +} + +func TestToolConfigToToolSchemaWithInputSchema(t *testing.T) { + // Test ToolConfig that combines Args with explicit InputSchema + toolConfig := ToolConfig{ + Name: "mixed_tool", + Title: "Mixed Tool", + Description: "A tool with both Args and InputSchema", + Method: "POST", + Endpoint: "https://api.example.com/mixed", + Args: []ArgConfig{ + { + Name: "arg1", + Type: "string", + Description: "First argument", + Required: true, + }, + }, + InputSchema: map[string]any{ + "custom_field": map[string]any{ + "type": "number", + "description": "Custom field from InputSchema", + }, + }, + OutputSchema: map[string]any{ + "result": map[string]any{ + "type": "string", + }, + }, + } + + toolSchema := toolConfig.ToToolSchema() + + // Verify both arg-generated and explicit InputSchema properties are present + assert.Equal(t, "object", toolSchema.InputSchema.Type) + assert.Len(t, toolSchema.InputSchema.Properties, 2) + assert.Contains(t, toolSchema.InputSchema.Properties, "arg1") + assert.Contains(t, toolSchema.InputSchema.Properties, "custom_field") + assert.Equal(t, []string{"arg1"}, toolSchema.InputSchema.Required) + + // Verify OutputSchema + assert.NotNil(t, toolSchema.OutputSchema) + assert.Equal(t, "object", toolSchema.OutputSchema.Type) + assert.Len(t, toolSchema.OutputSchema.Properties, 1) + assert.Contains(t, toolSchema.OutputSchema.Properties, "result") +} \ No newline at end of file diff --git a/internal/core/mcpproxy/sse.go b/internal/core/mcpproxy/sse.go index b55533df..5d99ffc5 100644 --- a/internal/core/mcpproxy/sse.go +++ b/internal/core/mcpproxy/sse.go @@ -129,9 +129,11 @@ func (t *SSETransport) FetchTools(ctx context.Context) ([]mcp.ToolSchema, error) } tools[i] = mcp.ToolSchema{ - Name: schema.Name, - Description: schema.Description, - InputSchema: inputSchema, + Name: schema.Name, + Title: "", // Will be supported when mcpgo library updates + Description: schema.Description, + InputSchema: inputSchema, + OutputSchema: nil, // Will be supported when mcpgo library updates } } diff --git a/internal/core/mcpproxy/stdio.go b/internal/core/mcpproxy/stdio.go index 495e6b28..ebf7904d 100644 --- a/internal/core/mcpproxy/stdio.go +++ b/internal/core/mcpproxy/stdio.go @@ -145,9 +145,11 @@ func (t *StdioTransport) FetchTools(ctx context.Context) ([]mcp.ToolSchema, erro } tools[i] = mcp.ToolSchema{ - Name: schema.Name, - Description: schema.Description, - InputSchema: inputSchema, + Name: schema.Name, + Title: "", // Will be supported when mcpgo library updates + Description: schema.Description, + InputSchema: inputSchema, + OutputSchema: nil, // Will be supported when mcpgo library updates } } diff --git a/internal/core/mcpproxy/streamable.go b/internal/core/mcpproxy/streamable.go index 9942a004..f7fdc9d9 100644 --- a/internal/core/mcpproxy/streamable.go +++ b/internal/core/mcpproxy/streamable.go @@ -129,9 +129,11 @@ func (t *StreamableTransport) FetchTools(ctx context.Context) ([]mcp.ToolSchema, } tools[i] = mcp.ToolSchema{ - Name: schema.Name, - Description: schema.Description, - InputSchema: inputSchema, + Name: schema.Name, + Title: "", // Will be supported when mcpgo library updates + Description: schema.Description, + InputSchema: inputSchema, + OutputSchema: nil, // Will be supported when mcpgo library updates } } diff --git a/internal/core/sse.go b/internal/core/sse.go index f617378b..91e969b8 100644 --- a/internal/core/sse.go +++ b/internal/core/sse.go @@ -364,9 +364,11 @@ func (s *Server) handlePostMessage(c *gin.Context, conn session.Connection) { toolSchemas := make([]mcp.ToolSchema, len(tools)) for i, tool := range tools { toolSchemas[i] = mcp.ToolSchema{ - Name: tool.Name, - Description: tool.Description, - InputSchema: tool.InputSchema, + Name: tool.Name, + Title: tool.Title, + Description: tool.Description, + InputSchema: tool.InputSchema, + OutputSchema: tool.OutputSchema, } } diff --git a/pkg/mcp/server_types.go b/pkg/mcp/server_types.go index 1723cd93..f3d122b6 100644 --- a/pkg/mcp/server_types.go +++ b/pkg/mcp/server_types.go @@ -52,10 +52,14 @@ type ( ToolSchema struct { // The name of the tool Name string `json:"name"` + // A human-readable title for the tool (optional) + Title string `json:"title,omitempty"` // A human-readable description of the tool Description string `json:"description"` // A JSON Schema object defining the expected parameters for the tool InputSchema ToolInputSchema `json:"inputSchema"` + // A JSON Schema object defining the expected output of the tool (optional) + OutputSchema *ToolInputSchema `json:"outputSchema,omitempty"` } ToolInputSchema struct { diff --git a/pkg/mcp/server_types_test.go b/pkg/mcp/server_types_test.go new file mode 100644 index 00000000..2ea64618 --- /dev/null +++ b/pkg/mcp/server_types_test.go @@ -0,0 +1,149 @@ +package mcp + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestToolSchemaWithOutputSchema(t *testing.T) { + // Test ToolSchema with new OutputSchema field + outputSchema := &ToolInputSchema{ + Type: "object", + Properties: map[string]any{ + "temperature": map[string]any{ + "type": "number", + "description": "Temperature in celsius", + }, + "conditions": map[string]any{ + "type": "string", + "description": "Weather conditions description", + }, + "humidity": map[string]any{ + "type": "number", + "description": "Humidity percentage", + }, + }, + Required: []string{"temperature", "conditions", "humidity"}, + } + + toolSchema := ToolSchema{ + Name: "get_weather_data", + Title: "Weather Data Retriever", + Description: "Get current weather data for a location", + InputSchema: ToolInputSchema{ + Type: "object", + Properties: map[string]any{ + "location": map[string]any{ + "type": "string", + "description": "City name or zip code", + }, + }, + Required: []string{"location"}, + }, + OutputSchema: outputSchema, + } + + // Test JSON serialization + jsonData, err := json.Marshal(toolSchema) + assert.NoError(t, err) + + // Test JSON deserialization + var unmarshaled ToolSchema + err = json.Unmarshal(jsonData, &unmarshaled) + assert.NoError(t, err) + + // Verify fields are preserved + assert.Equal(t, "get_weather_data", unmarshaled.Name) + assert.Equal(t, "Weather Data Retriever", unmarshaled.Title) + assert.Equal(t, "Get current weather data for a location", unmarshaled.Description) + assert.NotNil(t, unmarshaled.OutputSchema) + assert.Equal(t, "object", unmarshaled.OutputSchema.Type) + assert.Len(t, unmarshaled.OutputSchema.Properties, 3) + assert.Contains(t, unmarshaled.OutputSchema.Properties, "temperature") + assert.Contains(t, unmarshaled.OutputSchema.Properties, "conditions") + assert.Contains(t, unmarshaled.OutputSchema.Properties, "humidity") + assert.Equal(t, []string{"temperature", "conditions", "humidity"}, unmarshaled.OutputSchema.Required) +} + +func TestToolSchemaWithoutOutputSchema(t *testing.T) { + // Test ToolSchema without OutputSchema (backward compatibility) + toolSchema := ToolSchema{ + Name: "simple_tool", + Description: "A simple tool without output schema", + InputSchema: ToolInputSchema{ + Type: "object", + Properties: map[string]any{ + "param": map[string]any{ + "type": "string", + }, + }, + }, + } + + // Test JSON serialization + jsonData, err := json.Marshal(toolSchema) + assert.NoError(t, err) + + // Test JSON deserialization + var unmarshaled ToolSchema + err = json.Unmarshal(jsonData, &unmarshaled) + assert.NoError(t, err) + + // Verify fields are preserved + assert.Equal(t, "simple_tool", unmarshaled.Name) + assert.Equal(t, "", unmarshaled.Title) // Should be empty when not provided + assert.Equal(t, "A simple tool without output schema", unmarshaled.Description) + assert.Nil(t, unmarshaled.OutputSchema) // Should be nil when not provided +} + +func TestToolSchemaJSONFormat(t *testing.T) { + // Test that JSON output matches expected format from the issue + toolSchema := ToolSchema{ + Name: "get_weather_data", + Title: "Weather Data Retriever", + Description: "Get current weather data for a location", + InputSchema: ToolInputSchema{ + Type: "object", + Properties: map[string]any{ + "location": map[string]any{ + "type": "string", + "description": "City name or zip code", + }, + }, + Required: []string{"location"}, + }, + OutputSchema: &ToolInputSchema{ + Type: "object", + Properties: map[string]any{ + "temperature": map[string]any{ + "type": "number", + "description": "Temperature in celsius", + }, + "conditions": map[string]any{ + "type": "string", + "description": "Weather conditions description", + }, + "humidity": map[string]any{ + "type": "number", + "description": "Humidity percentage", + }, + }, + Required: []string{"temperature", "conditions", "humidity"}, + }, + } + + jsonData, err := json.Marshal(toolSchema) + assert.NoError(t, err) + + // Verify that the JSON contains the expected fields + jsonStr := string(jsonData) + assert.Contains(t, jsonStr, `"name":"get_weather_data"`) + assert.Contains(t, jsonStr, `"title":"Weather Data Retriever"`) + assert.Contains(t, jsonStr, `"inputSchema"`) + assert.Contains(t, jsonStr, `"outputSchema"`) + assert.Contains(t, jsonStr, `"temperature"`) + assert.Contains(t, jsonStr, `"conditions"`) + assert.Contains(t, jsonStr, `"humidity"`) +} \ No newline at end of file