Skip to content

fix: isolate instrumentation data per request #668

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
koic merged 1 commit into from
Nov 7, 2026
Merged

fix: isolate instrumentation data per request #668

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
5 changes: 5 additions & 0 deletions docs/_server/configuration.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,11 @@ If no exception reporter is configured, a default no-op reporter is used that si
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 SDK keeps separate instrumentation data for each request. Both hooks receive that request's hash.

`add_instrumentation_data` is deprecated. Update the data hash passed to `around_request` instead.
The helper retains shared-record behavior and is not safe for concurrent calls.

The hook receives a `data` hash and a `request_handler` block. You must call `request_handler.call` to execute the request:

**Signature:**
Expand Down
19 changes: 12 additions & 7 deletions lib/mcp/instrumentation.rb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,16 +5,17 @@ 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)
instrumentation_data = {}
@instrumentation_data = instrumentation_data
instrumentation_data.merge!(method: method)
# `self.` is required: the `server_context:` keyword above shadows the
# reader, and the value we want here is the user-defined hash passed to
# `Server.new`, not the per-call reporter context.
if configuration.instrument_server_context? && respond_to?(:server_context)
add_instrumentation_data(server_context: self.server_context)
instrumentation_data.merge!(server_context: self.server_context)
end

result = configuration.around_request.call(@instrumentation_data, &block)
result = configuration.around_request.call(instrumentation_data) { block.call(instrumentation_data) }

result
rescue => e
Expand All @@ -28,22 +29,26 @@ def instrument_call(method, server_context: {}, exception_already_reported: nil,
end

unless already_reported
add_instrumentation_data(error: :internal_error) unless @instrumentation_data.key?(:error)
instrumentation_data.merge!(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)
instrumentation_data.merge!(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)
configuration.instrumentation_callback.call(instrumentation_data)
end
end

# @deprecated Use the data hash passed to the {#instrument_call} block or
# `around_request` instead. `add_instrumentation_data` uses a shared record
# and is not safe for concurrent calls.
# @see #instrument_call
def add_instrumentation_data(**kwargs)
@instrumentation_data.merge!(kwargs)
end
Expand Down
81 changes: 48 additions & 33 deletions lib/mcp/server.rb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -733,9 +733,9 @@ def handle_request(request, method, session: nil, related_request_id: nil, auth_

handler = @handlers[method]
unless handler
instrument_call("unsupported_method", server_context: { request: request }) do
instrument_call("unsupported_method", server_context: { request: request }) do |instrumentation_data|
client = session&.client || @client
add_instrumentation_data(client: client) if client
instrumentation_data.merge!(client: client) if client
end
return
end
Expand Down Expand Up @@ -798,7 +798,7 @@ def handle_request(request, method, session: nil, related_request_id: nil, auth_
method,
server_context: { request: request },
exception_already_reported: ->(e) { reported_exception.equal?(e) },
) do
) do |instrumentation_data|
refuse_non_object_meta!(params)
envelope = lift_request_envelope(params, method: method, session: session)

Expand All @@ -817,7 +817,7 @@ def handle_request(request, method, session: nil, related_request_id: nil, auth_
when Methods::INITIALIZE
init(params, session: session)
when Methods::RESOURCES_READ
contents = read_resource_contents(params, session: session, related_request_id: related_request_id, cancellation: cancellation, envelope: envelope, auth_info: auth_info)
contents = read_resource_contents(params, instrumentation_data: instrumentation_data, session: session, related_request_id: related_request_id, cancellation: cancellation, envelope: envelope, auth_info: auth_info)

# An SEP-2322 `input_required` result must not be wrapped as `contents` or stamped with SEP-2549 cache hints.
contents.is_a?(InputRequiredResult) ? contents : build_read_resource_result(contents)
Expand All @@ -827,9 +827,9 @@ def handle_request(request, method, session: nil, related_request_id: nil, auth_

subscription_result(handler_result)
when Methods::TOOLS_CALL
call_tool(params, session: session, related_request_id: related_request_id, cancellation: cancellation, envelope: envelope, auth_info: auth_info)
call_tool(params, instrumentation_data: instrumentation_data, session: session, related_request_id: related_request_id, cancellation: cancellation, envelope: envelope, auth_info: auth_info)
when Methods::PROMPTS_GET
get_prompt(params, session: session, related_request_id: related_request_id, cancellation: cancellation, envelope: envelope, auth_info: auth_info)
get_prompt(params, instrumentation_data: instrumentation_data, session: session, related_request_id: related_request_id, cancellation: cancellation, envelope: envelope, auth_info: auth_info)
when Methods::COMPLETION_COMPLETE
complete(params, session: session, related_request_id: related_request_id, cancellation: cancellation, envelope: envelope, auth_info: auth_info)
when Methods::LOGGING_SET_LEVEL
Expand All @@ -838,10 +838,10 @@ def handle_request(request, method, session: nil, related_request_id: nil, auth_
dispatch_optional_context_handler(@handlers[method], params, session: session, related_request_id: related_request_id, cancellation: cancellation, envelope: envelope, auth_info: auth_info)
end
client = session&.client || @client
add_instrumentation_data(client: client) if client
instrumentation_data.merge!(client: client) if client

if cancellation&.cancelled?
add_instrumentation_data(cancelled: true, cancellation_reason: cancellation.reason)
instrumentation_data.merge!(cancelled: true, cancellation_reason: cancellation.reason)
next JsonRpcHandler::NO_RESPONSE
end

Expand All @@ -851,14 +851,15 @@ def handle_request(request, method, session: nil, related_request_id: nil, auth_
result = if envelope.nil? && @input_required_legacy_shim && session
run_legacy_input_required_shim(
result,
instrumentation_data: instrumentation_data,
method: method,
params: params,
session: session,
related_request_id: related_request_id,
cancellation: cancellation,
)
else
serialize_input_required_result(result, envelope: envelope, request: params, method: method)
serialize_input_required_result(result, instrumentation_data: instrumentation_data, envelope: envelope, request: params, method: method)
end
end

Expand All @@ -884,25 +885,25 @@ def handle_request(request, method, session: nil, related_request_id: nil, auth_

result
rescue CancelledError => e
add_instrumentation_data(cancelled: true, cancellation_reason: e.reason)
instrumentation_data.merge!(cancelled: true, cancellation_reason: e.reason)
next JsonRpcHandler::NO_RESPONSE
rescue OAuth::InsufficientScopeError => e
# A handler's `require_scopes!` failure is a client authorization error, not a server fault:
# surface it as an invalid-request JSON-RPC error whose message names the missing scopes.
# HTTP-level 403 challenges remain the transport `required_scopes` gate's job.
report_exception(e, { request: request })
add_instrumentation_data(error: :insufficient_scope)
instrumentation_data.merge!(error: :insufficient_scope)
converted = RequestHandlerError.new(e.message, request, error_type: :invalid_request, original_error: e)
reported_exception = converted
raise converted
rescue RequestHandlerError => e
report_exception(e.original_error || e, { request: request })
add_instrumentation_data(error: e.error_type)
instrumentation_data.merge!(error: e.error_type)
reported_exception = e
raise e
rescue => e
report_exception(e, { request: request })
add_instrumentation_data(error: :internal_error)
instrumentation_data.merge!(error: :internal_error)
wrapped = RequestHandlerError.new("Internal error handling #{method} request", request, original_error: e)
reported_exception = wrapped
raise wrapped
Expand Down Expand Up @@ -998,7 +999,7 @@ def lift_request_envelope(params, method:, session:)
# a final result. The capability gate enforces the SEP-2575 rule that servers MUST NOT rely on
# (or embed requests for) capabilities the client did not declare, and reports every missing capability at
# once so the client sees the full set.
def serialize_input_required_result(result, envelope:, request:, method:)
def serialize_input_required_result(result, instrumentation_data:, envelope:, request:, method:)
if envelope.nil?
raise RequestHandlerError.new(
"input_required results require the 2026-07-28 stateless lifecycle (SEP-2322)",
Expand All @@ -1010,7 +1011,7 @@ def serialize_input_required_result(result, envelope:, request:, method:)
missing = result.missing_client_capabilities(envelope.client_capabilities)
raise MissingRequiredClientCapabilityError.new(missing, request) unless missing.empty?

add_instrumentation_data(input_required: true)
instrumentation_data.merge!(input_required: true)
serialized = result.to_h

if @request_state_security && serialized[:requestState]
Expand Down Expand Up @@ -1039,7 +1040,7 @@ def serialize_input_required_result(result, envelope:, request:, method:)
# re-runs with `inputResponses`/`requestState` merged into the original params - the same deterministic replay
# contract the modern client driver follows. The `requestState` round-trips in-process as the raw value
# the handler wrote; `RequestStateSecurity` sealing is wire hardening and does not apply.
def run_legacy_input_required_shim(result, method:, params:, session:, related_request_id:, cancellation:)
def run_legacy_input_required_shim(result, instrumentation_data:, method:, params:, session:, related_request_id:, cancellation:)
rounds = 0

loop do
Expand Down Expand Up @@ -1071,6 +1072,7 @@ def run_legacy_input_required_shim(result, method:, params:, session:, related_r
result = redispatch_mrtr_method(
method,
retry_params,
instrumentation_data: instrumentation_data,
session: session,
related_request_id: related_request_id,
cancellation: cancellation,
Expand Down Expand Up @@ -1104,14 +1106,14 @@ def legacy_leg_params(entry)

# Re-runs the handler of one of the three MRTR-capable methods for
# the legacy shim. Legacy wire, so no envelope is threaded.
def redispatch_mrtr_method(method, params, session:, related_request_id:, cancellation:)
def redispatch_mrtr_method(method, params, instrumentation_data:, session:, related_request_id:, cancellation:)
case method
when Methods::TOOLS_CALL
call_tool(params, session: session, related_request_id: related_request_id, cancellation: cancellation)
call_tool(params, instrumentation_data: instrumentation_data, session: session, related_request_id: related_request_id, cancellation: cancellation)
when Methods::PROMPTS_GET
get_prompt(params, session: session, related_request_id: related_request_id, cancellation: cancellation)
get_prompt(params, instrumentation_data: instrumentation_data, session: session, related_request_id: related_request_id, cancellation: cancellation)
when Methods::RESOURCES_READ
contents = read_resource_contents(params, session: session, related_request_id: related_request_id, cancellation: cancellation)
contents = read_resource_contents(params, instrumentation_data: instrumentation_data, session: session, related_request_id: related_request_id, cancellation: cancellation)
contents.is_a?(InputRequiredResult) ? contents : build_read_resource_result(contents)
else
raise RequestHandlerError.new(
Expand Down Expand Up @@ -1394,21 +1396,21 @@ def list_tools(request)
apply_cache_metadata({ tools: page[:items], nextCursor: page[:next_cursor] }.compact)
end

def call_tool(request, session: nil, related_request_id: nil, cancellation: nil, envelope: nil, auth_info: nil)
def call_tool(request, instrumentation_data:, session: nil, related_request_id: nil, cancellation: nil, envelope: nil, auth_info: nil)
tool_name = request[:name]

tool = tools[tool_name]
unless tool
add_instrumentation_data(tool_name: tool_name, error: :tool_not_found)
instrumentation_data.merge!(tool_name: tool_name, error: :tool_not_found)

raise RequestHandlerError.new("Tool not found: #{tool_name}", request, error_type: :invalid_params)
end

arguments = request[:arguments] || {}
add_instrumentation_data(tool_name: tool_name, tool_arguments: arguments)
instrumentation_data.merge!(tool_name: tool_name, tool_arguments: arguments)

if tool.input_schema&.missing_required_arguments?(arguments)
add_instrumentation_data(error: :missing_required_arguments)
instrumentation_data.merge!(error: :missing_required_arguments)

missing = tool.input_schema.missing_required_arguments(arguments).join(", ")
return error_tool_response("Missing required arguments: #{missing}")
Expand All @@ -1418,7 +1420,7 @@ def call_tool(request, session: nil, related_request_id: nil, cancellation: nil,
begin
tool.input_schema.validate_arguments(arguments)
rescue Tool::InputSchema::ValidationError => e
add_instrumentation_data(error: :invalid_schema)
instrumentation_data.merge!(error: :invalid_schema)

return error_tool_response(e.message)
end
Expand Down Expand Up @@ -1473,11 +1475,11 @@ def list_prompts(request)
apply_cache_metadata({ prompts: page[:items], nextCursor: page[:next_cursor] }.compact)
end

def get_prompt(request, session: nil, related_request_id: nil, cancellation: nil, envelope: nil, auth_info: nil)
def get_prompt(request, instrumentation_data:, session: nil, related_request_id: nil, cancellation: nil, envelope: nil, auth_info: nil)
prompt_name = request[:name]
prompt = @prompts[prompt_name]
unless prompt
add_instrumentation_data(error: :prompt_not_found)
instrumentation_data.merge!(error: :prompt_not_found)
# The explicit `error_code` maps an unknown prompt to Invalid Params (-32602) rather than
# the default Internal Error (-32603), matching the `tools/call`, `resources/read`, and
# `completion/complete` siblings for the same not-found condition, while `error_type:
Expand All @@ -1490,7 +1492,7 @@ def get_prompt(request, session: nil, related_request_id: nil, cancellation: nil
)
end

add_instrumentation_data(prompt_name: prompt_name)
instrumentation_data.merge!(prompt_name: prompt_name)

prompt_args = request[:arguments]
prompt.validate_arguments!(prompt_args)
Expand Down Expand Up @@ -1766,9 +1768,9 @@ def skills_extension_keys(extensions)
# Default `resources/read` handler: routes to class-based resources and resource templates.
# Fully replaced when `resources_read_handler` is set. When no class-based resource or template is registered,
# unknown URIs keep the historical no-op `[]` response instead of raising.
def read_resource(request, server_context: nil)
def read_resource(request, instrumentation_data:, server_context: nil)
uri = request[:uri]
add_instrumentation_data(resource_uri: uri)
instrumentation_data.merge!(resource_uri: uri)

# Files of directory-backed skills are looked up by exact string equality in a map built at construction,
# so request data never reaches the filesystem; a miss falls through to the routing below.
Expand Down Expand Up @@ -1798,7 +1800,7 @@ def read_resource(request, server_context: nil)
params = template.match_uri(uri)
next unless params

add_instrumentation_data(resource_template: template.uri_template)
instrumentation_data.merge!(resource_template: template.uri_template)
return normalize_resource_contents(call_template_contents(template, params, server_context))
end

Expand Down Expand Up @@ -1884,9 +1886,22 @@ def complete(params, session: nil, related_request_id: nil, cancellation: nil, e
# Invokes `resources/read` via the registered handler. If the handler block opts in to `server_context:`,
# pass an `MCP::ServerContext` so the handler can observe cancellation via `server_context.cancelled?` or
# `server_context.raise_if_cancelled!`.
def read_resource_contents(request, session: nil, related_request_id: nil, cancellation: nil, envelope: nil, auth_info: nil)
def read_resource_contents(request, instrumentation_data:, session: nil, related_request_id: nil, cancellation: nil, envelope: nil, auth_info: nil)
handler = @handlers[Methods::RESOURCES_READ]
if handler == method(:read_resource)
server_context = build_server_context(
request: request,
session: session,
related_request_id: related_request_id,
cancellation: cancellation,
envelope: envelope,
auth_info: auth_info,
)
return read_resource(request, instrumentation_data: instrumentation_data, server_context: server_context)
end

dispatch_optional_context_handler(
@handlers[Methods::RESOURCES_READ],
handler,
request,
session: session,
related_request_id: related_request_id,
Expand Down
Loading
Loading