Merge pull request #54 from julianojulio/make_server_context_optional_and_sorbet_friendly

Make `server_context` optional for Tools and Prompts

Author: topherbullock

lib/mcp/prompt.rb

@@ -9,7 +9,7 @@ class << self
       attr_reader :description_value
       attr_reader :arguments_value
 
-      def template(args, server_context:)
+      def template(args, server_context: nil)
         raise NotImplementedError, "Subclasses must implement template"
       end
 
@@ -57,7 +57,7 @@ def define(name: nil, description: nil, arguments: [], &block)
           prompt_name name
           description description
           arguments arguments
-          define_singleton_method(:template) do |args, server_context:|
+          define_singleton_method(:template) do |args, server_context: nil|
             instance_exec(args, server_context:, &block)
           end
         end

lib/mcp/server.rb

@@ -243,13 +243,7 @@ def call_tool(request)
       end
 
       begin
-        call_params = tool_call_parameters(tool)
-
-        if call_params.include?(:server_context)
-          tool.call(**arguments.transform_keys(&:to_sym), server_context:).to_h
-        else
-          tool.call(**arguments.transform_keys(&:to_sym)).to_h
-        end
+        call_tool_with_args(tool, arguments)
       rescue => e
         raise RequestHandlerError.new("Internal error calling tool #{tool_name}", request, original_error: e)
       end
@@ -272,7 +266,7 @@ def get_prompt(request)
       prompt_args = request[:arguments]
       prompt.validate_arguments!(prompt_args)
 
-      prompt.template(prompt_args, server_context:).to_h
+      call_prompt_template_with_args(prompt, prompt_args)
     end
 
     def list_resources(request)
@@ -299,22 +293,29 @@ def index_resources_by_uri(resources)
       end
     end
 
-    def tool_call_parameters(tool)
-      method_def = tool_call_method_def(tool)
-      method_def.parameters.flatten
+    def accepts_server_context?(method_object)
+      parameters = method_object.parameters
+      accepts_server_context = parameters.any? { |_type, name| name == :server_context }
+      has_kwargs = parameters.any? { |type, _| type == :keyrest }
+
+      accepts_server_context || has_kwargs
     end
 
-    def tool_call_method_def(tool)
-      method = tool.method(:call)
+    def call_tool_with_args(tool, arguments)
+      args = arguments.transform_keys(&:to_sym)
 
-      if defined?(T::Utils) && T::Utils.respond_to?(:signature_for_method)
-        sorbet_typed_method_definition = T::Utils.signature_for_method(method)&.method
+      if accepts_server_context?(tool.method(:call))
+        tool.call(**args, server_context: server_context).to_h
+      else
+        tool.call(**args).to_h
+      end
+    end
 
-        # Return the Sorbet typed method definition if it exists, otherwise fallback to original method
-        # definition if Sorbet is defined but not used by this tool.
-        sorbet_typed_method_definition || method
+    def call_prompt_template_with_args(prompt, args)
+      if accepts_server_context?(prompt.method(:template))
+        prompt.template(args, server_context: server_context).to_h
       else
-        method
+        prompt.template(args).to_h
       end
     end
   end

lib/mcp/tool.rb

@@ -9,7 +9,7 @@ class << self
       attr_reader :input_schema_value
       attr_reader :annotations_value
 
-      def call(*args, server_context:)
+      def call(*args, server_context: nil)
         raise NotImplementedError, "Subclasses must implement call"
       end