Add define_custom_method to define custom methods

Author: atesgoral

lib/mcp/server.rb

@@ -20,6 +20,15 @@ def initialize(message, request, error_type: :internal_error, original_error: ni
       end
     end
 
+    class MethodAlreadyDefinedError < StandardError
+      attr_reader :method_name
+
+      def initialize(method_name)
+        super("Method #{method_name} already defined")
+        @method_name = method_name
+      end
+    end
+
     include Instrumentation
 
     attr_accessor :name, :version, :tools, :prompts, :resources, :server_context, :configuration, :capabilities, :transport
@@ -89,6 +98,14 @@ def define_prompt(name: nil, description: nil, arguments: [], &block)
       @prompts[prompt.name_value] = prompt
     end
 
+    def define_custom_method(method_name:, &block)
+      if @handlers.key?(method_name)
+        raise MethodAlreadyDefinedError, method_name
+      end
+
+      @handlers[method_name] = block
+    end
+
     def notify_tools_list_changed
       return unless @transport
 

test/mcp/server_test.rb

@@ -653,6 +653,38 @@ def call(message:, server_context: nil)
       assert_instrumentation_data({ method: "unsupported_method" })
     end
 
+    test "#handle handles custom methods" do
+      @server.define_custom_method(method_name: "add") do |params|
+        params[:a] + params[:b]
+      end
+
+      request = {
+        jsonrpc: "2.0",
+        id: 1,
+        method: "add",
+        params: { a: 1, b: 2 },
+      }
+
+      response = @server.handle(request)
+      assert_equal 3, response[:result]
+      assert_instrumentation_data({ method: "add" })
+    end
+
+    test "#handle handles custom notifications" do
+      @server.define_custom_method(method_name: "notify") do
+        nil
+      end
+
+      request = {
+        jsonrpc: "2.0",
+        method: "notify",
+      }
+
+      response = @server.handle(request)
+      assert_nil response
+      assert_instrumentation_data({ method: "notify" })
+    end
+
     test "the global configuration is used if no configuration is passed to the server" do
       server = Server.new(name: "test_server")
       assert_equal MCP.configuration.instrumentation_callback,