aia 1.1.0 → 2.0.0.0.pre.alpha

data/lib/aia/mcp_connection_manager.rb

@@ -0,0 +1,305 @@
+# 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
+  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]
+    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
+      AIA.config.connected_mcp_servers  = connected_keys
+      AIA.config.mcp_server_tool_counts = tool_counts
+      AIA.config.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]
+    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 = @mutex.synchronize { @connected_clients.keys.to_set }
+      failed_names    = @mutex.synchronize { @failed_servers.to_set { |f| f[:name] } }
+
+      servers.each 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)
+        @mutex.synchronize { @failed_servers << { name: srv_name, error: "global connection timeout (#{global_limit}s)" } }
+      end
+    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.
+    # rubocop:disable Metrics/MethodLength
+    def connect_one(server_config, name, spinner, logger)
+      timeout = server_timeout(server_config)
+      spinner.auto_spin
+
+      logger.info("MCP: connecting to '#{name}'...")
+      # rubocop:disable Metrics/BlockLength
+      Timeout.timeout(timeout) do
+        client = RobotLab::MCP::Client.new(server_config)
+        client.connect
+
+        if client.connected?
+          tools = client.list_tools
+          built_tools = tools.map do |tool_def|
+            tool_name  = tool_def[:name]
+            mcp_client = client
+            RobotLab::Tool.create(
+              name:        tool_name,
+              description: tool_def[:description],
+              parameters:  tool_def[:inputSchema],
+              mcp:         name
+            ) { |args| mcp_client.call_tool(tool_name, args) }
+          end
+
+          @mutex.synchronize do
+            @connected_clients[name]  = client
+            @server_tool_counts[name] = tools.size
+            add_tools_deduped(built_tools, name, logger)
+          end
+
+          logger.info("MCP: '#{name}' connected (#{tools.size} tools)")
+          spinner.success("(#{tools.size} tools)")
+        else
+          @mutex.synchronize do
+            @failed_servers << { name: name, error: "connection failed" }
+          end
+          logger.warn("MCP: '#{name}' failed to connect")
+          spinner.error("(connection failed)")
+        end
+      end
+      # rubocop:enable Metrics/BlockLength
+    rescue Timeout::Error
+      @mutex.synchronize do
+        @failed_servers << { name: name, error: "timed out after #{timeout}s" }
+      end
+      logger.warn("MCP: '#{name}' timed out after #{timeout}s")
+      spinner.error("(timed out)")
+    rescue StandardError => e
+      @mutex.synchronize do
+        @failed_servers << { name: name, error: e.message }
+      end
+      logger.warn("MCP: '#{name}' error: #{e.message}")
+      spinner.error("(#{e.message})")
+    end
+    # rubocop:enable Metrics/MethodLength
+
+    # 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_utility.rb

@@ -0,0 +1,57 @@
+# 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
+
+    def effective_mcp_server_names
+      return [] if AIA.config&.flags&.no_mcp
+      servers = AIA.config&.mcp_servers || []
+      return [] if servers.empty?
+
+      names     = servers.map { |s| server_name(s) }.compact
+      use_list  = Array(AIA.config.mcp_use)
+      skip_list = Array(AIA.config.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

data/lib/aia/mention_router.rb

@@ -0,0 +1,260 @@
+# frozen_string_literal: true
+
+# lib/aia/mention_router.rb
+#
+# Handles @mention routing in multi-model networks.
+# Scans prompts for @name tokens, matches them to robots in the
+# network, and sends the prompt only to mentioned robots.
+
+module AIA
+  class MentionRouter
+    include ContentExtractor
+    include HandlerProtocol
+
+    # Reserved @mention that broadcasts the message to every crew member.
+    BROADCAST_TOKEN = 'crew'
+
+    def initialize(ui_presenter:, tracker:, streaming_runner:)
+      @ui_presenter = ui_presenter
+      @tracker = tracker
+      @streaming_runner = streaming_runner
+    end
+
+    # Scan prompt for @mentions and route to matching robots.
+    # Returns true if mentions were handled, false otherwise.
+    #
+    # @param context [HandlerContext] — reads context.robot and context.prompt
+    # @return [Boolean]
+    def handle(context) # rubocop:disable Naming/PredicateMethod
+      robot  = context.robot
+      prompt = context.prompt
+      return false unless robot.network?
+
+      mention_tokens = prompt.scan(/@(\w+)/i).flatten
+      return false if mention_tokens.empty?
+
+      all_robots = robot.crew
+      matched = resolve_mentions(mention_tokens, all_robots)
+
+      if matched.empty?
+        report_unknown_mentions(mention_tokens, all_robots)
+        return true
+      end
+
+      clean_prompt = strip_addressing_prefix(prompt)
+
+      report_unknown_mentions(mention_tokens, all_robots)
+      run_mentioned_robots(matched, clean_prompt, concurrent: leading_address_only?(prompt))
+      true
+    end
+
+    private
+
+    # Remove only the leading run of @mentions (the addressing prefix), so
+    # "@alice summarize this" becomes "summarize this" while @mentions inside the
+    # message body — "hello @hemie have you met @joker" — are preserved as
+    # content and reach the routed robots intact.
+    def strip_addressing_prefix(prompt)
+      prompt.sub(/\A(?:\s*@\w+)+\s*/i, '').strip
+    end
+
+    # True when every @mention is part of the leading address — i.e. nothing but
+    # @mentions before the message — so it's a broadcast to its addressees and
+    # runs concurrently. A @mention woven into the body makes it sequential.
+    def leading_address_only?(prompt)
+      !strip_addressing_prefix(prompt).match?(/@\w+/)
+    end
+
+    # Resolve @mention tokens to robots. The reserved @crew token broadcasts to
+    # the entire crew; otherwise each token is matched to a robot by name.
+    #
+    # @return [Array<Robot>]
+    def resolve_mentions(tokens, all_robots)
+      return all_robots if tokens.any? { |token| token.casecmp?(BROADCAST_TOKEN) }
+
+      tokens.filter_map { |token| all_robots.find { |r| r.name.casecmp?(token) } }.uniq(&:name)
+    end
+
+    def report_unknown_mentions(mention_tokens, all_robots)
+      known_names = all_robots.map { |r| r.name.downcase }
+      unknown = mention_tokens.reject { |t| t.casecmp?(BROADCAST_TOKEN) || known_names.include?(t.downcase) }
+      return unless unknown.any?
+
+      available = all_robots.map(&:name).join(', ')
+      unknown.each do |name|
+        @ui_presenter.display_info("Unknown robot: @#{name}  (available: #{available})")
+      end
+    end
+
+    def run_mentioned_robots(robots, prompt, concurrent:)
+      if robots.size == 1
+        run_streamed(robots.first, prompt)
+      elsif concurrent
+        run_concurrently(robots, prompt)
+      else
+        run_sequentially(robots, prompt)
+      end
+      @ui_presenter.display_separator
+    end
+
+    # Robots mentioned in the body run one at a time, each streaming its reply.
+    # After each reply, it is injected into every *other* member's conversation
+    # so later robots build on it and earlier ones carry it forward — the
+    # sequential pipeline becomes a shared, accumulating conversation.
+    def run_sequentially(robots, prompt)
+      robots.each do |bot|
+        content = run_streamed(bot, prompt)
+        share_with_peers(robots, bot, content) if content
+      end
+    end
+
+    # Add one robot's reply to every other pipeline member's chat history.
+    def share_with_peers(robots, author, content)
+      note = "[#{author.name}]: #{content}"
+      (robots - [author]).each do |peer|
+        next unless peer.respond_to?(:chat) && peer.chat
+
+        peer.chat.add_message({ role: :user, content: note })
+      end
+    end
+
+    # A single addressed robot streams its reply token-by-token.
+    def run_streamed(bot, prompt)
+      result, streamed_content, elapsed = @streaming_runner.run(
+        bot, prompt,
+        header: "\nAI (#{bot.name}):\n   ",
+        spinner_message: "#{bot.name} processing..."
+      )
+      content = streamed_content || extract_content(result)
+
+      if streamed_content
+        puts
+        @ui_presenter.display_info("(#{format_duration(elapsed)})")
+      else
+        @ui_presenter.display_ai_response(reply_block(bot, content, elapsed))
+      end
+
+      finalize_reply(bot, prompt, result, content, elapsed)
+      content
+    rescue StandardError => e
+      @ui_presenter.display_info("Error from #{bot.name}: #{e.class}: #{e.message}")
+      nil
+    end
+
+    # A broadcast to multiple robots runs them concurrently — no token streaming,
+    # which would interleave on one console. Each robot runs in its own thread and
+    # pushes its result to a queue; the main thread renders each reply as it lands
+    # (finish order), so wall-clock is the slowest robot, not the sum. Rendering
+    # stays single-threaded, so display/tracking are never touched concurrently.
+    def run_concurrently(robots, prompt)
+      @ui_presenter.display_info("Broadcasting to #{robots.size} robots concurrently…")
+
+      done = Queue.new
+      robots.each { |bot| Thread.new { done << run_member(bot, prompt) } }
+      robots.size.times { render_member(prompt, *done.pop) }
+    end
+
+    # Run one crew member to completion (off the streaming path), in a worker
+    # thread. Returns [bot, result, elapsed, error] — no shared state touched.
+    def run_member(bot, prompt)
+      started = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      result  = bot.run(prompt, mcp: :inherit, tools: :inherit)
+      [bot, result, Process.clock_gettime(Process::CLOCK_MONOTONIC) - started, nil]
+    rescue StandardError => e
+      [bot, nil, Process.clock_gettime(Process::CLOCK_MONOTONIC) - started, e]
+    end
+
+    # Render one crew member's reply on the main thread.
+    def render_member(prompt, bot, result, elapsed, error)
+      if error
+        @ui_presenter.display_info("Error from #{bot.name}: #{error.class}: #{error.message}")
+        return
+      end
+
+      content = extract_content(result)
+      @ui_presenter.display_ai_response(reply_block(bot, content, elapsed))
+      finalize_reply(bot, prompt, result, content, elapsed)
+    end
+
+    def reply_block(bot, content, elapsed)
+      "**#{bot.name}** [#{bot.model || 'unknown'}] (#{format_duration(elapsed)}):\n#{content}"
+    end
+
+    def finalize_reply(bot, prompt, result, content, elapsed)
+      @tracker.record_turn(model: bot.model || 'unknown', input: prompt, result: result, elapsed: elapsed)
+      output_to_file(content)
+      display_metrics(result, elapsed: elapsed)
+      speak(content)
+    end
+
+    def display_metrics(result, elapsed: nil)
+      return unless AIA.config.flags.tokens
+
+      raw = result.respond_to?(:raw) ? result.raw : nil
+      return unless raw && raw.respond_to?(:input_tokens) && raw.input_tokens
+
+      model_id = extract_model_id(raw)
+      model_id ||= result.respond_to?(:robot_name) ? result.robot_name : "unknown"
+      metrics = {
+        model_id:      model_id,
+        input_tokens:  raw.input_tokens,
+        output_tokens: raw.output_tokens,
+        elapsed:       elapsed
+      }
+      @ui_presenter.display_token_metrics(metrics)
+    end
+
+    def extract_model_id(message)
+      return message.model_id if message.respond_to?(:model_id) && message.model_id
+      return message.model    if message.respond_to?(:model)    && message.model
+      nil
+    end
+
+    def speak(content)
+      return unless AIA.speak?
+
+      audio   = AIA.config.audio
+      command = audio.speak_command || 'say'
+      env     = {}
+      env['SPEECH_MODEL'] = audio.speech_model if audio.speech_model
+
+      if command == 'say'
+        run_with_spinner("Speaking...") do
+          if audio.voice && !audio.voice.to_s.strip.empty?
+            system(env, command, '-v', audio.voice, content.to_s)
+          else
+            system(env, command, content.to_s)
+          end
+        end
+      else
+        require 'tempfile'
+        tmpfile = Tempfile.new(['aia-tts-', '.mp3'])
+        tmpfile.close
+        begin
+          run_with_spinner("Converting to audio...") do
+            system(env, command, content.to_s, tmpfile.path)
+          end
+          if File.size?(tmpfile.path)
+            run_with_spinner("Speaking...") do
+              system('afplay', tmpfile.path)
+            end
+          end
+        ensure
+          tmpfile.unlink
+        end
+      end
+    rescue StandardError => e
+      $stderr.puts "Warning: Speech failed: #{e.message}"
+    end
+
+    def run_with_spinner(message)
+      spinner = TTY::Spinner.new("[:spinner] #{message}", format: :bouncing_ball, output: $stderr)
+      spinner.auto_spin
+      begin
+        yield
+      ensure
+        spinner.stop
+      end
+    end
+  end
+end