Internal Release 0.3.0

Add configurable exception reporter
Avoid leaking tool exceptions by sending them to a configurable exception reporter
Add optional context to pass to tools
Add basic instrumentation data/labels

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

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

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
 

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

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

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

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
 

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

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

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

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

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

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