Validate ToolCall Args against Schema
(cherry picked from commit eeb287117535dfb43785704171625323f01d5bda)
This commit is contained in:
parent
a601558a21
commit
600c96f66a
8 changed files with 275 additions and 3 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -150,7 +150,7 @@ module MCP
|
|||
report_exception(e, { request: request })
|
||||
if e.is_a?(RequestHandlerError)
|
||||
add_instrumentation_data(error: e.error_type)
|
||||
raise e
|
||||
raise StandardError, e.message
|
||||
end
|
||||
|
||||
add_instrumentation_data(error: :internal_error)
|
||||
|
|
@ -213,6 +213,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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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,32 @@ 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!
|
||||
schema = to_h
|
||||
begin
|
||||
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
|
||||
rescue StandardError => e
|
||||
raise ArgumentError, "Invalid JSON Schema: #{e.message}"
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ Gem::Specification.new do |spec|
|
|||
spec.require_paths = ["lib"]
|
||||
|
||||
spec.add_dependency("json_rpc_handler", "~> 0.1")
|
||||
spec.add_dependency("json-schema", "~> 4.1")
|
||||
spec.add_development_dependency("activesupport")
|
||||
spec.add_development_dependency("sorbet-static-and-runtime")
|
||||
end
|
||||
|
|
|
|||
|
|
@ -61,5 +61,45 @@ module MCP
|
|||
merged = config3.merge(config1)
|
||||
assert_equal "2025-03-27", merged.protocol_version
|
||||
end
|
||||
|
||||
test "defaults validate_tool_call_arguments to true" do
|
||||
config = Configuration.new
|
||||
assert config.validate_tool_call_arguments
|
||||
end
|
||||
|
||||
test "can set validate_tool_call_arguments to false" do
|
||||
config = Configuration.new(validate_tool_call_arguments: false)
|
||||
refute config.validate_tool_call_arguments
|
||||
end
|
||||
|
||||
test "validate_tool_call_arguments? returns false when set" do
|
||||
config = Configuration.new(validate_tool_call_arguments: false)
|
||||
refute config.validate_tool_call_arguments?
|
||||
end
|
||||
|
||||
test "validate_tool_call_arguments? returns true when not set" do
|
||||
config = Configuration.new
|
||||
assert config.validate_tool_call_arguments?
|
||||
end
|
||||
|
||||
test "merge preserves validate_tool_call_arguments from other config" do
|
||||
config1 = Configuration.new(validate_tool_call_arguments: false)
|
||||
config2 = Configuration.new
|
||||
merged = config1.merge(config2)
|
||||
assert merged.validate_tool_call_arguments?
|
||||
end
|
||||
|
||||
test "merge preserves validate_tool_call_arguments from self when other not set" do
|
||||
config1 = Configuration.new(validate_tool_call_arguments: false)
|
||||
config2 = Configuration.new
|
||||
merged = config2.merge(config1)
|
||||
refute merged.validate_tool_call_arguments
|
||||
end
|
||||
|
||||
test "raises ArgumentError when validate_tool_call_arguments is not a boolean" do
|
||||
assert_raises(ArgumentError) do
|
||||
Configuration.new(validate_tool_call_arguments: "true")
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
|
|
|
|||
|
|
@ -839,5 +839,123 @@ module MCP
|
|||
|
||||
refute_includes server_without_resources.capabilities, :resources
|
||||
end
|
||||
|
||||
test "tools/call validates arguments against input schema when validate_tool_call_arguments is true" 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: 123 },
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
assert_equal "2.0", response[:jsonrpc]
|
||||
assert_equal 1, response[:id]
|
||||
assert_equal(-32603, response[:error][:code])
|
||||
assert_includes response[:error][:data], "Invalid arguments"
|
||||
end
|
||||
|
||||
test "tools/call skips argument validation when validate_tool_call_arguments is false" do
|
||||
server = Server.new(
|
||||
tools: [TestTool],
|
||||
configuration: Configuration.new(validate_tool_call_arguments: false),
|
||||
)
|
||||
|
||||
response = server.handle(
|
||||
{
|
||||
jsonrpc: "2.0",
|
||||
id: 1,
|
||||
method: "tools/call",
|
||||
params: {
|
||||
name: "test_tool",
|
||||
arguments: { message: 123 },
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
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 validates arguments with complex types" do
|
||||
server = Server.new(
|
||||
tools: [ComplexTypesTool],
|
||||
configuration: Configuration.new(validate_tool_call_arguments: true),
|
||||
)
|
||||
|
||||
response = server.handle(
|
||||
{
|
||||
jsonrpc: "2.0",
|
||||
id: 1,
|
||||
method: "tools/call",
|
||||
params: {
|
||||
name: "complex_types_tool",
|
||||
arguments: {
|
||||
numbers: [1, 2, 3],
|
||||
strings: ["a", "b", "c"],
|
||||
objects: [{ name: "test" }],
|
||||
},
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
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
|
||||
|
||||
class TestTool < Tool
|
||||
tool_name "test_tool"
|
||||
description "a test tool for testing"
|
||||
input_schema({ properties: { message: { type: "string" } }, required: ["message"] })
|
||||
|
||||
class << self
|
||||
def call(message:, server_context: nil)
|
||||
Tool::Response.new([{ type: "text", content: "OK" }])
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
class ComplexTypesTool < Tool
|
||||
tool_name "complex_types_tool"
|
||||
description "a test tool with complex types"
|
||||
input_schema({
|
||||
properties: {
|
||||
numbers: { type: "array", items: { type: "number" } },
|
||||
strings: { type: "array", items: { type: "string" } },
|
||||
objects: {
|
||||
type: "array",
|
||||
items: {
|
||||
type: "object",
|
||||
properties: {
|
||||
name: { type: "string" },
|
||||
},
|
||||
required: ["name"],
|
||||
},
|
||||
},
|
||||
},
|
||||
required: ["numbers", "strings", "objects"],
|
||||
})
|
||||
|
||||
class << self
|
||||
def call(numbers:, strings:, objects:, server_context: nil)
|
||||
Tool::Response.new([{ type: "text", content: "OK" }])
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# frozen_string_literal: true
|
||||
|
||||
require "test_helper"
|
||||
require "mcp/tool/input_schema"
|
||||
|
||||
module MCP
|
||||
class Tool
|
||||
|
|
@ -27,6 +28,29 @@ module MCP
|
|||
input_schema = InputSchema.new(properties: { message: { type: "string" } }, required: [:message])
|
||||
assert_empty input_schema.missing_required_arguments({ message: "Hello, world!" })
|
||||
end
|
||||
|
||||
test "valid schema initialization" do
|
||||
schema = InputSchema.new(properties: { foo: { type: "string" } }, required: [:foo])
|
||||
assert_equal({ type: "object", properties: { foo: { type: "string" } }, required: [:foo] }, schema.to_h)
|
||||
end
|
||||
|
||||
test "invalid schema raises argument error" do
|
||||
assert_raises(ArgumentError) do
|
||||
InputSchema.new(properties: { foo: { type: "invalid_type" } }, required: [:foo])
|
||||
end
|
||||
end
|
||||
|
||||
test "validate arguments with valid data" do
|
||||
schema = InputSchema.new(properties: { foo: { type: "string" } }, required: [:foo])
|
||||
assert_nil(schema.validate_arguments({ foo: "bar" }))
|
||||
end
|
||||
|
||||
test "validate arguments with invalid data" do
|
||||
schema = InputSchema.new(properties: { foo: { type: "string" } }, required: [:foo])
|
||||
assert_raises(InputSchema::ValidationError) do
|
||||
schema.validate_arguments({ foo: 123 })
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
|
|
|
|||
|
|
@ -82,6 +82,23 @@ module MCP
|
|||
assert_equal expected, tool.input_schema.to_h
|
||||
end
|
||||
|
||||
test "raises detailed error message for invalid schema" do
|
||||
error = assert_raises(ArgumentError) do
|
||||
Class.new(MCP::Tool) do
|
||||
input_schema(
|
||||
properties: {
|
||||
count: { type: "integer", minimum: "not a number" },
|
||||
},
|
||||
required: [:count],
|
||||
)
|
||||
end
|
||||
end
|
||||
|
||||
assert_includes error.message, "Invalid JSON Schema"
|
||||
assert_includes error.message, "#/properties/count/minimum"
|
||||
assert_includes error.message, "string did not match the following type: number"
|
||||
end
|
||||
|
||||
test ".define allows definition of simple tools with a block" do
|
||||
tool = Tool.define(name: "mock_tool", description: "a mock tool for testing") do |_|
|
||||
Tool::Response.new([{ type: "text", content: "OK" }])
|
||||
|
|
@ -226,5 +243,23 @@ module MCP
|
|||
assert_equal response.content, [{ type: "text", content: "OK" }]
|
||||
assert_equal response.is_error, false
|
||||
end
|
||||
|
||||
test "input_schema rejects any $ref in schema" do
|
||||
schema_with_ref = {
|
||||
properties: {
|
||||
foo: { "$ref" => "#/definitions/bar" },
|
||||
},
|
||||
required: [],
|
||||
definitions: {
|
||||
bar: { type: "string" },
|
||||
},
|
||||
}
|
||||
error = assert_raises(ArgumentError) do
|
||||
Class.new(MCP::Tool) do
|
||||
input_schema schema_with_ref
|
||||
end
|
||||
end
|
||||
assert_match(/Invalid JSON Schema/, error.message)
|
||||
end
|
||||
end
|
||||
end
|
||||
|
|
|
|||
Loading…
Reference in a new issue