From 2121369ad644bca9fddaca2b2d78d86deaf957e9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ate=C5=9F=20G=C3=B6ral?= Date: Fri, 6 Nov 2026 03:09:48 -0500 Subject: [PATCH] fix: isolate instrumentation data per request --- docs/_server/configuration.md | 5 ++ lib/mcp/instrumentation.rb | 19 ++++--- lib/mcp/server.rb | 81 +++++++++++++++----------- test/mcp/instrumentation_test.rb | 98 ++++++++++++++++++++++++++++++++ test/mcp/server_test.rb | 83 +++++++++++++++++++++++++++ 5 files changed, 246 insertions(+), 40 deletions(-) diff --git a/docs/_server/configuration.md b/docs/_server/configuration.md index 7b8a900f..dd42344f 100644 --- a/docs/_server/configuration.md +++ b/docs/_server/configuration.md @@ -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:** diff --git a/lib/mcp/instrumentation.rb b/lib/mcp/instrumentation.rb index 5f1142fb..2b7d7b97 100644 --- a/lib/mcp/instrumentation.rb +++ b/lib/mcp/instrumentation.rb @@ -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 @@ -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 diff --git a/lib/mcp/server.rb b/lib/mcp/server.rb index 7ae234b7..eddcaaa3 100644 --- a/lib/mcp/server.rb +++ b/lib/mcp/server.rb @@ -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 @@ -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) @@ -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) @@ -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 @@ -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 @@ -851,6 +851,7 @@ 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, @@ -858,7 +859,7 @@ def handle_request(request, method, session: nil, related_request_id: nil, auth_ 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 @@ -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 @@ -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)", @@ -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] @@ -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 @@ -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, @@ -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( @@ -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}") @@ -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 @@ -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: @@ -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) @@ -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. @@ -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 @@ -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, diff --git a/test/mcp/instrumentation_test.rb b/test/mcp/instrumentation_test.rb index 8fb06af3..b5f0d2a3 100644 --- a/test/mcp/instrumentation_test.rb +++ b/test/mcp/instrumentation_test.rb @@ -455,6 +455,104 @@ def initialize(server_context) assert_equal({ request: { id: 42 } }, reported_context) end + test "#instrument_call isolates overlapping fibers and keeps both hooks on the same record" do + subject = Subject.new + completed = [] + around = [] + subject.configuration.instrumentation_callback = ->(data) { completed << data } + subject.configuration.around_request = ->(data, &request_handler) do + data.merge!(tool_name: "#{data[:method]}_tool", tool_arguments: { value: data[:method] }) + request_handler.call + data.merge!(client: { name: "slow_client" }, error: :custom_error) if data[:method] == "slow" + ensure + around << data + end + + slow = Fiber.new do + subject.instrument_call("slow") { Fiber.yield } + end + slow.resume + subject.instrument_call("fast") {} + slow.resume + + assert_equal(["fast", "slow"], completed.map { |data| data[:method] }) + assert_equal(["fast_tool", "slow_tool"], completed.map { |data| data[:tool_name] }) + assert_equal([{ value: "fast" }, { value: "slow" }], completed.map { |data| data[:tool_arguments] }) + refute(completed.first.key?(:client)) + refute(completed.first.key?(:error)) + assert_equal({ name: "slow_client" }, completed.last[:client]) + assert_equal(:custom_error, completed.last[:error]) + completed.zip(around).each do |callback_data, around_data| + assert_same(around_data, callback_data) + assert_kind_of(Numeric, callback_data[:duration]) + end + end + + test "#instrument_call keeps the outer record separate after a nested call raises" do + subject = Subject.new + completed = [] + subject.configuration.instrumentation_callback = ->(data) { completed << data } + subject.configuration.around_request = ->(data, &request_handler) do + data.merge!(tool_name: "outer_tool") if data[:method] == "outer" + request_handler.call + data.merge!(client: { name: "outer_client" }) if data[:method] == "outer" + end + + subject.instrument_call("outer") do + assert_raises(RuntimeError) do + subject.instrument_call("inner") { raise "inner failure" } + end + end + + assert_equal(["inner", "outer"], completed.map { |data| data[:method] }) + assert_equal(:internal_error, completed.first[:error]) + assert_equal("outer_tool", completed.last[:tool_name]) + assert_equal({ name: "outer_client" }, completed.last[:client]) + refute(completed.last.key?(:error)) + refute(completed.first.key?(:client)) + end + + test "#instrument_call keeps the outer record separate after a nested callback raises" do + subject = Subject.new + completed = [] + subject.configuration.instrumentation_callback = ->(data) do + completed << data + raise "callback failure" if completed.one? + end + subject.configuration.around_request = ->(data, &request_handler) do + data.merge!(tool_name: "outer_tool") if data[:method] == "outer" + request_handler.call + data.merge!(client: { name: "outer_client" }) if data[:method] == "outer" + end + + subject.instrument_call("outer") do + error = assert_raises(RuntimeError) { subject.instrument_call("inner") {} } + assert_equal("callback failure", error.message) + end + + assert_equal(["inner", "outer"], completed.map { |data| data[:method] }) + assert_equal("outer", completed.last[:method]) + assert_equal("outer_tool", completed.last[:tool_name]) + assert_equal({ name: "outer_client" }, completed.last[:client]) + refute(completed.first.key?(:tool_name)) + refute(completed.first.key?(:client)) + end + + test "#add_instrumentation_data preserves the legacy shared-record helper" do + subject = Subject.new + + subject.instrument_call("legacy") do + subject.add_instrumentation_data(additional_data: "legacy") + end + data = subject.instrumentation_data_received + + assert_equal("legacy", data[:method]) + assert_equal("legacy", data[:additional_data]) + assert_kind_of(Numeric, data[:duration]) + assert_same(data, subject.add_instrumentation_data(after_request: "updated")) + assert_equal("updated", data[:after_request]) + end + test "#instrument_call resets the instrumentation data between calls" do subject = Subject.new diff --git a/test/mcp/server_test.rb b/test/mcp/server_test.rb index 73827b93..b9f4aa08 100644 --- a/test/mcp/server_test.rb +++ b/test/mcp/server_test.rb @@ -1,6 +1,7 @@ # frozen_string_literal: true require "test_helper" +require "timeout" module MCP class ServerTest < ActiveSupport::TestCase @@ -2368,6 +2369,88 @@ def read_resource_request(uri) assert_equal(10.times.map { |i| "thread #{i}" }.sort, messages.sort) end + [false, true].each do |fail_slow_tool| + test "#handle concurrent tools keep instrumentation per request when failure is #{fail_slow_tool}" do + entered = Queue.new + release = Queue.new + completed = Queue.new + around = Queue.new + configuration = MCP::Configuration.new + configuration.instrumentation_callback = ->(data) { completed << data } + configuration.around_request = ->(data, &request_handler) do + request_handler.call + ensure + around << data + end + slow_tool = Tool.define(name: "slow_tool") do |value:| + entered << true + release.pop + raise "slow failure" if fail_slow_tool + + Tool::Response.new([{ type: "text", text: value }]) + end + fast_tool = Tool.define(name: "fast_tool") do |value:| + Tool::Response.new([{ type: "text", text: value }]) + end + server = Server.new(name: "test_server", tools: [slow_tool, fast_tool], configuration: configuration) + sessions = ["slow_client", "fast_client"].map do |name| + session = ServerSession.new(server: server, transport: mock) + server.handle({ + jsonrpc: "2.0", + method: "initialize", + params: initialize_params(clientInfo: { name: name, version: "1.0" }), + id: name, + }, session: session) + completed.pop + around.pop + session + end + call = ->(id, name, value, session) do + server.handle({ + jsonrpc: "2.0", + method: "tools/call", + params: { name: name, arguments: { value: value } }, + id: id, + }, session: session) + end + + slow_thread = Thread.new { call.call(1, "slow_tool", "slow", sessions.first) } + Timeout.timeout(5) { entered.pop } + fast_response = call.call(2, "fast_tool", "fast", sessions.last) + release << true + assert(slow_thread.join(5), "slow request did not finish") + slow_response = slow_thread.value + + assert_equal("fast", fast_response.dig(:result, :content, 0, :text)) + if fail_slow_tool + assert_equal(-32603, slow_response.dig(:error, :code)) + else + assert_equal("slow", slow_response.dig(:result, :content, 0, :text)) + end + records = 2.times.map { completed.pop(true) } + assert_equal(["fast_tool", "slow_tool"], records.map { |data| data[:tool_name] }) + assert_equal([{ value: "fast" }, { value: "slow" }], records.map { |data| data[:tool_arguments] }) + assert_equal("fast_client", records.first.dig(:client, :name)) + refute(records.first.key?(:error)) + if fail_slow_tool + assert_equal(:internal_error, records.last[:error]) + refute(records.last.key?(:client)) + else + assert_equal("slow_client", records.last.dig(:client, :name)) + refute(records.last.key?(:error)) + end + records.each do |data| + assert_equal("tools/call", data[:method]) + assert_kind_of(Numeric, data[:duration]) + assert_same(around.pop(true), data) + end + ensure + release << true + slow_thread&.kill + slow_thread&.join + end + end + test "#define_custom_method raises an error if the method is already defined" do assert_raises(Server::MethodAlreadyDefinedError) do @server.define_custom_method(method_name: "tools/call") do