aia 1.1.1 → 2.0.0.0.pre.beta2

data/lib/aia/mcp_config_normalizer.rb

@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+# lib/aia/mcp_config_normalizer.rb
+#
+# Normalizes a single MCP server config from AIA flat format to the nested
+# transport format expected by robot_lab. Server selection/filtering is
+# the responsibility of MCPDiscovery — this class only transforms shape.
+
+module AIA
+  class MCPConfigNormalizer
+    class << self
+      # Normalize a single MCP server config to robot_lab's nested transport format.
+      #
+      # @param server [Hash]
+      # @return [Hash]
+      def normalize(server)
+        server = server.is_a?(Hash) ? server.transform_keys(&:to_sym) : server.to_h.transform_keys(&:to_sym)
+
+        # Already in robot_lab format — pass through
+        return server if server[:transport]
+
+        # Legacy flat format: wrap command/args/env into transport
+        name      = server[:name]
+        transport = { type: server[:type] || 'stdio' }
+        transport[:command] = server[:command] if server[:command]
+        transport[:args]    = Array(server[:args]) if server[:args]
+        transport[:env]     = server[:env] if server[:env]
+
+        result = { name: name, transport: transport }
+        result[:timeout] = server[:timeout] if server[:timeout]
+        result
+      end
+    end
+  end
+end

data/lib/aia/mcp_connection_manager.rb

@@ -0,0 +1,315 @@
+# frozen_string_literal: true
+
+# lib/aia/mcp_connection_manager.rb
+#
+# Single owner for MCP connection state. Handles connecting to configured
+# MCP servers concurrently, tracking successes/failures, and injecting
+# results into robots.
+
+require 'timeout'
+require 'tty-spinner'
+
+module AIA
+  # :reek:TooManyMethods -- one small helper per connection-lifecycle concern (connect, wrap tools, record outcomes)
+  class MCPConnectionManager
+    # Maximum seconds to wait for a single MCP server to connect.
+    DEFAULT_TIMEOUT = 30
+
+    attr_reader :connected_clients, :connected_tools, :failed_servers
+
+    def initialize
+      @connected_clients    = {}
+      @connected_tools      = []
+      @failed_servers       = []
+      @server_tool_counts   = {}
+      @connected            = false
+      @mutex                = Mutex.new
+    end
+
+    # Connect all configured MCP servers concurrently.
+    # Each server gets its own spinner line and thread.
+    #
+    # @param servers [Array<Hash>] MCP server configurations
+    # @return [self]
+    # :reek:TooManyStatements -- connection ceremony: spinner setup, thread launch, global timeout guard, summary logging
+    def connect_all(servers)
+      return self unless servers.is_a?(Array) && servers.any?
+
+      reset_connection_state
+      logger = AIA::LoggerManager.mcp_logger
+      server_names = servers.map { |s| s.is_a?(Hash) ? s[:name] : s.to_s }.compact
+      logger.info("MCP initialization: connecting #{servers.size} server(s): #{server_names.join(', ')}")
+
+      multi = TTY::Spinner::Multi.new("[:spinner] Connecting MCP servers", format: :dots, output: $stderr)
+      threads = launch_server_threads(servers, multi, logger)
+
+      # The global cap must never be smaller than the slowest single server's
+      # timeout, or concurrent connections get killed before they can finish.
+      global_limit = [global_connection_timeout, *servers.map { |s| server_timeout(s) }].max
+      Timeout.timeout(global_limit) { threads.each(&:join) }
+
+      @connected = true
+      logger.info("MCP initialization complete: #{@connected_clients.size} connected, #{@failed_servers.size} failed")
+      self
+    rescue Timeout::Error
+      handle_global_timeout(servers, global_limit)
+    end
+
+    # Inject connected MCP clients and tools into robot(s).
+    #
+    # @param robot [RobotLab::Robot, RobotLab::Network] the robot or network
+    # @return [self]
+    def inject_into(robot)
+      clients, tools = @mutex.synchronize { [@connected_clients.dup, @connected_tools.dup] }
+
+      targets = robot.crew
+
+      targets.each do |target|
+        target.inject_mcp!(clients: clients, tools: tools)
+      end
+
+      self
+    end
+
+    # Update AIA config with connection results.
+    #
+    # @return [self]
+    def update_config
+      connected_keys, tool_counts, failed = @mutex.synchronize do
+        [@connected_clients.keys.dup, @server_tool_counts.dup, @failed_servers.dup]
+      end
+      cfg = AIA.config
+      cfg.connected_mcp_servers  = connected_keys
+      cfg.mcp_server_tool_counts = tool_counts
+      cfg.failed_mcp_servers     = failed
+      self
+    end
+
+    # Whether connect_all has been called.
+    def connected?
+      @connected
+    end
+
+    # Whether any tools were collected (from config servers or Ruby require).
+    def any_tools?
+      @mutex.synchronize { @connected_tools.any? }
+    end
+
+    # Absorb MCP clients registered in RubyLLM::MCP (e.g., via --require loading
+    # shared_tools/mcp/* files) that are not already tracked by this manager.
+    # Starts any unstarted clients, extracts their tools, and merges them into
+    # the connected state so inject_into will include them in the robot.
+    #
+    # @return [self]
+    # :reek:TooManyStatements -- per-client absorb loop: start if needed, extract tools, merge into connected state with error capture
+    def absorb_ruby_llm_mcp_clients
+      return self unless defined?(RubyLLM::MCP)
+
+      logger = AIA::LoggerManager.mcp_logger
+
+      RubyLLM::MCP.clients.each do |name, client|
+        next if @connected_clients.key?(name)
+
+        begin
+          client.start unless client.alive?
+
+          if client.alive?
+            tools = client.tools rescue []
+            @mutex.synchronize do
+              @connected_clients[name]  = client
+              @server_tool_counts[name] = tools.size
+              @connected_tools.concat(tools)
+            end
+            logger.info("MCP: absorbed RubyLLM::MCP client '#{name}' (#{tools.size} tools)")
+          else
+            logger.warn("MCP: RubyLLM::MCP client '#{name}' not alive, skipping")
+          end
+        rescue StandardError => e
+          logger.warn("MCP: error absorbing RubyLLM::MCP client '#{name}': #{e.message}")
+        end
+      end
+
+      @connected = true unless RubyLLM::MCP.clients.empty?
+      self
+    end
+
+    # Find an MCP client by server name.
+    #
+    # @param name [String] server name
+    # @return [RobotLab::MCP::Client, nil]
+    def client(name)
+      @connected_clients[name]
+    end
+
+    # Connected server names.
+    #
+    # @return [Array<String>]
+    def connected_server_names
+      @mutex.synchronize { @connected_clients.keys.dup }
+    end
+
+    # Failed server names.
+    #
+    # @return [Array<String>]
+    def failed_server_names
+      @mutex.synchronize { @failed_servers.map { |f| f[:name] } }
+    end
+
+    # Close all connected MCP clients and reset connection state.
+    # Errors from individual client closes are silently rescued —
+    # best-effort cleanup on process exit.
+    def close_all
+      clients = @mutex.synchronize { @connected_clients.dup }
+      clients.each_value do |client|
+        client.close
+      rescue StandardError
+        # best-effort — client may already be closed
+      end
+      @mutex.synchronize do
+        @connected_clients  = {}
+        @connected_tools    = []
+        @server_tool_counts = {}
+        @connected          = false
+      end
+    end
+
+    private
+
+    def reset_connection_state
+      @connected_clients = {}
+      @connected_tools   = []
+      @failed_servers    = []
+    end
+
+    def launch_server_threads(servers, multi, logger)
+      servers.map do |server_config|
+        name = server_config.is_a?(Hash) ? (server_config[:name] || server_config['name']) : server_config.to_s
+        spinner = multi.register("[:spinner] #{name}")
+        Thread.new(server_config, name, spinner) { |cfg, srv_name, sp| connect_one(cfg, srv_name, sp, logger) }
+      end
+    end
+
+    def handle_global_timeout(servers, global_limit)
+      logger = AIA::LoggerManager.mcp_logger
+      logger.warn("MCP: Global connection timeout (#{global_limit}s) reached")
+
+      connected_names, failed_names = @mutex.synchronize do
+        [@connected_clients.keys.to_set, @failed_servers.to_set { |f| f[:name] }]
+      end
+
+      stragglers = servers.filter_map do |server_config|
+        srv_name = server_config.is_a?(Hash) ? (server_config[:name] || server_config['name']) : server_config.to_s
+        next if connected_names.include?(srv_name) || failed_names.include?(srv_name)
+
+        { name: srv_name, error: "global connection timeout (#{global_limit}s)" }
+      end
+      @mutex.synchronize { @failed_servers.concat(stragglers) }
+    end
+
+    # Read global MCP connection timeout from config.
+    # Falls back to 120s if config is unavailable.
+    def global_connection_timeout
+      AIA.config&.concurrency&.mcp_timeout.to_i.then { |v| v.positive? ? v : 120 }
+    rescue StandardError
+      120
+    end
+
+    # Connect a single MCP server, updating the spinner on completion.
+    # :reek:TooManyStatements -- one connect attempt with per-outcome reporting (success, refused, timeout, error)
+    def connect_one(server_config, name, spinner, logger)
+      timeout = server_timeout(server_config)
+      spinner.auto_spin
+
+      logger.info("MCP: connecting to '#{name}'...")
+      Timeout.timeout(timeout) do
+        client = RobotLab::MCP::Client.new(server_config)
+        client.connect
+
+        if client.connected?
+          record_connected(name, client, build_mcp_tools(client, name), spinner, logger)
+        else
+          record_failed_server(name, "connection failed")
+          logger.warn("MCP: '#{name}' failed to connect")
+          spinner.error("(connection failed)")
+        end
+      end
+    rescue Timeout::Error
+      record_failed_server(name, "timed out after #{timeout}s")
+      logger.warn("MCP: '#{name}' timed out after #{timeout}s")
+      spinner.error("(timed out)")
+    rescue StandardError => e
+      msg = e.message
+      record_failed_server(name, msg)
+      logger.warn("MCP: '#{name}' error: #{msg}")
+      spinner.error("(#{msg})")
+    end
+
+    # Wrap each remote tool definition in a RobotLab::Tool proxying to the client.
+    def build_mcp_tools(client, server_name)
+      client.list_tools.map do |tool_def|
+        tool_name = tool_def[:name]
+        RobotLab::Tool.create(
+          name:        tool_name,
+          description: tool_def[:description],
+          parameters:  tool_def[:inputSchema],
+          mcp:         server_name
+        ) { |args| client.call_tool(tool_name, args) }
+      end
+    end
+
+    def record_connected(name, client, tools, spinner, logger)
+      tool_count = tools.size
+
+      @mutex.synchronize do
+        @connected_clients[name]  = client
+        @server_tool_counts[name] = tool_count
+        add_tools_deduped(tools, name, logger)
+      end
+
+      logger.info("MCP: '#{name}' connected (#{tool_count} tools)")
+      spinner.success("(#{tool_count} tools)")
+    end
+
+    def record_failed_server(name, error)
+      @mutex.synchronize { @failed_servers << { name: name, error: error } }
+    end
+
+    # Add tools to @connected_tools, skipping any whose name is already present.
+    # Logs a warning for each duplicate so the user knows which server wins.
+    # MUST be called inside @mutex.synchronize.
+    #
+    # @param new_tools [Array] tools to add
+    # @param server_name [String] name of the server providing these tools (for warning)
+    # @param logger [Logger] MCP logger
+    def add_tools_deduped(new_tools, server_name, logger)
+      new_tools.each do |tool|
+        tool_name = tool.respond_to?(:name) ? tool.name : tool.to_s
+        existing = @connected_tools.find do |t|
+          (t.respond_to?(:name) ? t.name : t.to_s) == tool_name
+        end
+
+        if existing
+          logger.warn("MCP: Tool '#{tool_name}' already registered from another server. " \
+                      "Skipping duplicate from '#{server_name}'.")
+        else
+          @connected_tools << tool
+        end
+      end
+    end
+
+    # Extract timeout in seconds from server config.
+    # Config values >= 1000 are treated as milliseconds.
+    def server_timeout(server_config)
+      raw = server_config.is_a?(Hash) ? server_config[:timeout] : nil
+      return DEFAULT_TIMEOUT if raw.nil?
+
+      seconds = raw.to_f
+      seconds /= 1000.0 if seconds >= 1000 # values >= 1000 are milliseconds
+      # Honor the configured per-server timeout. DEFAULT_TIMEOUT is only the
+      # fallback when no timeout is set (handled above) — it must NOT cap a
+      # larger configured value, or slow-starting servers (npx/uvx cold starts)
+      # get dropped and their tools never reach the robot or the tool filter.
+      seconds.positive? ? seconds : DEFAULT_TIMEOUT
+    end
+  end
+end

data/lib/aia/mcp_discovery.rb

@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+# lib/aia/mcp_discovery.rb
+#
+# Determines which MCP servers should be activated for a given prompt.
+# Two strategies: explicit (user directive) or all configured (default).
+
+module AIA
+  class MCPDiscovery
+    def initialize; end
+
+    # Discover relevant MCP servers for the current prompt.
+    # Priority: explicit --mcp-use list > all configured servers.
+    # --mcp-skip is always applied last.
+    #
+    # @param config the AIA configuration
+    # @return [Array<Hash>] relevant MCP server configs
+    def discover(config)
+      return [] if config.flags&.no_mcp
+
+      servers = select_servers(config)
+      apply_skip_filter(servers, config)
+    end
+
+    private
+
+    def select_servers(config)
+      return select_by_names(config.mcp_servers || [], Array(config.mcp_use)) if Array(config.mcp_use).any?
+
+      config.mcp_servers || []
+    end
+
+    def apply_skip_filter(servers, config)
+      skip_list = Array(config.mcp_skip)
+      return servers if skip_list.empty?
+
+      servers.reject { |s| skip_list.include?(AIA::Utility.server_name(s)) }
+    end
+
+    def select_by_names(all_servers, names)
+      all_servers.select { |s| names.include?(AIA::Utility.server_name(s)) }
+    end
+  end
+end

data/lib/aia/mcp_grouper.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+# lib/aia/mcp_grouper.rb
+#
+# Groups MCP servers into independent sets that can run concurrently.
+# Servers with the same `group` value must run sequentially.
+# Servers without a group are each independent.
+
+module AIA
+  class MCPGrouper
+    # Group MCP servers into independent sets.
+    #
+    # @param servers [Array<Hash>] MCP server configs
+    # @return [Array<Array<Hash>>] groups of servers that can run concurrently
+    def group(servers)
+      return [] if servers.nil? || servers.empty?
+
+      independent = []
+      groups = Hash.new { |h, k| h[k] = [] }
+
+      servers.each do |server|
+        group_name = server[:group] || server["group"]
+        if group_name
+          groups[group_name] << server
+        else
+          independent << [server]  # Each independent server is its own group
+        end
+      end
+
+      independent + groups.values
+    end
+  end
+end

data/lib/aia/mcp_server_config.rb

@@ -0,0 +1,30 @@
+# lib/aia/mcp_server_config.rb
+
+module AIA
+  # Normalized view of one MCP server hash. Server configs arrive with
+  # symbol or string keys, and with connection fields either nested under
+  # :transport or at the top level; this object owns all of that fallback
+  # logic so callers never touch the raw hash shape.
+  McpServerConfig = Data.define(:name, :command, :args, :env, :timeout_raw) do
+    def self.from_hash(server)
+      transport = fetch_key(server, :transport) || {}
+      new(
+        name:        AIA::Utility.server_name(server),
+        command:     fetch_key(transport, :command) || fetch_key(server, :command),
+        args:        Array(fetch_key(transport, :args) || fetch_key(server, :args)),
+        env:         fetch_key(transport, :env) || fetch_key(server, :env) || {},
+        timeout_raw: fetch_key(server, :timeout)
+      )
+    end
+
+    # Read a hash entry by symbol or string key.
+    def self.fetch_key(hash, key) = hash[key] || hash[key.to_s]
+
+    # Timeout in milliseconds; raw values below 1000 are treated as seconds.
+    def timeout_ms(default: 8_000, cap: 30_000)
+      raw = (timeout_raw || default).to_i
+      ms  = raw < 1000 ? raw * 1000 : raw
+      [ms, cap].min
+    end
+  end
+end

data/lib/aia/mcp_utility.rb

@@ -0,0 +1,60 @@
+# frozen_string_literal: true
+
+# lib/aia/mcp_utility.rb
+#
+# Pure MCP state-query methods mixed into AIA::Utility via `class << self include`.
+# No display or formatting concerns — those live in Utility directly.
+
+module AIA
+  module MCPUtility
+    def mcp_servers?
+      effective_mcp_server_names.any?
+    end
+
+    def mcp_server_names
+      connected = AIA.config&.connected_mcp_servers
+      return connected unless connected.nil?
+      return RubyLLM::MCP.clients.keys if defined?(RubyLLM::MCP) && RubyLLM::MCP.clients.any?
+      effective_mcp_server_names
+    end
+
+    def connected_mcp_servers?
+      mcp_server_names.any?
+    end
+
+    def failed_mcp_servers
+      AIA.config&.failed_mcp_servers || []
+    end
+
+    # :reek:TooManyStatements -- server-name resolution applying no_mcp, use-list, and skip-list filters in precedence order
+    def effective_mcp_server_names
+      cfg = AIA.config
+      return [] if cfg&.flags&.no_mcp
+
+      servers = cfg&.mcp_servers || []
+      return [] if servers.empty?
+
+      names     = servers.map { |s| server_name(s) }.compact
+      use_list  = Array(cfg.mcp_use)
+      skip_list = Array(cfg.mcp_skip)
+
+      if use_list.any?
+        names.select { |n| use_list.include?(n) }
+      elsif skip_list.any?
+        names.reject { |n| skip_list.include?(n) }
+      else
+        names
+      end
+    end
+
+    def server_name(s)
+      if s.is_a?(Hash)
+        s[:name] || s['name']
+      elsif s.respond_to?(:name)
+        s.name
+      else
+        s.to_s
+      end
+    end
+  end
+end