Merge pull request #30 from patvice/roots-and-sampling

Client Features: Support for Roots and Sampling

Author: patvice

README.md

@@ -12,6 +12,7 @@ This project is a Ruby client for the [Model Context Protocol (MCP)](https://mod
 - 🛠️ **Tool Integration**: Automatically converts MCP tools into RubyLLM-compatible tools
 - 📄 **Resource Management**: Access and include MCP resources (files, data) and resource templates in conversations
 - 🎯 **Prompt Integration**: Use predefined MCP prompts with arguments for consistent interactions
+- 🎛️ **Client Features**: Support for sampling and roots
 - 🎨 **Enhanced Chat Interface**: Extended RubyLLM chat methods for seamless MCP integration
 - 📚 **Simple API**: Easy-to-use interface that integrates seamlessly with RubyLLM
 
@@ -462,6 +463,81 @@ puts result
 # Result: { status: "success", data: "Processed data" }
 ```
 
+## Client Features
+
+The RubyLLM::MCP client provides support functionality that can be exposed to MCP servers. These features must be explicitly configured before creating client objects to ensure you're opting into this functionality.
+
+### Roots
+
+Roots provide MCP servers with access to underlying file system information. The implementation starts with a lightweight approach due to the MCP specification's current limitations on root usage.
+
+When roots are configured, the client will:
+
+- Expose roots as a supported capability to MCP servers
+- Support dynamic addition and removal of roots during the client lifecycle
+- Fire `notifications/roots/list_changed` events when roots are modified
+
+#### Configuration
+
+```ruby
+RubyLLM::MCP.config do |config|
+  config.roots = ["to/a/path", Rails.root]
+end
+
+client = RubyLLM::MCP::Client.new(...)
+```
+
+#### Usage
+
+```ruby
+# Access current root paths
+client.roots.paths
+# => ["to/a/path", #<Pathname:/to/rails/root/path>]
+
+# Add a new root (fires list_changed notification)
+client.roots.add("new/path")
+client.roots.paths
+# => ["to/a/path", #<Pathname:/to/rails/root/path>, "new/path"]
+
+# Remove a root (fires list_changed notification)
+client.roots.remove("to/a/path")
+client.roots.paths
+# => [#<Pathname:/to/rails/root/path>, "new/path"]
+```
+
+### Sampling
+
+Sampling allows MCP servers to offload LLM requests to the MCP client rather than making them directly from the server. This enables MCP servers to optionally use LLM connections through the client.
+
+#### Configuration
+
+```ruby
+RubyLLM::MCP.configure do |config|
+  config.sampling.enabled = true
+  config.sampling.preferred_model = "gpt-4.1"
+
+  # Optional: Use a block for dynamic model selection
+  config.sampling.preferred_model do |model_preferences|
+    model_preferences.hints.first
+  end
+
+  # Optional: Add guards to filter sampling requests
+  config.sampling.guard do |sample|
+    sample.message.include("Hello")
+  end
+end
+```
+
+#### How It Works
+
+With the above configuration:
+
+- Clients will respond to all incoming sample requests using the specified model (`gpt-4.1`)
+- Sample messages will only be approved if they contain the word "Hello" (when using the guard)
+- The `preferred_model` can be a string or a proc that provides dynamic model selection based on MCP server characteristics
+
+The `preferred_model` proc receives model preferences from the MCP server, allowing you to make intelligent model selection decisions based on the server's requirements for success.
+
 ## Transport Types
 
 ### SSE (Server-Sent Events)

lib/ruby_llm/mcp/client.rb

@@ -7,7 +7,7 @@ module MCP
     class Client
       extend Forwardable
 
-      attr_reader :name, :config, :transport_type, :request_timeout, :log_level, :on
+      attr_reader :name, :config, :transport_type, :request_timeout, :log_level, :on, :roots
 
       def initialize(name:, transport_type:, start: true, request_timeout: MCP.config.request_timeout, config: {})
         @name = name
@@ -25,10 +25,13 @@ def initialize(name:, transport_type:, start: true, request_timeout: MCP.config.
 
         @log_level = nil
 
+        setup_roots
+        setup_sampling
+
         @coordinator.start_transport if start
       end
 
-      def_delegators :@coordinator, :alive?, :capabilities, :ping
+      def_delegators :@coordinator, :alive?, :capabilities, :ping, :client_capabilities
 
       def start
         @coordinator.start_transport
@@ -165,6 +168,15 @@ def on_logging(level: Logging::WARNING, logger: nil, &block)
         self
       end
 
+      def sampling_callback_enabled?
+        @on.key?(:sampling) && !@on[:sampling].nil?
+      end
+
+      def on_sampling(&block)
+        @on[:sampling] = block
+        self
+      end
+
       private
 
       def setup_coordinator
@@ -187,6 +199,14 @@ def build_map(raw_data, klass)
           acc[instance.name] = instance
         end
       end
+
+      def setup_roots
+        @roots = Roots.new(paths: MCP.config.roots, coordinator: @coordinator)
+      end
+
+      def setup_sampling
+        @on[:sampling] = MCP.config.sampling.guard
+      end
     end
   end
 end

lib/ruby_llm/mcp/configuration.rb

@@ -3,12 +3,48 @@
 module RubyLLM
   module MCP
     class Configuration
-      attr_accessor :request_timeout, :log_file, :log_level, :has_support_complex_parameters
+      class Sampling
+        attr_accessor :enabled
+        attr_writer :preferred_model
+
+        def initialize
+          set_defaults
+        end
+
+        def reset!
+          set_defaults
+        end
+
+        def guard(&block)
+          @guard = block if block_given?
+          @guard
+        end
+
+        def preferred_model(&block)
+          @preferred_model = block if block_given?
+          @preferred_model
+        end
+
+        def enabled?
+          @enabled
+        end
+
+        private
+
+        def set_defaults
+          @enabled = false
+          @preferred_model = nil
+          @guard = nil
+        end
+      end
+
+      attr_accessor :request_timeout, :log_file, :log_level, :has_support_complex_parameters, :roots, :sampling
       attr_writer :logger
 
       REQUEST_TIMEOUT_DEFAULT = 8000
 
       def initialize
+        @sampling = Sampling.new
         set_defaults
       end
 
@@ -60,6 +96,9 @@ def set_defaults
         @log_level = ENV["RUBYLLM_MCP_DEBUG"] ? Logger::DEBUG : Logger::INFO
         @has_support_complex_parameters = false
         @logger = nil
+        @roots = []
+
+        @sampling.reset!
       end
     end
   end

lib/ruby_llm/mcp/coordinator.rb

@@ -23,6 +23,10 @@ def initialize(client, transport_type:, config: {})
         @capabilities = nil
       end
 
+      def name
+        client.name
+      end
+
       def request(body, **options)
         @transport.request(body, **options)
       rescue RubyLLM::MCP::Errors::TimeoutError => e
@@ -83,39 +87,11 @@ def ping
 
       def process_notification(result)
         notification = result.notification
-
-        case notification.type
-        when "notifications/tools/list_changed"
-          client.reset_tools!
-        when "notifications/resources/list_changed"
-          client.reset_resources!
-        when "notifications/resources/updated"
-          uri = notification.params["uri"]
-          resource = client.resources.find { |r| r.uri == uri }
-          resource&.reset_content!
-        when "notifications/prompts/list_changed"
-          client.reset_prompts!
-        when "notifications/message"
-          process_logging_message(notification)
-        when "notifications/progress"
-          process_progress_message(notification)
-        when "notifications/cancelled"
-          # TODO: - do nothing at the moment until we support client operations
-        else
-          message = "Unknown notification type: #{notification.type} params:#{notification.params.to_h}"
-          raise Errors::UnknownNotification.new(message: message)
-        end
+        NotificationHandler.new(self).execute(notification)
       end
 
       def process_request(result)
-        if result.ping?
-          ping_response(id: result.id)
-          return
-        end
-
-        # Handle server-initiated requests
-        # Currently, we do not support any client operations but will
-        raise RubyLLM::MCP::Errors::UnknownRequest.new(message: "Unknown request type: #{result.inspect}")
+        ResponseHandler.new(self).execute(result)
       end
 
       def initialize_request
@@ -190,20 +166,56 @@ def completion_prompt(**args)
         RubyLLM::MCP::Requests::CompletionPrompt.new(self, **args).call
       end
 
+      def set_logging(**args)
+        RubyLLM::MCP::Requests::LoggingSetLevel.new(self, **args).call
+      end
+
+      ## Notifications
+      #
       def initialize_notification
-        RubyLLM::MCP::Requests::InitializeNotification.new(self).call
+        RubyLLM::MCP::Notifications::Initialize.new(self).call
       end
 
       def cancelled_notification(**args)
-        RubyLLM::MCP::Requests::CancelledNotification.new(self, **args).call
+        RubyLLM::MCP::Notifications::Cancelled.new(self, **args).call
+      end
+
+      def roots_list_change_notification
+        RubyLLM::MCP::Notifications::RootsListChange.new(self).call
+      end
+
+      ## Responses
+      #
+      def ping_response(**args)
+        RubyLLM::MCP::Responses::Ping.new(self, **args).call
+      end
+
+      def roots_list_response(**args)
+        RubyLLM::MCP::Responses::RootsList.new(self, **args).call
+      end
+
+      def sampling_create_message_response(**args)
+        RubyLLM::MCP::Responses::SamplingCreateMessage.new(self, **args).call
       end
 
-      def ping_response(id: nil)
-        RubyLLM::MCP::Requests::PingResponse.new(self, id: id).call
+      def error_response(**args)
+        RubyLLM::MCP::Responses::Error.new(self, **args).call
       end
 
-      def set_logging(level:)
-        RubyLLM::MCP::Requests::LoggingSetLevel.new(self, level: level).call
+      def client_capabilities
+        capabilities = {}
+
+        if client.roots.active?
+          capabilities[:roots] = {
+            listChanged: true
+          }
+        end
+
+        if sampling_enabled?
+          capabilities[:sampling] = {}
+        end
+
+        capabilities
       end
 
       def build_transport
@@ -230,46 +242,10 @@ def build_transport
         end
       end
 
-      def process_logging_message(notification)
-        if client.logging_handler_enabled?
-          client.on[:logging].call(notification)
-        else
-          default_process_logging_message(notification)
-        end
-      end
-
-      def default_process_logging_message(notification, logger: RubyLLM::MCP.logger)
-        level = notification.params["level"]
-        logger_message = notification.params["logger"]
-        message = notification.params["data"]
-
-        message = "#{logger_message}: #{message}"
-
-        case level
-        when "debug"
-          logger.debug(message["message"])
-        when "info", "notice"
-          logger.info(message["message"])
-        when "warning"
-          logger.warn(message["message"])
-        when "error", "critical"
-          logger.error(message["message"])
-        when "alert", "emergency"
-          logger.fatal(message["message"])
-        end
-      end
-
-      def name
-        client.name
-      end
-
       private
 
-      def process_progress_message(notification)
-        progress_obj = RubyLLM::MCP::Progress.new(self, client.on[:progress], notification.params)
-        if client.tracking_progress?
-          progress_obj.execute_progress_handler
-        end
+      def sampling_enabled?
+        MCP.config.sampling.enabled?
       end
     end
   end

lib/ruby_llm/mcp/errors.rb

@@ -17,6 +17,8 @@ class CompletionNotAvailable < BaseError; end
         class ResourceSubscribeNotAvailable < BaseError; end
       end
 
+      class InvalidFormatError < BaseError; end
+
       class InvalidProtocolVersionError < BaseError; end
 
       class InvalidTransportType < BaseError; end