Skip to content

Commit c10c0dd

Browse files
committed
Refactor MCP tool and server to handle optional server_context parameter
- Updated `call` method in `MCP::Tool` and `MCP::Prompt` to accept `server_context` as an optional parameter. - Refactored `MCP::Server` to streamline argument handling for tool and prompt calls, introducing helper methods `call_tool_with_args` and `call_prompt_template_with_args`. - Adjusted tests to reflect changes in method signatures and ensure compatibility with optional `server_context` parameter.
1 parent 04710ac commit c10c0dd

File tree

6 files changed

+477
-24
lines changed

6 files changed

+477
-24
lines changed

lib/mcp/prompt.rb

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ class << self
99
attr_reader :description_value
1010
attr_reader :arguments_value
1111

12-
def template(args, server_context:)
12+
def template(args, server_context: nil)
1313
raise NotImplementedError, "Subclasses must implement template"
1414
end
1515

@@ -57,7 +57,7 @@ def define(name: nil, description: nil, arguments: [], &block)
5757
prompt_name name
5858
description description
5959
arguments arguments
60-
define_singleton_method(:template) do |args, server_context:|
60+
define_singleton_method(:template) do |args, server_context: nil|
6161
instance_exec(args, server_context:, &block)
6262
end
6363
end

lib/mcp/server.rb

Lines changed: 20 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -211,13 +211,7 @@ def call_tool(request)
211211
end
212212

213213
begin
214-
call_params = tool_call_parameters(tool)
215-
216-
if call_params.include?(:server_context)
217-
tool.call(**arguments.transform_keys(&:to_sym), server_context:).to_h
218-
else
219-
tool.call(**arguments.transform_keys(&:to_sym)).to_h
220-
end
214+
call_tool_with_args(tool, arguments)
221215
rescue => e
222216
raise RequestHandlerError.new("Internal error calling tool #{tool_name}", request, original_error: e)
223217
end
@@ -242,7 +236,7 @@ def get_prompt(request)
242236
prompt_args = request[:arguments]
243237
prompt.validate_arguments!(prompt_args)
244238

245-
prompt.template(prompt_args, server_context:).to_h
239+
call_prompt_template_with_args(prompt, prompt_args)
246240
end
247241

248242
def list_resources(request)
@@ -274,22 +268,29 @@ def index_resources_by_uri(resources)
274268
end
275269
end
276270

277-
def tool_call_parameters(tool)
278-
method_def = tool_call_method_def(tool)
279-
method_def.parameters.flatten
271+
def accepts_server_context?(method_object)
272+
parameters = method_object.parameters
273+
accepts_server_context = parameters.any? { |_type, name| name == :server_context }
274+
has_kwargs = parameters.any? { |type, _| type == :keyrest }
275+
276+
accepts_server_context || has_kwargs
280277
end
281278

282-
def tool_call_method_def(tool)
283-
method = tool.method(:call)
279+
def call_tool_with_args(tool, arguments)
280+
args = arguments.transform_keys(&:to_sym)
284281

285-
if defined?(T::Utils) && T::Utils.respond_to?(:signature_for_method)
286-
sorbet_typed_method_definition = T::Utils.signature_for_method(method)&.method
282+
if accepts_server_context?(tool.method(:call))
283+
tool.call(**args, server_context: server_context).to_h
284+
else
285+
tool.call(**args).to_h
286+
end
287+
end
287288

288-
# Return the Sorbet typed method definition if it exists, otherwise fallback to original method
289-
# definition if Sorbet is defined but not used by this tool.
290-
sorbet_typed_method_definition || method
289+
def call_prompt_template_with_args(prompt, args)
290+
if accepts_server_context?(prompt.method(:template))
291+
prompt.template(args, server_context: server_context).to_h
291292
else
292-
method
293+
prompt.template(args).to_h
293294
end
294295
end
295296
end

lib/mcp/tool.rb

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ class << self
99
attr_reader :input_schema_value
1010
attr_reader :annotations_value
1111

12-
def call(*args, server_context:)
12+
def call(*args, server_context: nil)
1313
raise NotImplementedError, "Subclasses must implement call"
1414
end
1515

0 commit comments

Comments
 (0)