Merge pull request #43 from tylerrowsell/tool-arg-validation

Validate ToolCall Args against Schema

lib/mcp/configuration.rb

@@ -4,12 +4,18 @@ module MCP
   class Configuration
     DEFAULT_PROTOCOL_VERSION = "2024-11-05"
 
-    attr_writer :exception_reporter, :instrumentation_callback, :protocol_version
+    attr_writer :exception_reporter, :instrumentation_callback, :protocol_version, :validate_tool_call_arguments
 
-    def initialize(exception_reporter: nil, instrumentation_callback: nil, protocol_version: nil)
+    def initialize(exception_reporter: nil, instrumentation_callback: nil, protocol_version: nil,
+      validate_tool_call_arguments: true)
       @exception_reporter = exception_reporter
       @instrumentation_callback = instrumentation_callback
       @protocol_version = protocol_version
+      unless validate_tool_call_arguments.is_a?(TrueClass) || validate_tool_call_arguments.is_a?(FalseClass)
+        raise ArgumentError, "validate_tool_call_arguments must be a boolean"
+      end
+
+      @validate_tool_call_arguments = validate_tool_call_arguments
     end
 
     def protocol_version
@@ -36,6 +42,12 @@ module MCP
       !@instrumentation_callback.nil?
     end
 
+    attr_reader :validate_tool_call_arguments
+
+    def validate_tool_call_arguments?
+      !!@validate_tool_call_arguments
+    end
+
     def merge(other)
       return self if other.nil?
 
@@ -54,11 +66,13 @@ module MCP
       else
         @protocol_version
       end
+      validate_tool_call_arguments = other.validate_tool_call_arguments
 
       Configuration.new(
         exception_reporter:,
         instrumentation_callback:,
         protocol_version:,
+        validate_tool_call_arguments:,
       )
     end
 

lib/mcp/server.rb

@@ -233,6 +233,15 @@ module MCP
         )
       end
 
+      if configuration.validate_tool_call_arguments && tool.input_schema
+        begin
+          tool.input_schema.validate_arguments(arguments)
+        rescue Tool::InputSchema::ValidationError => e
+          add_instrumentation_data(error: :invalid_schema)
+          raise RequestHandlerError.new(e.message, request, error_type: :invalid_schema)
+        end
+      end
+
       begin
         call_params = tool_call_parameters(tool)
 

lib/mcp/tool/input_schema.rb

@@ -1,13 +1,18 @@
 # frozen_string_literal: true
 
+require "json-schema"
+
 module MCP
   class Tool
     class InputSchema
+      class ValidationError < StandardError; end
+
       attr_reader :properties, :required
 
       def initialize(properties: {}, required: [])
         @properties = properties
         @required = required.map(&:to_sym)
+        validate_schema!
       end
 
       def to_h
@@ -21,6 +26,42 @@ module MCP
       def missing_required_arguments(arguments)
         (required - arguments.keys.map(&:to_sym))
       end
+
+      def validate_arguments(arguments)
+        errors = JSON::Validator.fully_validate(to_h, arguments)
+        if errors.any?
+          raise ValidationError, "Invalid arguments: #{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 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