14e02da475
APM integrations need to tag a span with the subject of the request, but the
user-defined `server_context` hash passed to `Server.new` never reaches
`around_request` or `instrumentation_callback`. The exception reporter's
`server_context` argument is a different value — it describes where a failure
occurred (`{ request: ... }`, `{ notification: ... }`), not who made the call —
so there is currently no way to read the request subject from a hook.
Add `Configuration#instrument_server_context`, off by default. When enabled,
`instrument_call` puts the host's `server_context` into the instrumentation data
before invoking `around_request`, so it is readable both before and after
`request_handler.call`.
The flag is opt-in rather than unconditional because the hash is
application-supplied and may hold values a tracing backend should not receive.
Hosts that do not expose a `server_context` reader are unaffected.
`instrument_call` reads it through `self.server_context`: the existing
`server_context:` keyword shadows the reader, and the value wanted here is the
user-defined hash, not the reporter context.
Follows the same shape as the earlier additions of `tool_arguments` (#218) and
`client` (#221) to this hash.
473 lines
17 KiB
Ruby
473 lines
17 KiB
Ruby
# frozen_string_literal: true
|
|
|
|
require "test_helper"
|
|
require "mcp/instrumentation"
|
|
|
|
module MCP
|
|
class InstrumentationTest < ActiveSupport::TestCase
|
|
class Subject
|
|
include Instrumentation
|
|
attr_reader :instrumentation_data_received, :configuration
|
|
|
|
def initialize
|
|
@configuration = MCP::Configuration.new
|
|
@configuration.instrumentation_callback = ->(data) { @instrumentation_data_received = data }
|
|
end
|
|
|
|
def instrumented_method
|
|
instrument_call("instrumented_method") do
|
|
# nothing to do
|
|
end
|
|
end
|
|
|
|
def instrumented_method_with_additional_data
|
|
instrument_call("instrumented_method_with_additional_data") do
|
|
add_instrumentation_data(additional_data: "test")
|
|
end
|
|
end
|
|
|
|
def instrumented_method_with_server_context(context)
|
|
instrument_call("instrumented_method_with_server_context", server_context: context) do
|
|
# nothing to do
|
|
end
|
|
end
|
|
|
|
def instrumented_method_that_raises_with_error_set
|
|
instrument_call("instrumented_method_that_raises_with_error_set") do
|
|
add_instrumentation_data(error: :custom_error)
|
|
raise "block error"
|
|
end
|
|
end
|
|
end
|
|
|
|
# Mirrors `MCP::Server`, which exposes the user-defined hash through
|
|
# `attr_accessor :server_context`.
|
|
class SubjectWithServerContext < Subject
|
|
attr_accessor :server_context
|
|
|
|
def initialize(server_context)
|
|
super()
|
|
@server_context = server_context
|
|
end
|
|
end
|
|
|
|
test "#instrument_call omits the user-defined server_context by default" do
|
|
subject = SubjectWithServerContext.new({ user_id: 42 })
|
|
|
|
subject.instrumented_method
|
|
|
|
refute_includes subject.instrumentation_data_received.keys, :server_context
|
|
end
|
|
|
|
test "#instrument_call exposes the user-defined server_context when instrument_server_context is enabled" do
|
|
subject = SubjectWithServerContext.new({ user_id: 42 })
|
|
subject.configuration.instrument_server_context = true
|
|
|
|
subject.instrumented_method
|
|
|
|
assert_equal({ user_id: 42 }, subject.instrumentation_data_received[:server_context])
|
|
end
|
|
|
|
test "#instrument_call exposes the user-defined server_context to around_request" do
|
|
subject = SubjectWithServerContext.new({ user_id: 42 })
|
|
subject.configuration.instrument_server_context = true
|
|
seen = nil
|
|
subject.configuration.around_request = ->(data, &handler) {
|
|
seen = data[:server_context]
|
|
handler.call
|
|
}
|
|
|
|
subject.instrumented_method
|
|
|
|
assert_equal({ user_id: 42 }, seen)
|
|
end
|
|
|
|
test "#instrument_call reports the user-defined server_context as nil when the host does not set one" do
|
|
subject = SubjectWithServerContext.new(nil)
|
|
subject.configuration.instrument_server_context = true
|
|
|
|
subject.instrumented_method
|
|
|
|
assert_includes subject.instrumentation_data_received.keys, :server_context
|
|
assert_nil subject.instrumentation_data_received[:server_context]
|
|
end
|
|
|
|
test "#instrument_call does not confuse the reporter context with the user-defined server_context" do
|
|
# `instrument_call(server_context:)` is the exception-reporter context, a
|
|
# different value from the user-defined hash. Enabling the flag must
|
|
# surface the latter, not the former.
|
|
subject = SubjectWithServerContext.new({ user_id: 42 })
|
|
subject.configuration.instrument_server_context = true
|
|
|
|
subject.instrumented_method_with_server_context({ request: "reporter-context" })
|
|
|
|
assert_equal({ user_id: 42 }, subject.instrumentation_data_received[:server_context])
|
|
end
|
|
|
|
test "#instrument_call skips server_context for hosts without the reader even when enabled" do
|
|
subject = Subject.new
|
|
subject.configuration.instrument_server_context = true
|
|
|
|
subject.instrumented_method
|
|
|
|
refute_includes subject.instrumentation_data_received.keys, :server_context
|
|
end
|
|
|
|
test "#instrument_call adds the method name to the instrumentation data" do
|
|
subject = Subject.new
|
|
|
|
subject.instrumented_method
|
|
assert_equal({ method: "instrumented_method" }, subject.instrumentation_data_received.tap do |data|
|
|
data.delete(:duration)
|
|
end)
|
|
end
|
|
|
|
test "#instrument_call exposes data added via add_instrumentation_data" do
|
|
subject = Subject.new
|
|
|
|
subject.instrumented_method_with_additional_data
|
|
assert_equal(
|
|
{ method: "instrumented_method_with_additional_data", additional_data: "test" },
|
|
subject.instrumentation_data_received.tap { |data| data.delete(:duration) },
|
|
)
|
|
end
|
|
|
|
test "#instrument_call invokes around_request wrapping execution" do
|
|
call_log = []
|
|
subject = Subject.new
|
|
subject.configuration.around_request = ->(_data, &request_handler) {
|
|
call_log << :before
|
|
request_handler.call
|
|
call_log << :after
|
|
}
|
|
|
|
subject.instrumented_method
|
|
|
|
assert_equal([:before, :after], call_log)
|
|
end
|
|
|
|
test "#instrument_call around_request receives method before request_handler.call" do
|
|
received_method = nil
|
|
subject = Subject.new
|
|
subject.configuration.around_request = ->(data, &request_handler) {
|
|
received_method = data[:method]
|
|
request_handler.call
|
|
}
|
|
|
|
subject.instrumented_method
|
|
|
|
assert_equal("instrumented_method", received_method)
|
|
end
|
|
|
|
test "#instrument_call around_request data is populated after request_handler.call" do
|
|
data_after_call = nil
|
|
subject = Subject.new
|
|
subject.configuration.around_request = ->(data, &request_handler) {
|
|
request_handler.call
|
|
data_after_call = data.dup
|
|
}
|
|
|
|
subject.instrumented_method_with_additional_data
|
|
|
|
assert_equal("instrumented_method_with_additional_data", data_after_call[:method])
|
|
assert_equal("test", data_after_call[:additional_data])
|
|
end
|
|
|
|
test "#instrument_call default around_request is pass-through" do
|
|
subject = Subject.new
|
|
|
|
subject.instrumented_method_with_additional_data
|
|
|
|
assert_equal(
|
|
{ method: "instrumented_method_with_additional_data", additional_data: "test" },
|
|
subject.instrumentation_data_received.tap { |data| data.delete(:duration) },
|
|
)
|
|
end
|
|
|
|
test "#instrument_call reports exception and sets error when around_request raises before request_handler.call" do
|
|
reported_exception = nil
|
|
reported_context = nil
|
|
subject = Subject.new
|
|
subject.configuration.exception_reporter = ->(e, server_context) {
|
|
reported_exception = e
|
|
reported_context = server_context
|
|
}
|
|
subject.configuration.around_request = ->(_data, &_request_handler) { raise "boom before" }
|
|
|
|
error = assert_raises(RuntimeError) { subject.instrumented_method }
|
|
assert_equal("boom before", error.message)
|
|
assert_same(error, reported_exception)
|
|
assert_equal({}, reported_context)
|
|
assert_equal(:internal_error, subject.instrumentation_data_received[:error])
|
|
assert(subject.instrumentation_data_received.key?(:duration))
|
|
end
|
|
|
|
test "#instrument_call reports exception and sets error when around_request raises after request_handler.call" do
|
|
reported_exception = nil
|
|
subject = Subject.new
|
|
subject.configuration.exception_reporter = ->(e, _server_context) { reported_exception = e }
|
|
subject.configuration.around_request = ->(_data, &request_handler) {
|
|
request_handler.call
|
|
raise "boom after"
|
|
}
|
|
|
|
error = assert_raises(RuntimeError) { subject.instrumented_method }
|
|
assert_equal("boom after", error.message)
|
|
assert_same(error, reported_exception)
|
|
assert_equal(:internal_error, subject.instrumentation_data_received[:error])
|
|
end
|
|
|
|
test "#instrument_call preserves block's custom error tag and reports the exception" do
|
|
report_count = 0
|
|
subject = Subject.new
|
|
subject.configuration.exception_reporter = ->(_e, _server_context) { report_count += 1 }
|
|
|
|
assert_raises(RuntimeError) { subject.instrumented_method_that_raises_with_error_set }
|
|
assert_equal(1, report_count)
|
|
assert_equal(:custom_error, subject.instrumentation_data_received[:error])
|
|
end
|
|
|
|
test "#instrument_call skips reporting when exception_already_reported returns true" do
|
|
reported = []
|
|
subject = Subject.new
|
|
subject.configuration.exception_reporter = ->(e, _server_context) { reported << e.message }
|
|
|
|
error = StandardError.new("inner failure")
|
|
subject.define_singleton_method(:instrumented_with_skip) do
|
|
instrument_call("instrumented_with_skip", exception_already_reported: ->(e) { error.equal?(e) }) do
|
|
raise error
|
|
end
|
|
end
|
|
|
|
assert_raises(StandardError) { subject.instrumented_with_skip }
|
|
assert_equal([], reported)
|
|
end
|
|
|
|
test "#instrument_call reports when exception_already_reported returns false for a new exception" do
|
|
reported = []
|
|
subject = Subject.new
|
|
subject.configuration.exception_reporter = ->(e, _server_context) { reported << e.message }
|
|
|
|
target = StandardError.new("inner failure")
|
|
subject.configuration.around_request = ->(_data, &request_handler) do
|
|
request_handler.call
|
|
ensure
|
|
raise "ensure boom"
|
|
end
|
|
|
|
subject.define_singleton_method(:instrumented_with_skip) do
|
|
instrument_call("instrumented_with_skip", exception_already_reported: ->(e) { target.equal?(e) }) do
|
|
raise target
|
|
end
|
|
end
|
|
|
|
assert_raises(RuntimeError) { subject.instrumented_with_skip }
|
|
assert_equal(["ensure boom"], reported)
|
|
end
|
|
|
|
test "#instrument_call falls back to reporting when exception_already_reported itself raises" do
|
|
reported = []
|
|
subject = Subject.new
|
|
subject.configuration.exception_reporter = ->(e, _server_context) { reported << e.message }
|
|
|
|
subject.define_singleton_method(:instrumented_with_broken_predicate) do
|
|
instrument_call(
|
|
"instrumented_with_broken_predicate",
|
|
exception_already_reported: ->(_e) { raise "predicate blew up" },
|
|
) do
|
|
raise "original failure"
|
|
end
|
|
end
|
|
|
|
error = assert_raises(RuntimeError) { subject.instrumented_with_broken_predicate }
|
|
assert_equal("original failure", error.message)
|
|
assert_equal(["original failure"], reported)
|
|
end
|
|
|
|
test "#instrument_call falls back to reporting when exception_already_reported raises a non-StandardError" do
|
|
reported = []
|
|
subject = Subject.new
|
|
subject.configuration.exception_reporter = ->(e, _server_context) { reported << e.message }
|
|
|
|
subject.define_singleton_method(:instrumented_with_system_exit_predicate) do
|
|
instrument_call(
|
|
"instrumented_with_system_exit_predicate",
|
|
exception_already_reported: ->(_e) { raise SystemExit, 9 },
|
|
) do
|
|
raise "original failure"
|
|
end
|
|
end
|
|
|
|
error = assert_raises(RuntimeError) { subject.instrumented_with_system_exit_predicate }
|
|
assert_equal("original failure", error.message)
|
|
assert_equal(["original failure"], reported)
|
|
end
|
|
|
|
test "#instrument_call normalizes non-boolean truthy return values from exception_already_reported" do
|
|
reported = []
|
|
subject = Subject.new
|
|
subject.configuration.exception_reporter = ->(e, _server_context) { reported << e.message }
|
|
|
|
subject.define_singleton_method(:instrumented_with_truthy_predicate) do
|
|
instrument_call(
|
|
"instrumented_with_truthy_predicate",
|
|
exception_already_reported: ->(_e) { "non-boolean truthy" },
|
|
) do
|
|
raise "original failure"
|
|
end
|
|
end
|
|
|
|
assert_raises(RuntimeError) { subject.instrumented_with_truthy_predicate }
|
|
assert_equal([], reported)
|
|
end
|
|
|
|
test "#instrument_call reports when exception_already_reported returns nil" do
|
|
reported = []
|
|
subject = Subject.new
|
|
subject.configuration.exception_reporter = ->(e, _server_context) { reported << e.message }
|
|
|
|
subject.define_singleton_method(:instrumented_with_nil_predicate) do
|
|
instrument_call(
|
|
"instrumented_with_nil_predicate",
|
|
exception_already_reported: ->(_e) { nil },
|
|
) do
|
|
raise "original failure"
|
|
end
|
|
end
|
|
|
|
assert_raises(RuntimeError) { subject.instrumented_with_nil_predicate }
|
|
assert_equal(["original failure"], reported)
|
|
end
|
|
|
|
test "#instrument_call does not carry reported state across invocations when the same exception is reused" do
|
|
reported = []
|
|
subject = Subject.new
|
|
subject.configuration.exception_reporter = ->(e, _server_context) { reported << e.message }
|
|
|
|
shared_error = RuntimeError.new("reused")
|
|
subject.configuration.around_request = ->(_data, &_request_handler) { raise shared_error }
|
|
|
|
2.times do
|
|
assert_raises(RuntimeError) { subject.instrumented_method }
|
|
end
|
|
assert_equal(["reused", "reused"], reported)
|
|
end
|
|
|
|
test "#instrument_call reports frozen exceptions without mutating them" do
|
|
reported = []
|
|
subject = Subject.new
|
|
subject.configuration.exception_reporter = ->(e, _server_context) { reported << e }
|
|
|
|
frozen_error = RuntimeError.new("frozen").freeze
|
|
subject.configuration.around_request = ->(_data, &_request_handler) { raise frozen_error }
|
|
|
|
error = assert_raises(RuntimeError) { subject.instrumented_method }
|
|
assert_same(frozen_error, error)
|
|
assert_equal([frozen_error], reported)
|
|
assert_predicate(frozen_error, :frozen?)
|
|
end
|
|
|
|
test "#instrument_call nested calls sharing an exception_already_reported predicate report only once" do
|
|
reported = []
|
|
subject = Subject.new
|
|
subject.configuration.exception_reporter = ->(e, _server_context) { reported << e.message }
|
|
|
|
reported_exception = nil
|
|
predicate = ->(e) { reported_exception.equal?(e) }
|
|
|
|
subject.define_singleton_method(:outer) do
|
|
instrument_call("outer", exception_already_reported: predicate) { inner }
|
|
end
|
|
|
|
subject.define_singleton_method(:inner) do
|
|
instrument_call("inner", exception_already_reported: predicate) do
|
|
raise "nested boom"
|
|
rescue => e
|
|
configuration.exception_reporter.call(e, {})
|
|
reported_exception = e
|
|
raise
|
|
end
|
|
end
|
|
|
|
assert_raises(RuntimeError) { subject.outer }
|
|
assert_equal(["nested boom"], reported)
|
|
end
|
|
|
|
test "#instrument_call keeps reported state isolated across concurrent threads on a shared subject" do
|
|
subject = Subject.new
|
|
reported = Queue.new
|
|
subject.configuration.exception_reporter = ->(e, _server_context) { reported << e.message }
|
|
|
|
threads = 10.times.map do |i|
|
|
Thread.new do
|
|
subject.instrument_call("threaded_#{i}") { raise "thread #{i}" }
|
|
rescue RuntimeError
|
|
# Swallow the intentional raise so the thread finishes cleanly.
|
|
# The test verifies side effects via the shared `reported` queue.
|
|
end
|
|
end
|
|
threads.each(&:join)
|
|
|
|
messages = []
|
|
messages << reported.pop until reported.empty?
|
|
|
|
assert_equal(10.times.map { |i| "thread #{i}" }, messages.sort)
|
|
end
|
|
|
|
test "#instrument_call reports each invocation when the same exception instance is raised from multiple threads" do
|
|
subject = Subject.new
|
|
reported = Queue.new
|
|
subject.configuration.exception_reporter = ->(e, _server_context) { reported << e }
|
|
|
|
shared_error = RuntimeError.new("shared across threads")
|
|
start_gate = Queue.new
|
|
|
|
threads = 10.times.map do
|
|
Thread.new do
|
|
start_gate.pop
|
|
begin
|
|
subject.instrument_call("shared_error_method") { raise shared_error }
|
|
rescue RuntimeError
|
|
# Swallow the intentional raise so the thread finishes cleanly.
|
|
# The test verifies side effects via the shared `reported` queue.
|
|
end
|
|
end
|
|
end
|
|
10.times { start_gate << :go }
|
|
threads.each(&:join)
|
|
|
|
collected = []
|
|
collected << reported.pop until reported.empty?
|
|
|
|
assert_equal(10, collected.size)
|
|
assert(collected.all? { |e| e.equal?(shared_error) })
|
|
end
|
|
|
|
test "#instrument_call forwards server_context to exception_reporter on around_request failure" do
|
|
reported_context = nil
|
|
subject = Subject.new
|
|
subject.configuration.exception_reporter = ->(_e, server_context) { reported_context = server_context }
|
|
subject.configuration.around_request = ->(_data, &_request_handler) { raise "boom" }
|
|
|
|
assert_raises(RuntimeError) do
|
|
subject.instrumented_method_with_server_context(request: { id: 42 })
|
|
end
|
|
assert_equal({ request: { id: 42 } }, reported_context)
|
|
end
|
|
|
|
test "#instrument_call resets the instrumentation data between calls" do
|
|
subject = Subject.new
|
|
|
|
subject.instrumented_method_with_additional_data
|
|
assert_equal(
|
|
{ method: "instrumented_method_with_additional_data", additional_data: "test" },
|
|
subject.instrumentation_data_received.tap { |data| data.delete(:duration) },
|
|
)
|
|
|
|
subject.instrumented_method
|
|
assert_equal({ method: "instrumented_method" }, subject.instrumentation_data_received.tap do |data|
|
|
data.delete(:duration)
|
|
end)
|
|
end
|
|
end
|
|
end
|