ruby-sdk/lib/mcp/instrumentation.rb
Koichi ITO fab6f81261 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.
2026-04-11 13:54:47 +09:00

47 lines
1.7 KiB
Ruby

# frozen_string_literal: true
module MCP
module Instrumentation
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 = 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
def add_instrumentation_data(**kwargs)
@instrumentation_data.merge!(kwargs)
end
end
end