Refactor schema validation to reduce duplication and add additionalProperties support

This refactoring extracts common schema validation logic into a base Schema
class, eliminating ~100 lines of duplicate code between InputSchema and
OutputSchema. It also adds support for the JSON Schema `additionalProperties`
keyword to allow fine-grained control over extra properties in tool arguments.

Previously `additionalParameters` could be set only on OutputSchema.

Key changes:

- Created MCP::Tool::Schema base class consolidating:
  - Schema validation against JSON Schema Draft 4 metaschema
  - Deep key transformation (string to symbol conversion)
  - $ref disallowance checking
  - Common validation methods

- Refactored InputSchema to inherit from Schema:
  - Simplified constructor to accept full schema hash
  - Changed required fields storage from symbols to strings
  - Removed ~40 lines of duplicated validation logic

- Refactored OutputSchema to inherit from Schema:
  - Removed ~54 lines of duplicated validation logic
  - Now shares all common functionality with InputSchema

- Updated Tool.input_schema setter:
  - Simplified to pass full hash directly to InputSchema

- Added additionalProperties support:
  - By default, tools now allow additional properties (JSON Schema default)
  - Can explicitly set `additionalProperties: false` to disallow extras
  - Added comprehensive test coverage for both behaviors

Author: ebeigarts

lib/mcp/tool.rb

@@ -75,9 +75,7 @@ def input_schema(value = NOT_SET)
         if value == NOT_SET
           input_schema_value
         elsif value.is_a?(Hash)
-          properties = value[:properties] || value["properties"] || {}
-          required = value[:required] || value["required"] || []
-          @input_schema_value = InputSchema.new(properties:, required:)
+          @input_schema_value = InputSchema.new(value)
         elsif value.is_a?(InputSchema)
           @input_schema_value = value
         end

lib/mcp/tool/input_schema.rb

@@ -1,78 +1,28 @@
 # frozen_string_literal: true
 
-require "json-schema"
+require_relative "schema"
 
 module MCP
   class Tool
-    class InputSchema
+    class InputSchema < Schema
       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?(InputSchema) && 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 missing_required_arguments?(arguments)
         missing_required_arguments(arguments).any?
       end
 
       def missing_required_arguments(arguments)
-        (required - arguments.keys.map(&:to_sym))
+        return [] unless schema[:required].is_a?(Array)
+
+        (schema[:required] - arguments.keys.map(&:to_s))
       end
 
       def validate_arguments(arguments)
-        errors = JSON::Validator.fully_validate(to_h, arguments)
+        errors = fully_validate(arguments)
         if errors.any?
           raise ValidationError, "Invalid arguments: #{errors.join(", ")}"
         end
       end
-
-      private
-
-      def validate_schema!
-        check_for_refs!
-        schema = to_h
-        gem_path = File.realpath(Gem.loaded_specs["json-schema"].full_gem_path)
-        schema_reader = JSON::Schema::Reader.new(
-          accept_uri: false,
-          accept_file: ->(path) { File.realpath(path.to_s).start_with?(gem_path) },
-        )
-        metaschema_path = Pathname.new(JSON::Validator.validator_for_name("draft4").metaschema)
-        # Converts metaschema to a file URI for cross-platform compatibility
-        metaschema_uri = JSON::Util::URI.file_uri(metaschema_path.expand_path.cleanpath.to_s.tr("\\", "/"))
-        metaschema = metaschema_uri.to_s
-        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 input schemas"
-          end
-
-          obj.each_value { |value| check_for_refs!(value) }
-        when Array
-          obj.each { |item| check_for_refs!(item) }
-        end
-      end
     end
   end
 end

lib/mcp/tool/output_schema.rb

@@ -1,70 +1,18 @@
 # frozen_string_literal: true
 
-require "json-schema"
+require_relative "schema"
 
 module MCP
   class Tool
-    class OutputSchema
+    class OutputSchema < Schema
       class ValidationError < StandardError; end
 
-      attr_reader :schema
-
-      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) && schema == other.schema
-      end
-
-      def to_h
-        @schema
-      end
-
       def validate_result(result)
-        errors = JSON::Validator.fully_validate(to_h, result)
+        errors = fully_validate(result)
         if errors.any?
           raise ValidationError, "Invalid result: #{errors.join(", ")}"
         end
       end
-
-      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!
-        schema = to_h
-        gem_path = File.realpath(Gem.loaded_specs["json-schema"].full_gem_path)
-        schema_reader = JSON::Schema::Reader.new(
-          accept_uri: false,
-          accept_file: ->(path) { File.realpath(path.to_s).start_with?(gem_path) },
-        )
-        metaschema_path = Pathname.new(JSON::Validator.validator_for_name("draft4").metaschema)
-        # Converts metaschema to a file URI for cross-platform compatibility
-        metaschema_uri = JSON::Util::URI.file_uri(metaschema_path.expand_path.cleanpath.to_s.tr("\\", "/"))
-        metaschema = metaschema_uri.to_s
-        errors = JSON::Validator.fully_validate(metaschema, schema, schema_reader: schema_reader)
-        if errors.any?
-          raise ArgumentError, "Invalid JSON Schema: #{errors.join(", ")}"
-        end
-      end
     end
   end
 end

lib/mcp/tool/schema.rb

@@ -0,0 +1,65 @@
+# frozen_string_literal: true
+
+require "json-schema"
+
+module MCP
+  class Tool
+    class Schema
+      attr_reader :schema
+
+      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?(self.class) && schema == other.schema
+      end
+
+      def to_h
+        @schema
+      end
+
+      private
+
+      def fully_validate(data)
+        JSON::Validator.fully_validate(to_h, data)
+      end
+
+      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 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!
+        schema = to_h
+        gem_path = File.realpath(Gem.loaded_specs["json-schema"].full_gem_path)
+        schema_reader = JSON::Schema::Reader.new(
+          accept_uri: false,
+          accept_file: ->(path) { File.realpath(path.to_s).start_with?(gem_path) },
+        )
+        metaschema_path = Pathname.new(JSON::Validator.validator_for_name("draft4").metaschema)
+        # Converts metaschema to a file URI for cross-platform compatibility
+        metaschema_uri = JSON::Util::URI.file_uri(metaschema_path.expand_path.cleanpath.to_s.tr("\\", "/"))
+        metaschema = metaschema_uri.to_s
+        errors = JSON::Validator.fully_validate(metaschema, schema, schema_reader: schema_reader)
+        if errors.any?
+          raise ArgumentError, "Invalid JSON Schema: #{errors.join(", ")}"
+        end
+      end
+    end
+  end
+end

test/mcp/server_test.rb

@@ -1096,6 +1096,62 @@ def call(message:, server_context: nil)
       assert_equal "OK", response[:result][:content][0][:content]
     end
 
+    test "tools/call allows additional properties by default" do
+      server = Server.new(
+        tools: [TestTool],
+        configuration: Configuration.new(validate_tool_call_arguments: true),
+      )
+
+      response = server.handle(
+        {
+          jsonrpc: "2.0",
+          id: 1,
+          method: "tools/call",
+          params: {
+            name: "test_tool",
+            arguments: {
+              message: "Hello, world!",
+              other_property: "I am allowed",
+            },
+          },
+        },
+      )
+
+      assert_equal "2.0", response[:jsonrpc]
+      assert_equal 1, response[:id]
+      assert response[:result], "Expected result key in response"
+      assert_equal "text", response[:result][:content][0][:type]
+      assert_equal "OK", response[:result][:content][0][:content]
+    end
+
+    test "tools/call disallows additional properties when additionalProperties set to false" do
+      server = Server.new(
+        tools: [TestToolWithAdditionalPropertiesSetToFalse],
+        configuration: Configuration.new(validate_tool_call_arguments: true),
+      )
+
+      response = server.handle(
+        {
+          jsonrpc: "2.0",
+          id: 1,
+          method: "tools/call",
+          params: {
+            name: "test_tool_with_additional_properties_set_to_false",
+            arguments: {
+              message: "Hello, world!",
+              extra: 123,
+            },
+          },
+        },
+      )
+
+      assert_equal "2.0", response[:jsonrpc]
+      assert_equal 1, response[:id]
+      assert_nil response[:error], "Expected no JSON-RPC error"
+      assert response[:result][:isError]
+      assert_includes response[:result][:content][0][:text], "Invalid arguments"
+    end
+
     test "tools/call with no args" do
       server = Server.new(tools: [@tool_with_no_args])
 
@@ -1123,7 +1179,19 @@ class TestTool < Tool
       input_schema({ properties: { message: { type: "string" } }, required: ["message"] })
 
       class << self
-        def call(message:, server_context: nil)
+        def call(server_context: nil, **kwargs)
+          Tool::Response.new([{ type: "text", content: "OK" }])
+        end
+      end
+    end
+
+    class TestToolWithAdditionalPropertiesSetToFalse < Tool
+      tool_name "test_tool_with_additional_properties_set_to_false"
+      description "a test tool with additionalProperties set to false for testing"
+      input_schema({ properties: { message: { type: "string" } }, required: ["message"], additionalProperties: false })
+
+      class << self
+        def call(server_context: nil, **kwargs)
           Tool::Response.new([{ type: "text", content: "OK" }])
         end
       end