Merge pull request #668 from atesgoral/fix/instrumentation-data-per-request

fix: isolate instrumentation data per request

Author: koic

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:**

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

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