ruby_llm-mcp 0.4.1 → 0.5.0

data/lib/ruby_llm/mcp/transport.rb

@@ -0,0 +1,58 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module MCP
+    class Transport
+      class << self
+        def transports
+          @transports ||= {}
+        end
+
+        def register_transport(transport_type, transport_class)
+          transports[transport_type] = transport_class
+        end
+      end
+
+      extend Forwardable
+
+      register_transport(:sse, RubyLLM::MCP::Transports::SSE)
+      register_transport(:stdio, RubyLLM::MCP::Transports::Stdio)
+      register_transport(:streamable, RubyLLM::MCP::Transports::StreamableHTTP)
+      register_transport(:streamable_http, RubyLLM::MCP::Transports::StreamableHTTP)
+
+      attr_reader :transport_type, :coordinator, :config, :pid
+
+      def initialize(transport_type, coordinator, config:)
+        @transport_type = transport_type
+        @coordinator = coordinator
+        @config = config
+        @pid = Process.pid
+      end
+
+      def_delegators :transport_protocol, :request, :alive?, :close, :start, :set_protocol_version
+
+      def transport_protocol
+        if @pid != Process.pid
+          @pid = Process.pid
+          @transport = build_transport
+          coordinator.restart_transport
+        end
+
+        @transport_protocol ||= build_transport
+      end
+
+      private
+
+      def build_transport
+        unless RubyLLM::MCP::Transport.transports.key?(transport_type)
+          supported_types = RubyLLM::MCP::Transport.transports.keys.join(", ")
+          message = "Invalid transport type: :#{transport_type}. Supported types are #{supported_types}"
+          raise Errors::InvalidTransportType.new(message: message)
+        end
+
+        transport_klass = RubyLLM::MCP::Transport.transports[transport_type]
+        transport_klass.new(coordinator: coordinator, **config)
+      end
+    end
+  end
+end

data/lib/ruby_llm/mcp/transports/http_client.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+require "httpx"
+
+module RubyLLM
+  module MCP
+    module Transports
+      class HTTPClient
+        CONNECTION_KEY = :ruby_llm_mcp_client_connection
+
+        def self.connection
+          Thread.current[CONNECTION_KEY] ||= build_connection
+        end
+
+        def self.build_connection
+          HTTPX.with(
+            pool_options: {
+              max_connections: RubyLLM::MCP.config.max_connections,
+              pool_timeout: RubyLLM::MCP.config.pool_timeout
+            }
+          )
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/mcp/{transport → transports}/sse.rb

@@ -8,11 +8,13 @@ require "securerandom"
 
 module RubyLLM
   module MCP
-    module Transport
+    module Transports
       class SSE
+        include Timeout
+
         attr_reader :headers, :id, :coordinator
 
-        def initialize(url, coordinator:, request_timeout:, headers: {})
+        def initialize(url:, coordinator:, request_timeout:, headers: {})
           @event_url = url
           @messages_url = nil
           @coordinator = coordinator
@@ -36,13 +38,10 @@ module RubyLLM
           @pending_requests = {}
           @pending_mutex = Mutex.new
           @connection_mutex = Mutex.new
-          @running = true
+          @running = false
           @sse_thread = nil
 
           RubyLLM::MCP.logger.info "Initializing SSE transport to #{@event_url} with client ID #{@client_id}"
-
-          # Start the SSE listener thread
-          start_sse_listener
         end
 
         def request(body, add_id: true, wait_for_response: true) # rubocop:disable Metrics/MethodLength
@@ -60,7 +59,8 @@ module RubyLLM
           end
 
           begin
-            http_client = HTTPX.with(timeout: { request_timeout: @request_timeout / 1000 }, headers: @headers)
+            http_client = HTTPClient.connection.with(timeout: { request_timeout: @request_timeout / 1000 },
+                                                     headers: @headers)
             response = http_client.post(@messages_url, body: JSON.generate(body))
 
             unless response.status == 200
@@ -83,16 +83,13 @@ module RubyLLM
           return unless wait_for_response
 
           begin
-            Timeout.timeout(@request_timeout / 1000) do
+            with_timeout(@request_timeout / 1000, request_id: request_id) do
               response_queue.pop
             end
-          rescue Timeout::Error
+          rescue RubyLLM::MCP::Errors::TimeoutError => e
             @pending_mutex.synchronize { @pending_requests.delete(request_id.to_s) }
             RubyLLM::MCP.logger.error "SSE request timeout (ID: #{request_id}) after #{@request_timeout / 1000} seconds"
-            raise Errors::TimeoutError.new(
-              message: "Request timed out after #{@request_timeout / 1000} seconds",
-              request_id: request_id
-            )
+            raise e
           end
         end
 
@@ -100,6 +97,13 @@ module RubyLLM
           @running
         end
 
+        def start
+          return if @running
+
+          @running = true
+          start_sse_listener
+        end
+
         def close
           RubyLLM::MCP.logger.info "Closing SSE transport connection"
           @running = false
@@ -107,6 +111,10 @@ module RubyLLM
           @sse_thread = nil
         end
 
+        def set_protocol_version(version)
+          @protocol_version = version
+        end
+
         private
 
         def start_sse_listener
@@ -125,7 +133,7 @@ module RubyLLM
             end
             @sse_thread.abort_on_exception = true
 
-            Timeout.timeout(100) do
+            with_timeout(@request_timeout / 1000) do
               endpoint = response_queue.pop
               set_message_endpoint(endpoint)
             end
@@ -179,7 +187,7 @@ module RubyLLM
           sleep 1
         end
 
-        def process_event(raw_event) # rubocop:disable Metrics/MethodLength
+        def process_event(raw_event)
           # Return if we believe that are getting a partial event
           return if raw_event[:data].nil?
 
@@ -209,15 +217,8 @@ module RubyLLM
             request_id = event["id"]&.to_s
             result = RubyLLM::MCP::Result.new(event)
 
-            if result.notification?
-              coordinator.process_notification(result)
-              return
-            end
-
-            if result.request?
-              coordinator.process_request(result) if coordinator.alive?
-              return
-            end
+            result = @coordinator.process_result(result)
+            return if result.nil?
 
             @pending_mutex.synchronize do
               # You can receieve duplicate events for the same request id, and we will ignore thoses

data/lib/ruby_llm/mcp/{transport → transports}/stdio.rb

@@ -7,11 +7,13 @@ require "securerandom"
 
 module RubyLLM
   module MCP
-    module Transport
+    module Transports
       class Stdio
+        include Timeout
+
         attr_reader :command, :stdin, :stdout, :stderr, :id, :coordinator
 
-        def initialize(command, request_timeout:, coordinator:, args: [], env: {})
+        def initialize(command:, request_timeout:, coordinator:, args: [], env: {})
           @request_timeout = request_timeout
           @command = command
           @coordinator = coordinator
@@ -23,11 +25,9 @@ module RubyLLM
           @id_mutex = Mutex.new
           @pending_requests = {}
           @pending_mutex = Mutex.new
-          @running = true
+          @running = false
           @reader_thread = nil
           @stderr_thread = nil
-
-          start_process
         end
 
         def request(body, add_id: true, wait_for_response: true)
@@ -58,15 +58,14 @@ module RubyLLM
           return unless wait_for_response
 
           begin
-            Timeout.timeout(@request_timeout / 1000) do
+            with_timeout(@request_timeout / 1000, request_id: request_id) do
               response_queue.pop
             end
-          rescue Timeout::Error
+          rescue RubyLLM::MCP::Errors::TimeoutError => e
             @pending_mutex.synchronize { @pending_requests.delete(request_id.to_s) }
-            raise RubyLLM::MCP::Errors::TimeoutError.new(
-              message: "Request timed out after #{@request_timeout / 1000} seconds",
-              request_id: request_id
-            )
+            log_message = "Stdio request timeout (ID: #{request_id}) after #{@request_timeout / 1000} seconds"
+            RubyLLM::MCP.logger.error(log_message)
+            raise e
           end
         end
 
@@ -74,6 +73,11 @@ module RubyLLM
           @running
         end
 
+        def start
+          start_process unless @running
+          @running = true
+        end
+
         def close # rubocop:disable Metrics/MethodLength
           @running = false
 
@@ -121,6 +125,10 @@ module RubyLLM
           @stderr_thread = nil
         end
 
+        def set_protocol_version(version)
+          @protocol_version = version
+        end
+
         private
 
         def start_process
@@ -199,22 +207,16 @@ module RubyLLM
           response = JSON.parse(line)
           request_id = response["id"]&.to_s
           result = RubyLLM::MCP::Result.new(response)
-
           RubyLLM::MCP.logger.debug "Result Received: #{result.inspect}"
-          # Handle notifications (process but don't return - continue processing other responses)
-          if result.notification?
-            coordinator.process_notification(result)
-            # Don't return here - continue to process potential tool responses
-          elsif result.request?
-            coordinator.process_request(result)
-            nil
-          else
-            # Handle regular responses (tool calls, etc.)
-            @pending_mutex.synchronize do
-              if result.matching_id?(request_id) && @pending_requests.key?(request_id)
-                response_queue = @pending_requests.delete(request_id)
-                response_queue&.push(result)
-              end
+
+          result = @coordinator.process_result(result)
+          return if result.nil?
+
+          # Handle regular responses (tool calls, etc.)
+          @pending_mutex.synchronize do
+            if result.matching_id?(request_id) && @pending_requests.key?(request_id)
+              response_queue = @pending_requests.delete(request_id)
+              response_queue&.push(result)
             end
           end
         rescue JSON::ParserError => e

data/lib/ruby_llm/mcp/{transport → transports}/streamable_http.rb

@@ -8,7 +8,7 @@ require "securerandom"
 
 module RubyLLM
   module MCP
-    module Transport
+    module Transports
       # Configuration options for reconnection behavior
       class ReconnectionOptions
         attr_reader :max_reconnection_delay, :initial_reconnection_delay,
@@ -40,10 +40,12 @@ module RubyLLM
 
       # Main StreamableHTTP transport class
       class StreamableHTTP
+        include Timeout
+
         attr_reader :session_id, :protocol_version, :coordinator
 
         def initialize( # rubocop:disable Metrics/ParameterLists
-          url,
+          url:,
           request_timeout:,
           coordinator:,
           headers: {},
@@ -110,6 +112,12 @@ module RubyLLM
           @abort_controller = false
         end
 
+        def set_protocol_version(version)
+          @protocol_version = version
+        end
+
+        private
+
         def terminate_session
           return unless @session_id
 
@@ -139,12 +147,6 @@ module RubyLLM
           end
         end
 
-        def set_protocol_version(version)
-          @protocol_version = version
-        end
-
-        private
-
         def handle_httpx_error_response!(response, context:, allow_eof_for_sse: false)
           return false unless response.is_a?(HTTPX::ErrorResponse)
 
@@ -200,7 +202,7 @@ module RubyLLM
         end
 
         def create_connection
-          client = HTTPX.with(
+          client = HTTPClient.connection.with(
             timeout: {
               connect_timeout: 10,
               read_timeout: @request_timeout / 1000,
@@ -257,7 +259,7 @@ module RubyLLM
         def create_connection_with_streaming_callbacks(request_id)
           buffer = +""
 
-          client = HTTPX.plugin(:callbacks).on_response_body_chunk do |request, _response, chunk|
+          client = HTTPClient.connection.plugin(:callbacks).on_response_body_chunk do |request, _response, chunk|
             next unless @running && !@abort_controller
 
             RubyLLM::MCP.logger.debug "Received chunk: #{chunk.bytesize} bytes for #{request.uri}"
@@ -562,19 +564,14 @@ module RubyLLM
             result = RubyLLM::MCP::Result.new(event_data, session_id: @session_id)
             RubyLLM::MCP.logger.debug "SSE Result Received: #{result.inspect}"
 
-            # Handle different types of messages
-            if result.notification?
-              @coordinator.process_notification(result)
-            elsif result.request?
-              @coordinator.process_request(result)
-            elsif result.response?
-              # Handle response to client request
-              request_id = result.id&.to_s
-              if request_id
-                @pending_mutex.synchronize do
-                  response_queue = @pending_requests.delete(request_id)
-                  response_queue&.push(result)
-                end
+            result = @coordinator.process_result(result)
+            return if result.nil?
+
+            request_id = result.id&.to_s
+            if request_id
+              @pending_mutex.synchronize do
+                response_queue = @pending_requests.delete(request_id)
+                response_queue&.push(result)
               end
             end
           rescue JSON::ParserError => e
@@ -591,15 +588,14 @@ module RubyLLM
         end
 
         def wait_for_response_with_timeout(request_id, response_queue)
-          Timeout.timeout(@request_timeout / 1000) do
+          with_timeout(@request_timeout / 1000, request_id: request_id) do
             response_queue.pop
           end
-        rescue Timeout::Error
+        rescue RubyLLM::MCP::Errors::TimeoutError => e
+          log_message = "StreamableHTTP request timeout (ID: #{request_id}) after #{@request_timeout / 1000} seconds"
+          RubyLLM::MCP.logger.error(log_message)
           @pending_mutex.synchronize { @pending_requests.delete(request_id.to_s) }
-          raise Errors::TimeoutError.new(
-            message: "Request timed out after #{@request_timeout / 1000} seconds",
-            request_id: request_id
-          )
+          raise e
         end
 
         def cleanup_sse_resources

data/lib/ruby_llm/mcp/transports/timeout.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module MCP
+    module Transports
+      module Timeout
+        def with_timeout(seconds, request_id: nil)
+          result = nil
+          exception = nil
+
+          worker = Thread.new do
+            result = yield
+          rescue StandardError => e
+            exception = e
+          end
+
+          if worker.join(seconds)
+            raise exception if exception
+
+            result
+          else
+            worker.kill # stop the thread (can still have some risk if shared resources)
+            raise RubyLLM::MCP::Errors::TimeoutError.new(
+              message: "Request timed out after #{seconds} seconds",
+              request_id: request_id
+            )
+          end
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/mcp/version.rb

@@ -2,6 +2,6 @@
 
 module RubyLLM
   module MCP
-    VERSION = "0.4.1"
+    VERSION = "0.5.0"
   end
 end

data/lib/ruby_llm/mcp.rb

@@ -4,19 +4,49 @@ require "ruby_llm"
 require "zeitwerk"
 require_relative "chat"
 
-loader = Zeitwerk::Loader.for_gem_extension(RubyLLM)
-loader.inflector.inflect("mcp" => "MCP")
-loader.inflector.inflect("sse" => "SSE")
-loader.inflector.inflect("openai" => "OpenAI")
-loader.inflector.inflect("streamable_http" => "StreamableHTTP")
-loader.setup
-
 module RubyLLM
   module MCP
     module_function
 
-    def client(*args, **kwargs)
-      @client ||= Client.new(*args, **kwargs)
+    def clients(config = RubyLLM::MCP.config.mcp_configuration)
+      @clients ||= {}
+      config.map do |options|
+        @clients[options[:name]] ||= Client.new(**options)
+      end
+    end
+
+    def add_client(options)
+      @clients ||= {}
+      @clients[options[:name]] ||= Client.new(**options)
+    end
+
+    def remove_client(name)
+      @clients ||= {}
+      client = @clients.delete(name)
+      client&.stop
+      client
+    end
+
+    def client(...)
+      Client.new(...)
+    end
+
+    def establish_connection(&)
+      clients.each(&:start)
+      yield clients
+    ensure
+      clients.each do |client|
+        client.stop if client.alive?
+      end
+    end
+
+    def tools(blacklist: [], whitelist: [])
+      tools = @clients.values.map(&:tools)
+                      .flatten
+                      .reject { |tool| blacklist.include?(tool.name) }
+
+      tools = tools.select { |tool| whitelist.include?(tool.name) } if whitelist.any?
+      tools.uniq(&:name)
     end
 
     def support_complex_parameters!
@@ -41,3 +71,14 @@ module RubyLLM
     end
   end
 end
+
+require_relative "mcp/railtie" if defined?(Rails::Railtie)
+
+loader = Zeitwerk::Loader.for_gem_extension(RubyLLM)
+loader.inflector.inflect("mcp" => "MCP")
+loader.inflector.inflect("sse" => "SSE")
+loader.inflector.inflect("openai" => "OpenAI")
+loader.inflector.inflect("streamable_http" => "StreamableHTTP")
+loader.inflector.inflect("http_client" => "HTTPClient")
+
+loader.setup

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: ruby_llm-mcp
 version: !ruby/object:Gem::Version
-  version: 0.4.1
+  version: 0.5.0
 platform: ruby
 authors:
 - Patrick Vice
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2025-07-02 00:00:00.000000000 Z
+date: 2025-07-07 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: httpx
@@ -68,7 +68,6 @@ files:
 - lib/ruby_llm/chat.rb
 - lib/ruby_llm/mcp.rb
 - lib/ruby_llm/mcp/attachment.rb
-- lib/ruby_llm/mcp/capabilities.rb
 - lib/ruby_llm/mcp/client.rb
 - lib/ruby_llm/mcp/completion.rb
 - lib/ruby_llm/mcp/configuration.rb
@@ -77,39 +76,51 @@ files:
 - lib/ruby_llm/mcp/error.rb
 - lib/ruby_llm/mcp/errors.rb
 - lib/ruby_llm/mcp/logging.rb
+- lib/ruby_llm/mcp/notification_handler.rb
+- lib/ruby_llm/mcp/notifications/cancelled.rb
+- lib/ruby_llm/mcp/notifications/initialize.rb
+- lib/ruby_llm/mcp/notifications/roots_list_change.rb
 - lib/ruby_llm/mcp/parameter.rb
 - lib/ruby_llm/mcp/progress.rb
 - lib/ruby_llm/mcp/prompt.rb
 - lib/ruby_llm/mcp/providers/anthropic/complex_parameter_support.rb
 - lib/ruby_llm/mcp/providers/gemini/complex_parameter_support.rb
 - lib/ruby_llm/mcp/providers/openai/complex_parameter_support.rb
-- lib/ruby_llm/mcp/requests/base.rb
-- lib/ruby_llm/mcp/requests/cancelled_notification.rb
+- lib/ruby_llm/mcp/railtie.rb
 - lib/ruby_llm/mcp/requests/completion_prompt.rb
 - lib/ruby_llm/mcp/requests/completion_resource.rb
 - lib/ruby_llm/mcp/requests/initialization.rb
-- lib/ruby_llm/mcp/requests/initialize_notification.rb
 - lib/ruby_llm/mcp/requests/logging_set_level.rb
-- lib/ruby_llm/mcp/requests/meta.rb
 - lib/ruby_llm/mcp/requests/ping.rb
-- lib/ruby_llm/mcp/requests/ping_response.rb
 - lib/ruby_llm/mcp/requests/prompt_call.rb
 - lib/ruby_llm/mcp/requests/prompt_list.rb
 - lib/ruby_llm/mcp/requests/resource_list.rb
 - lib/ruby_llm/mcp/requests/resource_read.rb
 - lib/ruby_llm/mcp/requests/resource_template_list.rb
 - lib/ruby_llm/mcp/requests/resources_subscribe.rb
+- lib/ruby_llm/mcp/requests/shared/meta.rb
+- lib/ruby_llm/mcp/requests/shared/pagination.rb
 - lib/ruby_llm/mcp/requests/tool_call.rb
 - lib/ruby_llm/mcp/requests/tool_list.rb
 - lib/ruby_llm/mcp/resource.rb
 - lib/ruby_llm/mcp/resource_template.rb
+- lib/ruby_llm/mcp/response_handler.rb
+- lib/ruby_llm/mcp/responses/error.rb
+- lib/ruby_llm/mcp/responses/ping.rb
+- lib/ruby_llm/mcp/responses/roots_list.rb
+- lib/ruby_llm/mcp/responses/sampling_create_message.rb
 - lib/ruby_llm/mcp/result.rb
+- lib/ruby_llm/mcp/roots.rb
+- lib/ruby_llm/mcp/sample.rb
+- lib/ruby_llm/mcp/server_capabilities.rb
 - lib/ruby_llm/mcp/tool.rb
-- lib/ruby_llm/mcp/transport/sse.rb
-- lib/ruby_llm/mcp/transport/stdio.rb
-- lib/ruby_llm/mcp/transport/streamable_http.rb
+- lib/ruby_llm/mcp/transport.rb
+- lib/ruby_llm/mcp/transports/http_client.rb
+- lib/ruby_llm/mcp/transports/sse.rb
+- lib/ruby_llm/mcp/transports/stdio.rb
+- lib/ruby_llm/mcp/transports/streamable_http.rb
+- lib/ruby_llm/mcp/transports/timeout.rb
 - lib/ruby_llm/mcp/version.rb
-- lib/tasks/release.rake
 homepage: https://github.com/patvice/ruby_llm-mcp
 licenses:
 - MIT

data/lib/ruby_llm/mcp/requests/base.rb

@@ -1,31 +0,0 @@
-# frozen_string_literal: true
-
-require "json"
-
-module RubyLLM
-  module MCP
-    module Requests
-      class Base
-        attr_reader :coordinator
-
-        def initialize(coordinator)
-          @coordinator = coordinator
-        end
-
-        def call
-          raise "Not implemented"
-        end
-
-        private
-
-        def validate_response!(response, body)
-          # TODO: Implement response validation
-        end
-
-        def raise_error(error)
-          raise "MCP Error: code: #{error['code']} message: #{error['message']} data: #{error['data']}"
-        end
-      end
-    end
-  end
-end

data/lib/ruby_llm/mcp/requests/meta.rb

@@ -1,30 +0,0 @@
-# frozen_string_literal: true
-
-require "securerandom"
-
-module RubyLLM
-  module MCP
-    module Requests
-      module Meta
-        def merge_meta(body)
-          meta = {}
-          meta.merge!(progress_token) if @coordinator.client.tracking_progress?
-
-          body[:params] ||= {}
-          body[:params].merge!({ _meta: meta }) unless meta.empty?
-          body
-        end
-
-        private
-
-        def progress_token
-          { progressToken: generate_progress_token }
-        end
-
-        def generate_progress_token
-          SecureRandom.uuid
-        end
-      end
-    end
-  end
-end