Merge pull request #144 from nilcolor/output_schema_arrays

Allow output schema to be array of objects

Author: koic

README.md

@@ -546,6 +546,27 @@ class DataTool < MCP::Tool
 end
 ```
 
+Output schema may also describe an array of objects:
+
+```ruby
+class WeatherTool < MCP::Tool
+  output_schema(
+    type: "array",
+    item: {
+      properties: {
+        temperature: { type: "number" },
+        condition: { type: "string" },
+        humidity: { type: "integer" }
+      },
+      required: ["temperature", "condition", "humidity"]
+    }
+  )
+end
+```
+
+Please note: in this case, you must provide `type: "array"`. The default type
+for output schemas is `object`.
+
 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

lib/mcp/tool.rb

@@ -84,9 +84,7 @@ 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:)
+          @output_schema_value = OutputSchema.new(value)
         elsif value.is_a?(OutputSchema)
           @output_schema_value = value
         end

lib/mcp/tool/output_schema.rb

@@ -7,23 +7,20 @@ class Tool
     class OutputSchema
       class ValidationError < StandardError; end
 
-      attr_reader :properties, :required
+      attr_reader :schema
 
-      def initialize(properties: {}, required: [])
-        @properties = properties
-        @required = required.map(&:to_sym)
+      def initialize(schema = {})
+        @schema = deep_transform_keys(JSON.parse(JSON.dump(schema)), &:to_sym)
+        @schema[:type] ||= "object"
         validate_schema!
       end
 
       def ==(other)
-        other.is_a?(OutputSchema) && properties == other.properties && required == other.required
+        other.is_a?(OutputSchema) && schema == other.schema
       end
 
       def to_h
-        { type: "object" }.tap do |hsh|
-          hsh[:properties] = properties if properties.any?
-          hsh[:required] = required if required.any?
-        end
+        @schema
       end
 
       def validate_result(result)
@@ -35,8 +32,24 @@ def validate_result(result)
 
       private
 
+      def deep_transform_keys(schema, &block)
+        case schema
+        when Hash
+          schema.each_with_object({}) do |(key, value), result|
+            if key.casecmp?("$ref")
+              raise ArgumentError, "Invalid JSON Schema: $ref is not allowed in tool output schemas"
+            end
+
+            result[yield(key)] = deep_transform_keys(value, &block)
+          end
+        when Array
+          schema.map { |e| deep_transform_keys(e, &block) }
+        else
+          schema
+        end
+      end
+
       def validate_schema!
-        check_for_refs!
         schema = to_h
         schema_reader = JSON::Schema::Reader.new(
           accept_uri: false,
@@ -48,19 +61,6 @@ def validate_schema!
           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/tool/output_schema_test.rb

@@ -5,15 +5,10 @@
 module MCP
   class Tool
     class OutputSchemaTest < ActiveSupport::TestCase
-      test "required arguments are converted to symbols" do
-        output_schema = OutputSchema.new(properties: { result: { type: "string" } }, required: ["result"])
-        assert_equal [:result], output_schema.required
-      end
-
       test "to_h returns a hash representation of the output schema" do
         output_schema = OutputSchema.new(properties: { result: { type: "string" } }, required: [:result])
         assert_equal(
-          { type: "object", properties: { result: { type: "string" } }, required: [:result] },
+          { type: "object", properties: { result: { type: "string" } }, required: ["result"] },
           output_schema.to_h,
         )
       end
@@ -41,7 +36,7 @@ class OutputSchemaTest < ActiveSupport::TestCase
 
       test "valid schema initialization" do
         schema = OutputSchema.new(properties: { foo: { type: "string" } }, required: [:foo])
-        assert_equal({ type: "object", properties: { foo: { type: "string" } }, required: [:foo] }, schema.to_h)
+        assert_equal({ type: "object", properties: { foo: { type: "string" } }, required: ["foo"] }, schema.to_h)
       end
 
       test "invalid schema raises argument error" do
@@ -135,6 +130,26 @@ class OutputSchemaTest < ActiveSupport::TestCase
           schema.validate_result(invalid_result)
         end
       end
+
+      test "allow to declare array schemas" do
+        schema = OutputSchema.new({
+          type: "array",
+          items: {
+            properties: { foo: { type: "string" } },
+            required: [:foo],
+          },
+        })
+        assert_equal(
+          {
+            type: "array",
+            items: {
+              properties: { foo: { type: "string" } },
+              required: ["foo"],
+            },
+          },
+          schema.to_h,
+        )
+      end
     end
   end
 end

test/mcp/tool_test.rb

@@ -267,7 +267,7 @@ def call(message, server_context: nil)
         title: "Mock Tool",
         description: "a mock tool for testing",
         inputSchema: { type: "object" },
-        outputSchema: { type: "object", properties: { result: { type: "string" } }, required: [:result] },
+        outputSchema: { type: "object", properties: { result: { type: "string" } }, required: ["result"] },
       }
       assert_equal expected, tool.to_h
     end
@@ -292,7 +292,7 @@ class HashOutputSchemaTool < Tool
       end
 
       tool = HashOutputSchemaTool
-      expected = { type: "object", properties: { result: { type: "string" } }, required: [:result] }
+      expected = { type: "object", properties: { result: { type: "string" } }, required: ["result"] }
       assert_equal expected, tool.output_schema.to_h
     end
 
@@ -302,7 +302,7 @@ class OutputSchemaObjectTool < Tool
       end
 
       tool = OutputSchemaObjectTool
-      expected = { type: "object", properties: { result: { type: "string" } }, required: [:result] }
+      expected = { type: "object", properties: { result: { type: "string" } }, required: ["result"] }
       assert_equal expected, tool.output_schema.to_h
     end
 
@@ -354,7 +354,7 @@ class OutputSchemaObjectTool < Tool
       assert_equal "mock_tool", tool.name_value
       assert_equal "a mock tool for testing", tool.description
       assert_instance_of Tool::OutputSchema, tool.output_schema
-      expected_output_schema = { type: "object", properties: { result: { type: "string" } }, required: [:result] }
+      expected_output_schema = { type: "object", properties: { result: { type: "string" } }, required: ["result"] }
       assert_equal expected_output_schema, tool.output_schema.to_h
     end
 
@@ -379,7 +379,7 @@ def call(message:, server_context: nil)
       expected_input = { type: "object", properties: { message: { type: "string" } }, required: [:message] }
       assert_equal expected_input, tool.input_schema.to_h
 
-      expected_output = { type: "object", properties: { result: { type: "string" }, success: { type: "boolean" } }, required: [:result, :success] }
+      expected_output = { type: "object", properties: { result: { type: "string" }, success: { type: "boolean" } }, required: ["result", "success"] }
       assert_equal expected_output, tool.output_schema.to_h
     end
   end