Merge pull request #133 from codenamev/main

Add Tool output_schema support with comprehensive validation

Author: koic

README.md

@@ -385,6 +385,14 @@ class MyTool < MCP::Tool
     },
     required: ["message"]
   )
+  output_schema(
+    properties: {
+      result: { type: "string" },
+      success: { type: "boolean" },
+      timestamp: { type: "string", format: "date-time" }
+    },
+    required: ["result", "success", "timestamp"]
+  )
   annotations(
     read_only_hint: true,
     destructive_hint: false,
@@ -451,6 +459,102 @@ Annotations can be set either through the class definition using the `annotation
 > [!NOTE]
 > This **Tool Annotations** feature is supported starting from `protocol_version: '2025-03-26'`.
 
+### Tool Output Schemas
+
+Tools can optionally define an `output_schema` to specify the expected structure of their results. This works similarly to how `input_schema` is defined and can be used in three ways:
+
+1. **Class definition with output_schema:**
+
+```ruby
+class WeatherTool < MCP::Tool
+  tool_name "get_weather"
+  description "Get current weather for a location"
+
+  input_schema(
+    properties: {
+      location: { type: "string" },
+      units: { type: "string", enum: ["celsius", "fahrenheit"] }
+    },
+    required: ["location"]
+  )
+
+  output_schema(
+    properties: {
+      temperature: { type: "number" },
+      condition: { type: "string" },
+      humidity: { type: "integer" }
+    },
+    required: ["temperature", "condition", "humidity"]
+  )
+
+  def self.call(location:, units: "celsius", server_context:)
+    # Call weather API and structure the response
+    api_response = WeatherAPI.fetch(location, units)
+    weather_data = {
+      temperature: api_response.temp,
+      condition: api_response.description,
+      humidity: api_response.humidity_percent
+    }
+
+    output_schema.validate_result(weather_data)
+
+    MCP::Tool::Response.new([{
+      type: "text",
+      text: weather_data.to_json
+    }])
+  end
+end
+```
+
+2. **Using Tool.define with output_schema:**
+
+```ruby
+tool = MCP::Tool.define(
+  name: "calculate_stats",
+  description: "Calculate statistics for a dataset",
+  input_schema: {
+    properties: {
+      numbers: { type: "array", items: { type: "number" } }
+    },
+    required: ["numbers"]
+  },
+  output_schema: {
+    properties: {
+      mean: { type: "number" },
+      median: { type: "number" },
+      count: { type: "integer" }
+    },
+    required: ["mean", "median", "count"]
+  }
+) do |args, server_context|
+  # Calculate statistics and validate against schema
+  MCP::Tool::Response.new([{ type: "text", text: "Statistics calculated" }])
+end
+```
+
+3. **Using OutputSchema objects:**
+
+```ruby
+class DataTool < MCP::Tool
+  output_schema MCP::Tool::OutputSchema.new(
+    properties: {
+      success: { type: "boolean" },
+      data: { type: "object" }
+    },
+    required: ["success"]
+  )
+end
+```
+
+MCP spec for the [Output Schema](https://modelcontextprotocol.io/specification/2025-06-18/server/tools#output-schema) specifies that:
+
+- **Server Validation**: Servers MUST provide structured results that conform to the output schema
+- **Client Validation**: Clients SHOULD validate structured results against the output schema
+- **Better Integration**: Enables strict schema validation, type information, and improved developer experience
+- **Backward Compatibility**: Tools returning structured content SHOULD also include serialized JSON in a TextContent block
+
+The output schema follows standard JSON Schema format and helps ensure consistent data exchange between MCP servers and clients.
+
 ### Prompts
 
 MCP spec includes [Prompts](https://modelcontextprotocol.io/specification/2025-06-18/server/prompts), which enable servers to define reusable prompt templates and workflows that clients can easily surface to users and LLMs.
@@ -758,7 +862,7 @@ The client provides a wrapper class for tools returned by the server:
 
 - `MCP::Client::Tool` - Represents a single tool with its metadata
 
-This class provide easy access to tool properties like name, description, and input schema.
+This class provides easy access to tool properties like name, description, input schema, and output schema.
 
 ## Releases
 

lib/mcp.rb

@@ -18,6 +18,7 @@
 require_relative "mcp/string_utils"
 require_relative "mcp/tool"
 require_relative "mcp/tool/input_schema"
+require_relative "mcp/tool/output_schema"
 require_relative "mcp/tool/response"
 require_relative "mcp/tool/annotations"
 require_relative "mcp/transport"

lib/mcp/client/tool.rb

@@ -3,12 +3,13 @@
 module MCP
   class Client
     class Tool
-      attr_reader :name, :description, :input_schema
+      attr_reader :name, :description, :input_schema, :output_schema
 
-      def initialize(name:, description:, input_schema:)
+      def initialize(name:, description:, input_schema:, output_schema: nil)
         @name = name
         @description = description
         @input_schema = input_schema
+        @output_schema = output_schema
       end
     end
   end

lib/mcp/tool.rb

@@ -14,14 +14,14 @@ def call(*args, server_context: nil)
       end
 
       def to_h
-        result = {
+        {
           name: name_value,
           title: title_value,
           description: description_value,
           inputSchema: input_schema_value.to_h,
+          outputSchema: @output_schema_value&.to_h,
+          annotations: annotations_value&.to_h,
         }.compact
-        result[:annotations] = annotations_value.to_h if annotations_value
-        result
       end
 
       def inherited(subclass)
@@ -30,6 +30,7 @@ def inherited(subclass)
         subclass.instance_variable_set(:@title_value, nil)
         subclass.instance_variable_set(:@description_value, nil)
         subclass.instance_variable_set(:@input_schema_value, nil)
+        subclass.instance_variable_set(:@output_schema_value, nil)
         subclass.instance_variable_set(:@annotations_value, nil)
       end
 
@@ -49,6 +50,8 @@ def input_schema_value
         @input_schema_value || InputSchema.new
       end
 
+      attr_reader :output_schema_value
+
       def title(value = NOT_SET)
         if value == NOT_SET
           @title_value
@@ -77,6 +80,18 @@ def input_schema(value = NOT_SET)
         end
       end
 
+      def output_schema(value = NOT_SET)
+        if value == NOT_SET
+          output_schema_value
+        elsif value.is_a?(Hash)
+          properties = value[:properties] || value["properties"] || {}
+          required = value[:required] || value["required"] || []
+          @output_schema_value = OutputSchema.new(properties:, required:)
+        elsif value.is_a?(OutputSchema)
+          @output_schema_value = value
+        end
+      end
+
       def annotations(hash = NOT_SET)
         if hash == NOT_SET
           @annotations_value
@@ -85,12 +100,13 @@ def annotations(hash = NOT_SET)
         end
       end
 
-      def define(name: nil, title: nil, description: nil, input_schema: nil, annotations: nil, &block)
+      def define(name: nil, title: nil, description: nil, input_schema: nil, output_schema: nil, annotations: nil, &block)
         Class.new(self) do
           tool_name name
           title title
           description description
           input_schema input_schema
+          output_schema output_schema
           self.annotations(annotations) if annotations
           define_singleton_method(:call, &block) if block
         end

lib/mcp/tool/output_schema.rb

@@ -0,0 +1,66 @@
+# frozen_string_literal: true
+
+require "json-schema"
+
+module MCP
+  class Tool
+    class OutputSchema
+      class ValidationError < StandardError; end
+
+      attr_reader :properties, :required
+
+      def initialize(properties: {}, required: [])
+        @properties = properties
+        @required = required.map(&:to_sym)
+        validate_schema!
+      end
+
+      def ==(other)
+        other.is_a?(OutputSchema) && properties == other.properties && required == other.required
+      end
+
+      def to_h
+        { type: "object" }.tap do |hsh|
+          hsh[:properties] = properties if properties.any?
+          hsh[:required] = required if required.any?
+        end
+      end
+
+      def validate_result(result)
+        errors = JSON::Validator.fully_validate(to_h, result)
+        if errors.any?
+          raise ValidationError, "Invalid result: #{errors.join(", ")}"
+        end
+      end
+
+      private
+
+      def validate_schema!
+        check_for_refs!
+        schema = to_h
+        schema_reader = JSON::Schema::Reader.new(
+          accept_uri: false,
+          accept_file: ->(path) { path.to_s.start_with?(Gem.loaded_specs["json-schema"].full_gem_path) },
+        )
+        metaschema = JSON::Validator.validator_for_name("draft4").metaschema
+        errors = JSON::Validator.fully_validate(metaschema, schema, schema_reader: schema_reader)
+        if errors.any?
+          raise ArgumentError, "Invalid JSON Schema: #{errors.join(", ")}"
+        end
+      end
+
+      def check_for_refs!(obj = properties)
+        case obj
+        when Hash
+          if obj.key?("$ref") || obj.key?(:$ref)
+            raise ArgumentError, "Invalid JSON Schema: $ref is not allowed in tool output schemas"
+          end
+
+          obj.each_value { |value| check_for_refs!(value) }
+        when Array
+          obj.each { |item| check_for_refs!(item) }
+        end
+      end
+    end
+  end
+end

test/mcp/client/tool_test.rb

@@ -28,6 +28,38 @@ def test_input_schema_returns_input_schema
           @tool.input_schema,
         )
       end
+
+      def test_output_schema_returns_nil_when_not_provided
+        assert_nil(@tool.output_schema)
+      end
+
+      def test_output_schema_returns_output_schema_when_provided
+        tool_with_output = Tool.new(
+          name: "test_tool_with_output",
+          description: "A test tool with output schema",
+          input_schema: { "type" => "object", "properties" => { "foo" => { "type" => "string" } } },
+          output_schema: { "type" => "object", "properties" => { "result" => { "type" => "string" } } },
+        )
+
+        assert_equal(
+          { "type" => "object", "properties" => { "result" => { "type" => "string" } } },
+          tool_with_output.output_schema,
+        )
+      end
+
+      def test_initialization_with_all_parameters
+        tool = Tool.new(
+          name: "full_tool",
+          description: "A tool with all parameters",
+          input_schema: { "type" => "object" },
+          output_schema: { "type" => "object", "properties" => { "status" => { "type" => "boolean" } } },
+        )
+
+        assert_equal("full_tool", tool.name)
+        assert_equal("A tool with all parameters", tool.description)
+        assert_equal({ "type" => "object" }, tool.input_schema)
+        assert_equal({ "type" => "object", "properties" => { "status" => { "type" => "boolean" } } }, tool.output_schema)
+      end
     end
   end
 end