Skip to content

Refactor schema validation to reduce duplication and add additionalProperties support #198

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ebeigarts wants to merge 1 commit into
base: main
Choose a base branch
from
+240 −162
Open

Refactor schema validation to reduce duplication and add additionalProperties support #198

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 1 addition & 3 deletions lib/mcp/tool.rb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
62 changes: 6 additions & 56 deletions lib/mcp/tool/input_schema.rb
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
58 changes: 3 additions & 55 deletions lib/mcp/tool/output_schema.rb
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
65 changes: 65 additions & 0 deletions lib/mcp/tool/schema.rb
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
70 changes: 69 additions & 1 deletion test/mcp/server_test.rb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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])

Expand Down Expand Up @@ -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
Expand Down
Loading