ruby_llm-mcp 0.4.1 → 0.5.1

data/lib/ruby_llm/mcp/sample.rb

@@ -0,0 +1,148 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module MCP
+    class Sample
+      class Hint
+        attr_reader :model, :cost_priority, :speed_priority, :intelligence_priority
+
+        def initialize(model, model_preferences)
+          @model = model
+          @model_preferences = model_preferences
+
+          @hints = model_preferences&.fetch("hints", [])
+          @cost_priority = model_preferences&.fetch("costPriority", nil)
+          @speed_priority = model_preferences&.fetch("speedPriority", nil)
+          @intelligence_priority = model_preferences&.fetch("intelligencePriority", nil)
+        end
+
+        def hints
+          @hints.map { |hint| hint["name"] }
+        end
+
+        def to_h
+          {
+            model: model,
+            hints: hints,
+            cost_priority: @cost_priority,
+            speed_priority: @speed_priority,
+            intelligence_priority: @intelligence_priority
+          }
+        end
+      end
+
+      REJECTED_MESSAGE = "Sampling request was rejected"
+
+      attr_reader :model_preferences, :system_prompt, :max_tokens, :raw_messages
+
+      def initialize(result, coordinator)
+        params = result.params
+        @id = result.id
+        @coordinator = coordinator
+
+        @raw_messages = params["messages"] || []
+        @model_preferences = Hint.new(params["model"], params["modelPreferences"])
+        @system_prompt = params["systemPrompt"]
+        @max_tokens = params["maxTokens"]
+      end
+
+      def execute
+        return unless callback_guard_success?
+
+        model = preferred_model
+        return unless model
+
+        chat_message = chat(model)
+        @coordinator.sampling_create_message_response(
+          id: @id, message: chat_message, model: model
+        )
+      end
+
+      def message
+        @message ||= raw_messages.map { |message| message.fetch("content")&.fetch("text") }.join("\n")
+      end
+
+      def to_h
+        {
+          id: @id,
+          model_preferences: @model_preferences.to_h,
+          system_prompt: @system_prompt,
+          max_tokens: @max_tokens
+        }
+      end
+
+      alias to_json to_h
+
+      private
+
+      def callback_guard_success?
+        return true unless @coordinator.client.sampling_callback_enabled?
+
+        unless @coordinator.client.on[:sampling].call(self)
+          @coordinator.error_response(id: @id, message: REJECTED_MESSAGE)
+          return false
+        end
+
+        true
+      rescue StandardError => e
+        RubyLLM::MCP.logger.error("Error in callback guard: #{e.message}, #{e.backtrace.join("\n")}")
+        @coordinator.error_response(id: @id, message: "Error executing sampling request")
+        false
+      end
+
+      def chat(model)
+        chat = RubyLLM::Chat.new(
+          model: model
+        )
+        if system_prompt
+          formated_system_message = create_message(system_message)
+          chat.add_message(formated_system_message)
+        end
+        raw_messages.each { |message| chat.add_message(create_message(message)) }
+
+        chat.complete
+      end
+
+      def preferred_model
+        @preferred_model ||= begin
+          model = RubyLLM::MCP.config.sampling.preferred_model
+          if model.respond_to?(:call)
+            model.call(model_preferences)
+          else
+            model
+          end
+        end
+      rescue StandardError => e
+        RubyLLM::MCP.logger.error("Error in preferred model: #{e.message}, #{e.backtrace.join("\n")}")
+        @coordinator.error_response(id: @id, message: "Failed to determine preferred model: #{e.message}")
+        false
+      end
+
+      def create_message(message)
+        role = message["role"]
+        content = create_content_for_message(message["content"])
+
+        RubyLLM::Message.new({ role: role, content: content })
+      end
+
+      def create_content_for_message(content)
+        case content["type"]
+        when "text"
+          MCP::Content.new(text: content["text"])
+        when "image", "audio"
+          attachment = MCP::Attachment.new(content["data"], content["mimeType"])
+          MCP::Content.new(text: nil, attachments: [attachment])
+        else
+          raise RubyLLM::MCP::Errors::InvalidFormatError.new(message: "Invalid content type: #{content['type']}")
+        end
+      end
+
+      def system_message
+        {
+          "role" => "system",
+          "content" => { "type" => "text", "text" => system_prompt }
+        }
+      end
+    end
+  end
+end

data/lib/ruby_llm/mcp/{capabilities.rb → server_capabilities.rb}

@@ -2,7 +2,7 @@
 
 module RubyLLM
   module MCP
-    class Capabilities
+    class ServerCapabilities
       attr_accessor :capabilities
 
       def initialize(capabilities = {})

data/lib/ruby_llm/mcp/tool.rb

@@ -12,16 +12,28 @@ module RubyLLM
         @idempotent_hint = annotation["idempotentHint"] || false
         @open_world_hint = annotation["openWorldHint"] || true
       end
+
+      def to_h
+        {
+          title: @title,
+          readOnlyHint: @read_only_hint,
+          destructiveHint: @destructive_hint,
+          idempotentHint: @idempotent_hint,
+          openWorldHint: @open_world_hint
+        }
+      end
     end
 
     class Tool < RubyLLM::Tool
-      attr_reader :name, :description, :parameters, :coordinator, :tool_response
+      attr_reader :name, :description, :parameters, :coordinator, :tool_response, :with_prefix
 
-      def initialize(coordinator, tool_response)
+      def initialize(coordinator, tool_response, with_prefix: false)
         super()
         @coordinator = coordinator
 
-        @name = tool_response["name"]
+        @with_prefix = with_prefix
+        @name = format_name(tool_response["name"])
+        @mcp_name = tool_response["name"]
         @description = tool_response["description"].to_s
         @parameters = create_parameters(tool_response["inputSchema"])
         @annotations = tool_response["annotations"] ? Annotation.new(tool_response["annotations"]) : nil
@@ -33,7 +45,7 @@ module RubyLLM
 
       def execute(**params)
         result = @coordinator.execute_tool(
-          name: @name,
+          name: @mcp_name,
           parameters: params
         )
 
@@ -54,6 +66,17 @@ module RubyLLM
         end
       end
 
+      def to_h
+        {
+          name: @name,
+          description: @description,
+          parameters: @parameters.to_h,
+          annotations: @annotations&.to_h
+        }
+      end
+
+      alias to_json to_h
+
       private
 
       def create_parameters(input_schema)
@@ -136,6 +159,14 @@ module RubyLLM
           resource.to_content
         end
       end
+
+      def format_name(name)
+        if @with_prefix
+          "#{@coordinator.name}_#{name}"
+        else
+          name
+        end
+      end
     end
   end
 end

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.1"
   end
 end