Merge pull request #19 from patvice/coordinator-refactor

Coordinator refactor

Author: patvice

lib/ruby_llm/mcp/client.rb

@@ -1,223 +1,103 @@
 # frozen_string_literal: true
 
+require "forwardable"
+
 module RubyLLM
   module MCP
     class Client
-      PROTOCOL_VERSION = "2025-03-26"
-      PV_2024_11_05 = "2024-11-05"
+      extend Forwardable
 
-      attr_reader :name, :config, :transport_type, :transport, :request_timeout, :reverse_proxy_url, :protocol_version,
-                  :capabilities
+      attr_reader :name, :config, :transport_type, :request_timeout
 
-      def initialize(name:, transport_type:, start: true, request_timeout: 8000, reverse_proxy_url: nil, config: {}) # rubocop:disable Metrics/ParameterLists
+      def initialize(name:, transport_type:, start: true, request_timeout: 8000, config: {})
         @name = name
-        @config = config
-        @protocol_version = PROTOCOL_VERSION
-        @headers = config[:headers] || {}
-
+        @config = config.merge(request_timeout: request_timeout)
         @transport_type = transport_type.to_sym
-        @transport = nil
-
-        @capabilities = nil
-
         @request_timeout = request_timeout
-        @reverse_proxy_url = reverse_proxy_url
 
-        if start
-          self.start
-        end
-      end
+        @coordinator = Coordinator.new(self, transport_type: @transport_type, config: @config)
 
-      def request(body, **options)
-        @transport.request(body, **options)
+        start_transport if start
       end
 
-      def start
-        case @transport_type
-        when :sse
-          @transport = RubyLLM::MCP::Transport::SSE.new(@config[:url], request_timeout: @request_timeout,
-                                                                       headers: @headers)
-        when :stdio
-          @transport = RubyLLM::MCP::Transport::Stdio.new(@config[:command], request_timeout: @request_timeout,
-                                                                             args: @config[:args], env: @config[:env])
-        when :streamable
-          @transport = RubyLLM::MCP::Transport::Streamable.new(@config[:url], request_timeout: @request_timeout,
-                                                                              headers: @headers)
-        else
-          raise "Invalid transport type: #{transport_type}"
-        end
+      def_delegators :@coordinator, :start_transport, :stop_transport, :restart_transport, :alive?, :capabilities
 
-        @initialize_response = initialize_request
-        @capabilities = RubyLLM::MCP::Capabilities.new(@initialize_response["result"]["capabilities"])
-        notification_request
-      end
-
-      def stop
-        @transport&.close
-        @transport = nil
-      end
-
-      def restart!
-        stop
-        start
-      end
-
-      def alive?
-        !!@transport&.alive?
-      end
+      alias start start_transport
+      alias stop stop_transport
+      alias restart! restart_transport
 
       def tools(refresh: false)
-        @tools = nil if refresh
-        @tools ||= fetch_and_create_tools
+        fetch(:tools, refresh) do
+          tools_data = @coordinator.tool_list.dig("result", "tools")
+          build_map(tools_data, MCP::Tool)
+        end
+
         @tools.values
       end
 
       def tool(name, refresh: false)
-        @tools = nil if refresh
-        @tools ||= fetch_and_create_tools
+        tools(refresh: refresh)
 
         @tools[name]
       end
 
       def resources(refresh: false)
-        @resources = nil if refresh
-        @resources ||= fetch_and_create_resources
+        fetch(:resources, refresh) do
+          resources_data = @coordinator.resource_list.dig("result", "resources")
+          build_map(resources_data, MCP::Resource)
+        end
+
         @resources.values
       end
 
       def resource(name, refresh: false)
-        @resources = nil if refresh
-        @resources ||= fetch_and_create_resources
+        resources(refresh: refresh)
 
         @resources[name]
       end
 
       def resource_templates(refresh: false)
-        @resource_templates = nil if refresh
-        @resource_templates ||= fetch_and_create_resource_templates
+        fetch(:resource_templates, refresh) do
+          templates_data = @coordinator.resource_template_list.dig("result", "resourceTemplates")
+          build_map(templates_data, MCP::ResourceTemplate)
+        end
+
         @resource_templates.values
       end
 
       def resource_template(name, refresh: false)
-        @resource_templates = nil if refresh
-        @resource_templates ||= fetch_and_create_resource_templates
+        resource_templates(refresh: refresh)
 
         @resource_templates[name]
       end
 
       def prompts(refresh: false)
-        @prompts = nil if refresh
-        @prompts ||= fetch_and_create_prompts
+        fetch(:prompts, refresh) do
+          prompts_data = @coordinator.prompt_list.dig("result", "prompts")
+          build_map(prompts_data, MCP::Prompt)
+        end
+
         @prompts.values
       end
 
       def prompt(name, refresh: false)
-        @prompts = nil if refresh
-        @prompts ||= fetch_and_create_prompts
+        prompts(refresh: refresh)
 
         @prompts[name]
       end
 
-      def execute_tool(**args)
-        RubyLLM::MCP::Requests::ToolCall.new(self, **args).call
-      end
-
-      def resource_read_request(**args)
-        RubyLLM::MCP::Requests::ResourceRead.new(self, **args).call
-      end
-
-      def completion_resource(**args)
-        RubyLLM::MCP::Requests::CompletionResource.new(self, **args).call
-      end
-
-      def completion_prompt(**args)
-        RubyLLM::MCP::Requests::CompletionPrompt.new(self, **args).call
-      end
-
-      def execute_prompt(**args)
-        RubyLLM::MCP::Requests::PromptCall.new(self, **args).call
-      end
-
       private
 
-      def initialize_request
-        RubyLLM::MCP::Requests::Initialization.new(self).call
-      end
-
-      def notification_request
-        RubyLLM::MCP::Requests::Notification.new(self).call
-      end
-
-      def tool_list_request
-        RubyLLM::MCP::Requests::ToolList.new(self).call
-      end
-
-      def resources_list_request
-        RubyLLM::MCP::Requests::ResourceList.new(self).call
-      end
-
-      def resource_template_list_request
-        RubyLLM::MCP::Requests::ResourceTemplateList.new(self).call
-      end
-
-      def prompt_list_request
-        RubyLLM::MCP::Requests::PromptList.new(self).call
-      end
-
-      def fetch_and_create_tools
-        tools_response = tool_list_request
-        tools_response = tools_response["result"]["tools"]
-
-        tools = {}
-        tools_response.each do |tool|
-          new_tool = RubyLLM::MCP::Tool.new(self, tool)
-          tools[new_tool.name] = new_tool
-        end
-
-        tools
-      end
-
-      def fetch_and_create_resources
-        resources_response = resources_list_request
-        resources_response = resources_response["result"]["resources"]
-
-        resources = {}
-        resources_response.each do |resource|
-          new_resource = RubyLLM::MCP::Resource.new(self, resource)
-          resources[new_resource.name] = new_resource
-        end
-
-        resources
+      def fetch(cache_key, refresh)
+        instance_variable_set("@#{cache_key}", nil) if refresh
+        instance_variable_get("@#{cache_key}") || instance_variable_set("@#{cache_key}", yield)
       end
 
-      def fetch_and_create_resource_templates
-        resource_templates_response = resource_template_list_request
-        resource_templates_response = resource_templates_response["result"]["resourceTemplates"]
-
-        resource_templates = {}
-        resource_templates_response.each do |resource_template|
-          new_resource_template = RubyLLM::MCP::ResourceTemplate.new(self, resource_template)
-          resource_templates[new_resource_template.name] = new_resource_template
+      def build_map(raw_data, klass)
+        raw_data.each_with_object({}) do |item, acc|
+          instance = klass.new(@coordinator, item)
+          acc[instance.name] = instance
         end
-
-        resource_templates
-      end
-
-      def fetch_and_create_prompts
-        prompts_response = prompt_list_request
-        prompts_response = prompts_response["result"]["prompts"]
-
-        prompts = {}
-        prompts_response.each do |prompt|
-          new_prompt = RubyLLM::MCP::Prompt.new(self,
-                                                name: prompt["name"],
-                                                description: prompt["description"],
-                                                arguments: prompt["arguments"])
-
-          prompts[new_prompt.name] = new_prompt
-        end
-
-        prompts
       end
     end
   end

lib/ruby_llm/mcp/coordinator.rb

@@ -0,0 +1,112 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module MCP
+    class Coordinator
+      PROTOCOL_VERSION = "2025-03-26"
+      PV_2024_11_05 = "2024-11-05"
+
+      attr_reader :client, :transport_type, :config, :request_timeout, :headers, :transport, :initialize_response,
+                  :capabilities, :protocol_version
+
+      def initialize(client, transport_type:, config: {})
+        @client = client
+        @transport_type = transport_type
+        @config = config
+
+        @protocol_version = PROTOCOL_VERSION
+        @headers = config[:headers] || {}
+
+        @transport = nil
+        @capabilities = nil
+      end
+
+      def request(body, **options)
+        @transport.request(body, **options)
+      end
+
+      def start_transport
+        case @transport_type
+        when :sse
+          @transport = RubyLLM::MCP::Transport::SSE.new(@config[:url],
+                                                        request_timeout: @config[:request_timeout],
+                                                        headers: @headers)
+        when :stdio
+          @transport = RubyLLM::MCP::Transport::Stdio.new(@config[:command],
+                                                          request_timeout: @config[:request_timeout],
+                                                          args: @config[:args],
+                                                          env: @config[:env])
+        when :streamable
+          @transport = RubyLLM::MCP::Transport::Streamable.new(@config[:url],
+                                                               request_timeout: @config[:request_timeout],
+                                                               headers: @headers)
+        else
+          message = "Invalid transport type: :#{transport_type}. Supported types are :sse, :stdio, :streamable"
+          raise Errors::InvalidTransportType.new(message: message)
+        end
+
+        @initialize_response = initialize_request
+        @capabilities = RubyLLM::MCP::Capabilities.new(@initialize_response["result"]["capabilities"])
+        initialize_notification
+      end
+
+      def stop_transport
+        @transport&.close
+        @transport = nil
+      end
+
+      def restart_transport
+        stop_transport
+        start_transport
+      end
+
+      def alive?
+        !!@transport&.alive?
+      end
+
+      def execute_tool(**args)
+        RubyLLM::MCP::Requests::ToolCall.new(self, **args).call
+      end
+
+      def resource_read(**args)
+        RubyLLM::MCP::Requests::ResourceRead.new(self, **args).call
+      end
+
+      def completion_resource(**args)
+        RubyLLM::MCP::Requests::CompletionResource.new(self, **args).call
+      end
+
+      def completion_prompt(**args)
+        RubyLLM::MCP::Requests::CompletionPrompt.new(self, **args).call
+      end
+
+      def execute_prompt(**args)
+        RubyLLM::MCP::Requests::PromptCall.new(self, **args).call
+      end
+
+      def initialize_request
+        RubyLLM::MCP::Requests::Initialization.new(self).call
+      end
+
+      def initialize_notification
+        RubyLLM::MCP::Requests::InitializeNotification.new(self).call
+      end
+
+      def tool_list
+        RubyLLM::MCP::Requests::ToolList.new(self).call
+      end
+
+      def resource_list
+        RubyLLM::MCP::Requests::ResourceList.new(self).call
+      end
+
+      def resource_template_list
+        RubyLLM::MCP::Requests::ResourceTemplateList.new(self).call
+      end
+
+      def prompt_list
+        RubyLLM::MCP::Requests::PromptList.new(self).call
+      end
+    end
+  end
+end