Support Sorbet typed tools

Author: chrisbutcher

Gemfile.lock

@@ -91,6 +91,13 @@ GEM
     ruby-progressbar (1.13.0)
     ruby2_keywords (0.0.5)
     securerandom (0.4.1)
+    sorbet (0.5.11495)
+      sorbet-static (= 0.5.11495)
+    sorbet-runtime (0.5.11495)
+    sorbet-static (0.5.11495-universal-darwin)
+    sorbet-static-and-runtime (0.5.11495)
+      sorbet (= 0.5.11495)
+      sorbet-runtime (= 0.5.11495)
     stringio (3.1.7)
     tzinfo (2.0.6)
       concurrent-ruby (~> 1.0)
@@ -112,6 +119,7 @@ DEPENDENCIES
   model_context_protocol!
   rake (~> 13.0)
   rubocop-shopify
+  sorbet-static-and-runtime
 
 BUNDLED WITH
    2.5.9

lib/model_context_protocol/server.rb

@@ -193,7 +193,8 @@ def call_tool(request)
       end
 
       begin
-        call_params = tool.method(:call).parameters.flatten
+        call_params = tool_call_parameters(tool)
+
         if call_params.include?(:server_context)
           tool.call(**arguments.transform_keys(&:to_sym), server_context:).to_h
         else
@@ -254,5 +255,24 @@ def index_resources_by_uri(resources)
         hash[resource.uri] = resource
       end
     end
+
+    def tool_call_parameters(tool)
+      method_def = tool_call_method_def(tool)
+      method_def.parameters.flatten
+    end
+
+    def tool_call_method_def(tool)
+      method = tool.method(:call)
+
+      if defined?(T::Utils) && T::Utils.respond_to?(:signature_for_method)
+        sorbet_typed_method_definition = T::Utils.signature_for_method(method)&.method
+
+        # 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
+      else
+        method
+      end
+    end
   end
 end

model_context_protocol.gemspec

@@ -29,4 +29,5 @@ Gem::Specification.new do |spec|
 
   spec.add_dependency("json_rpc_handler", "~> 0.1")
   spec.add_development_dependency("activesupport")
+  spec.add_development_dependency("sorbet-static-and-runtime")
 end

test/model_context_protocol/server_test.rb

@@ -1,3 +1,4 @@
+# typed: true
 # frozen_string_literal: true
 
 require "test_helper"
@@ -255,6 +256,43 @@ class ServerTest < ActiveSupport::TestCase
       assert_instrumentation_data({ method: "tools/call", tool_name: })
     end
 
+    test "#handle_json tools/call executes tool and returns result, when the tool is typed with Sorbet" do
+      class TypedTestTool < Tool
+        tool_name "test_tool"
+        description "a test tool for testing"
+        input_schema({ properties: { message: { type: "string" } }, required: ["message"] })
+
+        class << self
+          extend T::Sig
+
+          sig { params(message: String, server_context: T.nilable(T.untyped)).returns(Tool::Response) }
+          def call(message:, server_context: nil)
+            Tool::Response.new([{ type: "text", content: "OK" }])
+          end
+        end
+      end
+
+      request = JSON.generate({
+        jsonrpc: "2.0",
+        method: "tools/call",
+        params: { name: "test_tool", arguments: { message: "Hello, world!" } },
+        id: 1,
+      })
+
+      server = Server.new(
+        name: @server_name,
+        tools: [TypedTestTool],
+        prompts: [@prompt],
+        resources: [@resource],
+        resource_templates: [@resource_template],
+      )
+
+      raw_response = server.handle_json(request)
+      response = JSON.parse(raw_response, symbolize_names: true) if raw_response
+
+      assert_equal({ content: [{ type: "text", content: "OK" }], isError: false }, response[:result])
+    end
+
     test "#handle tools/call returns internal error and reports exception if the tool raises an error" do
       @server.configuration.exception_reporter.expects(:call).with do |exception, server_context|
         assert_not_nil exception

test/model_context_protocol/tool_test.rb

@@ -1,3 +1,4 @@
+# typed: true
 # frozen_string_literal: true
 
 require "test_helper"
@@ -203,5 +204,34 @@ class UpdatableAnnotationsTool < Tool
       tool.annotations(title: "Updated")
       assert_equal tool.annotations_value.title, "Updated"
     end
+
+    test "#call with Sorbet typed tools invokes the tool block and returns the response" do
+      class TypedTestTool < Tool
+        tool_name "test_tool"
+        description "a test tool for testing"
+        input_schema({ properties: { message: { type: "string" } }, required: ["message"] })
+        annotations(
+          title: "Test Tool",
+          read_only_hint: true,
+          destructive_hint: false,
+          idempotent_hint: true,
+          open_world_hint: false,
+        )
+
+        class << self
+          extend T::Sig
+
+          sig { params(message: String, server_context: T.nilable(T.untyped)).returns(Tool::Response) }
+          def call(message, server_context: nil)
+            Tool::Response.new([{ type: "text", content: "OK" }])
+          end
+        end
+      end
+
+      tool = TypedTestTool
+      response = tool.call("test")
+      assert_equal response.content, [{ type: "text", content: "OK" }]
+      assert_equal response.is_error, false
+    end
   end
 end

test/test_helper.rb

@@ -13,6 +13,8 @@
 require "active_support"
 require "active_support/test_case"
 
+require "sorbet-runtime"
+
 require_relative "instrumentation_test_helper"
 
 Minitest::Reporters.use!(Minitest::Reporters::ProgressReporter.new)