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 #302.
This commit is contained in:
Koichi ITO 2026-04-10 05:05:08 +09:00
commit fab6f81261
3 changed files with 790 additions and 33 deletions

View file

@ -7,11 +7,18 @@ module MCP
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 @@ module MCP
!@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 @@ module MCP
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 @@ module MCP
@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

View file

@ -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

View file

@ -375,7 +375,7 @@ module MCP
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 @@ module MCP
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 @@ module MCP
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