Merge pull request #309 from koic/add_around_request_hook_for_request_instrumentation

Add `around_request` hook for request instrumentation

lib/mcp/configuration.rb

@@ -7,11 +7,18 @@ module MCP
       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
@@ -50,10 +57,24 @@ module MCP
       !@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
@@ -72,20 +93,30 @@ module MCP
       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,
@@ -111,6 +142,11 @@ module MCP
       @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

lib/mcp/instrumentation.rb

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

lib/mcp/server.rb

@@ -375,7 +375,7 @@ module MCP
     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
@@ -385,7 +385,12 @@ module MCP
       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)
@@ -415,11 +420,14 @@ module MCP
         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