Lightweight ralph loop orchestrator for claude code

mini-arborist.rb

#!/usr/bin/env ruby
# frozen_string_literal: true

# mini-arborist: A lightweight iterative agent orchestrator
#
# Runs Claude Code agents in a loop where each agent performs a small, focused
# piece of work then exits.
#
# Usage:
#   Create a plan with claude:
#
#     mini-arborist plan my-feature --smart
#
#   When it's ready ask it to run mini-arborist to execute the plan and it'll iteratively spawn agents to implement it
#
require 'json'
require 'fileutils'
require 'shellwords'

DEFAULT_ITERATIONS = 5

def session_dir(plan_name = nil)
  plan_name ? ".mini-arborist/#{plan_name}" : '.mini-arborist'
end

def status_file(plan_name = nil)
  "#{session_dir(plan_name)}/status"
end

def scratch_file(plan_name = nil)
  "#{session_dir(plan_name)}/scratch.md"
end

def todo_file(plan_name = nil)
  "#{session_dir(plan_name)}/todo.md"
end

def plan_file(plan_name = nil)
  plan_name ? "#{session_dir(plan_name)}/plan.md" : 'plan.md'
end

def task_prompt(plan_name = nil, smart_mode = false)
  plan = plan_file(plan_name)
  scratch = scratch_file(plan_name)
  todo = todo_file(plan_name)
  status = status_file(plan_name)

  status_signals = if smart_mode
    <<~SIGNALS
      ## Status Signals

      Write ONE of these to `#{status}`:

      - `CONTINUE` or `CONTINUE:sonnet` - More work to do. Next agent uses sonnet (for coding/implementation).
      - `CONTINUE:opus` - More work to do. Next agent uses opus (for planning/analysis/exploration).
      - `DONE` - All work in #{plan} is complete and verified.
      - `BLOCKED: <reason>` - Need human input. Explain what you need.

      **IMPORTANT**: You choose the model for the NEXT agent only. Do NOT tell future agents what
      status to write in `#{scratch}` - each agent decides independently based on their work.
    SIGNALS
  else
    <<~SIGNALS
      ## Status Signals

      Write ONE of these to `#{status}`:

      - `CONTINUE` - More work to do. Next agent will be launched.
      - `DONE` - All work in #{plan} is complete and verified.
      - `BLOCKED: <reason>` - Need human input. Explain what you need.
    SIGNALS
  end

  <<~TASK
    # Agent Methodology

    You are one agent in a series. Do a small, focused piece of work, then exit.

    ## Communication Files

    - `#{plan}` - The goal (written by human, read-only)
    - `#{scratch}` - Instructions from the previous agent. Execute on these, then rewrite for the next agent.
    - `#{todo}` - Persistent checklist. Update as you complete items.
    - `#{status}` - Signal when done (see below)

    ## Iteration Flow

    1. Read `#{plan}` for overall context and design decisions
    2. Read `#{scratch}` for instructions from the last agent (empty = you're first)
    3. Do the work described (or start planning if first agent)
    4. Clear `#{scratch}` and write instructions for the next agent
    5. Update `#{todo}` with progress
    6. Write status to `#{status}`
    7. Exit immediately

    ## Key Rules

    - **If you write code, don't do extensive planning** - note what needs planning in `#{scratch}`
    - **If you do planning, don't edit code files** - write plans to `#{scratch}` for the next agent
    - Each iteration should be small and focused
    - Delegate complex work to future agents by writing clear instructions in `#{scratch}`

    #{status_signals}
    ## Git

    After code changes, commit with an appropriate message.

    ## Start

    Read #{plan} and #{scratch}, then begin your iteration.
  TASK
end

def initial_scratch(plan_name = nil)
  plan = plan_file(plan_name)
  todo = todo_file(plan_name)
  <<~SCRATCH
    First agent - read #{plan} and create an implementation plan. Break down the work into small steps. Write the plan to #{todo}, then write instructions for the next agent here.
  SCRATCH
end

def initial_todo
  <<~TODO
    # Todo

    (First agent will populate this)
  TODO
end

def planning_prompt(plan_name = nil, smart_mode = false)
  plan = plan_file(plan_name)
  scratch = scratch_file(plan_name)
  todo = todo_file(plan_name)
  run_cmd = plan_name ? "mini-arborist -f #{plan_name}" : "mini-arborist"
  run_cmd += " --smart" if smart_mode

  <<~PROMPT
    # Planning Agent

    You are a planning agent. Explore, analyze, and plan - do NOT implement.

    ## File Rules (CRITICAL)

    You may ONLY edit these files:
    - `#{plan}` - Your analysis and implementation plans
    - `#{todo}` - Task breakdown for implementation agents

    DO NOT edit any other files. If you need to test an idea, describe it in #{plan}
    or ask the user to run it. Implementation is done by later agents, not you.

    ## Allowed Actions
    - Read any files to understand the codebase
    - Run read-only commands (ls, grep, git log, gh, etc.)
    - Run exploratory commands to gather information
    - Write analysis and plans to the three files above

    ## Your Job
    Explore the codebase, then write a high-level strategic plan to #{plan},
    and task list to #{todo}.

    Focus on:
    - Overall approach and strategy
    - Key architectural decisions
    - High-level steps and sequencing
    - Important context and constraints

    Do NOT specify technical implementation details like:
    - Specific function names or variable names
    - Exact code changes or file modifications
    - Detailed algorithms or data structures

    The iteration loop will figure out the technical details. Your job is to provide
    strategic direction and context, not detailed implementation instructions.

    ## Running Implementation

    When planning is complete, run the implementation agents:

    ```
    #{run_cmd} -i <iterations> -n
    ```

    Guidelines:
    - Use at least `-i 5` iterations (more for larger tasks)
    - Estimate iterations based on the number of distinct implementation steps
    - Set Bash timeout to at least 2 minutes per iteration (e.g., 5 iterations = 600000ms timeout)
    - Always use `-n`

    After the implementation run completes, perform a code review of all changes
    made by the implementation agents. Check for correctness, edge cases, and
    adherence to the plan.
  PROMPT
end

def setup_files(plan_name = nil)
  # Create session directory first (plan file is inside it for named plans)
  FileUtils.mkdir_p(session_dir(plan_name))

  plan = plan_file(plan_name)
  scratch = scratch_file(plan_name)
  todo = todo_file(plan_name)
  status = status_file(plan_name)

  # Check for plan file
  unless File.exist?(plan)
    puts "Error: #{plan} not found"
    puts "Create #{plan} describing what you want to build, then run mini-arborist"
    exit 1
  end

  # Create scratch.md if it doesn't exist
  unless File.exist?(scratch)
    File.write(scratch, initial_scratch(plan_name))
    puts "Created #{scratch}"
  end

  # Create todo.md if it doesn't exist
  unless File.exist?(todo)
    File.write(todo, initial_todo)
    puts "Created #{todo}"
  end

  # Clear status file
  File.write(status, 'CONTINUE')
end

def read_status(plan_name = nil)
  status = status_file(plan_name)
  return ['CONTINUE', nil] unless File.exist?(status)

  content = File.read(status).strip
  status_part, model_hint = content.split(':', 2)
  [status_part, model_hint]
end

def parse_model_hint(model_hint)
  case model_hint&.downcase&.strip
  when 'sonnet', nil, ''
    'sonnet'
  when 'opus'
    'opus'
  else
    # Unknown/invalid -> fail safe to opus
    'opus'
  end
end

def format_tool_use(name, input)
  case name
  when 'Read'
    "📖 Read: #{input['file_path']}"
  when 'Edit'
    "✏️  Edit: #{input['file_path']}"
  when 'Write'
    "📝 Write: #{input['file_path']}"
  when 'Bash'
    cmd = input['command']&.split("\n")&.first || ''
    cmd = "#{cmd[0, 80]}..." if cmd.length > 80
    "💻 Bash: #{cmd}"
  when 'Glob'
    "🔍 Glob: #{input['pattern']}"
  when 'Grep'
    "🔍 Grep: #{input['pattern']}"
  when 'TodoWrite'
    "📋 Todo"
  when 'Task'
    "🤖 Task: #{input['description']}"
  else
    "🔧 #{name}"
  end
end

def format_assistant_message(data)
  message = data['message']
  return nil unless message

  parts = []
  (message['content'] || []).each do |block|
    case block['type']
    when 'text'
      text = block['text']&.strip
      parts << text if text && !text.empty?
    when 'tool_use'
      parts << format_tool_use(block['name'], block['input'] || {})
    end
  end

  parts.empty? ? nil : parts.join("\n")
end

def format_line(line)
  stripped = line.strip
  return nil if stripped.empty?

  data = JSON.parse(stripped)
  case data['type']
  when 'assistant'
    format_assistant_message(data)
  when 'result'
    result = data['result']
    result ? "✅ Done: #{result[0, 200]}#{'...' if result && result.length > 200}" : nil
  when 'error'
    "❌ Error: #{data['error'] || data['message'] || stripped}"
  end
rescue JSON::ParserError
  # Not JSON - probably an error message, show it raw
  stripped
end

def stream_formatter_available?
  @stream_formatter_available ||= system('which claude-stream-format > /dev/null 2>&1')
end

def run_claude(iteration, max_iterations, plan_name = nil, smart_mode = false, model = nil)
  puts "\n#{'=' * 60}"
  puts "🚀 Iteration #{iteration}/#{max_iterations}"
  puts "   Model: #{model}" if model
  puts '=' * 60

  # Clear status before run
  File.write(status_file(plan_name), '')

  if stream_formatter_available?
    run_claude_with_external_formatter(plan_name, smart_mode, model)
  else
    run_claude_with_builtin_formatter(plan_name, smart_mode, model)
  end
end

def run_claude_with_external_formatter(plan_name = nil, smart_mode = false, model = nil)
  prompt = task_prompt(plan_name, smart_mode)
  model_arg = model ? "--model #{Shellwords.escape(model)}" : ''
  cmd = "claude -p #{Shellwords.escape(prompt)} " \
        "#{model_arg} " \
        '--permission-mode acceptEdits --verbose --output-format stream-json 2>&1 | ' \
        'claude-stream-format'

  system(cmd)
  $?.success?
end

def run_claude_with_builtin_formatter(plan_name = nil, smart_mode = false, model = nil)
  prompt = task_prompt(plan_name, smart_mode)
  cmd = [
    'claude',
    '-p', prompt,
    '--permission-mode', 'acceptEdits',
    '--verbose',
    '--output-format', 'stream-json'
  ]
  cmd.concat(['--model', model]) if model

  IO.popen(cmd, err: [:child, :out]) do |io|
    io.each_line do |line|
      formatted = format_line(line)
      puts formatted if formatted
    end
  end

  status = $?
  unless status.success?
    puts "Exit code: #{status.exitstatus}"
  end
  status.success?
end

def run_loop(max_iterations, plan_name = nil, smart_mode = false)
  iteration = 0
  scratch = scratch_file(plan_name)
  run_cmd = plan_name ? "mini-arborist -f #{plan_name}" : "mini-arborist"

  loop do
    iteration += 1

    if iteration > max_iterations
      puts "\n⚠️  Reached maximum iterations (#{max_iterations})"
      puts "Run #{run_cmd} again to continue, or check #{scratch} for progress"
      break
    end

    # First iteration uses opus in smart mode, otherwise use hint from previous agent
    model = if smart_mode
              iteration == 1 ? 'opus' : parse_model_hint(read_status(plan_name)[1])
            else
              nil
            end

    success = run_claude(iteration, max_iterations, plan_name, smart_mode, model)

    unless success
      puts "\n❌ Claude exited with error"
      break
    end

    status, model_hint = read_status(plan_name)

    case status
    when 'DONE'
      puts "\n✅ Agent signaled DONE"
      puts "Check the results and #{scratch} for summary"
      break
    when 'CONTINUE', ''
      puts "\n🔄 Agent signaled CONTINUE"
      puts "   Next model: #{parse_model_hint(model_hint)}" if smart_mode
      # Loop continues
    when /^BLOCKED/
      puts "\n🚫 Agent is blocked: #{status.sub('BLOCKED:', '').strip}"
      puts "Address the issue and run #{run_cmd} again"
      break
    when /^ERROR/
      puts "\n❌ Agent error: #{status.sub('ERROR:', '').strip}"
      break
    else
      puts "\n⚠️  Unknown status: #{status}"
      puts "Treating as CONTINUE..."
    end
  end
end

def reset_for_new_task(plan_name = nil)
  scratch = scratch_file(plan_name)
  status = status_file(plan_name)

  puts "Resetting for new task..."
  File.write(scratch, initial_scratch(plan_name))
  File.write(status, 'CONTINUE')
  puts "Cleared #{scratch} - first agent will plan from #{plan_name}"
end

def show_help
  puts <<~HELP
    Usage: mini-arborist [options]

    Options:
      -f, --file NAME      Use NAME.md as the plan file (session state in .mini-arborist/NAME/)
      -i, --iterations N   Maximum iterations (default: #{DEFAULT_ITERATIONS})
      -n, --new            Reset state to start a new task (reads plan file fresh)
      -b, --branch         Start work in a new git branch (auto-generates name from plan)
      -P, --plan [NAME]    Interactive planning mode (explore codebase, no code changes)
      -s, --smart          Smart model selection (opus for planning, sonnet for coding)
      -h, --help           Show this help

    Commands:
      ls                   List available plans
      show [NAME]          Show plan file with glow
      rm NAME              Remove a plan and its state
      run NAME [opts]      Run a plan (alias for -f NAME, supports -i, --branch, etc.)
      plan NAME [-- args]  Interactive planning mode (alias for -P NAME)
                           Arguments after -- are passed to claude command

    Workflow:
      1. Create plan.md (or NAME.md) with your goal
      2. Run: mini-arborist [-f NAME]
      3. Re-run if it hits iteration limit
      4. For new work: update plan file, then run: mini-arborist [-f NAME] --new

    Examples:
      mini-arborist                    # Uses plan.md, state in .mini-arborist/
      mini-arborist -f auth-feature    # Uses auth-feature.md, state in .mini-arborist/auth-feature/
      mini-arborist -P auth-feature    # Planning mode for auth-feature.md
      mini-arborist --smart            # Smart model selection (first iteration uses opus)
      mini-arborist plan my-feature -- --model opus  # Planning mode with custom claude args
      mini-arborist ls                 # List all plans
      mini-arborist rm auth-feature    # Remove auth-feature plan
  HELP
end

def list_plans
  base_dir = '.mini-arborist'

  unless Dir.exist?(base_dir)
    puts "No plans found (#{base_dir} does not exist)"
    return
  end

  plans = Dir.children(base_dir)
    .select { |name| File.directory?(File.join(base_dir, name)) }
    .sort

  if plans.empty?
    puts "No plans found in #{base_dir}"
    return
  end

  puts "Available plans:"
  plans.each do |name|
    plan_path = plan_file(name)
    status_path = File.join(base_dir, name, 'status')
    status = File.exist?(status_path) ? File.read(status_path).strip : 'UNKNOWN'
    plan_exists = File.exist?(plan_path) ? '✓' : '✗'
    puts "  #{name} (plan: #{plan_exists}, status: #{status})"
  end
end

def remove_plan(name)
  if name.nil? || name.empty?
    puts "Error: rm requires a plan name"
    puts "Usage: mini-arborist rm NAME"
    exit 1
  end

  name = name.delete_suffix('.md')
  state_dir = File.join('.mini-arborist', name)

  unless Dir.exist?(state_dir)
    puts "Error: Plan '#{name}' not found (#{state_dir} does not exist)"
    exit 1
  end

  FileUtils.rm_rf(state_dir)
  puts "Removed #{state_dir}"
end

def show_plan(name)
  plan_name = name&.delete_suffix('.md')
  plan = plan_file(plan_name)

  unless File.exist?(plan)
    puts "Error: #{plan} not found"
    exit 1
  end

  exec('glow', '--pager', plan)
end

def parse_iterations
  ['-i', '--iterations'].each do |flag|
    idx = ARGV.index(flag)
    next unless idx

    value = ARGV[idx + 1]
    if value.nil? || value.start_with?('-')
      puts "Error: #{flag} requires a number"
      exit 1
    end
    return value.to_i
  end
  DEFAULT_ITERATIONS
end

def parse_branch
  ['-b', '--branch'].each do |flag|
    return true if ARGV.include?(flag)
  end
  nil
end

def parse_smart
  ['-s', '--smart'].each do |flag|
    return true if ARGV.include?(flag)
  end
  nil
end

def parse_plan_name
  ['-f', '--file'].each do |flag|
    idx = ARGV.index(flag)
    next unless idx

    value = ARGV[idx + 1]
    if value.nil? || value.start_with?('-')
      puts "Error: #{flag} requires a plan name"
      exit 1
    end
    # Strip .md extension if provided
    return value.delete_suffix('.md')
  end
  nil
end

def parse_plan_flag
  ['-P', '--plan'].each do |flag|
    idx = ARGV.index(flag)
    next unless idx

    # Check if next arg is a plan name (not another flag)
    value = ARGV[idx + 1]
    if value && !value.start_with?('-')
      # Strip .md extension if provided
      return value.delete_suffix('.md')
    end
    # Flag present but no name - return true to indicate default plan.md
    return true
  end
  nil
end

def branch_instruction(plan_name = nil)
  plan = plan_file(plan_name)
  <<~INSTRUCTION
    ## FIRST: Create a new branch

    Before doing anything else, create and switch to a new git branch.
    Derive a short, descriptive branch name from #{plan} (e.g., `feature/add-user-auth` or `fix/login-validation`).

  INSTRUCTION
end

def run_planning_mode(plan_name = nil, claude_args = [], smart_mode = false)
  # Create session directory first (plan file is inside it for named plans)
  FileUtils.mkdir_p(session_dir(plan_name))

  plan = plan_file(plan_name)

  # Create the plan file if it doesn't exist
  unless File.exist?(plan)
    File.write(plan, "# Plan\n\n(Describe your goal here)\n")
    puts "Created #{plan}"
  end

  setup_files(plan_name)

  puts "🌳 mini-arborist (planning mode)"
  puts "Plan file: #{plan}"
  puts "Working directory: #{Dir.pwd}"
  puts "Smart mode: enabled" if smart_mode
  puts "Launching interactive Claude session..."
  unless claude_args.empty?
    puts "Additional args: #{claude_args.join(' ')}"
  end
  puts

  # Build claude command with additional args
  # Note: In planning mode with --smart, we don't pass --model to let Claude choose naturally
  # Restrict all file writes to only .mini-arborist directory
  cmd = [
    'claude',
    '--allowed-tools', 'Edit(./.mini-arborist/**) Write(./.mini-arborist/**) Edit(./.mini-arborist/*) Write(./.mini-arborist/*) Bash(mini-arborist:*)',
    '--permission-mode', 'dontAsk',
    '--append-system-prompt', planning_prompt(plan_name, smart_mode)
  ]
  cmd.concat(claude_args)

  exec(*cmd)
end

def main
  if ARGV.include?('--help') || ARGV.include?('-h')
    show_help
    exit 0
  end

  # Handle commands
  if ARGV[0] == 'ls'
    list_plans
    exit 0
  end

  if ARGV[0] == 'rm'
    remove_plan(ARGV[1])
    exit 0
  end

  if ARGV[0] == 'show'
    show_plan(ARGV[1])
    # exec never returns
  end

  # Handle 'plan NAME' command (alias for -P NAME)
  if ARGV[0] == 'plan'
    plan_name = ARGV[1]&.delete_suffix('.md')
    if plan_name.nil? || plan_name.empty?
      puts "Error: plan requires a plan name"
      puts "Usage: mini-arborist plan NAME [-- CLAUDE_ARGS...]"
      exit 1
    end

    # Parse smart mode before removing from ARGV
    smart_mode = parse_smart

    # Parse claude args (everything after --)
    separator_idx = ARGV.index('--')
    claude_args = separator_idx ? ARGV[(separator_idx + 1)..-1] : []

    run_planning_mode(plan_name, claude_args, smart_mode)
    # exec never returns
  end

  # Handle 'run NAME' command (alias for -f NAME, with additional options)
  if ARGV[0] == 'run'
    run_plan_name = ARGV[1]&.delete_suffix('.md')
    if run_plan_name.nil? || run_plan_name.empty? || run_plan_name.start_with?('-')
      puts "Error: run requires a plan name"
      puts "Usage: mini-arborist run NAME [-i N] [--branch] [--new]"
      exit 1
    end
    # Remove 'run' and NAME from ARGV, inject -f NAME for normal parsing
    ARGV.shift(2)
    ARGV.unshift('-f', run_plan_name)
  end

  # Check for planning mode first (can have optional plan name)
  plan_flag = parse_plan_flag
  if plan_flag
    # plan_flag is either a plan name or true (for default)
    plan_name = plan_flag == true ? nil : plan_flag
    smart_mode = parse_smart
    run_planning_mode(plan_name, [], smart_mode)
    # exec never returns
  end

  # Parse options
  max_iterations = parse_iterations
  branch = parse_branch
  smart_mode = parse_smart
  plan_name = parse_plan_name

  plan = plan_file(plan_name)
  scratch = scratch_file(plan_name)

  puts "🌳 mini-arborist"
  puts "Plan file: #{plan}"
  puts "Working directory: #{Dir.pwd}"
  puts "Smart mode: enabled" if smart_mode
  puts

  setup_files(plan_name)

  if ARGV.include?('--new') || ARGV.include?('-n')
    reset_for_new_task(plan_name)
  end

  # Prepend branch instruction if --branch flag is set
  if branch
    current_scratch = File.read(scratch)
    File.write(scratch, branch_instruction(plan_name) + current_scratch)
    puts "Will create new branch (auto-generated from plan)"
  end

  puts "\nStarting agent loop (max #{max_iterations} iterations)..."
  puts "Press Ctrl+C to stop\n"

  run_loop(max_iterations, plan_name, smart_mode)
end

# Handle Ctrl+C gracefully
trap('INT') do
  puts "\n\n⏹️  Interrupted by user"
  exit 0
end

main