rails_console_ai 0.24.0 → 0.26.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 41b94c2b3a01cba1210242c6c8752510fc1976bf9b7195272f4a109dceadcdab
-  data.tar.gz: 6c813abfeeda804b115dfd2a385bc67e1796b0af80cc438e5f41d66886fa519e
+  metadata.gz: a2194416a93ce8522de376169eb627b90c8b2477ba253f0f1b877144251f9ee6
+  data.tar.gz: 65dcbeb6eef9dd2641181529aa671058e85a6cd1c5263acd51678c52c89d3567
 SHA512:
-  metadata.gz: 52281835764610528b039d12ac786ec7056d035e9055d5574bee3a3143c9833e8fa5943a5932980c7ee28ae77351a178186dc15b3afb2f0e9acb7c71b1104c5f
-  data.tar.gz: d1f56cd511ff3411b6e2783c0790bbc3d012a420971fe2cd400f7f573ffc168f5b6bca5c692618bf5c4ef2eaa9d0d0283958afcf1b7ac71d426dd4f2f4377425
+  metadata.gz: e4d5a1f2fe8ef6ae4593829ac1a8997eeeac652ea0d0a491f91a64b8e45a246d024a10aa269e853949e0d106779d4f1b9d671af21057bd6af802ff6a541ac15d
+  data.tar.gz: ec0e0d25f7bb2605a7c4180ebc56ab56d8f590d3632c98310ec65c3aecd987d04118e7199534a8ea5a1a8023a505de9ac467ce05d01393df667c2436bcebc68c

data/CHANGELOG.md

@@ -2,6 +2,23 @@
 
 All notable changes to this project will be documented in this file.
 
+## [0.26.0]
+
+- Add sub-agent support
+- Add integration tests
+- Increase max conversation rounds
+- Fix sub-agent model resolution
+- Improve plan step failure handling
+
+## [0.25.0]
+
+- Expand truncation limits
+- Allow any user on allow list to interact with Slack bot in a thread
+- Handle ctrl-c better in console
+- Fix stdout capture in Slack sessions
+- Improve Slack bot logging
+- Fix thread safety issues in Slack bot
+
 ## [0.24.0]
 
 - Refactor thinking text display and include in Slack with more technical detail

data/README.md

@@ -105,8 +105,69 @@ Say "think harder" in any query to auto-upgrade to the thinking model for that s
 - **History compaction** — `/compact` summarizes long conversations to reduce cost and latency
 - **Output trimming** — older execution outputs are automatically replaced with references; the LLM can recall them on demand via `recall_output`, and you can `/expand <id>` to see them
 - **Debug mode** — `/debug` shows context breakdown, token counts, and per-call cost estimates before and after each LLM call
+- **Sub-agents** — delegate multi-step investigations to a separate LLM context that returns only a concise summary, keeping the main conversation lean
 - **Safe mode** — configurable guards that block side effects (DB writes, HTTP mutations, email delivery) during AI code execution
 
+## Sub-Agents
+
+Sub-agents solve the context bloat problem. When the AI needs to investigate something (find a user's shard, explore model relationships, search code), those intermediate tool calls can inflate the main conversation to 90K+ tokens, causing the LLM to cut corners. Sub-agents fork the investigation into a separate LLM conversation and return only a concise summary.
+
+The AI decides when to use sub-agents via the `delegate_task` tool. It can target a custom agent by name or use a general-purpose investigation.
+
+### Custom Agents
+
+Define agents as markdown files in `.rails_console_ai/agents/`:
+
+```markdown
+---
+name: Find shard
+description: Given a user ID, determines which database shard they are on
+max_rounds: 5
+tools:
+  - execute_code
+  - recall_memory
+---
+
+You are a shard finder for a sharded Rails application.
+
+Steps:
+1. Find the user: User.find(id)
+2. Check user.shard
+3. Report: "User {username} (ID {id}) is on shard {shard}."
+```
+
+**Frontmatter fields:**
+
+| Field | Required | Description |
+|-------|----------|-------------|
+| `name` | yes | Agent name (shown in system prompt, used with `delegate_task`) |
+| `description` | yes | One-line description of what this agent does |
+| `max_rounds` | no | Max tool-use rounds (default: `sub_agent_max_rounds` config, default 15) |
+| `model` | no | Model override (e.g. use Haiku for simple lookups) |
+| `tools` | no | Array of tool names to include (default: all sub-agent tools) |
+
+The markdown body becomes additional system prompt instructions for the sub-agent.
+
+### How It Works
+
+1. Agent summaries appear in the AI's system prompt under `## Agents`
+2. The AI calls `delegate_task(task: "find user 56653's shard", agent: "Find shard")`
+3. A sub-agent spins up with its own context, tools, and provider
+4. It runs the investigation (up to `max_rounds` tool calls)
+5. The main conversation receives only: `"Sub-agent result: User 56653 is on shard 5."`
+
+Sub-agents have access to read-only memory tools (`recall_memory`, `recall_memories`), code execution (`execute_code`), and all schema/code introspection tools. They cannot ask the user questions, write memories, or spawn further sub-agents.
+
+### Configuration
+
+```ruby
+RailsConsoleAi.configure do |config|
+  config.sub_agent_max_rounds = 15         # default max rounds per sub-agent
+  config.sub_agent_model = nil             # nil = same model as main conversation
+  # config.sub_agent_model = 'claude-haiku-4-5-20251001'  # use a cheaper model
+end
+```
+
 ## Safety Guards
 
 Safety guards prevent AI-generated code from causing side effects. When a guard blocks an operation, the user is prompted to re-run with safe mode disabled.

data/lib/rails_console_ai/agent_loader.rb

@@ -0,0 +1,126 @@
+require 'yaml'
+
+module RailsConsoleAi
+  class AgentLoader
+    AGENTS_DIR = 'agents'
+    BUILTIN_DIR = File.expand_path('../agents', __FILE__)
+
+    def initialize(storage = nil)
+      @storage = storage || RailsConsoleAi.storage
+    end
+
+    def load_all_agents
+      agents = load_builtin_agents
+      app_agents = load_app_agents
+      # App-specific agents override built-ins with the same name
+      app_names = app_agents.map { |a| a['name'].to_s.downcase }.to_set
+      agents.reject! { |a| app_names.include?(a['name'].to_s.downcase) }
+      agents.concat(app_agents)
+      agents
+    rescue => e
+      RailsConsoleAi.logger.warn("RailsConsoleAi: failed to load agents: #{e.message}")
+      []
+    end
+
+    def agent_summaries
+      agents = load_all_agents
+      return nil if agents.empty?
+
+      agents.map { |a|
+        "- **#{a['name']}**: #{a['description']}"
+      }
+    end
+
+    def find_agent(name)
+      agents = load_all_agents
+      agents.find { |a| a['name'].to_s.downcase == name.to_s.downcase }
+    end
+
+    def save_agent(name:, description:, body:, max_rounds: nil, model: nil, tools: nil)
+      key = agent_key(name)
+      existing = find_agent(name)
+
+      frontmatter = {
+        'name' => name,
+        'description' => description
+      }
+      frontmatter['max_rounds'] = max_rounds if max_rounds
+      frontmatter['model'] = model if model
+      frontmatter['tools'] = Array(tools) if tools && !tools.empty?
+
+      content = "---\n#{YAML.dump(frontmatter).sub("---\n", '').strip}\n---\n\n#{body}\n"
+      @storage.write(key, content)
+
+      path = @storage.respond_to?(:root_path) ? File.join(@storage.root_path, key) : key
+      if existing
+        "Agent updated: \"#{name}\" (#{path})"
+      else
+        "Agent created: \"#{name}\" (#{path})"
+      end
+    rescue Storage::StorageError => e
+      "FAILED to save agent (#{e.message})."
+    end
+
+    def delete_agent(name:)
+      key = agent_key(name)
+      unless @storage.exists?(key)
+        found = load_all_agents.find { |a| a['name'].to_s.downcase == name.to_s.downcase }
+        return "No agent found: \"#{name}\"" unless found
+        key = agent_key(found['name'])
+      end
+
+      agent = load_agent(key)
+      @storage.delete(key)
+      "Agent deleted: \"#{agent ? agent['name'] : name}\""
+    rescue Storage::StorageError => e
+      "FAILED to delete agent (#{e.message})."
+    end
+
+    private
+
+    def load_builtin_agents
+      return [] unless File.directory?(BUILTIN_DIR)
+      Dir.glob(File.join(BUILTIN_DIR, '*.md')).sort.filter_map do |path|
+        content = File.read(path)
+        agent = parse_agent(content)
+        agent['builtin'] = true if agent
+        agent
+      end
+    rescue => e
+      RailsConsoleAi.logger.debug("RailsConsoleAi: failed to load built-in agents: #{e.message}")
+      []
+    end
+
+    def load_app_agents
+      keys = @storage.list("#{AGENTS_DIR}/*.md")
+      keys.filter_map { |key| load_agent(key) }
+    rescue => e
+      []
+    end
+
+    def agent_key(name)
+      slug = name.downcase.strip
+        .gsub(/[^a-z0-9\s-]/, '')
+        .gsub(/[\s]+/, '-')
+        .gsub(/-+/, '-')
+        .sub(/^-/, '').sub(/-$/, '')
+      "#{AGENTS_DIR}/#{slug}.md"
+    end
+
+    def load_agent(key)
+      content = @storage.read(key)
+      return nil if content.nil? || content.strip.empty?
+      parse_agent(content)
+    rescue => e
+      RailsConsoleAi.logger.warn("RailsConsoleAi: failed to load agent #{key}: #{e.message}")
+      nil
+    end
+
+    def parse_agent(content)
+      return nil unless content =~ /\A---\s*\n(.*?\n)---\s*\n(.*)/m
+      frontmatter = YAML.safe_load($1, permitted_classes: [Time, Date]) || {}
+      body = $2.strip
+      frontmatter.merge('body' => body)
+    end
+  end
+end

data/lib/rails_console_ai/agents/explore-data.md

@@ -0,0 +1,24 @@
+---
+name: Explore data
+description: Run queries to investigate records, follow associations, and gather specific data
+max_rounds: 20
+---
+
+You are a data investigator for a Rails application. Your job is to run queries, check
+records, follow associations, and gather specific data points.
+
+Strategy:
+1. Use describe_model to understand the model's columns and associations before querying
+2. Use execute_code to run ActiveRecord queries
+3. Follow associations to find related data (e.g., user -> booking_pages -> booking_requests)
+4. Use model methods to get computed values — NEVER fabricate URLs, tokens, or derived data manually
+
+Rules:
+- Always use describe_model before querying a model you haven't seen yet.
+- Prefer ActiveRecord over raw SQL.
+- When you find a model has methods that return what you need (check with
+  `.methods.grep(/pattern/)`), USE those methods instead of constructing values yourself.
+- Include specific IDs, values, and counts in your findings.
+
+End with a concise, factual summary of what you found. Include specific record IDs, values,
+and relationships. Keep your summary under 300 words.

data/lib/rails_console_ai/agents/investigate-code.md

@@ -0,0 +1,23 @@
+---
+name: Investigate code
+description: Search and read source code to understand how a feature, pattern, or method works
+max_rounds: 20
+---
+
+You are a code investigator for a Rails application. Your job is to search the codebase,
+read relevant files, and understand how a specific feature or pattern works.
+
+Strategy:
+1. Start with search_code to find where the relevant code lives
+2. Use read_file to examine the key files (read specific sections, not entire files)
+3. Follow the chain: find the method, check what it calls, trace the data flow
+4. Use describe_model if database models are involved
+
+Rules:
+- Be thorough but efficient. Search first, then read specific files.
+- If a search returns too many results, narrow it with a more specific query or directory.
+- Look at actual method implementations, not just where things are called.
+- Note file paths and line numbers in your findings.
+
+End with a concise summary of your findings: key files, classes, methods, and how they
+connect. Keep your summary under 500 words.

data/lib/rails_console_ai/channel/console.rb

@@ -127,8 +127,8 @@ module RailsConsoleAi
 
       # --- Omitted output tracking (shared with Executor) ---
 
-      MAX_DISPLAY_LINES = 10
-      MAX_DISPLAY_CHARS = 2000
+      MAX_DISPLAY_LINES = 20
+      MAX_DISPLAY_CHARS = 4000
 
       def init_omitted_tracking
         @omitted_outputs = {}
@@ -230,7 +230,16 @@ module RailsConsoleAi
           @interactive_console_capture.write("ai> #{input}\n")
           @engine.log_interactive_turn
 
-          status = @engine.send_and_execute
+          expected_stdout = $stdout
+          begin
+            status = @engine.send_and_execute
+          rescue Interrupt
+            $stdout = expected_stdout
+            $stdout.puts "\n\e[33m  Cancelled.\e[0m"
+            @engine.pop_last_message
+            @engine.log_interactive_turn
+            next
+          end
           if status == :interrupted
             @engine.pop_last_message
             @engine.log_interactive_turn
@@ -413,6 +422,7 @@ module RailsConsoleAi
         @real_stdout.puts "\e[2m    /debug       Toggle debug summaries (context stats, cost per call)\e[0m"
         @real_stdout.puts "\e[2m    /retry       Re-execute the last code block\e[0m"
         @real_stdout.puts "\e[2m    > code       Execute Ruby directly (skip LLM)\e[0m"
+        @real_stdout.puts "\e[2m    Ctrl-C       Cancel the current operation\e[0m"
         @real_stdout.puts "\e[2m    exit/quit    Leave interactive mode\e[0m"
       end
 

data/lib/rails_console_ai/channel/slack.rb

@@ -42,7 +42,7 @@ module RailsConsoleAi
           post(stripped)
         else
           @output_log.write("#{stripped}\n")
-          STDOUT.puts "#{@log_prefix} (status) #{stripped}"
+          log_prefixed("(status)", stripped)
         end
       end
 
@@ -56,7 +56,7 @@ module RailsConsoleAi
 
       def display_tool_call(text)
         @output_log.write("-> #{text}\n")
-        STDOUT.puts "#{@log_prefix} -> #{text}"
+        log_prefixed("->", text)
       end
 
       def display_code(code)
@@ -161,7 +161,7 @@ module RailsConsoleAi
       def post(text)
         return if text.nil? || text.strip.empty?
         @output_log.write("#{text}\n")
-        STDOUT.puts "#{@log_prefix} >> #{text}"
+        log_prefixed(">>", text)
         @slack_bot.send(:post_message,
           channel: @channel_id,
           thread_ts: @thread_ts,
@@ -171,6 +171,10 @@ module RailsConsoleAi
         RailsConsoleAi.logger.error("Slack post failed: #{e.message}")
       end
 
+      def log_prefixed(tag, text)
+        text.each_line { |line| STDOUT.puts "#{@log_prefix} #{tag} #{line.rstrip}" }
+      end
+
       def strip_ansi(text)
         text.to_s.gsub(ANSI_REGEX, '')
       end

data/lib/rails_console_ai/channel/sub_agent.rb

@@ -0,0 +1,84 @@
+require 'rails_console_ai/channel/base'
+
+module RailsConsoleAi
+  module Channel
+    class SubAgent < Base
+      def initialize(parent_channel:, task_label: nil)
+        @parent = parent_channel
+        @label = task_label
+      end
+
+      def display(text)
+        # Swallowed — sub-agent final text is returned as tool result, not displayed
+      end
+
+      def display_thinking(text)
+        @parent.display_thinking(text)
+      end
+
+      def display_status(text)
+        @parent.display_status(text)
+      end
+
+      def display_warning(text)
+        @parent.display_warning(text)
+      end
+
+      def display_error(text)
+        @parent.display_error(text)
+      end
+
+      def display_tool_call(text)
+        @parent.display_tool_call(text)
+      end
+
+      def display_code(code)
+        # Swallowed — sub-agent auto-executes, no need to show code
+      end
+
+      def display_result_output(output)
+        @parent.display_result_output(output)
+      end
+
+      def display_result(result)
+        # Swallowed — sub-agent return values aren't useful to show
+      end
+
+      def prompt(text)
+        '(sub-agent cannot ask user)'
+      end
+
+      def confirm(text)
+        'y'
+      end
+
+      def user_identity
+        @parent.user_identity
+      end
+
+      def mode
+        'sub_agent'
+      end
+
+      def cancelled?
+        @parent.cancelled?
+      end
+
+      def supports_danger?
+        false  # Sub-agents must never silently bypass safety guards
+      end
+
+      def supports_editing?
+        false
+      end
+
+      def wrap_llm_call(&block)
+        yield
+      end
+
+      def system_instructions
+        nil
+      end
+    end
+  end
+end

data/lib/rails_console_ai/configuration.rb

@@ -31,7 +31,9 @@ module RailsConsoleAi
                   :code_search_paths,
                   :channels,
                   :bypass_guards_for_methods,
-                  :user_extra_info
+                  :user_extra_info,
+                  :sub_agent_max_rounds,
+                  :sub_agent_model
 
     def initialize
       @provider     = :anthropic
@@ -64,6 +66,8 @@ module RailsConsoleAi
       @channels = {}
       @bypass_guards_for_methods = []
       @user_extra_info = {}
+      @sub_agent_max_rounds = 15
+      @sub_agent_model = nil
     end
 
     def resolve_user_extra_info(username)
@@ -135,7 +139,7 @@ module RailsConsoleAi
       end
 
       if allow
-        Array(allow).each { |key| safety_guards.allow(guard_name, key) }
+        Array(allow).each { |key| safety_guards.allow_global(guard_name, key) }
       end
     end
 

data/lib/rails_console_ai/context_builder.rb

@@ -20,6 +20,7 @@ module RailsConsoleAi
       parts << guide_context
       parts << trusted_methods_context
       parts << skills_context
+      parts << agents_context
       parts << user_extra_info_context
       parts << pinned_memory_context
       parts << memory_context
@@ -93,6 +94,14 @@ module RailsConsoleAi
           a new skill with a step-by-step recipe. Skills are procedures (how to do X); memories
           are facts (what you learned about X). Do NOT use save_memory when asked to create a skill.
 
+        You have a delegate_task tool to spawn sub-agents for investigation tasks:
+        - Use delegate_task when a task requires multiple tool calls to investigate
+          (e.g., finding a user's shard, exploring how a feature works in the code,
+          gathering data across multiple models).
+        - The sub-agent runs in a separate context and returns only a concise summary.
+        - This keeps your conversation context small and efficient.
+        - If a custom agent is available for the task (see Agents section), specify it by name.
+
         RULES:
         - Give ONE concise answer. Do not offer multiple alternatives or variations.
         - For multi-step tasks, use execute_plan to break the work into small, clear steps.
@@ -145,6 +154,19 @@ module RailsConsoleAi
       nil
     end
 
+    def agents_context
+      require 'rails_console_ai/agent_loader'
+      summaries = RailsConsoleAi::AgentLoader.new.agent_summaries
+      return nil if summaries.nil? || summaries.empty?
+
+      lines = ["## Agents (use delegate_task tool to invoke)"]
+      lines.concat(summaries)
+      lines.join("\n")
+    rescue => e
+      RailsConsoleAi.logger.debug("RailsConsoleAi: agents context failed: #{e.message}")
+      nil
+    end
+
     def user_extra_info_context
       info = @config.resolve_user_extra_info(@user_name)
       return nil if info.nil? || info.strip.empty?

data/lib/rails_console_ai/conversation_engine.rb

@@ -3,8 +3,8 @@ module RailsConsoleAi
     attr_reader :history, :total_input_tokens, :total_output_tokens,
                 :interactive_session_id, :session_name
 
-    LARGE_OUTPUT_THRESHOLD = 10_000      # chars — truncate tool results larger than this immediately
-    LARGE_OUTPUT_PREVIEW_CHARS = 8_000   # chars — how much of the output the LLM sees upfront
+    LARGE_OUTPUT_THRESHOLD = 20_000      # chars — truncate tool results larger than this immediately
+    LARGE_OUTPUT_PREVIEW_CHARS = 16_000  # chars — how much of the output the LLM sees upfront
     LOOP_WARN_THRESHOLD = 3              # same tool+args repeated → inject warning
     LOOP_BREAK_THRESHOLD = 5             # same tool+args repeated → break loop
 
@@ -882,6 +882,17 @@ module RailsConsoleAi
             @channel.display_status("     #{preview}#{cached_tag}")
           end
 
+          # Aggregate sub-agent token usage into parent's cost tracking
+          if tc[:name] == 'delegate_task' && tools.last_sub_agent_usage
+            sa = tools.last_sub_agent_usage
+            @total_input_tokens += sa[:input] || 0
+            @total_output_tokens += sa[:output] || 0
+            if sa[:model]
+              @token_usage[sa[:model]][:input] += sa[:input] || 0
+              @token_usage[sa[:model]][:output] += sa[:output] || 0
+            end
+          end
+
           if RailsConsoleAi.configuration.debug
             $stderr.puts "\e[35m[debug] tool result (#{tool_result.to_s.length} chars)\e[0m"
           end
@@ -1033,6 +1044,11 @@ module RailsConsoleAi
       when 'execute_plan'
         steps = args['steps']
         steps ? "(#{steps.length} steps)" : ''
+      when 'delegate_task'
+        task_preview = args['task'].to_s[0, 100]
+        task_preview += '...' if args['task'].to_s.length > 100
+        agent = args['agent'] ? ", agent: \"#{args['agent']}\"" : ''
+        "(\"#{task_preview}\"#{agent})"
       else ''
       end
     end
@@ -1104,6 +1120,9 @@ module RailsConsoleAi
       when 'execute_plan'
         steps_done = result.scan(/^Step \d+/).length
         steps_done > 0 ? "#{steps_done} steps executed" : truncate(result, 80)
+      when 'delegate_task'
+        # Show the full sub-agent result — this is the whole point of delegation
+        result
       else
         truncate(result, 80)
       end

data/lib/rails_console_ai/executor.rb

@@ -100,13 +100,18 @@ module RailsConsoleAi
       @last_safety_exception = nil
       captured_output = StringIO.new
       old_stdout = $stdout
-      # When a channel is present it handles display (with truncation), so capture only.
-      # Without a channel, tee so output appears live on the terminal.
-      $stdout = if @channel
-                  captured_output
-                else
-                  TeeIO.new(old_stdout, captured_output)
-                end
+      # Three capture strategies:
+      # 1. Slack mode (PrefixedIO active): thread-local capture to avoid cross-thread pollution
+      # 2. Console mode (channel present): capture-only, channel.display_result_output shows it after
+      # 3. No channel (tests/one-shot): TeeIO so output appears live AND is captured
+      use_thread_local = defined?(RailsConsoleAi::PrefixedIO) && $stdout.is_a?(RailsConsoleAi::PrefixedIO)
+      if use_thread_local
+        Thread.current[:capture_io] = captured_output
+      elsif @channel
+        $stdout = captured_output
+      else
+        $stdout = TeeIO.new(old_stdout, captured_output)
+      end
 
       RailsConsoleAi::SafetyError.clear!
 
@@ -114,7 +119,7 @@ module RailsConsoleAi
         binding_context.eval(code, "(rails_console_ai)", 1)
       end
 
-      $stdout = old_stdout
+      restore_stdout(use_thread_local, old_stdout)
 
       # Check if a SafetyError was raised but swallowed by a rescue inside the eval'd code
       if (swallowed = RailsConsoleAi::SafetyError.last_raised)
@@ -136,8 +141,12 @@ module RailsConsoleAi
 
       @last_output = captured_output.string
       result
+    rescue Interrupt
+      restore_stdout(use_thread_local, old_stdout)
+      @last_output = captured_output&.string
+      raise
     rescue RailsConsoleAi::SafetyError => e
-      $stdout = old_stdout if old_stdout
+      restore_stdout(use_thread_local, old_stdout)
       RailsConsoleAi::SafetyError.clear!
       @last_error = "SafetyError: #{e.message}"
       @last_safety_error = true
@@ -146,13 +155,13 @@ module RailsConsoleAi
       @last_output = captured_output&.string
       nil
     rescue SyntaxError => e
-      $stdout = old_stdout if old_stdout
+      restore_stdout(use_thread_local, old_stdout)
       @last_error = "SyntaxError: #{e.message}"
       log_execution_error(@last_error)
       @last_output = nil
       nil
     rescue => e
-      $stdout = old_stdout if old_stdout
+      restore_stdout(use_thread_local, old_stdout)
       # Check if a SafetyError is wrapped (e.g. ActiveRecord::StatementInvalid wrapping our error)
       if safety_error?(e)
         safety_exc = extract_safety_exception(e)
@@ -333,6 +342,14 @@ module RailsConsoleAi
 
     private
 
+    def restore_stdout(use_thread_local, old_stdout)
+      if use_thread_local
+        Thread.current[:capture_io] = nil
+      else
+        $stdout = old_stdout if old_stdout
+      end
+    end
+
     def danger_allowed?
       @channel.nil? || @channel.supports_danger?
     end

data/lib/rails_console_ai/prefixed_io.rb

@@ -5,6 +5,9 @@ module RailsConsoleAi
     end
 
     def write(str)
+      if (capture = Thread.current[:capture_io])
+        return capture.write(str)
+      end
       prefix = Thread.current[:log_prefix]
       if prefix && str.is_a?(String) && !str.strip.empty?
         prefixed = str.gsub(/^(?=.)/, "#{prefix} ")
@@ -15,6 +18,9 @@ module RailsConsoleAi
     end
 
     def puts(*args)
+      if (capture = Thread.current[:capture_io])
+        return capture.puts(*args)
+      end
       prefix = Thread.current[:log_prefix]
       if prefix
         args = [""] if args.empty?
@@ -32,6 +38,9 @@ module RailsConsoleAi
     end
 
     def print(*args)
+      if (capture = Thread.current[:capture_io])
+        return capture.print(*args)
+      end
       @io.print(*args)
     end