Merge pull request #221 from chattermill/md-client-instrumentation

Add client info to instrumentation callback

Author: koic

README.md

@@ -400,6 +400,7 @@ The instrumentation callback receives a hash with the following possible keys:
 - `resource_uri`: (String, optional) The URI of the resource called
 - `error`: (String, optional) Error code if a lookup failed
 - `duration`: (Float) Duration of the call in seconds
+- `client`: (Hash, optional) Client information with `name` and `version` keys, from the initialize request
 
 **Type:**
 

lib/mcp/server.rb

@@ -75,6 +75,7 @@ def initialize(
       @resource_index = index_resources_by_uri(resources)
       @server_context = server_context
       @configuration = MCP.configuration.merge(configuration)
+      @client = nil
 
       validate!
 
@@ -258,15 +259,17 @@ def validate_tool_name!
     def handle_request(request, method)
       handler = @handlers[method]
       unless handler
-        instrument_call("unsupported_method") {}
+        instrument_call("unsupported_method") do
+          add_instrumentation_data(client: @client) if @client
+        end
         return
       end
 
       Methods.ensure_capability!(method, capabilities)
 
       ->(params) {
         instrument_call(method) do
-          case method
+          result = case method
           when Methods::TOOLS_LIST
             { tools: @handlers[Methods::TOOLS_LIST].call(params) }
           when Methods::PROMPTS_LIST
@@ -280,6 +283,9 @@ def handle_request(request, method)
           else
             @handlers[method].call(params)
           end
+          add_instrumentation_data(client: @client) if @client
+
+          result
         rescue => e
           report_exception(e, { request: request })
           if e.is_a?(RequestHandlerError)
@@ -314,6 +320,7 @@ def server_info
     end
 
     def init(request)
+      @client = request[:clientInfo] if request
       {
         protocolVersion: configuration.protocol_version,
         capabilities: capabilities,

test/mcp/server_test.rb

@@ -159,6 +159,53 @@ class ServerTest < ActiveSupport::TestCase
       assert_instrumentation_data({ method: "initialize" })
     end
 
+    test "#handle initialize request with clientInfo includes client in instrumentation data" do
+      client_info = { name: "test_client", version: "1.0.0" }
+      request = {
+        jsonrpc: "2.0",
+        method: "initialize",
+        id: 1,
+        params: {
+          clientInfo: client_info,
+        },
+      }
+
+      @server.handle(request)
+      assert_instrumentation_data({ method: "initialize", client: client_info })
+    end
+
+    test "instrumentation data includes client info for subsequent requests after initialize" do
+      client_info = { name: "test_client", version: "1.0.0" }
+      initialize_request = {
+        jsonrpc: "2.0",
+        method: "initialize",
+        id: 1,
+        params: {
+          clientInfo: client_info,
+        },
+      }
+      @server.handle(initialize_request)
+
+      ping_request = {
+        jsonrpc: "2.0",
+        method: "ping",
+        id: 2,
+      }
+      @server.handle(ping_request)
+      assert_instrumentation_data({ method: "ping", client: client_info })
+    end
+
+    test "instrumentation data does not include client key when no clientInfo provided" do
+      request = {
+        jsonrpc: "2.0",
+        method: "ping",
+        id: 1,
+      }
+
+      @server.handle(request)
+      assert_instrumentation_data({ method: "ping" })
+    end
+
     test "#handle returns nil for notification requests" do
       request = {
         jsonrpc: "2.0",