aia 1.1.1 → 2.0.0.0.pre.alpha

data/lib/aia/special_mode_handler.rb

@@ -0,0 +1,302 @@
+# frozen_string_literal: true
+
+# lib/aia/special_mode_handler.rb
+#
+# Handles special execution modes triggered by directives:
+# /verify      — two independent answers + reconciliation
+# /decompose   — break into parallel sub-tasks (Async::Barrier for I/O concurrency)
+# /concurrent  — concurrent MCP server access
+# /orchestrate — 3-tier layered orchestration (orchestrator → leads → specialists)
+
+require 'async'
+
+module AIA
+  class SpecialModeHandler
+    include ContentExtractor
+
+    # Special-mode dispatch table. Each row is triggered by a TurnState flag and
+    # runs either a bespoke method (`run:`) or, for the uniform "sub-handler
+    # returns content" family, the shared #run_sub_handler (`sub:`). Adding a
+    # mode means adding a row here plus (for a bespoke mode) its handle_* method
+    # — never editing #handle.
+    MODES = [
+      { flag: :force_verify,         run: :handle_verification   },
+      { flag: :force_decompose,      run: :handle_decomposition  },
+      { flag: :force_concurrent_mcp, run: :handle_concurrent_mcp },
+      { flag: :force_debate,         sub: { handler: :debate_handler,   label: 'Debate'     } },
+      { flag: :force_delegate,       sub: { handler: :delegate_handler, label: 'Delegation' } },
+      { flag: :force_spawn,          sub: { handler: :spawn_handler,    label: 'Spawn',
+                                            context: lambda { |prompt, ts|
+                                              type = ts.spawn_type
+                                              spec = ts.spawn_spec
+                                              ts.spawn_type = nil
+                                              ts.spawn_spec = nil
+                                              { prompt: prompt, specialist_type: type, spawn_spec: spec }
+                                            } } },
+      { flag: :force_orchestrate,    run: :handle_orchestration  }
+    ].freeze
+
+    # Memoized sub-handler accessors, used to propagate a robot swap.
+    HANDLERS = %i[debate_handler delegate_handler spawn_handler layered_orchestrator].freeze
+
+    def initialize(robot:, ui_presenter:, tracker:)
+      @robot        = robot
+      @ui_presenter = ui_presenter
+      @tracker      = tracker
+      # Handlers are instantiated on first use to avoid loading their
+      # dependencies (async, trakflow, etc.) unless the mode is invoked.
+      @debate_handler       = nil
+      @delegate_handler     = nil
+      @spawn_handler        = nil
+      @layered_orchestrator = nil
+    end
+
+    # Update the robot reference (e.g., after a model switch).
+    def robot=(new_robot)
+      @robot = new_robot
+      HANDLERS.each do |name|
+        inst = instance_variable_get(:"@#{name}")
+        inst.robot = new_robot if inst
+      end
+    end
+
+    # Check TurnState flags and dispatch to the appropriate handler.
+    # Returns true if a special mode was handled.
+    #
+    # @param prompt [String] the user prompt
+    # @return [Boolean]
+    def handle(prompt)
+      turn_state = AIA.turn_state
+      mode = MODES.find { |m| turn_state.public_send(m[:flag]) }
+      return false unless mode
+
+      turn_state.public_send(:"#{mode[:flag]}=", false)
+      if mode[:sub]
+        run_sub_handler(mode[:sub], prompt, turn_state)
+      else
+        send(mode[:run], prompt)
+      end
+    end
+
+    private
+
+    # Run a "sub-handler returns content" mode: debate / delegate / spawn share
+    # this exact shape, differing only by handler, label, and (spawn) an extra
+    # context field built by spec[:context].
+    def run_sub_handler(spec, prompt, turn_state)
+      ctx_args = spec[:context] ? spec[:context].call(prompt, turn_state) : { prompt: prompt }
+      content  = send(spec[:handler]).handle(HandlerContext.new(**ctx_args))
+      return false unless content
+
+      display_and_save(content)
+      true
+    rescue StandardError => e
+      @ui_presenter.display_info("#{spec[:label]} failed: #{e.message}. Falling back to normal mode.")
+      false
+    end
+
+    def handle_verification(prompt)
+      require_relative 'verification_network'
+      @ui_presenter.display_info("Running verification (2 independent + reconciliation)...")
+
+      network = VerificationNetwork.build(AIA.config)
+      result = @ui_presenter.with_spinner("Verifying") { network.run(message: prompt) }
+
+      present_result(result, prompt: prompt, ui_presenter: @ui_presenter, tracker: @tracker)
+      true
+    rescue StandardError => e
+      @ui_presenter.display_info("Verification failed: #{e.message}. Falling back to normal mode.")
+      false
+    end
+
+    # rubocop:disable Metrics/AbcSize, Metrics/MethodLength
+    def handle_decomposition(prompt)
+      require_relative 'prompt_decomposer'
+      total_start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+
+      @ui_presenter.display_info("Decomposing prompt into sub-tasks...")
+      decomposer    = PromptDecomposer.new(@robot)
+      decompose_t0  = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      subtasks      = decomposer.decompose(prompt)
+      decompose_dur = Process.clock_gettime(Process::CLOCK_MONOTONIC) - decompose_t0
+
+      if subtasks.empty?
+        @ui_presenter.display_info("Prompt cannot be meaningfully decomposed. Running normally.")
+        return false
+      end
+
+      @ui_presenter.display_info("Decomposed into #{subtasks.size} sub-tasks:")
+      subtasks.each_with_index { |t, i| @ui_presenter.display_info("  #{i + 1}. #{t}") }
+
+      timings         = Array.new(subtasks.size, 0.0)
+      raw_subtasks    = Array.new(subtasks.size)
+      wall_start      = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+
+      results = Sync do
+        barrier = Async::Barrier.new
+        tasks = subtasks.each_with_index.map do |task, i|
+          barrier.async do
+            t0 = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+            @ui_presenter.display_info("Processing sub-task #{i + 1}...")
+            r = @robot.run(task, mcp: :inherit, tools: :inherit)
+            timings[i]      = Process.clock_gettime(Process::CLOCK_MONOTONIC) - t0
+            raw_subtasks[i] = r
+            extract_content(r)
+          rescue => e
+            timings[i] = Process.clock_gettime(Process::CLOCK_MONOTONIC) - t0
+            @ui_presenter.display_info("  Sub-task #{i + 1} failed: #{e.message}")
+            nil
+          end
+        end
+        barrier.wait
+        tasks.map(&:wait)
+      end
+      # rubocop:enable Metrics/AbcSize, Metrics/MethodLength
+
+      parallel_wall    = Process.clock_gettime(Process::CLOCK_MONOTONIC) - wall_start
+      serial_est       = timings.sum
+      parallel_speedup = serial_est / parallel_wall
+
+      raise DecomposeError, "All sub-tasks failed" if results.all?(&:nil?)
+
+      @ui_presenter.display_info("Synthesizing results...")
+      synth_t0  = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      final     = decomposer.synthesize(prompt, results.compact)
+      synth_dur = Process.clock_gettime(Process::CLOCK_MONOTONIC) - synth_t0
+
+      total_wall = Process.clock_gettime(Process::CLOCK_MONOTONIC) - total_start
+
+      lines = ["", "Benchmark:"]
+      lines << "  Decompose probe : #{format('%.2f', decompose_dur)}s"
+      timings.each_with_index { |t, i| lines << "  Sub-task #{i + 1}      : #{format('%.2f', t)}s" }
+      lines << "  Synthesis       : #{format('%.2f', synth_dur)}s"
+      lines << "  ─────────────────────────────────────────"
+      lines << "  Sub-tasks serial est. : #{format('%.2f', serial_est)}s  (sum of sub-task times)"
+      lines << "  Sub-tasks concurrent  : #{format('%.2f', parallel_wall)}s  (#{format('%.1f', parallel_speedup)}x parallel speedup)"
+      lines << "  Total flow wall time  : #{format('%.2f', total_wall)}s  (compare against --tokens Time for a normal run)"
+      @ui_presenter.display_info(lines.join("\n"))
+
+      display_decompose_token_metrics(raw_subtasks, final, total_wall)
+
+      present_result(final, prompt: prompt, ui_presenter: @ui_presenter, tracker: @tracker)
+      true
+    rescue StandardError => e
+      @ui_presenter.display_info("Decomposition failed: #{e.message}. Falling back to normal mode.")
+      false
+    end
+
+    def handle_concurrent_mcp(prompt)
+      return false unless (AIA.config.mcp_servers || []).size > 1
+
+      discovery = MCPDiscovery.new
+      relevant = discovery.discover(AIA.config)
+      return false if relevant.size <= 1
+
+      grouper = MCPGrouper.new
+      groups = grouper.group(relevant)
+      return false if groups.size < 2
+
+      @ui_presenter.display_info("Running concurrent MCP across #{groups.size} server groups...")
+
+      network = RobotFactory.build_concurrent_mcp_network(AIA.config, groups)
+      result = @ui_presenter.with_spinner("Processing (concurrent)") { network.run(message: prompt) }
+      extract_content(result)
+
+      present_result(result, prompt: prompt, ui_presenter: @ui_presenter, tracker: @tracker)
+      true
+    rescue StandardError => e
+      @ui_presenter.display_info("Concurrent MCP failed: #{e.message}. Falling back to normal mode.")
+      false
+    end
+
+    def handle_orchestration(prompt)
+      @ui_presenter.display_info("Starting 3-tier layered orchestration...")
+
+      content = layered_orchestrator.handle(HandlerContext.new(robot: @robot, prompt: prompt))
+
+      if content
+        display_and_save(content)
+      else
+        @ui_presenter.display_separator
+      end
+
+      true  # always consume the prompt; never fall through to normal mode
+    rescue StandardError => e
+      @ui_presenter.display_info("Orchestration failed: #{e.class}: #{e.message}")
+      @ui_presenter.display_separator
+      true
+    end
+
+    def debate_handler
+      @debate_handler ||= begin
+        require_relative 'debate_handler'
+        DebateHandler.new(robot: @robot, ui_presenter: @ui_presenter, tracker: @tracker)
+      end
+    end
+
+    def delegate_handler
+      @delegate_handler ||= begin
+        require_relative 'delegate_handler'
+        require_relative 'task_decomposer'
+        require_relative 'task_executor'
+        DelegateHandler.new(
+          robot: @robot, ui_presenter: @ui_presenter,
+          tracker: @tracker, task_coordinator: AIA.task_coordinator
+        )
+      end
+    end
+
+    def spawn_handler
+      @spawn_handler ||= begin
+        require_relative 'spawn_handler'
+        SpawnHandler.new(robot: @robot, ui_presenter: @ui_presenter, tracker: @tracker)
+      end
+    end
+
+    def layered_orchestrator
+      @layered_orchestrator ||= begin
+        require_relative 'layered_orchestrator'
+        LayeredOrchestrator.new(robot: @robot, ui_presenter: @ui_presenter, tracker: @tracker)
+      end
+    end
+
+    # Display content, write to output file, and print separator.
+    # Used by handlers that receive pre-extracted content from sub-handlers.
+    def display_and_save(content)
+      @ui_presenter.display_ai_response(content)
+      output_to_file(content)
+      @ui_presenter.display_separator
+    end
+
+    # Aggregate token counts from all sub-task raw results plus the synthesis
+    # result and display a combined metrics table. Only runs when --tokens is set.
+    def display_decompose_token_metrics(raw_subtasks, synthesis_result, total_elapsed)
+      return unless AIA.config.flags.tokens
+
+      total_input  = 0
+      total_output = 0
+
+      raw_subtasks.compact.each do |r|
+        raw = r.respond_to?(:raw) ? r.raw : nil
+        next unless raw.respond_to?(:input_tokens) && raw.input_tokens
+        total_input  += raw.input_tokens  || 0
+        total_output += raw.output_tokens || 0
+      end
+
+      synth_raw = synthesis_result.respond_to?(:raw) ? synthesis_result.raw : nil
+      if synth_raw.respond_to?(:input_tokens) && synth_raw.input_tokens
+        total_input  += synth_raw.input_tokens  || 0
+        total_output += synth_raw.output_tokens || 0
+      end
+
+      return if total_input.zero? && total_output.zero?
+
+      @ui_presenter.display_token_metrics({
+        model_id:      AIA.config.models.first.name,
+        input_tokens:  total_input,
+        output_tokens: total_output,
+        elapsed:       total_elapsed
+      })
+    end
+  end
+end

data/lib/aia/startup_coordinator.rb

@@ -0,0 +1,150 @@
+# frozen_string_literal: true
+
+# lib/aia/startup_coordinator.rb
+#
+# Handles all session startup tasks: MCP connection, tool loading,
+# filter initialization, task coordination, and bus attachment.
+# Extracted from Session to give Session a single responsibility.
+
+require "fileutils"
+require "tty-spinner"
+
+module AIA
+  class StartupCoordinator
+    attr_reader :filters, :mcp_manager
+
+    def initialize(robot:, ui_presenter:)
+      @robot       = robot
+      @ui          = ui_presenter
+      @filters     = {}
+      @mcp_manager = nil
+    end
+
+    # Run all startup coordination tasks.
+    #
+    # @param config [AIA::Config]
+    def run(config)
+      connect_mcp_servers(config)
+
+      tools = all_available_tools(config)
+
+      spinner = TTY::Spinner.new("  [:spinner] Initializing tool filters...", format: :dots)
+      spinner.auto_spin
+      @filters = ToolFilterRegistry.build_from_config(config, tools)
+      spinner.success
+
+      initialize_task_coordinator if trakflow_available?
+
+      attach_bus_if_network
+    end
+
+    private
+
+    # Connect ALL configured MCP servers eagerly at startup, then absorb any
+    # MCP clients registered via --require (e.g., shared_tools/mcp/* files).
+    # Uses MCPDiscovery to determine which servers to connect (respects
+    # --mcp-use, --mcp-skip, and KBS mcp_activate decisions), then normalizes
+    # each server config via MCPConfigNormalizer before handing off to
+    # MCPConnectionManager.
+    def connect_mcp_servers(config)
+      return if config.flags.no_mcp
+
+      discovered = MCPDiscovery.new.discover(config)
+      validate_mcp_use_names(config, discovered) if Array(config.mcp_use).any?
+
+      # Prefer the robot's pre-built mcp_config (already normalized) when
+      # present, but always apply --mcp-skip so that flag is honoured.
+      # Fall back to the MCPDiscovery result (which includes both filters).
+      servers = if @robot.respond_to?(:mcp_config) && @robot.mcp_config.is_a?(Array)
+                  skip_list = Array(config.mcp_skip)
+                  @robot.mcp_config.reject { |s| skip_list.include?(AIA::Utility.server_name(s)) }
+                else
+                  discovered.map { |s| MCPConfigNormalizer.normalize(s) }
+                end
+
+      @mcp_manager = MCPConnectionManager.new
+      @mcp_manager.connect_all(servers) if servers.is_a?(Array) && servers.any?
+      @mcp_manager.absorb_ruby_llm_mcp_clients
+      return unless @mcp_manager.connected?
+
+      @mcp_manager.update_config
+      @mcp_manager.inject_into(@robot) if @mcp_manager.any_tools?
+    end
+
+    # Initialize TaskCoordinator, auto-creating the TrakFlow project if needed.
+    def initialize_task_coordinator
+      require_relative 'task_coordinator'
+      require_relative 'trakflow_bridge'
+      require_relative 'tools/task_board_tool'
+      require_relative 'tools/delegate_to_foreman_tool'
+
+      ensure_trakflow_initialized unless TrakFlow.initialized?
+
+      AIA.task_coordinator = TaskCoordinator.new
+    rescue StandardError => e
+      AIA.debug_warn("TaskCoordinator initialization failed (non-fatal): #{e.message}", exc: e)
+    end
+
+    # Return true if the trak_flow gem is available; false if it cannot be loaded.
+    def trakflow_available?
+      defined?(TrakFlow) || begin
+        require 'trak_flow'
+        true
+      rescue LoadError
+        false
+      end
+    end
+
+    # Bootstrap a TrakFlow project so the task board is always available.
+    # Uses the default database path (~/.config/trak_flow/tf.db).
+    def ensure_trakflow_initialized
+      trakflow_dir = TrakFlow.trak_flow_dir
+      FileUtils.mkdir_p(trakflow_dir)
+
+      db = TrakFlow::Storage::Database.new
+      db.connect
+
+      TrakFlow.reset_root!
+    end
+
+    # Attach a shared TypedBus to multi-model networks.
+    def attach_bus_if_network
+      return unless @robot.network?
+
+      RobotFactory.attach_bus(@robot)
+    rescue StandardError => e
+      AIA.debug_warn("Bus attachment failed (non-fatal): #{e.message}", exc: e)
+    end
+
+    # Warn when --mcp-use names don't match any configured server.
+    # Helps users catch typos before a silent empty-tools session.
+    def validate_mcp_use_names(config, discovered_servers)
+      requested  = Array(config.mcp_use)
+      available  = Array(config.mcp_servers).map { |s| AIA::Utility.server_name(s) }
+      found      = discovered_servers.map { |s| AIA::Utility.server_name(s) }
+      missing    = requested - found
+
+      return if missing.empty?
+
+      $stderr.puts "WARNING: --mcp-use specified server(s) not found in config: #{missing.join(', ')}"
+      $stderr.puts "         Available servers: #{available.join(', ')}" if available.any?
+    end
+
+    # Collect all tools available to the robot: local + MCP.
+    def all_available_tools(config)
+      local = Array(config.loaded_tools)
+      mcp   = collect_mcp_tools
+      local + mcp
+    end
+
+    # Retrieve MCP tools from the connection manager.
+    # MCPConnectionManager builds RobotLab::Tool instances via connect_one and
+    # absorb_ruby_llm_mcp_clients — those are what get injected into the robot
+    # and what robot_lab filters by name. RubyLLM::MCP.clients is a different
+    # registry and does not contain the tools connected here.
+    def collect_mcp_tools
+      return [] unless @mcp_manager
+      @mcp_manager.connected_tools
+    end
+  end
+end

data/lib/aia/streaming_runner.rb

@@ -0,0 +1,169 @@
+# frozen_string_literal: true
+
+# lib/aia/streaming_runner.rb
+#
+# Encapsulates the streaming execution pattern used across ChatLoop,
+# MentionRouter, and expert routing. Manages spinner lifecycle and
+# streaming callback for any robot.
+
+require "tty-spinner"
+
+module AIA
+  class StreamingRunner
+    # Hard ceiling on tools sent to the provider. OpenAI and Anthropic both
+    # reject tool arrays longer than 128. Overridable via config.max_tools.
+    DEFAULT_MAX_TOOLS = 128
+
+    def initialize
+      @spinner = TTY::Spinner.new("[:spinner] Processing...", format: :bouncing_ball)
+    end
+
+    # Run a robot with streaming output.
+    # Spinner shows until the first chunk arrives, then stops and
+    # chunks are printed directly to stdout.
+    #
+    # @param robot [RobotLab::Robot, RobotLab::Network] the robot to run
+    # @param prompt [String] the prompt to send
+    # @param header [String] text printed before the first streamed chunk
+    # @param spinner_message [String] spinner label
+    # @param tools [Array<String>, nil] tool names to allow for this turn (nil = all)
+    # @return [Array(Object, String, Float)] [result, streamed_content_or_nil, elapsed_seconds]
+    def run(robot, prompt, header: "\nAI:\n   ", spinner_message: "Processing...", tools: nil)
+      @spinner.reset
+      @spinner.update(title: spinner_message)
+      @spinner.auto_spin
+      streamed = []
+      header_printed = false
+      in_think_block = false  # tracks position inside a <think>...</think> span across chunks
+      start_time = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+
+      streaming_block = proc do |chunk|
+        text = chunk.respond_to?(:content) ? chunk.content.to_s : chunk.to_s
+        next if text.empty?
+
+        # qwen3 and similar local reasoning models embed thinking inline in the
+        # content stream wrapped in <think>...</think> tags. Filter them out
+        # unless --thinking is on. Tag boundaries may split across chunks, so
+        # in_think_block persists across calls via the closure.
+        unless show_thinking?
+          text, in_think_block = filter_thinking(text, in_think_block)
+          next if text.empty?
+        end
+
+        unless header_printed
+          @spinner.stop
+          print header
+          header_printed = true
+        end
+
+        streamed << text
+        $stdout.print(text)
+      end
+
+      tools_param = resolve_tools_param(tools, robot)
+
+      begin
+        result = robot.run(prompt, mcp: :inherit, tools: tools_param, &streaming_block)
+      rescue Exception # rubocop:disable Lint/RescueException
+        @spinner.stop unless header_printed
+        raise
+      end
+
+      @spinner.stop unless header_printed
+      elapsed = Process.clock_gettime(Process::CLOCK_MONOTONIC) - start_time
+      content = streamed.empty? ? nil : streamed.join
+      [result, content, elapsed]
+    end
+
+    private
+
+    # Translate the resolved tool list into robot_lab's ToolConfig vocabulary,
+    # then clamp it to the provider's limit:
+    #   nil   -> :inherit  (no filter active / filter errored — use the full set)
+    #   []    -> :none     (filter ran and found nothing relevant — send no tools)
+    #   names -> names     (filtered subset; ToolConfig.filter_tools applies them)
+    # Distinguishing [] from nil is what lets a "nothing relevant" turn send zero
+    # tools instead of the whole build-time set.
+    def resolve_tools_param(tools, robot)
+      param =
+        if tools.nil?
+          :inherit
+        elsif tools.empty?
+          :none
+        else
+          tools
+        end
+
+      enforce_tool_cap(param, robot)
+    end
+
+    # Clamp the effective tool list to the provider's maximum. Returns the
+    # value unchanged when it's already within budget (including :none, and
+    # :inherit when the robot's own tool count is under the cap). When over,
+    # expands :inherit to the robot's tool names, trims to the cap, and warns —
+    # a reduced tool set beats a failed turn.
+    def enforce_tool_cap(tools_param, robot)
+      return tools_param if tools_param == :none
+
+      max   = max_tools
+      names = tools_param == :inherit ? robot_tool_names(robot) : tools_param
+      return tools_param if names.size <= max
+
+      dropped = names.size - max
+      $stderr.puts(
+        "⚠ Tool list (#{names.size}) exceeds the provider limit of #{max}; " \
+        "sending #{max}, dropping #{dropped}. Enable --auto-tool-filter or reduce " \
+        "the available tools to control which ones are sent."
+      )
+      names.first(max)
+    end
+
+    # The provider tool cap, from config.max_tools when set, else the default.
+    def max_tools
+      configured = AIA.respond_to?(:config) && AIA.config.respond_to?(:max_tools) ? AIA.config.max_tools : nil
+      configured&.positive? ? configured : DEFAULT_MAX_TOOLS
+    end
+
+    # The robot's full set of tool names (local + MCP). Used only to size and
+    # trim an over-limit :inherit set.
+    def robot_tool_names(robot)
+      local = robot.respond_to?(:local_tools) ? Array(robot.local_tools) : []
+      mcp   = robot.respond_to?(:mcp_tools)   ? Array(robot.mcp_tools)   : []
+      (local + mcp).map { |t| t.respond_to?(:name) ? t.name : t.class.name }
+    end
+
+    def show_thinking?
+      AIA.config&.flags&.thinking
+    end
+
+    # Strip <think>...</think> spans from a streaming chunk, handling the case
+    # where tag boundaries fall between chunk deliveries.
+    #
+    # @param text [String] raw chunk text
+    # @param in_think_block [Boolean] whether a <think> tag is currently open
+    # @return [Array(String, Boolean)] [filtered_text, updated_in_think_block]
+    def filter_thinking(text, in_think_block)
+      output = +''
+
+      until text.empty?
+        if in_think_block
+          if (idx = text.index('</think>'))
+            in_think_block = false
+            text = text[(idx + '</think>'.length)..]
+          else
+            break  # entire remaining chunk is thinking — discard
+          end
+        elsif (idx = text.index('<think>'))
+          output << text[0...idx]
+          in_think_block = true
+          text = text[(idx + '<think>'.length)..]
+        else
+          output << text
+          break
+        end
+      end
+
+      [output, in_think_block]
+    end
+  end
+end