diff --git a/client/http_test.go b/client/http_test.go index ad4d4ba55..23e7cf8c1 100644 --- a/client/http_test.go +++ b/client/http_test.go @@ -2,6 +2,7 @@ package client import ( "context" + "encoding/json" "fmt" "sync" "testing" @@ -10,9 +11,10 @@ import ( "github.com/mark3labs/mcp-go/client/transport" "github.com/mark3labs/mcp-go/mcp" "github.com/mark3labs/mcp-go/server" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) - func TestHTTPClient(t *testing.T) { hooks := &server.Hooks{} hooks.AddAfterCallTool(func(ctx context.Context, id any, message *mcp.CallToolRequest, result *mcp.CallToolResult) { @@ -192,6 +194,67 @@ func TestHTTPClient(t *testing.T) { }) } +func TestHTTPClient_ListTools_WithOutputSchema(t *testing.T) { + // 1. Setup Server + srv := server.NewMCPServer("test-server", "1.0.0") + + // Define a tool with a structured output type, including descriptions. + type WeatherData struct { + Temperature float64 `json:"temperature" jsonschema:"description=The temperature in Celsius."` + Conditions string `json:"conditions" jsonschema:"description=Weather conditions (e.g. Cloudy)."` + } + tool := mcp.NewTool("get_weather", mcp.WithOutputType[WeatherData]()) + srv.AddTool(tool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return nil, nil // Handler not needed for this test + }) + + // Use the dedicated test helper to create the server + httpServer := server.NewTestStreamableHTTPServer(srv) + defer httpServer.Close() + + // 2. Setup Client + // Use the correct client constructor + client, err := NewStreamableHttpClient(httpServer.URL) + require.NoError(t, err) + + // Initialize the client session before making other requests. + _, err = client.Initialize(context.Background(), mcp.InitializeRequest{ + Params: mcp.InitializeParams{ + ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, + ClientInfo: mcp.Implementation{ + Name: "test-client", + Version: "1.0.0", + }, + // Client does not need to declare tool capabilities, + // it's a server-side declaration. + Capabilities: mcp.ClientCapabilities{}, + }, + }) + require.NoError(t, err, "client not initialized") + + // 3. Client calls ListTools + result, err := client.ListTools(context.Background(), mcp.ListToolsRequest{}) + require.NoError(t, err) + require.Len(t, result.Tools, 1, "Should retrieve one tool") + + // 4. Assert on the received tool's OutputSchema + retrievedTool := result.Tools[0] + assert.Equal(t, "get_weather", retrievedTool.Name) + require.NotNil(t, retrievedTool.OutputSchema, "OutputSchema should be present") + + // Unmarshal and verify the content of the schema + var schemaData map[string]interface{} + err = json.Unmarshal(retrievedTool.OutputSchema, &schemaData) + require.NoError(t, err) + + properties := schemaData["properties"].(map[string]interface{}) + tempProp := properties["temperature"].(map[string]interface{}) + condProp := properties["conditions"].(map[string]interface{}) + + assert.Equal(t, "The temperature in Celsius.", tempProp["description"]) + assert.Equal(t, "Weather conditions (e.g. Cloudy).", condProp["description"]) +} + type SafeMap struct { mu sync.RWMutex data map[string]int diff --git a/go.mod b/go.mod index 9b9fe2d48..68ae1c1ce 100644 --- a/go.mod +++ b/go.mod @@ -4,13 +4,19 @@ go 1.23 require ( github.com/google/uuid v1.6.0 + github.com/invopop/jsonschema v0.13.0 + github.com/santhosh-tekuri/jsonschema v1.2.4 github.com/spf13/cast v1.7.1 github.com/stretchr/testify v1.9.0 github.com/yosida95/uritemplate/v3 v3.0.2 ) require ( + github.com/bahlo/generic-list-go v0.2.0 // indirect + github.com/buger/jsonparser v1.1.1 // indirect github.com/davecgh/go-spew v1.1.1 // indirect + github.com/mailru/easyjson v0.7.7 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect + github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect ) diff --git a/go.sum b/go.sum index 31ed86d18..bcc73d5aa 100644 --- a/go.sum +++ b/go.sum @@ -1,3 +1,7 @@ +github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk= +github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg= +github.com/buger/jsonparser v1.1.1 h1:2PnMjfWD7wBILjqQbt530v576A/cAbQvEW9gGIpYMUs= +github.com/buger/jsonparser v1.1.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8= @@ -6,18 +10,27 @@ github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38= github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/invopop/jsonschema v0.13.0 h1:KvpoAJWEjR3uD9Kbm2HWJmqsEaHt8lBUpd0qHcIi21E= +github.com/invopop/jsonschema v0.13.0/go.mod h1:ffZ5Km5SWWRAIN6wbDXItl95euhFz2uON45H2qjYt+0= +github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0= +github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/rogpeppe/go-internal v1.9.0 h1:73kH8U+JUqXU8lRuOHeVHaa/SZPifC7BkcraZVejAe8= github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs= +github.com/santhosh-tekuri/jsonschema v1.2.4 h1:hNhW8e7t+H1vgY+1QeEQpveR6D4+OwKPXCfD2aieJis= +github.com/santhosh-tekuri/jsonschema v1.2.4/go.mod h1:TEAUOeZSmIxTTuHatJzrvARHiuO9LYd+cIxzgEHCQI4= github.com/spf13/cast v1.7.1 h1:cuNEagBQEHWN1FnbGEjCXL2szYEXqfJPbP2HNUaca9Y= github.com/spf13/cast v1.7.1/go.mod h1:ancEpBxwJDODSW/UG4rDrAqiKolqNNh2DX3mk86cAdo= github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg= github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc= +github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw= github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4= github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= diff --git a/mcp/tools.go b/mcp/tools.go index 3e3931b09..fce263ea2 100644 --- a/mcp/tools.go +++ b/mcp/tools.go @@ -1,11 +1,16 @@ package mcp import ( + "bytes" "encoding/json" "errors" "fmt" "reflect" "strconv" + "sync" + + jsonschemaGenerator "github.com/invopop/jsonschema" + "github.com/santhosh-tekuri/jsonschema" ) var errToolSchemaConflict = errors.New("provide either InputSchema or RawInputSchema, not both") @@ -36,6 +41,8 @@ type ListToolsResult struct { type CallToolResult struct { Result Content []Content `json:"content"` // Can be TextContent, ImageContent, AudioContent, or EmbeddedResource + // Structured content that conforms to the tool's output schema + StructuredContent any `json:"structuredContent,omitempty"` // Whether the tool call ended in an error. // // If not set, this is assumed to be false (the call was successful). @@ -468,6 +475,13 @@ type ToolListChangedNotification struct { Notification } +// validatorState holds the thread-safe state for output validation +type validatorState struct { + validator *jsonschema.Schema + initErr error // Store initialization error to be shared across goroutines + once sync.Once +} + // Tool represents the definition for a tool the client can call. type Tool struct { // The name of the tool. @@ -476,10 +490,17 @@ type Tool struct { Description string `json:"description,omitempty"` // A JSON Schema object defining the expected parameters for the tool. InputSchema ToolInputSchema `json:"inputSchema"` + // A JSON Schema object defining the expected output for the tool. + OutputSchema json.RawMessage `json:"outputSchema,omitempty"` // Alternative to InputSchema - allows arbitrary JSON Schema to be provided RawInputSchema json.RawMessage `json:"-"` // Hide this from JSON marshaling // Optional properties describing tool behavior Annotations ToolAnnotation `json:"annotations"` + + // Internal fields for output validation, not serialized + outputType reflect.Type `json:"-"` + // Thread-safe validator state (pointer to avoid copy issues) + validatorState *validatorState `json:"-"` } // GetName returns the name of the tool. @@ -487,11 +508,78 @@ func (t Tool) GetName() string { return t.Name } +// HasOutputSchema returns true if the tool has an output schema defined. +// This indicates that the tool can return structured content. +func (t Tool) HasOutputSchema() bool { + return t.OutputSchema != nil +} + +// validateStructuredOutput performs the actual validation using the compiled schema +func (t Tool) validateStructuredOutput(result *CallToolResult) error { + if t.validatorState == nil || t.validatorState.validator == nil { + return fmt.Errorf("output validator not initialized for tool %s", t.Name) + } + return t.validatorState.validator.ValidateInterface(result.StructuredContent) +} + +// ensureOutputSchemaValidator compiles and caches the JSON schema validator if not already done +// This method is thread-safe and ensures the validator is initialized exactly once +func (t *Tool) ensureOutputSchemaValidator() error { + if t.OutputSchema == nil { + return nil + } + + // Use sync.Once to ensure initialization happens only once + t.validatorState.once.Do(func() { + // If validator is already set (e.g., by WithOutputType), return early + if t.validatorState.validator != nil { + return + } + + compiler := jsonschema.NewCompiler() + + if err := compiler.AddResource("output-schema", bytes.NewReader(t.OutputSchema)); err != nil { + t.validatorState.initErr = err + return + } + + compiledSchema, err := compiler.Compile("output-schema") + if err != nil { + t.validatorState.initErr = err + return + } + + t.validatorState.validator = compiledSchema + }) + + // Return the shared initialization error (if any) + return t.validatorState.initErr +} + +// ValidateStructuredOutput validates the structured content against the tool's output schema. +// Returns nil if the tool has no output schema or if validation passes. +// Returns an error if the tool has an output schema but the structured content is invalid. +func (t *Tool) ValidateStructuredOutput(result *CallToolResult) error { + if !t.HasOutputSchema() { + return nil + } + + if result.StructuredContent == nil { + return fmt.Errorf("tool %s has output schema but structuredContent is nil", t.Name) + } + + if err := t.ensureOutputSchemaValidator(); err != nil { + return err + } + + return t.validateStructuredOutput(result) +} + // MarshalJSON implements the json.Marshaler interface for Tool. // It handles marshaling either InputSchema or RawInputSchema based on which is set. func (t Tool) MarshalJSON() ([]byte, error) { // Create a map to build the JSON structure - m := make(map[string]any, 3) + m := make(map[string]any, 5) // Add the name and description m["name"] = t.Name @@ -510,6 +598,11 @@ func (t Tool) MarshalJSON() ([]byte, error) { m["inputSchema"] = t.InputSchema } + // Add output schema if defined + if t.OutputSchema != nil { + m["outputSchema"] = t.OutputSchema + } + m["annotations"] = t.Annotations return json.Marshal(m) @@ -522,17 +615,17 @@ type ToolInputSchema struct { } // MarshalJSON implements the json.Marshaler interface for ToolInputSchema. -func (tis ToolInputSchema) MarshalJSON() ([]byte, error) { +func (schema ToolInputSchema) MarshalJSON() ([]byte, error) { m := make(map[string]any) - m["type"] = tis.Type + m["type"] = schema.Type // Marshal Properties to '{}' rather than `nil` when its length equals zero - if tis.Properties != nil { - m["properties"] = tis.Properties + if schema.Properties != nil { + m["properties"] = schema.Properties } - if len(tis.Required) > 0 { - m["required"] = tis.Required + if len(schema.Required) > 0 { + m["required"] = schema.Required } return json.Marshal(m) @@ -574,6 +667,8 @@ func NewTool(name string, opts ...ToolOption) Tool { Properties: make(map[string]any), Required: nil, // Will be omitted from JSON if empty }, + OutputSchema: nil, + RawInputSchema: nil, Annotations: ToolAnnotation{ Title: "", ReadOnlyHint: ToBoolPtr(false), @@ -581,6 +676,8 @@ func NewTool(name string, opts ...ToolOption) Tool { IdempotentHint: ToBoolPtr(false), OpenWorldHint: ToBoolPtr(true), }, + outputType: nil, + validatorState: &validatorState{}, // Initialize immediately for thread safety } for _, opt := range opts { @@ -602,13 +699,13 @@ func NewToolWithRawSchema(name, description string, schema json.RawMessage) Tool Name: name, Description: description, RawInputSchema: schema, + validatorState: &validatorState{}, // Initialize for thread safety } return tool } -// WithDescription adds a description to the Tool. -// The description should provide a clear, human-readable explanation of what the tool does. +// WithDescription sets the description field of the Tool. func WithDescription(description string) ToolOption { return func(t *Tool) { t.Description = description @@ -1076,3 +1173,133 @@ func WithBooleanItems(opts ...PropertyOption) PropertyOption { schema["items"] = itemSchema } } + +// +// Output Schema Configuration Functions +// + +// WithOutputSchema sets the output schema for the Tool. +// This allows the tool to define the structure of its return data. +func WithOutputSchema(schema json.RawMessage) ToolOption { + return func(t *Tool) { + cleaned, err := ExtractMCPSchema(schema) + if err != nil { + // Fallback to original if cleaning fails + t.OutputSchema = schema + } else { + t.OutputSchema = cleaned + } + } +} + +// +// New Output Schema Functions +// + +// WithOutputType sets the output schema for the Tool using Go generics and struct tags. +// This replaces the builder pattern with a cleaner interface based on struct definitions. +func WithOutputType[T any]() ToolOption { + return func(t *Tool) { + var zero T + validator, schemaBytes, err := compileOutputSchema(zero) + if err != nil { + // Skip setting output schema if compilation fails + // This allows the tool to work without validation + return + } + + t.OutputSchema = schemaBytes + t.outputType = reflect.TypeOf(zero) + + // Initialize validatorState with pre-compiled validator + t.validatorState = &validatorState{ + validator: validator, + initErr: nil, // No error since compilation succeeded + } + } +} + +// compileOutputSchema generates JSON schema from a struct and compiles it for validation +func compileOutputSchema[T any](sample T) (*jsonschema.Schema, json.RawMessage, error) { + // Generate JSON Schema from struct + reflector := jsonschemaGenerator.Reflector{ + // Use Draft 7 which is widely supported + DoNotReference: true, + } + schema := reflector.Reflect(&sample) + + // Manually override the schema version to Draft-07. + // This is required because our validator, santhosh-tekuri/jsonschema, + // only supports up to Draft-07. This workaround ensures compatibility + // between the generator and the validator. + schema.Version = "http://json-schema.org/draft-07/schema#" + + // Serialize to JSON + schemaBytes, err := json.Marshal(schema) + if err != nil { + return nil, nil, fmt.Errorf("failed to marshal schema: %w", err) + } + + // Clean schema to MCP inline format + cleanedBytes, err := ExtractMCPSchema(schemaBytes) + if err != nil { + return nil, nil, err + } + + // Compile for validation using santhosh-tekuri/jsonschema + compiler := jsonschema.NewCompiler() + if err := compiler.AddResource("schema", bytes.NewReader(cleanedBytes)); err != nil { + return nil, nil, fmt.Errorf("failed to add schema resource: %w", err) + } + + validator, err := compiler.Compile("schema") + if err != nil { + return nil, nil, fmt.Errorf("failed to compile schema: %w", err) + } + + return validator, cleanedBytes, nil +} + +// ValidateOutput validates the structured content against the tool's output schema. +// Skips validation when IsError is true or the tool has no output schema defined. +func (t *Tool) ValidateOutput(result *CallToolResult) error { + // Skip validation if IsError is true or no output schema defined + if result.IsError || !t.HasOutputSchema() { + return nil + } + + if result.StructuredContent == nil { + return fmt.Errorf("tool %s requires structured output but got nil", t.Name) + } + + // Ensure the validator is compiled + if err := t.ensureOutputSchemaValidator(); err != nil { + return fmt.Errorf("failed to compile output schema for tool %s: %w", t.Name, err) + } + + // Convert structured content to JSON-compatible format for validation + var validationData any + + // If it's already a map, slice, or primitive, use it as-is + // If it's a struct, convert it to JSON-compatible format + switch result.StructuredContent.(type) { + case map[string]any, []any, string, int, int64, float64, bool, nil: + validationData = result.StructuredContent + default: + // Convert struct to JSON-compatible format via JSON marshaling/unmarshaling + jsonBytes, err := json.Marshal(result.StructuredContent) + if err != nil { + return fmt.Errorf("failed to marshal structured content for validation: %w", err) + } + + if err := json.Unmarshal(jsonBytes, &validationData); err != nil { + return fmt.Errorf("failed to unmarshal structured content for validation: %w", err) + } + } + + if t.validatorState == nil || t.validatorState.validator == nil { + return fmt.Errorf("output validator not initialized for tool %s", t.Name) + } + + return t.validatorState.validator.ValidateInterface(validationData) +} diff --git a/mcp/tools_test.go b/mcp/tools_test.go index 0cd71230e..b8545c58d 100644 --- a/mcp/tools_test.go +++ b/mcp/tools_test.go @@ -712,3 +712,469 @@ func TestNewItemsAPICompatibility(t *testing.T) { }) } } + +// Test HasOutputSchema method with empty schema +func TestHasOutputSchemaEmpty(t *testing.T) { + tool := NewTool("test-tool") + + // Empty schema should return false + assert.False(t, tool.HasOutputSchema()) +} + +// Test HasOutputSchema method with defined schema +func TestHasOutputSchemaWithSchema(t *testing.T) { + schema := json.RawMessage(`{"type": "object", "properties": {"result": {"type": "string"}}}`) + tool := NewTool("test-tool", WithOutputSchema(schema)) + + // Should return true when schema is defined + assert.True(t, tool.HasOutputSchema()) +} + +// Test Tool JSON marshaling includes output schema when defined +func TestToolMarshalWithOutputSchema(t *testing.T) { + schema := json.RawMessage(`{ + "type": "object", + "properties": { + "temperature": {"type": "string", "description": "Temperature value"}, + "condition": {"type": "string"} + }, + "required": ["temperature"] + }`) + + tool := NewTool("weather-tool", + WithDescription("Get weather information"), + WithString("location", Required()), + WithOutputSchema(schema), + ) + + // Marshal to JSON + data, err := json.Marshal(tool) + assert.NoError(t, err) + + // Unmarshal to verify structure + var result map[string]any + err = json.Unmarshal(data, &result) + assert.NoError(t, err) + + // Check that outputSchema is present + outputSchema, exists := result["outputSchema"] + assert.True(t, exists) + + schema2, ok := outputSchema.(map[string]any) + assert.True(t, ok) + assert.Equal(t, "object", schema2["type"]) + + // Check properties + properties, ok := schema2["properties"].(map[string]any) + assert.True(t, ok) + assert.Contains(t, properties, "temperature") + assert.Contains(t, properties, "condition") + + // Check required fields + required, ok := schema2["required"].([]any) + assert.True(t, ok) + assert.Contains(t, required, "temperature") +} + +// Test Tool JSON marshaling omits output schema when empty +func TestToolMarshalWithoutOutputSchema(t *testing.T) { + tool := NewTool("simple-tool", + WithDescription("Simple tool without output schema"), + WithString("input", Required()), + ) + + // Marshal to JSON + data, err := json.Marshal(tool) + assert.NoError(t, err) + + // Unmarshal to verify structure + var result map[string]any + err = json.Unmarshal(data, &result) + assert.NoError(t, err) + + // Check that outputSchema is not present + _, exists := result["outputSchema"] + assert.False(t, exists) +} + +// Test WithOutputSchema function +func TestWithOutputSchema(t *testing.T) { + schema := json.RawMessage(`{ + "type": "object", + "properties": { + "data": {"type": "string"} + }, + "required": ["data"] + }`) + + tool := NewTool("custom-tool", WithOutputSchema(schema)) + + assert.True(t, tool.HasOutputSchema()) + cleaned, _ := ExtractMCPSchema(schema) + assert.Equal(t, cleaned, tool.OutputSchema) +} + +// Test NewStructuredToolResult function +func TestNewStructuredToolResult(t *testing.T) { + // This test is removed as NewStructuredToolResult is deprecated + // Use NewToolResultStructured[T]() instead +} + +// Test NewStructuredToolError function +func TestNewStructuredToolError(t *testing.T) { + // This test is removed as NewStructuredToolError is deprecated + // Use NewToolResultErrorStructured[T]() instead +} + +// Test WithOutputType function with struct-based schema generation +func TestWithOutputType(t *testing.T) { + type WeatherOutput struct { + Temperature float64 `json:"temperature" jsonschema:"description=Temperature in Celsius"` + Condition string `json:"condition" jsonschema:"required"` + Humidity int `json:"humidity,omitempty" jsonschema:"minimum=0,maximum=100"` + } + + tool := NewTool("weather-tool", + WithDescription("Get weather information"), + WithString("location", Required()), + WithOutputType[WeatherOutput](), + ) + + assert.True(t, tool.HasOutputSchema()) + assert.NotNil(t, tool.OutputSchema) + + // Marshal and verify JSON structure + data, err := json.Marshal(tool) + assert.NoError(t, err) + + var result map[string]any + err = json.Unmarshal(data, &result) + assert.NoError(t, err) + + // Check that outputSchema is present and valid + outputSchema, exists := result["outputSchema"] + assert.True(t, exists) + + schema, ok := outputSchema.(map[string]any) + assert.True(t, ok) + assert.Equal(t, "object", schema["type"]) + + // Check properties exist + properties, ok := schema["properties"].(map[string]any) + assert.True(t, ok) + assert.Contains(t, properties, "temperature") + assert.Contains(t, properties, "condition") + assert.Contains(t, properties, "humidity") +} + +// Test new helper functions +func TestNewToolResultStructured(t *testing.T) { + type ResponseData struct { + Message string `json:"message"` + Status int `json:"status"` + } + + data := ResponseData{ + Message: "Operation successful", + Status: 200, + } + + result := NewToolResultStructured(data) + + assert.False(t, result.IsError) + assert.Equal(t, data, result.StructuredContent) + assert.Len(t, result.Content, 1) + + textContent, ok := result.Content[0].(TextContent) + assert.True(t, ok) + assert.Equal(t, "text", textContent.Type) + + // Verify the JSON content matches the structured content + var parsedJSON ResponseData + err := json.Unmarshal([]byte(textContent.Text), &parsedJSON) + assert.NoError(t, err) + assert.Equal(t, data.Message, parsedJSON.Message) + assert.Equal(t, data.Status, parsedJSON.Status) +} + +func TestNewToolResultWithStructured(t *testing.T) { + type ResponseData struct { + Value int `json:"value"` + } + + data := ResponseData{Value: 42} + textContent := NewTextContent("Custom text content") + imageContent := NewImageContent("iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==", "image/png") + + result := NewToolResultWithStructured([]Content{textContent, imageContent}, data) + + assert.False(t, result.IsError) + assert.Equal(t, data, result.StructuredContent) + assert.Len(t, result.Content, 2) + + // Verify content is preserved as-is + assert.Equal(t, textContent, result.Content[0]) + assert.Equal(t, imageContent, result.Content[1]) +} + +func TestNewToolResultErrorStructured(t *testing.T) { + type ErrorData struct { + Code string `json:"code"` + Message string `json:"message"` + } + + errorData := ErrorData{ + Code: "TIMEOUT", + Message: "Request timed out", + } + + result := NewToolResultErrorStructured(errorData) + + assert.True(t, result.IsError) + assert.Equal(t, errorData, result.StructuredContent) + assert.Len(t, result.Content, 1) + + textContent, ok := result.Content[0].(TextContent) + assert.True(t, ok) + assert.Equal(t, "text", textContent.Type) + + // Verify the JSON content matches the structured content + var parsedJSON ErrorData + err := json.Unmarshal([]byte(textContent.Text), &parsedJSON) + assert.NoError(t, err) + assert.Equal(t, errorData.Code, parsedJSON.Code) + assert.Equal(t, errorData.Message, parsedJSON.Message) +} + +func TestNewToolResultErrorWithStructured(t *testing.T) { + type ErrorData struct { + Code string `json:"code"` + } + + errorData := ErrorData{Code: "ERR_404"} + textContent := NewTextContent("Not found error occurred") + + result := NewToolResultErrorWithStructured([]Content{textContent}, errorData) + + assert.True(t, result.IsError) + assert.Equal(t, errorData, result.StructuredContent) + assert.Len(t, result.Content, 1) + assert.Equal(t, textContent, result.Content[0]) +} + +// Test validation with new API +func TestValidateOutputWithSchema(t *testing.T) { + type ValidOutput struct { + Message string `json:"message" jsonschema:"required"` + Count int `json:"count" jsonschema:"minimum=0"` + } + + tool := NewTool("test-tool", + WithString("input", Required()), + WithOutputType[ValidOutput](), + ) + + // Test valid structured content + validData := ValidOutput{Message: "Success", Count: 5} + validResult := NewToolResultStructured(validData) + + err := tool.ValidateOutput(validResult) + assert.NoError(t, err, "Valid structured content should pass validation") + + // Test invalid structured content (missing required field) + invalidData := map[string]any{"count": 10} // missing required "message" + invalidResult := &CallToolResult{ + IsError: false, + StructuredContent: invalidData, + Content: []Content{NewTextContent("test")}, + } + + err = tool.ValidateOutput(invalidResult) + assert.Error(t, err, "Invalid structured content should fail validation") +} + +// Test ValidateOutput behavior with different scenarios +func TestValidateOutput(t *testing.T) { + // Test case 1: Tool without output schema should not validate + toolNoSchema := NewTool("no-schema-tool", WithString("input", Required())) + result := NewToolResultText("Just text content") + + err := toolNoSchema.ValidateOutput(result) + assert.NoError(t, err, "Tool without output schema should not validate") + + // Test case 2: Error result should skip validation + schema := json.RawMessage(`{"type": "object", "properties": {"message": {"type": "string"}}, "required": ["message"]}`) + toolWithSchema := NewTool("schema-tool", + WithString("input", Required()), + WithOutputSchema(schema), + ) + errorResult := &CallToolResult{IsError: true, StructuredContent: map[string]any{"invalid": "data"}} + + err = toolWithSchema.ValidateOutput(errorResult) + assert.NoError(t, err, "Error result should skip validation") + + // Test case 3: Tool with output schema but no structured content should return error + resultNoStructured := NewToolResultText("Just text content") + + err = toolWithSchema.ValidateOutput(resultNoStructured) + assert.Error(t, err, "Tool with output schema but no structured content should return error") + assert.Contains(t, err.Error(), "requires structured output") +} + +func TestEnsureOutputSchemaValidatorThreadSafety(t *testing.T) { + // Create a tool with output schema but no pre-compiled validator + outputSchema := json.RawMessage(`{ + "type": "object", + "properties": { + "message": {"type": "string"}, + "status": {"type": "integer"} + }, + "required": ["message"] + }`) + + tool := NewTool("thread-safe-tool", + WithDescription("A tool to test thread safety"), + WithOutputSchema(outputSchema), + ) + + // Run multiple goroutines concurrently to call ensureOutputSchemaValidator + const numGoroutines = 100 + errors := make(chan error, numGoroutines) + + for range numGoroutines { + go func() { + err := tool.ensureOutputSchemaValidator() + errors <- err + }() + } + + // Collect all errors + for i := 0; i < numGoroutines; i++ { + err := <-errors + assert.NoError(t, err, "ensureOutputSchemaValidator should not return error") + } + + // Verify the validator was properly initialized + assert.NotNil(t, tool.validatorState, "validatorState should be initialized") + assert.NotNil(t, tool.validatorState.validator, "validator should be initialized") + assert.NoError(t, tool.validatorState.initErr, "initErr should be nil for successful compilation") + + // Test validation with the initialized validator + result := &CallToolResult{ + Content: []Content{NewTextContent("Success")}, + StructuredContent: map[string]any{ + "message": "test message", + "status": 200, + }, + IsError: false, + } + + err := tool.ValidateOutput(result) + assert.NoError(t, err, "ValidateOutput should succeed with valid data") +} + +func TestEnsureOutputSchemaValidatorWithOutputType(t *testing.T) { + type TestOutput struct { + Message string `json:"message" jsonschema:"required"` + Status int `json:"status" jsonschema:"minimum=100,maximum=599"` + } + + // Create a tool with WithOutputType (which pre-compiles the validator) + tool := NewTool("pre-compiled-tool", + WithDescription("A tool with pre-compiled validator"), + WithOutputType[TestOutput](), + ) + + // Verify the validator is already set + assert.NotNil(t, tool.validatorState, "validatorState should be initialized") + assert.NotNil(t, tool.validatorState.validator, "validator should be pre-compiled") + assert.NoError(t, tool.validatorState.initErr, "initErr should be nil for pre-compiled validator") + + // Call ensureOutputSchemaValidator multiple times concurrently + const numGoroutines = 50 + errors := make(chan error, numGoroutines) + + for range numGoroutines { + go func() { + err := tool.ensureOutputSchemaValidator() + errors <- err + }() + } + + // Collect all errors + for range numGoroutines { + err := <-errors + assert.NoError(t, err, "ensureOutputSchemaValidator should not return error") + } + + // Test validation still works correctly + result := &CallToolResult{ + Content: []Content{NewTextContent("Success")}, + StructuredContent: TestOutput{ + Message: "test message", + Status: 200, + }, + IsError: false, + } + + err := tool.ValidateOutput(result) + assert.NoError(t, err, "ValidateOutput should succeed with valid data") +} + +func TestEnsureOutputSchemaValidatorErrorConsistency(t *testing.T) { + // Create a tool with invalid output schema to test error handling consistency + invalidSchema := json.RawMessage(`{ + "type": "invalid-type-that-should-cause-error", + "$ref": "#/invalid/reference" + }`) + + tool := NewTool("error-test-tool", + WithDescription("A tool to test error consistency"), + WithOutputSchema(invalidSchema), + ) + + // Run multiple goroutines concurrently + const numGoroutines = 50 + errors := make(chan error, numGoroutines) + + for range numGoroutines { + go func() { + err := tool.ensureOutputSchemaValidator() + errors <- err + }() + } + + // Collect all errors - they should all be consistent + var firstErr error + var errorCount, successCount int + for i := range numGoroutines { + err := <-errors + if i == 0 { + firstErr = err + } + + if err != nil { + errorCount++ + } else { + successCount++ + } + + // All results should be consistent - either all errors or all success + if (firstErr == nil) != (err == nil) { + t.Errorf("Inconsistent error handling in goroutine %d: first error was %v, but got %v", i, firstErr, err) + } + } + + // Should either be all errors or all success + if errorCount > 0 && successCount > 0 { + t.Errorf("Mixed results: %d errors and %d successes - should be consistent", errorCount, successCount) + } + + // Verify that subsequent calls return the same result + for i := range 5 { + finalErr := tool.ensureOutputSchemaValidator() + if (firstErr == nil) != (finalErr == nil) { + t.Errorf("Error state changed after concurrent access: original was %v, subsequent call %d returned %v", firstErr, i, finalErr) + } + } +} diff --git a/mcp/utils.go b/mcp/utils.go index 3e652efd7..423f15ce0 100644 --- a/mcp/utils.go +++ b/mcp/utils.go @@ -3,6 +3,7 @@ package mcp import ( "encoding/json" "fmt" + "strings" "github.com/spf13/cast" ) @@ -353,6 +354,57 @@ func NewToolResultErrorf(format string, a ...any) *CallToolResult { } } +// marshalToContent converts structured data to JSON content for backwards compatibility +func marshalToContent(data any) []Content { + jsonBytes, err := json.Marshal(data) + var text string + if err != nil { + text = fmt.Sprintf("Error serializing structured content: %v", err) + } else { + text = string(jsonBytes) + } + + return []Content{NewTextContent(text)} +} + +// NewToolResultStructured creates a CallToolResult with only structured content. +// The helper automatically generates JSON content for backwards compatibility. +func NewToolResultStructured[T any](data T) *CallToolResult { + return &CallToolResult{ + Content: marshalToContent(data), + StructuredContent: data, + IsError: false, + } +} + +// NewToolResultWithStructured creates a CallToolResult with both explicit content and structured content. +func NewToolResultWithStructured[T any](content []Content, data T) *CallToolResult { + return &CallToolResult{ + Content: content, + StructuredContent: data, + IsError: false, + } +} + +// NewToolResultErrorStructured creates a CallToolResult for an error case with structured content. +// The helper automatically generates JSON content and sets IsError to true. +func NewToolResultErrorStructured[T any](data T) *CallToolResult { + return &CallToolResult{ + Content: marshalToContent(data), + StructuredContent: data, + IsError: true, + } +} + +// NewToolResultErrorWithStructured creates a CallToolResult for an error case with both explicit content and structured content. +func NewToolResultErrorWithStructured[T any](content []Content, data T) *CallToolResult { + return &CallToolResult{ + Content: content, + StructuredContent: data, + IsError: true, + } +} + // NewListResourcesResult creates a new ListResourcesResult func NewListResourcesResult( resources []Resource, @@ -817,3 +869,40 @@ func ParseStringMap(request CallToolRequest, key string, defaultValue map[string func ToBoolPtr(b bool) *bool { return &b } + +// ExtractMCPSchema converts a JSON Schema into the inline object schema format +// required by MCP. It removes top-level metadata and inlines local definitions. +// This is not a full JSON-Schema dereferencer. +func ExtractMCPSchema(raw json.RawMessage) (json.RawMessage, error) { + if len(raw) == 0 { + return raw, nil + } + + var m map[string]any + if err := json.Unmarshal(raw, &m); err != nil { + return raw, fmt.Errorf("failed to unmarshal schema: %w", err) + } + + // Inline local $ref into $defs if present. + if ref, ok := m["$ref"].(string); ok && strings.HasPrefix(ref, "#/$defs/") { + defName := strings.TrimPrefix(ref, "#/$defs/") + if defs, ok := m["$defs"].(map[string]any); ok { + if defSchema, ok := defs[defName].(map[string]any); ok { + m = defSchema + } + } + } + + // Remove metadata + delete(m, "$schema") + delete(m, "$id") + delete(m, "$defs") + delete(m, "$ref") + + cleaned, err := json.Marshal(m) + if err != nil { + return raw, fmt.Errorf("failed to marshal cleaned schema: %w", err) + } + + return cleaned, nil +} diff --git a/server/server.go b/server/server.go index 46e6d9c57..2039685db 100644 --- a/server/server.go +++ b/server/server.go @@ -1027,13 +1027,26 @@ func (s *MCPServer) handleToolCall( result, err := finalHandler(ctx, request) if err != nil { - return nil, &requestError{ - id: id, - code: mcp.INTERNAL_ERROR, - err: err, + return nil, &requestError{id, mcp.INTERNAL_ERROR, err} + } + + // Only validate if the result is not an error. + // The validation function itself also checks for IsError, but checking it + // here first prevents overwriting a legitimate error response from the tool + // with a validation error. + if !result.IsError { + if err := tool.Tool.ValidateOutput(result); err != nil { + // If validation fails, return a new error result. + return &mcp.CallToolResult{ + Content: []mcp.Content{ + mcp.NewTextContent(fmt.Sprintf("tool output schema validation failed: %v", err)), + }, + IsError: true, + }, nil } } + // If validation passes or was skipped, return the original result from the handler. return result, nil } diff --git a/server/server_test.go b/server/server_test.go index 1c81d18dd..4371df937 100644 --- a/server/server_test.go +++ b/server/server_test.go @@ -1,6 +1,7 @@ package server import ( + "bytes" "context" "encoding/base64" "encoding/json" @@ -2022,3 +2023,397 @@ func TestMCPServer_ProtocolNegotiation(t *testing.T) { }) } } + +// TestMCPServer_StructuredOutputValidation_Success tests that tools with valid structured output +// pass validation successfully. The tool returns structured content that matches the defined +// output schema, and the server should accept it without errors. +func TestMCPServer_StructuredOutputValidation_Success(t *testing.T) { + server := NewMCPServer("test-server", "1.0.0") + + // Define output schema for weather data + outputSchema := json.RawMessage(`{ + "type": "object", + "properties": { + "temperature": {"type": "number", "description": "Temperature in celsius"}, + "conditions": {"type": "string", "description": "Weather conditions description"}, + "humidity": {"type": "number", "description": "Humidity percentage"} + }, + "required": ["temperature", "conditions", "humidity"] + }`) + + // Create a weather tool with output schema that requires specific structure + tool := mcp.NewTool("get_weather_data", + mcp.WithDescription("Get current weather data for a location"), + mcp.WithString("location", mcp.Description("City name or zip code"), mcp.Required()), + mcp.WithOutputSchema(outputSchema), + ) + + // Add tool handler that returns valid structured content + server.AddTool(tool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return &mcp.CallToolResult{ + Content: []mcp.Content{ + mcp.TextContent{ + Type: "text", + Text: `{"temperature": 22.5, "conditions": "Partly cloudy", "humidity": 65}`, + }, + }, + StructuredContent: map[string]any{ + "temperature": 22.5, + "conditions": "Partly cloudy", + "humidity": 65, + }, + }, nil + }) + + message := `{ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": { + "name": "get_weather_data", + "arguments": { + "location": "Tokyo" + } + } + }` + + response := server.HandleMessage(context.Background(), []byte(message)) + assert.NotNil(t, response) + + // Verify we get a successful JSON-RPC response + resp, ok := response.(mcp.JSONRPCResponse) + assert.True(t, ok, "Expected JSONRPCResponse, got %T", response) + + // Verify the result is a CallToolResult with no error + result, ok := resp.Result.(mcp.CallToolResult) + assert.True(t, ok, "Expected CallToolResult, got %T", resp.Result) + + // Verify that IsError is false (validation passed) + assert.False(t, result.IsError, "Expected IsError to be false for valid structured output") + + // Verify structured content is present and correct + assert.NotNil(t, result.StructuredContent, "StructuredContent should be present") + structuredMap, ok := result.StructuredContent.(map[string]any) + assert.True(t, ok, "StructuredContent should be a map") + assert.Equal(t, 22.5, structuredMap["temperature"]) + assert.Equal(t, "Partly cloudy", structuredMap["conditions"]) + assert.Equal(t, 65, structuredMap["humidity"]) +} + +// TestMCPServer_StructuredOutputValidation_Failure tests that tools with invalid structured output +// fail validation properly. The tool returns structured content that doesn't match the defined +// output schema, and the server should return an error result with IsError set to true. +func TestMCPServer_StructuredOutputValidation_Failure(t *testing.T) { + server := NewMCPServer("test-server", "1.0.0") + + // Define output schema for weather data + outputSchema := json.RawMessage(`{ + "type": "object", + "properties": { + "temperature": {"type": "number", "description": "Temperature in celsius"}, + "conditions": {"type": "string", "description": "Weather conditions description"}, + "humidity": {"type": "number", "description": "Humidity percentage"} + }, + "required": ["temperature", "conditions", "humidity"] + }`) + + // Create a weather tool with output schema that requires specific structure + tool := mcp.NewTool("get_weather_data", + mcp.WithDescription("Get current weather data for a location"), + mcp.WithString("location", mcp.Description("City name or zip code"), mcp.Required()), + mcp.WithOutputSchema(outputSchema), + ) + + // Add tool handler that returns invalid structured content + server.AddTool(tool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { + // Return structured content that doesn't match the schema + // Schema expects: {temperature: number, conditions: string, humidity: number} + // But we return: {location: string, status: string} - missing required fields + return &mcp.CallToolResult{ + Content: []mcp.Content{ + mcp.TextContent{ + Type: "text", + Text: "Weather data retrieved", + }, + }, + StructuredContent: map[string]any{ + "location": "Tokyo", + "status": "success", // This doesn't match the required output schema + }, + }, nil + }) + + message := `{ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": { + "name": "get_weather_data", + "arguments": { + "location": "Tokyo" + } + } + }` + + response := server.HandleMessage(context.Background(), []byte(message)) + assert.NotNil(t, response) + + // Verify we get a successful JSON-RPC response (not a protocol error) + resp, ok := response.(mcp.JSONRPCResponse) + assert.True(t, ok, "Expected JSONRPCResponse, got %T", response) + + // Verify the result is a CallToolResult with error flag set + result, ok := resp.Result.(mcp.CallToolResult) + assert.True(t, ok, "Expected CallToolResult, got %T", resp.Result) + + // Verify that IsError is true due to validation failure + assert.True(t, result.IsError, "Expected IsError to be true due to validation failure") + + // Verify error message contains validation failure information + assert.Len(t, result.Content, 1, "Expected exactly one content item with error message") + textContent, ok := result.Content[0].(mcp.TextContent) + assert.True(t, ok, "Expected TextContent, got %T", result.Content[0]) + assert.Contains(t, textContent.Text, "tool output schema validation failed", + "Error message should indicate schema validation failure") +} + +// TestMCPServer_WithOutputType_Validation_Success tests that tools defined with WithOutputType +// and returning valid structured output pass validation. +func TestMCPServer_WithOutputType_Validation_Success(t *testing.T) { + server := NewMCPServer("test-server", "1.0.0") + + // NOTE: The `jsonschema` tag parser used by invopop/jsonschema has a limitation: + // it uses commas as delimiters and does not support escaping them. + // Therefore, descriptions should not contain commas. + type WeatherData struct { + Temperature float64 `json:"temperature" jsonschema:"description=The current temperature in Celsius."` + Conditions string `json:"conditions" jsonschema:"description=A brief description of weather conditions (e.g. Cloudy or Sunny)."` + Humidity int `json:"humidity" jsonschema:"description=The relative humidity percentage."` + } + + tool := mcp.NewTool("get_weather_data", + mcp.WithOutputType[WeatherData]()) + + // Print the generated schema for verification. + var prettyJSON bytes.Buffer + err := json.Indent(&prettyJSON, tool.OutputSchema, "", " ") + require.NoError(t, err) + t.Logf("Generated Output Schema:\n%s", prettyJSON.String()) + + // Basic sanity check: generated schema should not contain $schema keyword (it is already cleaned). + assert.NotContains(t, string(tool.OutputSchema), "$schema", "schema should not contain $schema keyword") + + // Verify that the generated schema contains the descriptions. + var schema map[string]any + err = json.Unmarshal(tool.OutputSchema, &schema) + require.NoError(t, err) + + properties := schema["properties"].(map[string]any) + tempProp := properties["temperature"].(map[string]any) + assert.Equal(t, "The current temperature in Celsius.", tempProp["description"]) + condProp := properties["conditions"].(map[string]any) + assert.Equal(t, "A brief description of weather conditions (e.g. Cloudy or Sunny).", condProp["description"]) + humProp := properties["humidity"].(map[string]any) + assert.Equal(t, "The relative humidity percentage.", humProp["description"]) + + // Add tool handler that returns valid structured content using the new helper. + server.AddTool(tool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { + weatherData := WeatherData{ + Temperature: 22.5, + Conditions: "Partly cloudy", + Humidity: 65, + } + return mcp.NewToolResultStructured(weatherData), nil + }) + + message := `{ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": { + "name": "get_weather_data", + "arguments": { + "location": "Tokyo" + } + } + }` + + response := server.HandleMessage(context.Background(), []byte(message)) + require.NotNil(t, response) + + resp, ok := response.(mcp.JSONRPCResponse) + require.True(t, ok, "Expected JSONRPCResponse, got %T", response) + + result, ok := resp.Result.(mcp.CallToolResult) + require.True(t, ok, "Expected CallToolResult, got %T", resp.Result) + + // Verify that IsError is false (validation passed). + assert.False(t, result.IsError, "Expected IsError to be false for valid structured output") + + // Verify structured content is present and correct. + require.NotNil(t, result.StructuredContent, "StructuredContent should be present") + weatherData, ok := result.StructuredContent.(WeatherData) + require.True(t, ok, "StructuredContent should be of type WeatherData, but was %T", result.StructuredContent) + + assert.Equal(t, 22.5, weatherData.Temperature) + assert.Equal(t, "Partly cloudy", weatherData.Conditions) + assert.Equal(t, 65, weatherData.Humidity) +} + +// TestMCPServer_WithOutputType_Validation_Failure tests that tools defined with WithOutputType +// fail validation when returning invalid structured output. +func TestMCPServer_WithOutputType_Validation_Failure(t *testing.T) { + server := NewMCPServer("test-server", "1.0.0") + + // This struct is the "correct" schema. + type WeatherData struct { + Temperature float64 `json:"temperature"` + Conditions string `json:"conditions"` + } + + // This struct is what the handler will return. It's missing 'Conditions'. + type InvalidWeatherData struct { + Temperature float64 `json:"temperature"` + } + + // Create a tool with an output type that requires 'temperature' and 'conditions'. + tool := mcp.NewTool("get_weather_data", + mcp.WithOutputType[WeatherData]()) + + // Add a handler that returns invalid content (missing 'conditions'). + server.AddTool(tool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { + invalidData := InvalidWeatherData{ + Temperature: 25.0, + } + return mcp.NewToolResultStructured(invalidData), nil + }) + + message := `{ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": { + "name": "get_weather_data", + "arguments": {} + } + }` + + response := server.HandleMessage(context.Background(), []byte(message)) + require.NotNil(t, response) + + resp, ok := response.(mcp.JSONRPCResponse) + require.True(t, ok, "Expected JSONRPCResponse, got %T", response) + + result, ok := resp.Result.(mcp.CallToolResult) + require.True(t, ok, "Expected CallToolResult, got %T", resp.Result) + + // Verify that IsError is true and the error message indicates a validation failure. + assert.True(t, result.IsError, "Expected IsError to be true for invalid structured output") + textContent, ok := result.Content[0].(mcp.TextContent) + require.True(t, ok) + assert.Contains(t, textContent.Text, "tool output schema validation failed") + // The specific error message from the validator can be complex, + // so we'll just check for the key part. + assert.Contains(t, textContent.Text, "missing properties: \"conditions\"") +} + +// TestMCPServer_WithOutputType_ErrorResponseSkipsValidation tests that if a tool handler +// returns a result with IsError=true, schema validation is skipped. +func TestMCPServer_WithOutputType_ErrorResponseSkipsValidation(t *testing.T) { + server := NewMCPServer("test-server", "1.0.0") + + type WeatherData struct { + Temperature float64 `json:"temperature"` + Conditions string `json:"conditions"` + } + + tool := mcp.NewTool("get_weather_data", + mcp.WithOutputType[WeatherData]()) + + // This data is invalid, but it shouldn't be validated because IsError is true. + invalidData := map[string]any{ + "temperature": "this-is-not-a-float", + } + + // Add a handler that returns an error response with invalid structured content. + server.AddTool(tool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return mcp.NewToolResultErrorStructured(invalidData), nil + }) + + message := `{ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": { + "name": "get_weather_data", + "arguments": {} + } + }` + + response := server.HandleMessage(context.Background(), []byte(message)) + require.NotNil(t, response) + + resp, ok := response.(mcp.JSONRPCResponse) + require.True(t, ok, "Expected JSONRPCResponse, got %T", response) + + result, ok := resp.Result.(mcp.CallToolResult) + require.True(t, ok, "Expected CallToolResult, got %T", resp.Result) + + // IsError should be true because the handler returned an error response. + // The content should not have been replaced by a validation error message. + assert.True(t, result.IsError, "Expected IsError to be true") + assert.Equal(t, invalidData, result.StructuredContent) + + // Check that the content is the marshaled version of invalidData, not a validation error. + textContent, ok := result.Content[0].(mcp.TextContent) + require.True(t, ok) + assert.NotContains(t, textContent.Text, "tool output schema validation failed") + + expectedJSON, _ := json.Marshal(invalidData) + assert.Equal(t, string(expectedJSON), textContent.Text) +} + +// TestMCPServer_ToolDefinition_IncludesOutputSchema verifies that when a tool is defined +// using WithOutputType, its OutputSchema is correctly generated with descriptions +// and stored within the server's tool registry. +func TestMCPServer_ToolDefinition_IncludesOutputSchema(t *testing.T) { + server := NewMCPServer("test-server", "1.0.0") + + // NOTE: The `jsonschema` tag parser used by invopop/jsonschema has a limitation: + // it uses commas as delimiters and does not support escaping them. + // Therefore, descriptions should not contain commas. + type WeatherData struct { + Temperature float64 `json:"temperature" jsonschema:"description=The current temperature in Celsius."` + Conditions string `json:"conditions" jsonschema:"description=A brief description of weather conditions (e.g. Cloudy or Sunny)."` + } + + tool := mcp.NewTool("get_weather_data", + mcp.WithOutputType[WeatherData]()) + + // Add the tool to the server. The handler is nil because we are not testing execution. + server.AddTool(tool, nil) + + // Retrieve the tool from the server's internal registry to inspect its state. + registeredTool, ok := server.tools[tool.Name] + require.True(t, ok, "Tool should be registered on the server") + require.NotNil(t, registeredTool.Tool.OutputSchema, "OutputSchema should not be nil") + + // Print the schema for verification during test runs. + var prettySchema bytes.Buffer + err := json.Indent(&prettySchema, registeredTool.Tool.OutputSchema, "", " ") + require.NoError(t, err) + t.Logf("Generated Output Schema:\n%s", prettySchema.String()) + + // Unmarshal the schema to verify its contents. + var schemaData map[string]interface{} + err = json.Unmarshal(registeredTool.Tool.OutputSchema, &schemaData) + require.NoError(t, err, "Failed to unmarshal output schema") + + // Verify that the descriptions from the struct tags are present. + properties := schemaData["properties"].(map[string]interface{}) + tempProp := properties["temperature"].(map[string]interface{}) + condProp := properties["conditions"].(map[string]interface{}) + + assert.Equal(t, "The current temperature in Celsius.", tempProp["description"]) + assert.Equal(t, "A brief description of weather conditions (e.g. Cloudy or Sunny).", condProp["description"]) +}