Add around_request hook for request instrumentation

## Motivation and Context

`instrumentation_callback` fires after request execution, making it impossible to wrap
requests with Application Performance Monitoring (APM) spans (e.g. Datadog tracing).
This adds `around_request`, which wraps request handling and allows executing code
before and after each request.

For example, to wrap requests with a Datadog trace:

```ruby
MCP.configure do |config|
  config.around_request = ->(data, &request_handler) {
    Datadog::Tracing.trace("mcp.#{data[:method]}") do
      request_handler.call
    end
  }
end
```

`instrumentation_callback` is soft-deprecated in favor of `around_request`.

## How Has This Been Tested?

Added tests for `around_request` in configuration, instrumentation,
and server test files. All existing tests continue to pass.

## Breaking Changes

`around_request` itself is optional and defaults to a pass-through,
and `instrumentation_callback` continues to work as before.

However, `instrument_call` now accepts a `server_context:` keyword and its block is invoked
without arguments (previously `yield block` passed the block Proc as an argument).
This affects anyone mixing `MCP::Instrumentation` into their own classes,
though the public `Server` API is unchanged.

Resolves modelcontextprotocol#302.

Author: koic

README.md

@@ -159,15 +159,17 @@ MCP.configure do |config|
     end
   }
 
-  config.instrumentation_callback = ->(data) {
-    puts "Got instrumentation data #{data.inspect}"
+  config.around_request = ->(data, &request_handler) {
+    logger.info("Start: #{data[:method]}")
+    request_handler.call
+    logger.info("Done: #{data[:method]}, tool: #{data[:tool_name]}")
   }
 end
 ```
 
 or by creating an explicit configuration and passing it into the server.
 This is useful for systems where an application hosts more than one MCP server but
-they might require different instrumentation callbacks.
+they might require different configurations.
 
 ```ruby
 configuration = MCP::Configuration.new
@@ -179,8 +181,10 @@ configuration.exception_reporter = ->(exception, server_context) {
   end
 }
 
-configuration.instrumentation_callback = ->(data) {
-  puts "Got instrumentation data #{data.inspect}"
+configuration.around_request = ->(data, &request_handler) {
+  logger.info("Start: #{data[:method]}")
+  request_handler.call
+  logger.info("Done: #{data[:method]}, tool: #{data[:tool_name]}")
 }
 
 server = MCP::Server.new(
@@ -193,7 +197,8 @@ server = MCP::Server.new(
 
 #### `server_context`
 
-The `server_context` is a user-defined hash that is passed into the server instance and made available to tools, prompts, and exception/instrumentation callbacks. It can be used to provide contextual information such as authentication state, user IDs, or request-specific data.
+The `server_context` is a user-defined hash that is passed into the server instance and made available to tool and prompt calls.
+It can be used to provide contextual information such as authentication state, user IDs, or request-specific data.
 
 **Type:**
 
@@ -210,7 +215,9 @@ server = MCP::Server.new(
 )
 ```
 
-This hash is then passed as the `server_context` argument to tool and prompt calls, and is included in exception and instrumentation callbacks.
+This hash is then passed as the `server_context` keyword argument to tool and prompt calls.
+Note that exception and instrumentation callbacks do not receive this user-defined hash.
+See the relevant sections below for the arguments they receive.
 
 #### Request-specific `_meta` Parameter
 
@@ -263,17 +270,77 @@ end
 The exception reporter receives:
 
 - `exception`: The Ruby exception object that was raised
-- `server_context`: The context hash provided to the server
+- `server_context`: A hash describing where the failure occurred (e.g., `{ request: <raw JSON-RPC request> }`
+  for request handling, `{ notification: "tools_list_changed" }` for notification delivery).
+  This is not the user-defined `server_context` passed to `Server.new`.
 
 **Signature:**
 
 ```ruby
 exception_reporter = ->(exception, server_context) { ... }
 ```
 
-##### Instrumentation Callback
+##### Around Request
 
-The instrumentation callback receives a hash with the following possible keys:
+The `around_request` hook wraps request handling, allowing you to execute code before and after each request.
+This is useful for Application Performance Monitoring (APM) tracing, logging, or other observability needs.
+
+The hook receives a `data` hash and a `request_handler` block. You must call `request_handler.call` to execute the request:
+
+**Signature:**
+
+```ruby
+around_request = ->(data, &request_handler) { request_handler.call }
+```
+
+**`data` availability by timing:**
+
+- Before `request_handler.call`: `method`
+- After `request_handler.call`: `tool_name`, `tool_arguments`, `prompt_name`, `resource_uri`, `error`, `client`
+- Not available inside `around_request`: `duration` (added after `around_request` returns)
+
+> [!NOTE]
+> `tool_name`, `prompt_name` and `resource_uri` may only be populated for the corresponding request methods
+> (`tools/call`, `prompts/get`, `resources/read`), and may not be set depending on how the request is handled
+> (for example, `prompt_name` is not recorded when the prompt is not found).
+> `duration` is added after `around_request` returns, so it is not visible from within the hook.
+
+**Example:**
+
+```ruby
+MCP.configure do |config|
+  config.around_request = ->(data, &request_handler) {
+    logger.info("Start: #{data[:method]}")
+    request_handler.call
+    logger.info("Done: #{data[:method]}, tool: #{data[:tool_name]}")
+  }
+end
+```
+
+##### Instrumentation Callback (soft-deprecated)
+
+> [!NOTE]
+> `instrumentation_callback` is soft-deprecated. Use `around_request` instead.
+>
+> To migrate, wrap the call in `begin/ensure` so the callback still runs when the request fails:
+>
+> ```ruby
+> # Before
+> config.instrumentation_callback = ->(data) { log(data) }
+>
+> # After
+> config.around_request = ->(data, &request_handler) do
+>   request_handler.call
+> ensure
+>   log(data)
+> end
+> ```
+>
+> Note that `data[:duration]` is not available inside `around_request`.
+> If you need it, measure elapsed time yourself within the hook, or keep using `instrumentation_callback`.
+
+The instrumentation callback is called after each request finishes, whether successfully or with an error.
+It receives a hash with the following possible keys:
 
 - `method`: (String) The protocol method called (e.g., "ping", "tools/list")
 - `tool_name`: (String, optional) The name of the tool called
@@ -284,25 +351,10 @@ The instrumentation callback receives a hash with the following possible keys:
 - `duration`: (Float) Duration of the call in seconds
 - `client`: (Hash, optional) Client information with `name` and `version` keys, from the initialize request
 
-> [!NOTE]
-> `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.
-
-**Type:**
+**Signature:**
 
 ```ruby
 instrumentation_callback = ->(data) { ... }
-# where data is a Hash with keys as described above
-```
-
-**Example:**
-
-```ruby
-MCP.configure do |config|
-  config.instrumentation_callback = ->(data) {
-    puts "Instrumentation: #{data.inspect}"
-  }
-end
 ```
 
 ### Server Protocol Version

lib/mcp/configuration.rb

@@ -7,11 +7,18 @@ class Configuration
       LATEST_STABLE_PROTOCOL_VERSION, "2025-06-18", "2025-03-26", "2024-11-05",
     ]
 
-    attr_writer :exception_reporter, :instrumentation_callback
+    attr_writer :exception_reporter, :around_request
 
-    def initialize(exception_reporter: nil, instrumentation_callback: nil, protocol_version: nil,
+    # @deprecated Use {#around_request=} instead. `instrumentation_callback`
+    #   fires only after a request completes and cannot wrap execution in a
+    #   surrounding block (e.g. for Application Performance Monitoring (APM) spans).
+    # @see #around_request=
+    attr_writer :instrumentation_callback
+
+    def initialize(exception_reporter: nil, around_request: nil, instrumentation_callback: nil, protocol_version: nil,
       validate_tool_call_arguments: true)
       @exception_reporter = exception_reporter
+      @around_request = around_request
       @instrumentation_callback = instrumentation_callback
       @protocol_version = protocol_version
       if protocol_version
@@ -50,10 +57,24 @@ def exception_reporter?
       !@exception_reporter.nil?
     end
 
+    def around_request
+      @around_request || default_around_request
+    end
+
+    def around_request?
+      !@around_request.nil?
+    end
+
+    # @deprecated Use {#around_request} instead. `instrumentation_callback`
+    #   fires only after a request completes and cannot wrap execution in a
+    #   surrounding block (e.g. for Application Performance Monitoring (APM) spans).
+    # @see #around_request
     def instrumentation_callback
       @instrumentation_callback || default_instrumentation_callback
     end
 
+    # @deprecated Use {#around_request?} instead.
+    # @see #around_request?
     def instrumentation_callback?
       !@instrumentation_callback.nil?
     end
@@ -72,20 +93,30 @@ def merge(other)
       else
         @exception_reporter
       end
+
+      around_request = if other.around_request?
+        other.around_request
+      else
+        @around_request
+      end
+
       instrumentation_callback = if other.instrumentation_callback?
         other.instrumentation_callback
       else
         @instrumentation_callback
       end
+
       protocol_version = if other.protocol_version?
         other.protocol_version
       else
         @protocol_version
       end
+
       validate_tool_call_arguments = other.validate_tool_call_arguments
 
       Configuration.new(
         exception_reporter: exception_reporter,
+        around_request: around_request,
         instrumentation_callback: instrumentation_callback,
         protocol_version: protocol_version,
         validate_tool_call_arguments: validate_tool_call_arguments,
@@ -111,6 +142,11 @@ def default_exception_reporter
       @default_exception_reporter ||= ->(exception, server_context) {}
     end
 
+    def default_around_request
+      @default_around_request ||= ->(_data, &request_handler) { request_handler.call }
+    end
+
+    # @deprecated Use {#default_around_request} instead.
     def default_instrumentation_callback
       @default_instrumentation_callback ||= ->(data) {}
     end

lib/mcp/instrumentation.rb

@@ -2,19 +2,40 @@
 
 module MCP
   module Instrumentation
-    def instrument_call(method, &block)
+    def instrument_call(method, server_context: {}, exception_already_reported: nil, &block)
       start_time = Process.clock_gettime(Process::CLOCK_MONOTONIC)
       begin
         @instrumentation_data = {}
         add_instrumentation_data(method: method)
 
-        result = yield block
+        result = configuration.around_request.call(@instrumentation_data, &block)
 
         result
+      rescue => e
+        already_reported = begin
+          !!exception_already_reported&.call(e)
+        # rubocop:disable Lint/RescueException
+        rescue Exception
+          # rubocop:enable Lint/RescueException
+          # The predicate is expected to be side-effect-free and return a boolean.
+          # Any exception raised from it (including non-StandardError such as SystemExit)
+          # must not shadow the original exception.
+          false
+        end
+
+        unless already_reported
+          add_instrumentation_data(error: :internal_error) unless @instrumentation_data.key?(:error)
+          configuration.exception_reporter.call(e, server_context)
+        end
+
+        raise
       ensure
         end_time = Process.clock_gettime(Process::CLOCK_MONOTONIC)
         add_instrumentation_data(duration: end_time - start_time)
 
+        # Backward compatibility: `instrumentation_callback` is soft-deprecated
+        # in favor of `around_request`, but existing callers still expect it
+        # to fire after every request until it is removed in a future version.
         configuration.instrumentation_callback.call(@instrumentation_data)
       end
     end

lib/mcp/server.rb

@@ -375,7 +375,7 @@ def schema_contains_ref?(schema)
     def handle_request(request, method, session: nil, related_request_id: nil)
       handler = @handlers[method]
       unless handler
-        instrument_call("unsupported_method") do
+        instrument_call("unsupported_method", server_context: { request: request }) do
           client = session&.client || @client
           add_instrumentation_data(client: client) if client
         end
@@ -385,7 +385,12 @@ def handle_request(request, method, session: nil, related_request_id: nil)
       Methods.ensure_capability!(method, capabilities)
 
       ->(params) {
-        instrument_call(method) do
+        reported_exception = nil
+        instrument_call(
+          method,
+          server_context: { request: request },
+          exception_already_reported: ->(e) { reported_exception.equal?(e) },
+        ) do
           result = case method
           when Methods::INITIALIZE
             init(params, session: session)
@@ -415,11 +420,14 @@ def handle_request(request, method, session: nil, related_request_id: nil)
         rescue RequestHandlerError => e
           report_exception(e.original_error || e, { request: request })
           add_instrumentation_data(error: e.error_type)
+          reported_exception = e
           raise e
         rescue => e
           report_exception(e, { request: request })
           add_instrumentation_data(error: :internal_error)
-          raise RequestHandlerError.new("Internal error handling #{method} request", request, original_error: e)
+          wrapped = RequestHandlerError.new("Internal error handling #{method} request", request, original_error: e)
+          reported_exception = wrapped
+          raise wrapped
         end
       }
     end

test/mcp/configuration_test.rb

@@ -123,6 +123,53 @@ class ConfigurationTest < ActiveSupport::TestCase
       refute merged.validate_tool_call_arguments
     end
 
+    test "initializes with a default pass-through around_request" do
+      config = Configuration.new
+      called = false
+      config.around_request.call({}) { called = true }
+      assert called
+    end
+
+    test "allows setting a custom around_request" do
+      config = Configuration.new
+      call_log = []
+      config.around_request = ->(_data, &request_handler) {
+        call_log << :before
+        request_handler.call
+        call_log << :after
+      }
+
+      config.around_request.call({}) { call_log << :execute }
+      assert_equal([:before, :execute, :after], call_log)
+    end
+
+    test "around_request? returns false by default" do
+      config = Configuration.new
+      refute config.around_request?
+    end
+
+    test "around_request? returns true when set" do
+      config = Configuration.new
+      config.around_request = ->(_data, &request_handler) { request_handler.call }
+      assert config.around_request?
+    end
+
+    test "merge preserves around_request from other config" do
+      custom = ->(_data, &request_handler) { request_handler.call }
+      config1 = Configuration.new
+      config2 = Configuration.new(around_request: custom)
+      merged = config1.merge(config2)
+      assert_equal custom, merged.around_request
+    end
+
+    test "merge preserves around_request from self when other not set" do
+      custom = ->(_data, &request_handler) { request_handler.call }
+      config1 = Configuration.new(around_request: custom)
+      config2 = Configuration.new
+      merged = config1.merge(config2)
+      assert_equal custom, merged.around_request
+    end
+
     test "raises ArgumentError when protocol_version is not a supported value" do
       exception = assert_raises(ArgumentError) do
         Configuration.new(protocol_version: "1999-12-31")