robot_lab 0.3.1 → 0.5.0

data/lib/robot_lab/mcp.rb

@@ -0,0 +1,259 @@
+# frozen_string_literal: true
+
+module RobotLab
+  # robot_lab's thin layer over ruby_llm's built-in MCP client.
+  #
+  # ruby_llm 2.1 ships +RubyLLM::MCP+: a class that describes one Model
+  # Context Protocol server (stdio or Streamable HTTP), connects lazily, and
+  # exposes the server's tools, resources, prompts, OAuth flow, elicitation
+  # and long-running tasks. robot_lab adds only what a robot needs on top:
+  #
+  # - {MCP.build} turns a robot_lab config Hash (the shape used in templates,
+  #   +RunConfig#mcp+ and +RobotLab.build(mcp: [...])+) into a +RubyLLM::MCP+
+  #   instance, and passes instances and subclasses through unchanged.
+  # - {MCP::Server} is the +RubyLLM::MCP+ subclass those hashes become; it
+  #   carries a +description+ for {MCP::ServerDiscovery} and a +connected?+ flag.
+  # - {MCP::Registry} holds a robot's servers by name so repeated runs reuse
+  #   one connection (and one stdio process) per server.
+  #
+  # @example A config hash
+  #   RobotLab::MCP.build(
+  #     name: "filesystem",
+  #     transport: { type: "stdio", command: "npx", args: ["-y", "@modelcontextprotocol/server-filesystem", "."] },
+  #     timeout: 30, only: %w[read_file list_directory]
+  #   )
+  #
+  # @example A RubyLLM::MCP subclass
+  #   class GitHub < RubyLLM::MCP
+  #     command "github-mcp-server", "stdio"
+  #     env GITHUB_PERSONAL_ACCESS_TOKEN: ENV["GITHUB_PERSONAL_ACCESS_TOKEN"]
+  #     requires_approval if: :destructive?
+  #   end
+  #   RobotLab.build(name: "dev", mcp: [GitHub])
+  module MCP
+    # Transport types that map onto ruby_llm's Streamable HTTP client.
+    HTTP_TYPES = %w[streamable-http streamable_http http https].freeze
+
+    # Transport types the old robot_lab client offered that ruby_llm does not.
+    UNSUPPORTED_TYPES = %w[sse ws websocket].freeze
+
+    # Keys consumers attach for their own routing/grouping; never errors.
+    IGNORED_KEYS = %i[topics independent group session_id auth_provider].freeze
+
+    # Keys that may appear nested under +transport:+ or flat at the top level.
+    TRANSPORT_KEYS = %i[type command args env url headers bearer_token oauth directory cwd].freeze
+
+    # Keys the factory understands at the top level of a config hash.
+    SETTING_KEYS = %i[name description transport timeout prefix only except
+                      requires_approval defer log_level input_requests extensions].freeze
+
+    # Errors a server can raise while connecting or listing tools. A robot
+    # rescues these per server so one dead server does not abort a run.
+    CONNECT_ERRORS = [RubyLLM::Error, RubyLLM::ConfigurationError, SystemCallError, IOError, ArgumentError].freeze
+
+    module_function
+
+    # Build a server from a config Hash, or pass through an instance or class.
+    #
+    # @param spec [Hash, RubyLLM::MCP, Class<RubyLLM::MCP>]
+    # @return [RubyLLM::MCP]
+    # @raise [MCPError] for an unusable spec
+    def build(spec)
+      case spec
+      when RubyLLM::MCP then spec
+      when Class
+        raise MCPError, "#{spec} is not a RubyLLM::MCP subclass" unless spec <= RubyLLM::MCP
+
+        spec.new
+      when Hash then define(spec).new
+      else
+        raise MCPError, "Cannot build an MCP server from #{spec.class}; " \
+                        'pass a Hash, a RubyLLM::MCP, or a RubyLLM::MCP subclass'
+      end
+    end
+
+    # Turn a config Hash into an anonymous {Server} subclass with ruby_llm's
+    # class-level DSL applied. Public so hosts can build classes once and
+    # instantiate them per request (ruby_llm's recommended pattern).
+    #
+    # @param config [Hash]
+    # @return [Class<Server>]
+    # @raise [MCPError] when the hash has no name, no usable transport, or an unsupported one
+    def define(config)
+      config = normalize(config)
+      klass  = Class.new(Server)
+      apply_settings(klass, config)
+      klass
+    end
+
+    # The name a spec will be registered under.
+    #
+    # @param spec [Hash, RubyLLM::MCP, Class<RubyLLM::MCP>]
+    # @return [String]
+    # @raise [MCPError] when the spec has no name
+    def name_of(spec)
+      name = case spec
+             when Hash then spec[:name] || spec['name']
+             when Class then spec.default_name
+             when RubyLLM::MCP then spec.name
+             end
+      raise MCPError, "MCP server spec has no name: #{spec.inspect}" if name.nil? || name.to_s.strip.empty?
+
+      name.to_s
+    end
+
+    # Normalize a timeout to seconds. Values of 1000 or more are taken to be
+    # milliseconds (the convention aia and older configs used); the result is
+    # never below one second.
+    #
+    # @param value [Numeric, String]
+    # @return [Float]
+    def seconds(value)
+      secs = value.to_f
+      secs /= 1000.0 if secs >= 1000
+      [secs, 1.0].max
+    end
+
+    # --- internals (module_function keeps them callable as MCP.xxx for the Class.new body) ---
+
+    # Symbolize keys, fold flat transport keys into +:transport+, require a name.
+    # :reek:TooManyStatements -- one linear normalization pass over the hash shape.
+    def normalize(config)
+      config = config.to_h.transform_keys(&:to_sym)
+      name_of(config)
+
+      raw = config[:transport]
+      transport = case raw
+                  when Hash           then raw.transform_keys(&:to_sym)
+                  when String, Symbol then { type: raw.to_s }
+                  when nil            then {}
+                  else                     raw # a transport object (duck type)
+                  end
+      if transport.is_a?(Hash)
+        TRANSPORT_KEYS.each { |key| transport[key] = config[key] if config.key?(key) && !transport.key?(key) }
+      end
+      config.merge(transport: transport)
+    end
+
+    def apply_settings(klass, config)
+      klass.default_name = config[:name].to_s
+      klass.description(config[:description]) if config[:description]
+      apply_transport(klass, config)
+      apply_tool_shaping(klass, config)
+      log_ignored_keys(config)
+    end
+
+    # :reek:TooManyStatements -- each optional setting is one guarded DSL call.
+    def apply_tool_shaping(klass, config)
+      klass.timeout(seconds(config[:timeout]))      if config[:timeout]
+      klass.prefix(config[:prefix])                 if config[:prefix]
+      klass.only(*Array(config[:only]))             if config[:only]
+      klass.except(*Array(config[:except]))         if config[:except]
+      klass.log_level(config[:log_level])           if config[:log_level]
+      klass.input_requests(*Array(config[:input_requests])) if config.key?(:input_requests)
+      apply_extensions(klass, config[:extensions])
+      apply_approvals(klass, config[:requires_approval])
+      apply_deferral(klass, config[:defer])
+    end
+
+    def apply_transport(klass, config)
+      transport = config[:transport]
+      return klass.transport(transport) if transport_object?(transport)
+
+      case (type = transport_type(config, transport))
+      when 'stdio'            then apply_stdio(klass, config, transport)
+      when *HTTP_TYPES        then apply_http(klass, config, transport)
+      when *UNSUPPORTED_TYPES then raise MCPError, unsupported_transport_message(config[:name], type)
+      else raise MCPError, "Unknown MCP transport type '#{type}' for server '#{config[:name]}'"
+      end
+    end
+
+    def transport_object?(transport)
+      !transport.is_a?(Hash) && transport.respond_to?(:request) && transport.respond_to?(:close)
+    end
+
+    def transport_type(config, transport)
+      type = transport[:type]&.to_s&.downcase
+      return type unless type.nil? || type.empty?
+      return 'stdio' if transport[:command]
+      return 'streamable-http' if transport[:url]
+
+      raise MCPError, "MCP server '#{config[:name]}' needs a command (stdio) or a url (Streamable HTTP)"
+    end
+
+    def apply_stdio(klass, config, transport)
+      command = transport[:command]
+      raise MCPError, "MCP server '#{config[:name]}' uses stdio but has no command" if command.nil?
+
+      klass.command(*(Array(command) + Array(transport[:args])).map(&:to_s))
+      env = transport[:env]
+      klass.env(**env.to_h.transform_keys(&:to_sym)) if env && !env.empty?
+      directory = transport[:directory] || transport[:cwd]
+      klass.directory(directory.to_s) if directory
+    end
+
+    # :reek:TooManyStatements -- url, headers, bearer token and oauth are four independent guarded calls.
+    def apply_http(klass, config, transport)
+      url = transport[:url]
+      raise MCPError, "MCP server '#{config[:name]}' uses Streamable HTTP but has no url" if url.nil?
+
+      klass.url(url.to_s)
+      transport[:headers].to_h.each { |header, value| klass.header(header.to_s, value) }
+      klass.bearer_token(transport[:bearer_token]) if transport[:bearer_token]
+      apply_oauth(klass, transport[:oauth])
+    end
+
+    def apply_oauth(klass, oauth)
+      case oauth
+      when nil, false then nil
+      when true       then klass.oauth
+      else                 klass.oauth(**oauth.to_h.transform_keys(&:to_sym))
+      end
+    end
+
+    def apply_extensions(klass, extensions)
+      case extensions
+      when nil then nil
+      when Hash
+        extensions.each { |name, settings| klass.extension(extension_name(name), **settings.to_h.transform_keys(&:to_sym)) }
+      else Array(extensions).each { |name| klass.extension(extension_name(name)) }
+      end
+    end
+
+    # ruby_llm takes a Symbol for the extensions it knows (:tasks, :apps) and
+    # a vendor-prefixed String for the rest; YAML configs give Strings for both.
+    def extension_name(name)
+      name.to_s.include?('/') ? name.to_s : name.to_sym
+    end
+
+    def apply_approvals(klass, value)
+      case value
+      when nil, false then nil
+      when true       then klass.requires_approval
+      when Hash       then klass.requires_approval(**value.transform_keys(&:to_sym))
+      else                 klass.requires_approval(*Array(value).map(&:to_s))
+      end
+    end
+
+    def apply_deferral(klass, value)
+      case value
+      when nil, false then nil
+      when true       then klass.defer
+      else                 klass.defer(*Array(value).map(&:to_s))
+      end
+    end
+
+    def log_ignored_keys(config)
+      unknown = config.keys - SETTING_KEYS - TRANSPORT_KEYS - IGNORED_KEYS
+      return if unknown.empty?
+
+      RobotLab.config.logger.debug { "MCP server '#{config[:name]}': ignoring unknown keys #{unknown.join(', ')}" }
+    end
+
+    def unsupported_transport_message(name, type)
+      "MCP server '#{name}' uses transport '#{type}', which ruby_llm 2.1 does not provide " \
+        '(only stdio and Streamable HTTP). Point the server at its Streamable HTTP endpoint, ' \
+        'or pass a transport: object that responds to request/notify/cancel/close.'
+    end
+  end
+end

data/lib/robot_lab/robot/mcp_management.rb

@@ -2,12 +2,108 @@
 
 module RobotLab
   class Robot < RubyLLM::Agent
-    # MCP client lifecycle and hierarchical tool/MCP resolution.
+    # MCP servers: the robot's registry, per-run server selection, tool
+    # collection, and connection failures.
     #
-    # Owns:    @mcp_clients, @mcp_tools, @mcp_initialized, @failed_mcp_configs
-    # Reads:   @mcp_config, @tools_config, @name, @chat, @local_tools
-    # Contract: ivars initialized by initialize_runtime_state; lazy init on first run
+    # Servers are +RubyLLM::MCP+ instances (see RobotLab::MCP). They connect
+    # lazily on the first +tools+ request and stay connected across runs
+    # until #disconnect. Each run attaches the selected servers' tools to the
+    # chat with +with_tools+, so the per-run +tools:+ allowlist, +max_tools+
+    # cap and +:none+ apply to MCP tools exactly as to local ones.
+    #
+    # Owns:    @mcp_registry, @mcp_failures, @mcp_tool_owners, @run_mcp_servers
+    # Reads:   @mcp_config, @tools_config, @name, @mcp_discovery, @bus
+    # Contract: ivars initialized by initialize_runtime_state
     module MCPManagement
+      # The robot's MCP servers, keyed by name. Each entry is a
+      # +RubyLLM::MCP+ with the full protocol surface: +tools+, +call+,
+      # +resources+, +resource(uri)+, +prompts+, +prompt(name)+,
+      # +instructions+, +close+, and the OAuth flow.
+      #
+      # @return [MCP::Registry]
+      #
+      # @example
+      #   robot.mcp[:github].resources.map(&:uri)
+      #   robot.run(robot.mcp[:github].prompt("code_review", code: diff))
+      def mcp
+        @mcp_registry
+      end
+
+      # Register MCP servers on this robot. Accepts config Hashes,
+      # +RubyLLM::MCP+ instances and +RubyLLM::MCP+ subclasses, in any mix.
+      # Registering a spec already known returns the same server, so hosts
+      # may call this on every request. +with_mcp(nil)+ closes and forgets
+      # every server.
+      #
+      # Registered servers are available to runs that select them through
+      # the +mcp:+ cascade (+run(..., mcp: :inherit)+ picks up the build-time
+      # list; a run-time Array picks exactly those).
+      #
+      # @param servers [Array<Hash, RubyLLM::MCP, Class<RubyLLM::MCP>>, nil]
+      # @return [self]
+      def with_mcp(*servers)
+        if servers == [nil]
+          @mcp_registry.clear
+        else
+          servers.flatten.compact.each { |spec| @mcp_registry.add(spec) }
+        end
+        self
+      end
+
+      # Eagerly register the build-time MCP servers and connect to each.
+      # Connections are otherwise made on the first run that selects a
+      # server. Failures are recorded in #mcp_failures, not raised.
+      #
+      # @return [self]
+      def connect_mcp!
+        collect_mcp_tools(servers_for(resolve_mcp_hierarchy(@mcp_config)))
+        self
+      end
+
+      # The MCP tools offered by the given servers (default: every
+      # registered server). Connects servers that are not connected yet and
+      # skips, with a warning, any that cannot be reached.
+      #
+      # @param servers [Enumerable<RubyLLM::MCP>]
+      # @return [Array<RubyLLM::MCP::Tool>]
+      def mcp_tools(servers = @mcp_registry)
+        collect_mcp_tools(servers.to_a)
+      end
+
+      # Names of servers whose last connection attempt failed. A server is
+      # retried on the next run that selects it and leaves this list once it
+      # answers.
+      #
+      # @return [Array<String>]
+      def failed_mcp_server_names
+        @mcp_failures.keys
+      end
+
+      # The last connection error per failed server name.
+      #
+      # @return [Hash{String => Exception}]
+      def mcp_failures
+        @mcp_failures.dup
+      end
+
+      # The server that owns +tool+, or nil for a local tool.
+      #
+      # @param tool [RubyLLM::Tool]
+      # @return [RubyLLM::MCP, nil]
+      def mcp_server_for(tool)
+        @mcp_tool_owners[tool]
+      end
+
+      # Close every MCP server (stopping stdio processes) and the bus channel.
+      # Servers stay registered and reconnect on their next use.
+      #
+      # @return [self]
+      def disconnect
+        @mcp_registry.each(&:close)
+        teardown_bus_channel if @bus
+        self
+      end
+
       private
 
       # Resolve MCP hierarchy: runtime -> robot build -> network -> config
@@ -32,121 +128,79 @@ module RobotLab
         network&.network
       end
 
-      # Ensure MCP clients are initialized for the given server configs.
-      # On subsequent calls, retries any servers that previously failed to connect.
-      # :reek:TooManyStatements -- first-run init vs retry paths share the needed-server list; linear either way.
-      def ensure_mcp_clients(mcp_servers)
-        return if mcp_servers.empty?
-
-        needed_servers = mcp_servers.map { |s| s.is_a?(Hash) ? s[:name] : s.to_s }.compact
-
-        if @mcp_initialized
-          # Already initialized — retry any servers that are needed but not connected
-          retry_failed_servers(mcp_servers, needed_servers)
-          return
-        end
+      # The registered server instances a run's resolved +mcp+ value names.
+      # Specs not registered yet are registered here. With +mcp_discovery+,
+      # the candidates are narrowed to those relevant to +message+ first.
+      #
+      # @param resolved [Array<Hash, RubyLLM::MCP, Class>] output of resolve_mcp_hierarchy
+      # @param message [String, RubyLLM::MCP::Prompt, nil] this run's request
+      # @return [Array<RubyLLM::MCP>]
+      def servers_for(resolved, message: nil)
+        specs = Array(resolved)
+        specs = MCP::ServerDiscovery.select(discovery_query(message), from: specs) if @mcp_discovery && message
+        specs.map { |spec| @mcp_registry.add(spec) }
+      end
 
-        @mcp_clients = {}
-        @mcp_tools = []
-        @failed_mcp_configs = {}
+      # The text ServerDiscovery scores servers against.
+      def discovery_query(message)
+        return message.to_s unless message.is_a?(RubyLLM::MCP::Prompt)
 
-        mcp_servers.each do |server_config|
-          init_mcp_client(server_config)
-        end
+        message.messages.map { |prompt_message| prompt_message.content.to_s }.join(' ')
+      end
 
-        @mcp_initialized = true
+      # The tools of +servers+, connecting each as needed. A server that
+      # cannot be reached is recorded in @mcp_failures, warned about, and
+      # contributes no tools; the run goes on with the rest.
+      #
+      # @param servers [Array<RubyLLM::MCP>]
+      # @return [Array<RubyLLM::MCP::Tool>]
+      def collect_mcp_tools(servers)
+        servers.flat_map { |server| tools_from(server) }
       end
 
-      # :reek:TooManyStatements -- connect success/failure bookkeeping plus the rescue path for one server.
-      def init_mcp_client(server_config)
-        client = MCP::Client.new(server_config)
-        client.connect
+      # One server's tools, or [] (recorded and warned) when it cannot be reached.
+      def tools_from(server)
+        tools = server.tools
+        @mcp_failures.delete(server.name)
+        tools.each { |tool| @mcp_tool_owners[tool] = server }
+        tools
+      rescue *MCP::CONNECT_ERRORS => e
+        record_mcp_failure(server, e)
+        []
+      end
 
-        if client.connected?
-          server_name = client.server.name
-          @mcp_clients[server_name] = client
-          discover_mcp_tools(client, server_name)
-        else
-          server_name = extract_server_name(server_config)
-          @failed_mcp_configs[server_name] = server_config
-          RobotLab.config.logger.warn(
-            "Robot '#{@name}' failed to connect to MCP server: #{server_name}"
-          )
-        end
-      rescue StandardError => e
-        server_name = extract_server_name(server_config)
-        @failed_mcp_configs[server_name] = server_config
+      def record_mcp_failure(server, error)
+        @mcp_failures[server.name] = error
         RobotLab.config.logger.warn(
-          "Robot '#{@name}' error connecting to MCP server '#{server_name}': #{e.message}"
+          "Robot '#{@name}' could not reach MCP server '#{server.name}': #{error.class}: #{error.message}"
         )
       end
 
-      # Retry connecting to servers that previously failed
-      # :reek:TooManyStatements -- per-server retry with success bookkeeping and a best-effort rescue.
-      def retry_failed_servers(_mcp_servers, needed_servers)
-        return if @failed_mcp_configs.nil? || @failed_mcp_configs.empty?
-
-        # Only retry servers that are still needed and still failed
-        to_retry = @failed_mcp_configs.slice(*needed_servers)
-        return if to_retry.empty?
-
-        logger = RobotLab.config.logger
-        to_retry.each do |name, server_config|
-          logger.info(
-            "Robot '#{@name}' retrying MCP server: #{name}"
-          )
-
-          client = MCP::Client.new(server_config)
-          client.connect
-
-          if client.connected?
-            @mcp_clients[name] = client
-            @failed_mcp_configs.delete(name)
-            discover_mcp_tools(client, name)
-            logger.info(
-              "Robot '#{@name}' successfully connected to MCP server '#{name}' on retry"
-            )
-          end
-        rescue StandardError => e
-          logger.warn(
-            "Robot '#{@name}' retry failed for MCP server '#{name}': #{e.message}"
-          )
-        end
-      end
-
-      # :reek:FeatureEnvy -- adapting each MCP tool definition into a local Tool is exactly this method's purpose.
-      # :reek:NestedIterators -- the inner block is the tool's execution closure, not an iteration.
-      def discover_mcp_tools(client, server_name)
-        tools = client.list_tools
-
-        tools.each do |tool_def|
-          tool_name = tool_def[:name]
-          mcp_client = client
+      # Whether +tool+ is an MCP tool its server declared deferred (offered to
+      # the provider's tool search instead of the model's context).
+      def deferred_mcp_tool?(tool)
+        server = mcp_server_for(tool)
+        return false unless server
 
-          tool = Tool.create(
-            name: tool_name,
-            description: tool_def[:description],
-            parameters: tool_def[:inputSchema],
-            mcp: server_name
-          ) { |args| mcp_client.call_tool(tool_name, args) }
-
-          @mcp_tools << tool
-        end
+        server.defers?(tool)
+      end
 
-        RobotLab.config.logger.info(
-          "Robot '#{@name}' discovered #{tools.size} tools from MCP server '#{server_name}'"
-        )
+      # Whether the chat currently holds MCP tools from an earlier run.
+      def mcp_tools_attached?
+        @chat.tools.each_value.any?(RubyLLM::MCP::Tool)
       end
 
-      def extract_server_name(server_config)
-        case server_config
-        when Hash
-          server_config[:name] || server_config['name'] || server_config.to_s
-        when MCP::Server
-          server_config.name
-        else
-          server_config.to_s
-        end
+      # Two tools with one name cannot both be offered; ruby_llm raises when
+      # it builds the tool list, so catch it here with the fix spelled out.
+      #
+      # @param tools [Array<RubyLLM::Tool>]
+      # @raise [MCPError]
+      def check_unique_tool_names!(tools)
+        duplicates = tools.group_by { |tool| tool.name.to_s }.select { |_, same| same.size > 1 }.keys
+        return if duplicates.empty?
+
+        raise MCPError, "Robot '#{@name}' has more than one tool named #{duplicates.join(', ')}. " \
+                        'Give one of the MCP servers a prefix:, or rename the local tool.'
       end
     end
   end