aia 1.1.1 → 2.0.0.0.pre.beta2

data/lib/aia/patches/ruby_llm_tool_error.rb

@@ -0,0 +1,96 @@
+# lib/aia/patches/ruby_llm_tool_error.rb
+#
+# Two ruby_llm hardening patches:
+#
+#   1. RubyLLM::Chat#execute_tool — rescue any exception raised during tool
+#      execution and return it as an error hash instead of raising. Without
+#      this, an exception inside a tool skips the add_message call in
+#      handle_tool_calls, leaving an orphaned tool_call in the conversation
+#      history with no matching tool_result. Most LLMs reject that state on the
+#      next request, killing the session. ruby_llm's own execute_tool does NOT
+#      rescue tool exceptions, so this is still required as of 1.16.
+#
+#   2. RubyLLM::Provider#complete — normalize OpenAI o-series / gpt-5 params
+#      (max_tokens -> max_completion_tokens) before the request.
+#
+# Both patches are applied via #prepend + super so they preserve ruby_llm's own
+# method body (including its instrumentation) and survive internal changes to
+# that body. They are guarded against drift: if a future ruby_llm renames or
+# removes the target method, the patch is skipped with a loud warning rather
+# than silently installing dead code or breaking every request.
+
+module AIA
+  class OpenAIParameterNormalizer
+    class << self
+      def normalize(params, model, provider)
+        normalized = params.dup
+        return normalized unless uses_completion_tokens_param?(model, provider)
+        return normalized unless normalized.key?(:max_tokens)
+
+        max_tokens = normalized.delete(:max_tokens)
+        normalized[:max_completion_tokens] ||= max_tokens
+        normalized
+      end
+
+      private
+
+      def uses_completion_tokens_param?(model, provider)
+        provider_slug(provider) == 'openai' && model_id(model).match?(/\A(o\d|gpt-5)/i)
+      end
+
+      def provider_slug(provider)
+        provider.respond_to?(:slug) ? provider.slug.to_s : nil
+      end
+
+      def model_id(model)
+        model.respond_to?(:id) ? model.id.to_s : model.to_s
+      end
+    end
+  end
+
+  # Prepended onto RubyLLM::Provider. Uses a pass-through (*args, **kwargs)
+  # signature so any kwargs ruby_llm adds in future are forwarded to super
+  # unchanged; we only rewrite params before delegating.
+  module RubyLLMOpenAIParameters
+    def complete(*, **kwargs, &)
+      kwargs[:params] = AIA::OpenAIParameterNormalizer.normalize(
+        kwargs[:params] || {}, kwargs[:model], self
+      )
+      super
+    end
+  end
+
+  # Prepended onto RubyLLM::Chat. Calls super (ruby_llm's own instrumented
+  # execute_tool) and only adds exception rescue, so tool-call instrumentation
+  # is preserved.
+  module RubyLLMChatToolError
+    private
+
+    def execute_tool(tool_call)
+      super
+    rescue SignalException, SystemExit
+      raise
+    rescue Exception => e # rubocop:disable Lint/RescueException
+      RubyLLM.logger.warn { "Tool #{tool_call.name} raised #{e.class}: #{e.message}" }
+      { error: "#{e.class}: #{e.message}" }
+    end
+  end
+end
+
+# --- Apply patches, guarded against ruby_llm API drift ---
+
+if RubyLLM::Provider.private_method_defined?(:complete) ||
+   RubyLLM::Provider.method_defined?(:complete)
+  RubyLLM::Provider.prepend(AIA::RubyLLMOpenAIParameters)
+else
+  $stderr.puts "WARNING: [AIA] ruby_llm patch skipped: RubyLLM::Provider#complete " \
+               "not found (ruby_llm #{RubyLLM::VERSION}). OpenAI param normalization is off."
+end
+
+if RubyLLM::Chat.private_method_defined?(:execute_tool) ||
+   RubyLLM::Chat.method_defined?(:execute_tool)
+  RubyLLM::Chat.prepend(AIA::RubyLLMChatToolError)
+else
+  $stderr.puts "WARNING: [AIA] ruby_llm patch skipped: RubyLLM::Chat#execute_tool " \
+               "not found (ruby_llm #{RubyLLM::VERSION}). Tool-exception handling is degraded."
+end

data/lib/aia/pipeline_orchestrator.rb

@@ -0,0 +1,272 @@
+# frozen_string_literal: true
+
+# lib/aia/pipeline_orchestrator.rb
+#
+# Handles per-prompt pipeline processing: building prompt text,
+# executing against the robot, tracking results, and displaying output.
+# Extracted from Session to separate startup concerns from per-turn concerns.
+
+require "json"
+require_relative 'skill_utils'
+
+module AIA
+  class PipelineOrchestrator
+    include ContentExtractor
+    include SkillUtils
+
+    def initialize(robot:, prompt_handler:, input_collector:, ui_presenter:, session_tracker:)
+      @robot           = robot
+      @prompt_handler  = prompt_handler
+      @input_collector = input_collector
+      @ui              = ui_presenter
+      @tracker         = session_tracker
+    end
+
+    # Process all prompts in the pipeline.
+    #
+    # @param config [AIA::Config]
+    # :reek:TooManyStatements -- shift-based loop so front-matter pipeline edits take effect; tracking interleaved
+    # :reek:DuplicateMethodCall -- config.pipeline is deliberately re-read each iteration because prompt front matter mutates it mid-loop
+    def process(config)
+      bridge   = TrakFlowBridge.new
+      tracking = bridge.available? && config.flags.track_pipeline
+
+      bridge.create_plan_from_pipeline(config.pipeline) if tracking
+
+      # Use shift-based loop so that prompt front matter which sets
+      # config.pipeline (via `next:` or `pipeline:`) is picked up on
+      # the next iteration rather than being silently ignored.
+      until config.pipeline.empty?
+        prompt_id = config.pipeline.shift
+        next if prompt_id.nil? || prompt_id.empty?
+
+        begin
+          bridge.update_step_status(prompt_id, :started) if tracking
+
+          prompt_text = build_prompt_text(prompt_id, config)
+          next if prompt_text.nil? || prompt_text.strip.empty?
+
+          result, content = compute_result(prompt_text, config)
+
+          bridge.update_step_status(prompt_id, :completed) if tracking
+
+          @tracker.record_turn(
+            model:  config.models.first.name,
+            input:  prompt_text,
+            result: result
+          )
+
+          @ui.display_ai_response(content)
+          output_to_file(content, config)
+          display_metrics(result, config)
+          @ui.display_separator
+        rescue StandardError => e
+          bridge.update_step_status(prompt_id, :failed, reason: e.message) if tracking
+          raise
+        end
+      end
+    end
+
+    # Allow robot to be updated (after model switch)
+    attr_writer :robot
+
+    private
+
+    # Returns [result, displayable_content] for one prompt — image when in paint
+    # mode, otherwise the LLM text response.
+    def compute_result(prompt_text, config)
+      if AIA.paint?
+        image = generate_image(prompt_text, config)
+        [image, image]
+      else
+        result = execute_prompt(prompt_text, config)
+        [result, extract_content(result)]
+      end
+    end
+
+    # :reek:TooManyStatements -- image request assembly, spinner-wrapped generation, then save/display of the result
+    def generate_image(prompt_text, config)
+      img_cfg = config.image
+      quality = img_cfg.quality
+      style   = img_cfg.style
+      params  = {}
+      params[:quality] = quality if quality && !quality.to_s.strip.empty?
+      params[:style]   = style   if style   && !style.to_s.strip.empty?
+
+      image = @ui.with_spinner("Generating image") do
+        RubyLLM.paint(
+          prompt_text,
+          model:  img_cfg.model || 'dall-e-3',
+          size:   img_cfg.size  || '1024x1024',
+          params: params
+        )
+      end
+
+      out_file = config.output.file
+      if out_file
+        ext  = '.png'
+        path = out_file.chomp(File.extname(out_file)) + ext
+        image.save(path)
+        "Image saved to: #{path}"
+      else
+        image.url || "(image returned no URL)"
+      end
+    end
+
+    # Execute a prompt, optionally using concurrent MCP
+    def execute_prompt(prompt_text, config)
+      concurrent_network = maybe_use_concurrent_mcp(prompt_text, config)
+
+      if concurrent_network
+        @ui.with_spinner("Processing (concurrent)") do
+          concurrent_network.run(prompt_text)
+        end
+      else
+        @ui.with_spinner("Processing") do
+          @robot.run(prompt_text, mcp: :inherit, tools: :inherit)
+        end
+      end
+    end
+
+    # Check if concurrent MCP mode should be used
+    # :reek:TooManyStatements -- ordered eligibility gates (server count, auto/forced, discovery, grouping threshold) before the build
+    def maybe_use_concurrent_mcp(prompt_text, config)
+      return nil unless (config.mcp_servers || []).size > 1
+
+      turn_state  = AIA.turn_state
+      concurrency = config.respond_to?(:concurrency) ? config.concurrency : nil
+      return nil unless concurrency&.auto || turn_state.force_concurrent_mcp
+
+      if turn_state.force_concurrent_mcp
+        turn_state.force_concurrent_mcp = false
+      end
+
+      discovery         = MCPDiscovery.new
+      relevant_servers  = discovery.discover(config)
+      return nil if relevant_servers.size <= 1
+
+      grouper = MCPGrouper.new
+      groups  = grouper.group(relevant_servers)
+
+      threshold = concurrency.respond_to?(:threshold) ? (concurrency.threshold || 2) : 2
+      return nil if groups.size < threshold
+
+      RobotFactory.build_concurrent_mcp_network(config, groups)
+    rescue StandardError => e
+      $stderr.puts "Warning: Concurrent MCP setup failed: #{e.message}"
+      nil
+    end
+
+    # Build prompt text from a prompt_id
+    # :reek:TooManyStatements -- layered prompt assembly: parameters, role prefix, skills, stdin, context files, in a fixed order
+    def build_prompt_text(prompt_id, config)
+      parsed = @prompt_handler.fetch_prompt(prompt_id)
+      return nil unless parsed
+
+      parameters = parsed.respond_to?(:parameters) ? parsed.parameters : nil
+      if parameters && !parameters.empty?
+        values = @input_collector.collect(parameters)
+        values.each { |k, v| parameters[k] = v }
+      end
+
+      prompt_text = parsed.to_s
+
+      role = config.prompts.role
+      if role && !role.empty?
+        role_parsed = @prompt_handler.fetch_role(role)
+        if role_parsed
+          prompt_text = "#{role_parsed}\n\n#{prompt_text}"
+        end
+      end
+
+      # In pipeline mode, append --skill content to each prompt text.
+      # In chat mode, skills are injected into the system prompt instead (system_prompt_assembler).
+      unless config.flags&.chat == true
+        skill_content = load_skills_content(
+          Array(config.prompts&.skills),
+          skills_base_dir(config)
+        )
+        prompt_text = "#{prompt_text}\n\n#{skill_content}" if skill_content
+      end
+
+      stdin_content = config.stdin_content
+      if stdin_content && !stdin_content.strip.empty?
+        prompt_text = "#{prompt_text}\n\n#{stdin_content}"
+        config.stdin_content = nil
+      end
+
+      add_context_files(prompt_text, config)
+    end
+
+    # Append context file contents to prompt
+    def add_context_files(prompt_text, config)
+      context_files = config.context_files
+      return prompt_text if context_files.nil? || context_files.empty?
+
+      context = context_files.map do |file|
+        File.read(file) rescue "Error reading file: #{file}"
+      end.join("\n\n")
+
+      return prompt_text if context.strip.empty?
+
+      "#{prompt_text}\n\n#{context}"
+    end
+
+    # Display token metrics if enabled.
+    # Uses result.raw (RubyLLM::Message) for token data, matching ChatLoop's approach.
+    def display_metrics(result, config)
+      return unless config.flags.tokens
+
+      if defined?(SimpleFlow::Result) && result.is_a?(SimpleFlow::Result)
+        display_network_metrics(result)
+        return
+      end
+
+      raw = result.respond_to?(:raw) ? result.raw : nil
+      return unless raw && raw.respond_to?(:input_tokens) && raw.input_tokens
+
+      model_id = (raw.respond_to?(:model_id) && raw.model_id) ||
+                 (raw.respond_to?(:model) && raw.model) ||
+                 config.models.first.name
+      @ui.display_token_metrics(
+        model_id:      model_id,
+        input_tokens:  raw.input_tokens,
+        output_tokens: raw.output_tokens
+      )
+    end
+
+    def display_network_metrics(flow_result)
+      metrics_list = []
+      flow_result.context.each do |task_name, robot_result|
+        next if task_name == :run_params
+        next unless robot_result.respond_to?(:raw)
+
+        raw = robot_result.raw
+        next unless raw && raw.respond_to?(:input_tokens) && raw.input_tokens
+
+        model_id = (raw.respond_to?(:model_id) && raw.model_id) ||
+                   (raw.respond_to?(:model) && raw.model) ||
+                   task_name.to_s
+        display_name = robot_result.respond_to?(:robot_name) ? robot_result.robot_name : task_name.to_s
+        metrics_list << {
+          model_id:      model_id,
+          display_name:  display_name,
+          input_tokens:  raw.input_tokens || 0,
+          output_tokens: raw.output_tokens || 0,
+          elapsed:       robot_result.respond_to?(:duration) ? robot_result.duration : nil
+        }
+      end
+      @ui.display_multi_model_metrics(metrics_list) unless metrics_list.empty?
+    end
+
+    # Write content to the output file if configured
+    def output_to_file(content, config)
+      out_file = config.output.file
+      return unless out_file
+
+      File.open(out_file, "a") do |f|
+        f.puts "AI: #{content}"
+      end
+    end
+  end
+end

data/lib/aia/plugin_loader.rb

@@ -0,0 +1,170 @@
+# frozen_string_literal: true
+
+module AIA
+  # Loads, reloads, and unloads top-level Ruby plugin files from the configured
+  # plugins directory. A plugin is any `*.rb` file directly under
+  # `config.paths.plugins_dir`; it may define top-level (Kernel) methods and/or
+  # top-level constants. The loader tracks exactly what each file defines so a
+  # changed plugin can be cleanly reloaded and a deleted plugin fully removed
+  # from the running process.
+  class PluginLoader
+    class << self
+      # Per-process registry of loaded plugins, keyed by basename.
+      # Each value: { path:, signature:, defined: { methods: [Symbol], constants: [Symbol] } }
+      def registry
+        @registry ||= {}
+      end
+
+      def mutex
+        @mutex ||= Mutex.new
+      end
+
+      # Load every `*.rb` file at the top level of the plugins directory,
+      # starting from a clean slate (any previously-tracked definitions are
+      # removed first). Returns the sorted basenames of the loaded plugins.
+      #
+      # @param config [#paths, #loaded_plugins=]
+      # @return [Array<String>]
+      def load!(config)
+        @config = config
+        reset!(config)
+
+        dir = plugins_dir(config)
+        return [] unless dir
+
+        Dir.glob(File.join(dir, '*.rb')).each { |file| load_file(file, config) }
+        loaded_names
+      end
+
+      # Load — or, if already loaded, reload — a single plugin file. On reload,
+      # the file's previous definitions are removed before re-executing so that
+      # renamed/deleted methods don't linger. Returns the basename, or nil on
+      # failure.
+      #
+      # @param path [String] absolute path to a `*.rb` plugin file
+      # @param config [Object, nil] config to keep `loaded_plugins` in sync with
+      # @return [String, nil]
+      def load_file(path, config = @config)
+        mutex.synchronize do
+          name = File.basename(path, '.rb')
+          remove_definitions(registry[name]) if registry.key?(name)
+
+          before = current_definitions
+          load(path)
+          defined = diff_definitions(before, current_definitions)
+
+          registry[name] = { path: path, signature: signature(path), defined: defined }
+          sync_config(config)
+          name
+        end
+      rescue ScriptError, StandardError => e
+        warn_failure(path, e)
+        nil
+      end
+
+      # Unload a plugin by basename: remove the methods and constants it defined
+      # and forget it. Returns the basename if it was loaded, else nil.
+      #
+      # @param name [String] plugin basename (no extension)
+      # @param config [Object, nil]
+      # @return [String, nil]
+      def unload_file(name, config = @config)
+        mutex.synchronize do
+          entry = registry.delete(name)
+          remove_definitions(entry)
+          sync_config(config)
+          entry && name
+        end
+      end
+
+      # @return [Array<String>] sorted basenames currently loaded
+      def loaded_names
+        registry.keys.sort
+      end
+
+      # Remove all tracked plugin definitions and clear the registry. Used by
+      # load! for a clean reload and by tests for isolation.
+      #
+      # @param config [Object, nil]
+      # @return [void]
+      def reset!(config = @config)
+        mutex.synchronize do
+          registry.each_value { |entry| remove_definitions(entry) }
+          registry.clear
+          sync_config(config)
+        end
+      end
+
+      private
+
+      def plugins_dir(config)
+        dir = config.paths&.plugins_dir
+        return nil if dir.nil? || dir.to_s.strip.empty?
+        return nil unless Dir.exist?(dir)
+
+        dir
+      end
+
+      # Snapshot of the top-level methods and constants currently defined on
+      # Object. Top-level `def` creates private instance methods on Object;
+      # top-level constants/classes/modules become constants on Object.
+      #
+      # @return [Hash{Symbol=>Array<Symbol>}]
+      def current_definitions
+        {
+          methods:   Object.private_instance_methods(false) + Object.instance_methods(false),
+          constants: Object.constants(false)
+        }
+      end
+
+      # What appeared between two snapshots.
+      def diff_definitions(before, after)
+        {
+          methods:   after[:methods]   - before[:methods],
+          constants: after[:constants] - before[:constants]
+        }
+      end
+
+      # Remove the methods and constants a plugin defined, if still present and
+      # owned directly by Object.
+      def remove_definitions(entry)
+        return unless entry
+
+        defined = entry[:defined] || {}
+        Array(defined[:methods]).each { |m| remove_object_method(m) }
+        Array(defined[:constants]).each { |c| remove_object_constant(c) }
+      end
+
+      def remove_object_method(name)
+        return unless Object.private_method_defined?(name, false) ||
+                      Object.method_defined?(name, false)
+
+        Object.send(:remove_method, name)
+      rescue NameError
+        # Already gone — nothing to do.
+      end
+
+      def remove_object_constant(name)
+        Object.send(:remove_const, name) if Object.const_defined?(name, false)
+      rescue NameError
+        # Already gone — nothing to do.
+      end
+
+      # A cheap change-signature for a file: [mtime, size]. Nil if unreadable.
+      def signature(path)
+        stat = File.stat(path)
+        [stat.mtime, stat.size]
+      rescue SystemCallError
+        nil
+      end
+
+      def sync_config(config)
+        config.loaded_plugins = loaded_names if config.respond_to?(:loaded_plugins=)
+      end
+
+      def warn_failure(path, error)
+        $stderr.puts "Warning: Failed to load plugin '#{path}': #{error.message}"
+      end
+    end
+  end
+end