ruby_llm-mcp 0.7.1 → 1.0.0

data/lib/ruby_llm/mcp/client.rb

@@ -7,17 +7,42 @@ module RubyLLM
     class Client
       extend Forwardable
 
-      attr_reader :name, :config, :transport_type, :request_timeout, :log_level, :on, :roots
+      attr_reader :name, :config, :transport_type, :request_timeout, :log_level, :on, :roots, :adapter,
+                  :on_logging_level
       attr_accessor :linked_resources
 
-      def initialize(name:, transport_type:, start: true, request_timeout: MCP.config.request_timeout, config: {})
+      def initialize(name:, transport_type:, sdk: nil, adapter: nil, start: true, # rubocop:disable Metrics/ParameterLists, Metrics/MethodLength
+                     request_timeout: MCP.config.request_timeout, config: {})
         @name = name
-        @with_prefix = config.delete(:with_prefix) || false
-        @config = config.merge(request_timeout: request_timeout)
         @transport_type = transport_type.to_sym
+        @adapter_type = adapter || sdk || MCP.config.default_adapter
+
+        # Validate early
+        MCP.config.adapter_config.validate!(
+          adapter: @adapter_type,
+          transport: @transport_type
+        )
+
+        @with_prefix = config.delete(:with_prefix) || false
+        @resolved_protocol_version = config[:protocol_version] ||
+                                     config["protocol_version"] ||
+                                     MCP.config.protocol_version
+        @resolved_extensions = Extensions::Registry.merge(
+          MCP.config.extensions.to_h,
+          config[:extensions] || config["extensions"]
+        )
+
+        @config = config.merge(
+          request_timeout: request_timeout,
+          protocol_version: @resolved_protocol_version,
+          extensions: @resolved_extensions
+        )
         @request_timeout = request_timeout
 
-        @coordinator = setup_coordinator
+        # Store OAuth config for later use
+        @oauth_config = config[:oauth] || config["oauth"]
+        @oauth_provider = nil
+        @oauth_storage = nil
 
         @on = {}
         @tools = {}
@@ -29,31 +54,73 @@ module RubyLLM
 
         @linked_resources = []
 
-        setup_roots
-        setup_sampling
+        # Build adapter based on configuration
+        @adapter = build_adapter
+
+        setup_roots if @adapter.supports?(:roots)
+        setup_sampling if @adapter.supports?(:sampling)
+        setup_event_handlers
+        sync_elicitation_handler_state
 
-        @coordinator.start_transport if start
+        @adapter.start if start
       end
 
-      def_delegators :@coordinator, :alive?, :capabilities, :ping, :client_capabilities
+      def_delegators :@adapter, :alive?, :capabilities, :ping, :client_capabilities,
+                     :register_in_flight_request, :unregister_in_flight_request,
+                     :cancel_in_flight_request
 
       def start
-        @coordinator.start_transport
+        @adapter.start
       end
 
       def stop
-        @coordinator.stop_transport
+        @adapter.stop
       end
 
       def restart!
-        @coordinator.restart_transport
+        @adapter.restart!
+      end
+
+      # Get or create OAuth provider for this client
+      # @param type [Symbol] OAuth provider type (:standard or :browser, defaults to :standard)
+      # @param options [Hash] additional options passed to provider
+      # @return [OAuthProvider, BrowserOAuthProvider] OAuth provider instance
+      def oauth(type: :standard, **options)
+        # Return existing provider if already created
+        return @oauth_provider if @oauth_provider
+
+        # Get provider from transport if it already exists
+        transport_oauth = transport_oauth_provider
+        return transport_oauth if transport_oauth
+
+        # Create new provider lazily
+        server_url = @config[:url] || @config["url"]
+        unless server_url
+          raise Errors::ConfigurationError.new(
+            message: "Cannot create OAuth provider without server URL in config"
+          )
+        end
+
+        oauth_options = {
+          server_url: server_url,
+          scope: @oauth_config&.dig(:scope) || @oauth_config&.dig("scope"),
+          storage: oauth_storage,
+          **options
+        }
+
+        @oauth_provider = Auth.create_oauth(
+          server_url,
+          type: type,
+          **oauth_options
+        )
       end
 
       def tools(refresh: false)
+        require_feature!(:tools)
         return [] unless capabilities.tools_list?
 
         fetch(:tools, refresh) do
-          tools = @coordinator.tool_list
+          tools = @adapter.tool_list
           build_map(tools, MCP::Tool, with_prefix: @with_prefix)
         end
 
@@ -71,10 +138,11 @@ module RubyLLM
       end
 
       def resources(refresh: false)
+        require_feature!(:resources)
         return [] unless capabilities.resources_list?
 
         fetch(:resources, refresh) do
-          resources = @coordinator.resource_list
+          resources = @adapter.resource_list
           resources = build_map(resources, MCP::Resource)
           include_linked_resources(resources)
         end
@@ -88,15 +156,30 @@ module RubyLLM
         @resources[name]
       end
 
+      def unsubscribe_from_resource(resource_or_uri) # rubocop:disable Naming/PredicateMethod
+        require_feature!(:subscriptions)
+
+        uri = if resource_or_uri.respond_to?(:uri)
+                resource_or_uri.uri
+              else
+                resource_or_uri.to_s
+              end
+        @adapter.resources_unsubscribe(uri: uri)
+        resource = @resources.values.find { |existing| existing.uri == uri }
+        resource&.instance_variable_set(:@subscribed, false)
+        true
+      end
+
       def reset_resources!
         @resources = {}
       end
 
       def resource_templates(refresh: false)
+        require_feature!(:resource_templates)
         return [] unless capabilities.resources_list?
 
         fetch(:resource_templates, refresh) do
-          resource_templates = @coordinator.resource_template_list
+          resource_templates = @adapter.resource_template_list
           build_map(resource_templates, MCP::ResourceTemplate)
         end
 
@@ -114,10 +197,11 @@ module RubyLLM
       end
 
       def prompts(refresh: false)
+        require_feature!(:prompts)
         return [] unless capabilities.prompt_list?
 
         fetch(:prompts, refresh) do
-          prompts = @coordinator.prompt_list
+          prompts = @adapter.prompt_list
           build_map(prompts, MCP::Prompt)
         end
 
@@ -134,11 +218,46 @@ module RubyLLM
         @prompts = {}
       end
 
+      def tasks_list
+        require_feature!(:tasks)
+        return [] unless capabilities.tasks? && capabilities.tasks_list?
+
+        @adapter.tasks_list.map { |task| MCP::Task.new(@adapter, task) }
+      end
+
+      def task_get(task_id)
+        require_feature!(:tasks)
+        result = @adapter.task_get(task_id: task_id)
+        MCP::Task.new(@adapter, result.value)
+      end
+
+      def task_result(task_id)
+        require_feature!(:tasks)
+        result = @adapter.task_result(task_id: task_id)
+        result.value
+      end
+
+      def task_cancel(task_id)
+        require_feature!(:tasks)
+        unless capabilities.tasks_cancel?
+          message = "Task cancellation is not available for this MCP server"
+          raise Errors::Capabilities::TaskCancelNotAvailable.new(message: message)
+        end
+
+        result = @adapter.task_cancel(task_id: task_id)
+        MCP::Task.new(@adapter, result.value)
+      end
+
       def tracking_progress?
         @on.key?(:progress) && !@on[:progress].nil?
       end
 
       def on_progress(&block)
+        require_feature!(:progress_tracking)
+        if alive?
+          @adapter.set_progress_tracking(enabled: true)
+        end
+
         @on[:progress] = block
         self
       end
@@ -147,8 +266,23 @@ module RubyLLM
         @on.key?(:human_in_the_loop) && !@on[:human_in_the_loop].nil?
       end
 
-      def on_human_in_the_loop(&block)
-        @on[:human_in_the_loop] = block
+      def on_human_in_the_loop(handler_class = nil, **options)
+        require_feature!(:human_in_the_loop)
+
+        if block_given?
+          raise ArgumentError, "Block-based human-in-the-loop callbacks are no longer supported. Use a handler class."
+        end
+
+        if handler_class
+          # Validate handler class
+          validate_handler_class!(handler_class, :execute)
+
+          @on[:human_in_the_loop] = { class: handler_class, options: options }
+        else
+          # Clear handler when called without arguments
+          @on[:human_in_the_loop] = nil
+        end
+
         self
       end
 
@@ -161,9 +295,10 @@ module RubyLLM
       end
 
       def on_logging(level: Logging::WARNING, &block)
+        require_feature!(:logging)
         @on_logging_level = level
         if alive?
-          @coordinator.set_logging(level: level)
+          @adapter.set_logging(level: level)
         end
 
         @on[:logging] = block
@@ -174,17 +309,69 @@ module RubyLLM
         @on.key?(:sampling) && !@on[:sampling].nil?
       end
 
-      def on_sampling(&block)
-        @on[:sampling] = block
+      def on_sampling(handler_class = nil, **options, &block)
+        require_feature!(:sampling)
+
+        if handler_class
+          # Validate handler class
+          validate_handler_class!(handler_class, :execute)
+
+          # Handler class provided
+          @on[:sampling] = if options.any?
+                             lambda do |sample|
+                               handler_class.new(sample: sample, coordinator: @adapter.native_client, **options).call
+                             end
+                           else
+                             handler_class
+                           end
+        elsif block_given?
+          # Block provided (backward compatible)
+          @on[:sampling] = block
+        else
+          # Clear handler when called without arguments
+          @on[:sampling] = nil
+        end
+
         self
       end
 
       def elicitation_enabled?
+        @adapter_type == :ruby_llm && MCP.config.elicitation.enabled?
+      end
+
+      def elicitation_callback_enabled?
         @on.key?(:elicitation) && !@on[:elicitation].nil?
       end
 
-      def on_elicitation(&block)
-        @on[:elicitation] = block
+      def on_elicitation(handler_class = nil, **options, &block)
+        require_feature!(:elicitation)
+
+        if handler_class
+          # Validate handler class
+          validate_handler_class!(handler_class, :execute)
+
+          # Handler class provided
+          @on[:elicitation] = if options.any?
+                                lambda do |elicitation|
+                                  handler_class.new(
+                                    elicitation: elicitation,
+                                    coordinator: @adapter.native_client,
+                                    **options
+                                  ).call
+                                end
+                              else
+                                handler_class
+                              end
+        elsif block_given?
+          # Block provided (backward compatible)
+          @on[:elicitation] = block
+        else
+          # Clear handler when called without arguments
+          @on[:elicitation] = nil
+        end
+
+        sync_elicitation_handler_state
+
         self
       end
 
@@ -212,10 +399,55 @@ module RubyLLM
 
       private
 
-      def setup_coordinator
-        Coordinator.new(self,
-                        transport_type: @transport_type,
-                        config: @config)
+      # Get OAuth provider from adapter's transport if available
+      # @return [OAuthProvider, BrowserOAuthProvider, nil] OAuth provider or nil
+      def transport_oauth_provider
+        return nil unless @adapter
+
+        # For RubyLLMAdapter
+        if @adapter.respond_to?(:native_client)
+          transport = @adapter.native_client.transport
+          transport_protocol = transport.transport_protocol
+          return transport_protocol.oauth_provider if transport_protocol.respond_to?(:oauth_provider)
+        end
+
+        # For MCPSdkAdapter with wrapped transports
+        if @adapter.respond_to?(:mcp_client) && @adapter.instance_variable_get(:@mcp_client)
+          mcp_client = @adapter.instance_variable_get(:@mcp_client)
+          if mcp_client&.transport.respond_to?(:native_transport)
+            return mcp_client.transport.native_transport.oauth_provider
+          end
+        end
+
+        nil
+      end
+
+      def build_adapter
+        case @adapter_type
+        when :ruby_llm
+          RubyLLM::MCP::Adapters::RubyLLMAdapter.new(self,
+                                                     transport_type: @transport_type,
+                                                     config: @config)
+        when :mcp_sdk
+          RubyLLM::MCP::Adapters::MCPSdkAdapter.new(self,
+                                                    transport_type: @transport_type,
+                                                    config: @config)
+        else
+          raise ArgumentError, "Unknown adapter type: #{@adapter_type}"
+        end
+      end
+
+      def require_feature!(feature)
+        unless @adapter.supports?(feature)
+          raise Errors::UnsupportedFeature.new(
+            message: <<~MSG.strip
+              Feature '#{feature}' is not supported by the #{@adapter_type} adapter.
+
+              This feature requires the :ruby_llm adapter.
+              Change your configuration to use adapter: :ruby_llm
+            MSG
+          )
+        end
       end
 
       def fetch(cache_key, refresh)
@@ -229,9 +461,9 @@ module RubyLLM
       def build_map(raw_data, klass, with_prefix: false)
         raw_data.each_with_object({}) do |item, acc|
           instance = if with_prefix
-                       klass.new(@coordinator, item, with_prefix: @with_prefix)
+                       klass.new(@adapter, item, with_prefix: @with_prefix)
                      else
-                       klass.new(@coordinator, item)
+                       klass.new(@adapter, item)
                      end
           acc[instance.name] = instance
         end
@@ -246,7 +478,7 @@ module RubyLLM
       end
 
       def setup_roots
-        @roots = Roots.new(paths: MCP.config.roots, coordinator: @coordinator)
+        @roots = Roots.new(paths: MCP.config.roots, adapter: @adapter)
       end
 
       def setup_sampling
@@ -254,11 +486,56 @@ module RubyLLM
       end
 
       def setup_event_handlers
-        @on[:progress] = MCP.config.on_progress
-        @on[:human_in_the_loop] = MCP.config.on_human_in_the_loop
-        @on[:logging] = MCP.config.on_logging
-        @on_logging_level = MCP.config.on_logging_level
-        @on[:elicitation] = MCP.config.on_elicitation
+        # Only setup handlers that are supported
+        if @adapter.supports?(:progress_tracking)
+          @on[:progress] = MCP.config.on_progress
+          if @on[:progress] && alive?
+            @adapter.set_progress_tracking(enabled: true)
+          end
+        end
+
+        if @adapter.supports?(:human_in_the_loop)
+          @on[:human_in_the_loop] = MCP.config.on_human_in_the_loop
+        end
+
+        if @adapter.supports?(:logging)
+          @on[:logging] = MCP.config.on_logging
+          @on_logging_level = MCP.config.on_logging_level
+        end
+
+        if @adapter.supports?(:elicitation)
+          @on[:elicitation] = MCP.config.on_elicitation
+        end
+      end
+
+      def sync_elicitation_handler_state
+        return unless @adapter.supports?(:elicitation)
+
+        @adapter.set_elicitation_enabled(enabled: elicitation_enabled?)
+      end
+
+      # Get or create OAuth storage shared with transport
+      def oauth_storage
+        # Try to get storage from transport's OAuth provider
+        transport_oauth = transport_oauth_provider
+        return transport_oauth.storage if transport_oauth
+
+        # Create new storage shared with client
+        @oauth_storage ||= Auth::MemoryStorage.new
+      end
+
+      # Validate that a handler class has required methods
+      # @param handler_class [Class] the handler class to validate
+      # @param required_method [Symbol] the method that must be defined
+      # @raise [ArgumentError] if validation fails
+      def validate_handler_class!(handler_class, required_method)
+        unless Handlers.handler_class?(handler_class)
+          raise ArgumentError, "Handler must be a class, got #{handler_class.class}"
+        end
+
+        unless handler_class.method_defined?(required_method)
+          raise ArgumentError, "Handler class #{handler_class} must define ##{required_method} method"
+        end
       end
     end
   end