Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
104 changes: 78 additions & 26 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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(
Expand All @@ -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:**

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

Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down
40 changes: 38 additions & 2 deletions lib/mcp/configuration.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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,
Expand All @@ -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
Expand Down
25 changes: 23 additions & 2 deletions lib/mcp/instrumentation.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
14 changes: 11 additions & 3 deletions lib/mcp/server.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
Expand Down Expand Up @@ -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
Expand Down
47 changes: 47 additions & 0 deletions test/mcp/configuration_test.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
Loading