diff --git a/Gemfile.lock b/Gemfile.lock index 06e342d..de021d3 100644 --- a/Gemfile.lock +++ b/Gemfile.lock @@ -1,7 +1,7 @@ PATH remote: . specs: - model_context_protocol (0.2.5) + model_context_protocol (0.3.0) json_rpc_handler (~> 0.1) GEM diff --git a/README.md b/README.md index 022b094..6920053 100644 --- a/README.md +++ b/README.md @@ -47,7 +47,8 @@ module ModelContextProtocol server = ModelContextProtocol::Server.new( name: "my_server", tools: [someTool, anotherTool], - prompts: [myPrompt] + prompts: [myPrompt], + context: nil, ) render(json: server.handle(request.body.read).to_h) end @@ -55,6 +56,41 @@ module ModelContextProtocol end ``` +## Configuration + +The gem can be configured using the `ModelContextProtocol.configure` block: + +```ruby +ModelContextProtocol.configure do |config| + config.exception_reporter = ->(exception, context) do + # Your exception reporting logic here + # For example with Bugsnag: + Bugsnag.notify(exception) do |report| + report.add_metadata(:model_context_protocol, context) + end + end + + config.instrumentation_callback = -> (data) { puts "Got instrumentation data #{data.inspect}" } +end +``` + +### Exception Reporting + +The exception reporter receives two arguments: +- `exception`: The Ruby exception object that was raised +- `context`: A hash containing contextual information about where the error occurred + +The context hash includes: +- For tool calls: `{ tool_name: "name", arguments: { ... } }` +- For general request handling: `{ request: { ... } }` + +When an exception occurs: +1. The exception is reported via the configured reporter +2. For tool calls, a generic error response is returned to the client: `{ error: "Internal error occurred", is_error: true }` +3. For other requests, the exception is re-raised after reporting + +If no exception reporter is configured, a default no-op reporter is used that silently ignores exceptions. + ## Tools MCP spec includes [Tools](https://modelcontextprotocol.io/docs/concepts/tools) which provide functionality to LLM apps. @@ -65,11 +101,10 @@ This gem provides a `ModelContextProtocol::Tool` class that can be used to creat ```ruby class MyTool < ModelContextProtocol::Tool - tool_name "my_tool" - tool_description "This tool performs specific functionality..." - tool_input_schema [{ type: "text", name: "message" }] + description "This tool performs specific functionality..." + input_schema [{ type: "text", name: "message" }] - def call(message) + def call(message, context:) Tool::Response.new([{ type: "text", content: "OK" }]) end end @@ -80,11 +115,14 @@ tool = MyTool.new 2. By using the `ModelContextProtocol::Tool.define` method with a block: ```ruby -tool = ModelContextProtocol::Tool.define(name: "my_tool", description: "This tool performs specific functionality...") do |args| +tool = ModelContextProtocol::Tool.define(name: "my_tool", description: "This tool performs specific functionality...") do |args, context| Tool::Response.new([{ type: "text", content: "OK" }]) end ``` +The context parameter is the context passed into the server and can be used to pass per request information, +e.g. around authentication state. + ## Prompts MCP spec includes [Prompts](https://modelcontextprotocol.io/docs/concepts/prompts), which enable servers to define reusable prompt templates and workflows that clients can easily surface to users and LLMs. @@ -167,7 +205,8 @@ Register prompts with the MCP server: ```ruby server = ModelContextProtocol::Server.new( name: "my_server", - prompts: [MyPrompt.new] + prompts: [MyPrompt.new], + context: nil, ) ``` @@ -176,6 +215,28 @@ The server will handle prompt listing and execution through the MCP protocol met - `prompts/list` - Lists all registered prompts and their schemas - `prompts/get` - Retrieves and executes a specific prompt with arguments +### Instrumentation + +The server allows registering a callback to receive information about instrumentation. +To register a handler pass a proc/lambda to as `instrumentation_callback` into the server constructor. + +```ruby +ModelContextProtocol.configure do |config| + config.instrumentation_callback = -> (data) { puts "Got instrumentation data #{data.inspect}" } +end +``` + +The data contains the following keys: +`method`: the metod called, e.g. `ping`, `tools/list`, `tools/call` etc +`tool_name`: the name of the tool called +`prompt_name`: the name of the prompt called +`resource_uri`: the uri of the resource called +`error`: if looking up tools/prompts etc failed, e.g. `tool_not_found` +`duration`: the duration of the call in seconds + +`tool_name`, `prompt_name` and `resource_uri` are only populated if a matching handler is registered. +This is to avoid potential issues with metric cardinality + ## Releases TODO diff --git a/lib/model_context_protocol.rb b/lib/model_context_protocol.rb index 6444991..c5be614 100644 --- a/lib/model_context_protocol.rb +++ b/lib/model_context_protocol.rb @@ -8,8 +8,19 @@ require "model_context_protocol/content" require "model_context_protocol/resource" require "model_context_protocol/prompt" require "model_context_protocol/version" +require "model_context_protocol/configuration" module ModelContextProtocol + class << self + def configure + yield(configuration) + end + + def configuration + @configuration ||= Configuration.new + end + end + class Annotations attr_reader :audience, :priority diff --git a/lib/model_context_protocol/configuration.rb b/lib/model_context_protocol/configuration.rb new file mode 100644 index 0000000..3792477 --- /dev/null +++ b/lib/model_context_protocol/configuration.rb @@ -0,0 +1,13 @@ +# frozen_string_literal: true + +module ModelContextProtocol + class Configuration + attr_accessor :exception_reporter + attr_accessor :instrumentation_callback + + def initialize + @exception_reporter = ->(exception, context) {} # Default no-op reporter + @instrumentation_callback = ->(data) {} # Default no-op callback + end + end +end diff --git a/lib/model_context_protocol/instrumentation.rb b/lib/model_context_protocol/instrumentation.rb new file mode 100644 index 0000000..044896a --- /dev/null +++ b/lib/model_context_protocol/instrumentation.rb @@ -0,0 +1,24 @@ +# frozen_string_literal: true + +module ModelContextProtocol + module Instrumentation + def instrument_call(method, &block) + start_time = Process.clock_gettime(Process::CLOCK_MONOTONIC) + @instrumentation_data = {} + add_instrumentation_data(method:) + + result = yield block + + end_time = Process.clock_gettime(Process::CLOCK_MONOTONIC) + add_instrumentation_data(duration: end_time - start_time) + + ModelContextProtocol.configuration.instrumentation_callback.call(@instrumentation_data) + + result + end + + def add_instrumentation_data(**kwargs) + @instrumentation_data.merge!(kwargs) + end + end +end diff --git a/lib/model_context_protocol/server.rb b/lib/model_context_protocol/server.rb index cf623a5..08f3ed3 100644 --- a/lib/model_context_protocol/server.rb +++ b/lib/model_context_protocol/server.rb @@ -1,20 +1,30 @@ # frozen_string_literal: true require "json_rpc_handler" +require_relative "instrumentation" module ModelContextProtocol class Server + class RequestHandlerError < StandardError + def initialize(message, request) + super(message) + @request = request + end + end + PROTOCOL_VERSION = "2024-11-05" - attr_accessor :name, :tools, :prompts, :resources + include Instrumentation - def initialize(name: "model_context_protocol", tools: [], prompts: [], resources: []) + attr_accessor :name, :tools, :prompts, :resources, :context + + def initialize(name: "model_context_protocol", tools: [], prompts: [], resources: [], context: nil) @name = name @tools = tools.to_h { |t| [t.name, t] } @prompts = prompts.to_h { |p| [p.name, p] } @resources = resources @resource_index = resources.index_by(&:uri) - + @context = context @handlers = { "resources/list" => method(:list_resources), "resources/read" => method(:read_resource), @@ -29,18 +39,21 @@ module ModelContextProtocol def handle(request) JsonRpcHandler.handle(request) do |method| - handler = case method - when "tools/list" - ->(params) { { tools: @handlers["tools/list"].call(params) } } - when "prompts/list" - ->(params) { { prompts: @handlers["prompts/list"].call(params) } } - when "resources/list" - ->(params) { { resources: @handlers["resources/list"].call(params) } } - else - @handlers[method] + instrument_call(method) do + case method + when "tools/list" + ->(params) { { tools: @handlers["tools/list"].call(params) } } + when "prompts/list" + ->(params) { { prompts: @handlers["prompts/list"].call(params) } } + when "resources/list" + ->(params) { { resources: @handlers["resources/list"].call(params) } } + else + @handlers[method] + end + rescue => e + report_exception(e, { request: request }) + raise RequestHandlerError.new("Internal error handling #{request[:method]} request", request) end - - handler end end @@ -86,6 +99,7 @@ module ModelContextProtocol end def init(request) + add_instrumentation_data(method: "initialize") { protocolVersion: PROTOCOL_VERSION, capabilities: capabilities, @@ -94,47 +108,76 @@ module ModelContextProtocol end def list_tools(request) + add_instrumentation_data(method: "tools/list") @tools.map { |_, tool| tool.to_h } end def call_tool(request) + add_instrumentation_data(method: "tools/call") tool_name = request[:name] tool = tools[tool_name] - raise "Tool not found #{tool_name}" unless tool + unless tool + add_instrumentation_data(error: :tool_not_found) + raise "Tool not found #{tool_name}" + end - result = tool.call(**request[:arguments]) - result.to_h + add_instrumentation_data(tool_name:) + + begin + result = tool.call(**request[:arguments], context:) + result.to_h + rescue => e + report_exception(e, { tool_name: tool_name, arguments: request[:arguments] }) + add_instrumentation_data(error: :internal_error) + raise RequestHandlerError.new("Internal error calling tool #{tool_name}", request) + end end def list_prompts(request) + add_instrumentation_data(method: "prompts/list") @prompts.map { |_, prompt| prompt.to_h } end def get_prompt(request) + add_instrumentation_data(method: "prompts/get") prompt_name = request[:name] prompt = @prompts[prompt_name] + unless prompt + add_instrumentation_data(error: :prompt_not_found) + raise "Prompt not found #{prompt_name}" + end - raise "Prompt not found #{prompt_name}" unless prompt + add_instrumentation_data(prompt_name:) prompt_args = request[:arguments] prompt.validate_arguments!(prompt_args) - result = prompt.template(prompt_args) - - result.to_h + prompt.template(prompt_args).to_h end def list_resources(request) + add_instrumentation_data(method: "resources/list") + @resources.map(&:to_h) end def read_resource(request) + add_instrumentation_data(method: "resources/read") resource_uri = request[:uri] resource = @resource_index[resource_uri] - raise "Resource not found #{resource_uri}" unless resource + unless resource + add_instrumentation_data(error: :resource_not_found) + raise "Resource not found #{resource_uri}" + end + + add_instrumentation_data(resource_uri:) resource.to_h end + + def report_exception(exception, context = {}) + ModelContextProtocol.configuration.exception_reporter.call(exception, context) + end end end diff --git a/lib/model_context_protocol/tool.rb b/lib/model_context_protocol/tool.rb index 2ccfbae..c6acfe9 100644 --- a/lib/model_context_protocol/tool.rb +++ b/lib/model_context_protocol/tool.rb @@ -44,8 +44,8 @@ module ModelContextProtocol def define(name: nil, description: nil, input_schema: nil, &block) new(name:, description:, input_schema:).tap do |tool| - tool.define_singleton_method(:call) do |*args| - instance_exec(*args, &block) + tool.define_singleton_method(:call) do |*args, context:| + instance_exec(*args, context:, &block) end end end @@ -59,7 +59,7 @@ module ModelContextProtocol @input_schema = input_schema || self.class.input_schema_value end - def call(*args) + def call(*args, context:) raise NotImplementedError, "Subclasses must implement call" end diff --git a/lib/model_context_protocol/version.rb b/lib/model_context_protocol/version.rb index 9bf8da1..b40e57b 100644 --- a/lib/model_context_protocol/version.rb +++ b/lib/model_context_protocol/version.rb @@ -1,5 +1,5 @@ # frozen_string_literal: true module ModelContextProtocol - VERSION = "0.2.5" + VERSION = "0.3.0" end diff --git a/test/instrumentation_test_helper.rb b/test/instrumentation_test_helper.rb new file mode 100644 index 0000000..de65be2 --- /dev/null +++ b/test/instrumentation_test_helper.rb @@ -0,0 +1,24 @@ +# frozen_string_literal: true + +module InstrumentationTestHelper + class Instrumentation + attr_reader :callback + attr_reader :data + + def initialize + @callback = ->(data) { @data = data } + end + end + + def instrumentation_helper + @instrumentation_helper ||= Instrumentation.new + end + + def assert_instrumentation_data(expected_data) + data = instrumentation_helper.data.dup + duration = data.delete(:duration) + assert_not_nil(duration, "Duration is not set") + assert_operator(duration, :>=, 0, "Duration is not positive or zero") + assert_equal(expected_data, data, "Instrumentation data does not match expected data") + end +end diff --git a/test/model_context_protocol/configuration_test.rb b/test/model_context_protocol/configuration_test.rb new file mode 100644 index 0000000..dfdde5c --- /dev/null +++ b/test/model_context_protocol/configuration_test.rb @@ -0,0 +1,38 @@ +# frozen_string_literal: true + +require "test_helper" + +module ModelContextProtocol + class ConfigurationTest < ActiveSupport::TestCase + test "initializes with a default no-op exception reporter" do + config = Configuration.new + assert_respond_to config, :exception_reporter + + # The default reporter should be callable but do nothing + exception = StandardError.new("test error") + context = { test: "context" } + + # Should not raise any errors + config.exception_reporter.call(exception, context) + end + + test "allows setting a custom exception reporter" do + config = Configuration.new + reported_exception = nil + reported_context = nil + + config.exception_reporter = ->(exception, context) do + reported_exception = exception + reported_context = context + end + + test_exception = StandardError.new("test error") + test_context = { foo: "bar" } + + config.exception_reporter.call(test_exception, test_context) + + assert_equal test_exception, reported_exception + assert_equal test_context, reported_context + end + end +end diff --git a/test/model_context_protocol/instrumentation_test.rb b/test/model_context_protocol/instrumentation_test.rb new file mode 100644 index 0000000..6b5ef99 --- /dev/null +++ b/test/model_context_protocol/instrumentation_test.rb @@ -0,0 +1,64 @@ +# frozen_string_literal: true + +require "test_helper" + +module ModelContextProtocol + class InstrumentationTest < ActiveSupport::TestCase + class Subject + include Instrumentation + attr_reader :instrumentation_data_received + + def initialize + ModelContextProtocol.configure do |config| + config.instrumentation_callback = ->(data) { @instrumentation_data_received = data } + end + end + + def instrumented_method + instrument_call("instrumented_method") do + # nothing to do + end + end + + def instrumented_method_with_additional_data + instrument_call("instrumented_method_with_additional_data") do + add_instrumentation_data(additional_data: "test") + end + end + end + + test "#instrument_call adds the method name to the instrumentation data" do + subject = Subject.new + + subject.instrumented_method + assert_equal({ method: "instrumented_method" }, subject.instrumentation_data_received.tap do |data| + data.delete(:duration) + end) + end + + test "#instrument_call exposes data added via add_instrumentation_data" do + subject = Subject.new + + subject.instrumented_method_with_additional_data + assert_equal( + { method: "instrumented_method_with_additional_data", additional_data: "test" }, + subject.instrumentation_data_received.tap { |data| data.delete(:duration) }, + ) + end + + test "#instrument_call resets the instrumentation data between calls" do + subject = Subject.new + + subject.instrumented_method_with_additional_data + assert_equal( + { method: "instrumented_method_with_additional_data", additional_data: "test" }, + subject.instrumentation_data_received.tap { |data| data.delete(:duration) }, + ) + + subject.instrumented_method + assert_equal({ method: "instrumented_method" }, subject.instrumentation_data_received.tap do |data| + data.delete(:duration) + end) + end + end +end diff --git a/test/model_context_protocol/server_test.rb b/test/model_context_protocol/server_test.rb index ae826bc..1b2ae44 100644 --- a/test/model_context_protocol/server_test.rb +++ b/test/model_context_protocol/server_test.rb @@ -4,6 +4,7 @@ require "test_helper" module ModelContextProtocol class ServerTest < ActiveSupport::TestCase + include InstrumentationTestHelper setup do @tool = Tool.new(name: "test_tool", description: "Test tool", input_schema: {}) @@ -31,7 +32,16 @@ module ModelContextProtocol ) @server_name = "test_server" - @server = Server.new(name: @server_name, tools: [@tool], prompts: [@prompt], resources: [@resource]) + ModelContextProtocol.configure do |config| + config.instrumentation_callback = instrumentation_helper.callback + end + + @server = Server.new( + name: @server_name, + tools: [@tool], + prompts: [@prompt], + resources: [@resource], + ) end # https://spec.modelcontextprotocol.io/specification/2024-11-05/basic/utilities/ping/#behavior-requirements @@ -51,6 +61,7 @@ module ModelContextProtocol }, response, ) + assert_instrumentation_data({ method: "ping" }) end test "#handle initialize request returns protocol info, server info, and capabilities" do @@ -81,6 +92,7 @@ module ModelContextProtocol } assert_equal expected_result, response + assert_instrumentation_data({ method: "initialize" }) end test "#handle returns nil for notification requests" do @@ -90,6 +102,7 @@ module ModelContextProtocol } assert_nil @server.handle(request) + assert_instrumentation_data({ method: "some_notification" }) end test "#handle tools/list returns available tools" do @@ -105,6 +118,7 @@ module ModelContextProtocol assert_equal "test_tool", result[:tools][0][:name] assert_equal "Test tool", result[:tools][0][:description] assert_equal({}, result[:tools][0][:inputSchema]) + assert_instrumentation_data({ method: "tools/list" }) end test "#tools_list_handler sets the tools/list handler" do @@ -121,6 +135,7 @@ module ModelContextProtocol response = @server.handle(request) result = response[:result] assert_equal({ tools: [{ name: "hammer", description: "Hammer time!" }] }, result) + assert_instrumentation_data({ method: "tools/list" }) end test "#handle tools/call executes tool and returns result" do @@ -128,7 +143,7 @@ module ModelContextProtocol tool_args = { "arg" => "value" } tool_response = Tool::Response.new([{ "result" => "success" }]) - @tool.expects(:call).with(**tool_args).returns(tool_response) + @tool.expects(:call).with(**tool_args, context: @server.context).returns(tool_response) request = { jsonrpc: "2.0", @@ -142,10 +157,20 @@ module ModelContextProtocol response = @server.handle(request) assert_equal tool_response.to_h, response[:result] + assert_instrumentation_data({ method: "tools/call", tool_name: }) end - test "#handle tools/call returns error if the tool raises an error" do - @tool.expects(:call).raises(StandardError.new("Tool error")) + test "#handle tools/call returns internal error and reports exception if the tool raises an error" do + exception = StandardError.new("Tool error") + @tool.expects(:call).raises(exception) + + ModelContextProtocol.configuration.exception_reporter.expects(:call).with( + exception, + { + tool_name: "test_tool", + arguments: {}, + }, + ) request = { jsonrpc: "2.0", @@ -160,7 +185,8 @@ module ModelContextProtocol response = @server.handle(request) assert_equal "Internal error", response[:error][:message] - assert_equal "Tool error", response[:error][:data] + assert_equal "Internal error calling tool test_tool", response[:error][:data] + assert_instrumentation_data({ method: "tools/call", tool_name: "test_tool", error: :internal_error }) end test "#handle tools/call returns error for unknown tool" do @@ -177,6 +203,7 @@ module ModelContextProtocol response = @server.handle(request) assert_equal "Internal error", response[:error][:message] assert_equal "Tool not found unknown_tool", response[:error][:data] + assert_instrumentation_data({ method: "tools/call", error: :tool_not_found }) end test "#tools_call_handler sets the tools/call handler" do @@ -194,6 +221,7 @@ module ModelContextProtocol response = @server.handle(request) assert_equal({ content: "my_tool called successfully", is_error: false }, response[:result]) + assert_instrumentation_data({ method: "tools/call", tool_name: "my_tool" }) end test "#handle prompts/list returns list of prompts" do @@ -205,6 +233,7 @@ module ModelContextProtocol response = @server.handle(request) assert_equal({ prompts: [@prompt.to_h] }, response[:result]) + assert_instrumentation_data({ method: "prompts/list" }) end test "#prompts_list_handler sets the prompts/list handler" do @@ -220,6 +249,7 @@ module ModelContextProtocol response = @server.handle(request) assert_equal({ prompts: [{ name: "foo_prompt", description: "Foo prompt" }] }, response[:result]) + assert_instrumentation_data({ method: "prompts/list" }) end test "#handle prompts/get returns templated prompt" do @@ -242,6 +272,7 @@ module ModelContextProtocol response = @server.handle(request) assert_equal(expected_result, response[:result]) + assert_instrumentation_data({ method: "prompts/get", prompt_name: "test_prompt" }) end test "#handle prompts/get returns error if prompt is not found" do @@ -257,6 +288,7 @@ module ModelContextProtocol response = @server.handle(request) assert_equal("Prompt not found unknown_prompt", response[:error][:data]) + assert_instrumentation_data({ method: "prompts/get", error: :prompt_not_found }) end test "#handle prompts/get returns error if prompt arguments are invalid" do @@ -272,6 +304,7 @@ module ModelContextProtocol response = @server.handle(request) assert_equal "Missing required arguments: test_argument", response[:error][:data] + assert_instrumentation_data({ method: "prompts/get", prompt_name: "test_prompt" }) end test "#prompts_get_handler sets the prompts/get handler" do @@ -297,6 +330,7 @@ module ModelContextProtocol { description: "foo_bar_prompt", messages: [{ role: "user", content: { text: "bar" } }] }, response[:result], ) + assert_instrumentation_data({ method: "prompts/get", prompt_name: "foo_bar_prompt" }) end test "#handle resources/list returns a list of resources" do @@ -308,6 +342,7 @@ module ModelContextProtocol response = @server.handle(request) assert_equal({ resources: [@resource.to_h] }, response[:result]) + assert_instrumentation_data({ method: "resources/list" }) end test "#resources_list_handler sets the resources/list handler" do @@ -326,6 +361,7 @@ module ModelContextProtocol { resources: [{ uri: "test_resource", name: "Test resource", description: "Test resource" }] }, response[:result], ) + assert_instrumentation_data({ method: "resources/list" }) end test "#handle resources/read returns a resource" do @@ -340,6 +376,7 @@ module ModelContextProtocol response = @server.handle(request) assert_equal(@resource.to_h, response[:result]) + assert_instrumentation_data({ method: "resources/read", resource_uri: @resource.uri }) end test "#handle resources/read returns error if resource is not found" do @@ -354,6 +391,7 @@ module ModelContextProtocol response = @server.handle(request) assert_equal "Resource not found unknown_resource", response[:error][:data] + assert_instrumentation_data({ method: "resources/read", error: :resource_not_found }) end test "#resources_read_handler sets the resources/read handler" do @@ -389,6 +427,7 @@ module ModelContextProtocol }, response[:result], ) + assert_instrumentation_data({ method: "resources/read" }) end test "#handle unknown method returns method not found error" do @@ -402,6 +441,7 @@ module ModelContextProtocol assert_equal "Method not found", response[:error][:message] assert_equal "unknown_method", response[:error][:data] + assert_instrumentation_data({ method: "unknown_method" }) end end end diff --git a/test/test_helper.rb b/test/test_helper.rb index b2fd0b1..e678470 100644 --- a/test/test_helper.rb +++ b/test/test_helper.rb @@ -13,6 +13,8 @@ require "mocha/minitest" require "active_support" require "active_support/test_case" +require_relative "instrumentation_test_helper" + Minitest::Reporters.use!(Minitest::Reporters::ProgressReporter.new) module ActiveSupport