Merge pull request #369 from koic/switch_to_json_schemer

Speed Up `Tool::Schema` Validation by 5x to 100x
This commit is contained in:
Koichi ITO 2026-06-13 20:13:00 +09:00 • committed by GitHub
commit 39676bc93c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 81 additions and 42 deletions

View file

@ -12,9 +12,9 @@ module MCP
end end
def missing_required_arguments(arguments) def missing_required_arguments(arguments)
return [] unless schema[:required].is_a?(Array) return [] unless @schema[:required].is_a?(Array)
(schema[:required] - arguments.keys.map(&:to_s)) (@schema[:required] - arguments.keys.map(&:to_s))
end end
def validate_arguments(arguments) def validate_arguments(arguments)

View file

@ -1,7 +1,7 @@
# frozen_string_literal: true # frozen_string_literal: true
require "digest" require "digest"
require "json-schema" require "json_schemer"
module MCP module MCP
class Tool class Tool
@ -38,11 +38,10 @@ module MCP
# JSON Schema 2020-12 is the default dialect for MCP schema definitions # JSON Schema 2020-12 is the default dialect for MCP schema definitions
# per MCP 2025-11-25 (SEP-1613). Note: emission only — runtime validation # per MCP 2025-11-25 (SEP-1613). Note: emission only — runtime validation
# is still performed against the JSON Schema draft-04 metaschema because # is still performed against the JSON Schema draft-04 metaschema.
# the `json-schema` gem does not yet support 2020-12.
JSON_SCHEMA_2020_12_URI = "https://json-schema.org/draft/2020-12/schema" JSON_SCHEMA_2020_12_URI = "https://json-schema.org/draft/2020-12/schema"
attr_reader :schema DRAFT4_META_SCHEMA_URI = "http://json-schema.org/draft-04/schema#"
def initialize(schema = {}) def initialize(schema = {})
@schema = JSON.parse(JSON.dump(schema), symbolize_names: true) @schema = JSON.parse(JSON.dump(schema), symbolize_names: true)
@ -51,7 +50,7 @@ module MCP
end end
def ==(other) def ==(other)
other.is_a?(self.class) && schema == other.schema other.is_a?(self.class) && @schema == other.instance_variable_get(:@schema)
end end
def to_h def to_h
@ -62,8 +61,38 @@ module MCP
private private
def stringify(obj)
case obj
when Hash
obj.each_with_object({}) { |(k, v), h| h[k.to_s] = stringify(v) }
when Array
obj.map { |v| stringify(v) }
when Symbol
obj.to_s
else
obj
end
end
# Lazily built so a cache hit in `validate_schema!` avoids the schemer construction cost.
# Memoized per Schema instance because schema content is fixed at construction,
# so the compiled schemer is reusable across many `fully_validate` calls.
#
# `format: false` preserves the legacy behavior of the previous `json-schema` based implementation,
# which did not enforce `format` keywords. `RegexpError` from a malformed `pattern` is re-raised as
# `ArgumentError` so callers see the same exception class they used to.
def schemer
@schemer ||= JSONSchemer.schema(
stringify(schema_for_validation),
meta_schema: DRAFT4_META_SCHEMA_URI,
format: false,
)
rescue RegexpError => e
raise ArgumentError, "Invalid JSON Schema: #{e.message}"
end
def fully_validate(data) def fully_validate(data)
JSON::Validator.fully_validate(schema_for_validation, data) schemer.validate(stringify(data)).map { |validation_error| validation_error.fetch("error") }
end end
def validate_schema! def validate_schema!
@ -75,16 +104,7 @@ module MCP
key = Digest::SHA256.hexdigest(JSON.generate(target, max_nesting: false)) key = Digest::SHA256.hexdigest(JSON.generate(target, max_nesting: false))
return if VALIDATION_CACHE.validated?(key) return if VALIDATION_CACHE.validated?(key)
gem_path = File.realpath(Gem.loaded_specs["json-schema"].full_gem_path) errors = schemer.validate_schema.map { |validation_error| validation_error.fetch("error") }
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, target, schema_reader: schema_reader)
if errors.any? if errors.any?
raise ArgumentError, "Invalid JSON Schema: #{errors.join(", ")}" raise ArgumentError, "Invalid JSON Schema: #{errors.join(", ")}"
end end
@ -92,9 +112,8 @@ module MCP
VALIDATION_CACHE.store(key) VALIDATION_CACHE.store(key)
end end
# The `json-schema` gem's draft-04 validator cannot resolve newer or unknown `$schema` # `json_schemer` is pinned to the draft-04 metaschema, so strip top-level `$schema` before validation:
# dialect URIs. Strip the top-level `$schema` before validation so a dialect URI # this preserves the legacy behavior of ignoring the advertised dialect URI when the SDK validates schemas.
# (whether SDK-injected by `to_h` or user-supplied) does not break the validator.
def schema_for_validation def schema_for_validation
return @schema unless @schema.key?(:"$schema") return @schema unless @schema.key?(:"$schema")

View file

@ -30,5 +30,5 @@ Gem::Specification.new do |spec|
spec.executables = spec.files.grep(%r{^exe/}) { |f| File.basename(f) } spec.executables = spec.files.grep(%r{^exe/}) { |f| File.basename(f) }
spec.require_paths = ["lib"] spec.require_paths = ["lib"]
spec.add_dependency("json-schema", ">= 4.1") spec.add_dependency("json_schemer", ">= 2.4")
end end

View file

@ -7,7 +7,7 @@ module MCP
class InputSchemaTest < ActiveSupport::TestCase class InputSchemaTest < ActiveSupport::TestCase
test "required arguments are converted to strings" do test "required arguments are converted to strings" do
input_schema = InputSchema.new(properties: { message: { type: "string" } }, required: [:message]) input_schema = InputSchema.new(properties: { message: { type: "string" } }, required: [:message])
assert_equal ["message"], input_schema.schema[:required] assert_equal ["message"], input_schema.to_h[:required]
end end
test "to_h returns a hash representation of the input schema" do test "to_h returns a hash representation of the input schema" do
@ -139,10 +139,10 @@ module MCP
test "unexpected errors bubble up from validate_arguments" do test "unexpected errors bubble up from validate_arguments" do
schema = InputSchema.new(properties: { foo: { type: "string" } }, required: ["foo"]) schema = InputSchema.new(properties: { foo: { type: "string" } }, required: ["foo"])
JSON::Validator.stub(:fully_validate, ->(*) { raise "unexpected error" }) do JSONSchemer::Schema.any_instance.stubs(:validate).raises("unexpected error")
assert_raises(RuntimeError) do assert_raises(RuntimeError) do
schema.validate_arguments({ foo: "bar" }) schema.validate_arguments(foo: "bar")
end
end end
end end
@ -200,6 +200,26 @@ module MCP
schema6 = InputSchema.new(properties: { foo: { type: "string" } }, required: ["foo"], additionalProperties: false) schema6 = InputSchema.new(properties: { foo: { type: "string" } }, required: ["foo"], additionalProperties: false)
refute_equal schema1, schema6 refute_equal schema1, schema6
end end
test "format keyword is not enforced (legacy behavior)" do
schema = InputSchema.new(
properties: { email: { type: "string", format: "email" } },
required: ["email"],
)
assert_nil(schema.validate_arguments(email: "not_an_email"))
end
test "invalid pattern raises ArgumentError, not RegexpError" do
error = assert_raises(ArgumentError) do
InputSchema.new(properties: { id: { type: "string", pattern: "[" } })
end
assert_includes error.message, "Invalid JSON Schema"
end
test "Symbol values in arguments are treated as strings" do
schema = InputSchema.new(properties: { foo: { type: "string" } }, required: ["foo"])
assert_nil(schema.validate_arguments(foo: :bar))
end
end end
end end
end end

View file

@ -110,10 +110,10 @@ module MCP
test "unexpected errors bubble up from validate_result" do test "unexpected errors bubble up from validate_result" do
schema = OutputSchema.new(properties: { foo: { type: "string" } }, required: ["foo"]) schema = OutputSchema.new(properties: { foo: { type: "string" } }, required: ["foo"])
JSON::Validator.stub(:fully_validate, ->(*) { raise "unexpected error" }) do JSONSchemer::Schema.any_instance.stubs(:validate).raises("unexpected error")
assert_raises(RuntimeError) do assert_raises(RuntimeError) do
schema.validate_result({ foo: "bar" }) schema.validate_result(foo: "bar")
end
end end
end end

View file

@ -10,7 +10,7 @@ module MCP
end end
test "validates a schema once and reuses the result for identical schemas" do test "validates a schema once and reuses the result for identical schemas" do
JSON::Validator.expects(:fully_validate).once.returns([]) JSONSchemer::Schema.any_instance.expects(:validate_schema).once.returns([])
schema = { properties: { validates_once: { type: "string" } } } schema = { properties: { validates_once: { type: "string" } } }
InputSchema.new(schema) InputSchema.new(schema)
@ -18,7 +18,7 @@ module MCP
end end
test "validates distinct schemas separately" do test "validates distinct schemas separately" do
JSON::Validator.expects(:fully_validate).twice.returns([]) JSONSchemer::Schema.any_instance.expects(:validate_schema).twice.returns([])
InputSchema.new(properties: { distinct_a: { type: "string" } }) InputSchema.new(properties: { distinct_a: { type: "string" } })
InputSchema.new(properties: { distinct_b: { type: "string" } }) InputSchema.new(properties: { distinct_b: { type: "string" } })
@ -64,13 +64,13 @@ module MCP
break break
end end
JSON::Validator.stub(:fully_validate, []) do JSONSchemer::Schema.any_instance.stubs(:validate_schema).returns([])
assert_nothing_raised do assert_nothing_raised do
InputSchema.new(schema) InputSchema.new(schema)
InputSchema.new(schema) InputSchema.new(schema)
end end
end end
end
test "ValidationCache evicts the oldest entry beyond its max size" do test "ValidationCache evicts the oldest entry beyond its max size" do
cache = Schema::ValidationCache.new(max_size: 2) cache = Schema::ValidationCache.new(max_size: 2)

View file

@ -167,8 +167,8 @@ module MCP
end end
assert_includes error.message, "Invalid JSON Schema" assert_includes error.message, "Invalid JSON Schema"
assert_includes error.message, "#/properties/count/minimum" assert_includes error.message, "properties/count/minimum"
assert_includes error.message, "string did not match the following type: number" assert_includes error.message, "number"
end end
test ".define allows definition of simple tools with a block" do test ".define allows definition of simple tools with a block" do
@ -431,8 +431,8 @@ module MCP
end end
assert_includes error.message, "Invalid JSON Schema" assert_includes error.message, "Invalid JSON Schema"
assert_includes error.message, "#/properties/count/minimum" assert_includes error.message, "properties/count/minimum"
assert_includes error.message, "string did not match the following type: number" assert_includes error.message, "number"
end end
test "output_schema accepts $ref in schema" do test "output_schema accepts $ref in schema" do