@hanzo/dev/src/lib/code-act-agent.ts

Hanzo Dev - Meta AI development CLI that manages and runs all LLMs and CLI tools

import { FileEditor, ChunkLocalizer } from './editor';
import { FunctionCallingSystem, FunctionCall } from './function-calling';
import { spawn } from 'child_process';
import * as fs from 'fs';
import * as path from 'path';
import chalk from 'chalk';

export interface AgentTask {
  id: string;
  description: string;
  status: 'pending' | 'running' | 'completed' | 'failed';
  result?: any;
  error?: string;
  retries: number;
}

export interface CodeActPlan {
  steps: string[];
  parallelizable: boolean[];
  currentStep: number;
}

export class CodeActAgent {
  private editor: FileEditor;
  private functionCalling: FunctionCallingSystem;
  private tasks: Map<string, AgentTask> = new Map();
  private maxRetries: number = 3;
  private parallelExecutor: ParallelExecutor;

  constructor() {
    this.editor = new FileEditor();
    this.functionCalling = new FunctionCallingSystem();
    this.parallelExecutor = new ParallelExecutor();
  }

  // Plan and execute a complex task
  async executeTask(description: string): Promise<void> {
    console.log(chalk.cyan(`\n🤖 CodeAct Agent: ${description}\n`));
    
    // Generate plan
    const plan = await this.generatePlan(description);
    console.log(chalk.yellow('📋 Execution Plan:'));
    plan.steps.forEach((step, i) => {
      const parallel = plan.parallelizable[i] ? ' [parallel]' : '';
      console.log(`  ${i + 1}. ${step}${parallel}`);
    });
    console.log();

    // Execute plan
    await this.executePlan(plan);
  }

  private async generatePlan(description: string): Promise<CodeActPlan> {
    // In a real implementation, this would use an LLM to generate the plan
    // For now, we'll use pattern matching
    const steps: string[] = [];
    const parallelizable: boolean[] = [];

    if (description.includes('refactor')) {
      steps.push('Analyze current code structure');
      parallelizable.push(false);
      steps.push('Identify refactoring opportunities');
      parallelizable.push(false);
      steps.push('Apply refactoring changes');
      parallelizable.push(true);
      steps.push('Run tests');
      parallelizable.push(false);
      steps.push('Fix any issues');
      parallelizable.push(false);
    } else if (description.includes('test')) {
      steps.push('Discover test files');
      parallelizable.push(true);
      steps.push('Run tests');
      parallelizable.push(true);
      steps.push('Analyze failures');
      parallelizable.push(false);
      steps.push('Fix failing tests');
      parallelizable.push(true);
      steps.push('Re-run tests');
      parallelizable.push(false);
    } else {
      // Default plan
      steps.push('Analyze requirements');
      parallelizable.push(false);
      steps.push('Implement changes');
      parallelizable.push(false);
      steps.push('Validate changes');
      parallelizable.push(false);
    }

    return { steps, parallelizable, currentStep: 0 };
  }

  private async executePlan(plan: CodeActPlan): Promise<void> {
    for (let i = 0; i < plan.steps.length; i++) {
      plan.currentStep = i;
      const step = plan.steps[i];
      
      console.log(chalk.blue(`\n▶ Step ${i + 1}: ${step}`));
      
      if (plan.parallelizable[i] && i + 1 < plan.steps.length && plan.parallelizable[i + 1]) {
        // Collect parallel tasks
        const parallelTasks: string[] = [step];
        while (i + 1 < plan.steps.length && plan.parallelizable[i + 1]) {
          i++;
          parallelTasks.push(plan.steps[i]);
        }
        
        // Execute in parallel
        await this.executeParallelSteps(parallelTasks);
      } else {
        // Execute single step
        await this.executeStep(step);
      }
    }
    
    console.log(chalk.green('\n✅ Task completed successfully!\n'));
  }

  private async executeStep(step: string): Promise<void> {
    const taskId = this.generateTaskId();
    const task: AgentTask = {
      id: taskId,
      description: step,
      status: 'running',
      retries: 0
    };
    
    this.tasks.set(taskId, task);
    
    try {
      // Execute with retry logic
      await this.executeWithRetry(task);
      task.status = 'completed';
      console.log(chalk.green(`  ✓ ${step}`));
    } catch (error) {
      task.status = 'failed';
      task.error = error instanceof Error ? error.message : String(error);
      console.log(chalk.red(`  ✗ ${step}: ${task.error}`));
      throw error;
    }
  }

  private async executeParallelSteps(steps: string[]): Promise<void> {
    console.log(chalk.yellow(`\n⚡ Executing ${steps.length} steps in parallel...`));
    
    const tasks = steps.map(step => {
      const taskId = this.generateTaskId();
      const task: AgentTask = {
        id: taskId,
        description: step,
        status: 'pending',
        retries: 0
      };
      this.tasks.set(taskId, task);
      return task;
    });

    const results = await this.parallelExecutor.executeTasks(tasks, async (task) => {
      task.status = 'running';
      await this.executeWithRetry(task);
      task.status = 'completed';
      console.log(chalk.green(`  ✓ ${task.description}`));
    });

    // Check for failures
    const failures = results.filter(r => r.status === 'failed');
    if (failures.length > 0) {
      throw new Error(`${failures.length} parallel tasks failed`);
    }
  }

  private async executeWithRetry(task: AgentTask): Promise<void> {
    while (task.retries < this.maxRetries) {
      try {
        // Map step description to actual actions
        await this.mapStepToActions(task.description);
        return;
      } catch (error) {
        task.retries++;
        if (task.retries >= this.maxRetries) {
          throw error;
        }
        
        console.log(chalk.yellow(`  ⚠ Retry ${task.retries}/${this.maxRetries}: ${error}`));
        
        // Attempt to fix the error
        await this.attemptErrorCorrection(error);
      }
    }
  }

  private async mapStepToActions(step: string): Promise<void> {
    // This would use LLM in real implementation
    // For now, use pattern matching
    
    if (step.includes('test')) {
      await this.runTests();
    } else if (step.includes('analyze') || step.includes('identify')) {
      await this.analyzeCode();
    } else if (step.includes('refactor') || step.includes('implement')) {
      await this.implementChanges();
    } else if (step.includes('fix')) {
      await this.fixIssues();
    }
    
    // Simulate some work
    await new Promise(resolve => setTimeout(resolve, 500));
  }

  private async runTests(): Promise<void> {
    const result = await this.functionCalling.callFunction({
      id: Date.now().toString(),
      name: 'run_command',
      arguments: { command: 'npm test' }
    });
    
    if (!result.result?.success) {
      throw new Error('Tests failed');
    }
  }

  private async analyzeCode(): Promise<void> {
    // Use search and analysis tools
    const result = await this.functionCalling.callFunction({
      id: Date.now().toString(),
      name: 'search_files',
      arguments: { pattern: '*.ts', path: '.' }
    });
    
    // Analyze results...
  }

  private async implementChanges(): Promise<void> {
    // Simulate making changes
    console.log(chalk.gray('    Making code changes...'));
  }

  private async fixIssues(): Promise<void> {
    // Simulate fixing issues
    console.log(chalk.gray('    Applying fixes...'));
  }

  private async attemptErrorCorrection(error: any): Promise<void> {
    console.log(chalk.yellow('  🔧 Attempting automatic error correction...'));
    
    const errorMessage = error instanceof Error ? error.message : String(error);
    
    // Pattern-based error correction
    if (errorMessage.includes('Tests failed')) {
      console.log(chalk.gray('    Analyzing test failures...'));
      // Would use LLM to analyze and fix test failures
    } else if (errorMessage.includes('compile')) {
      console.log(chalk.gray('    Fixing compilation errors...'));
      // Would use LLM to fix compilation errors
    } else if (errorMessage.includes('lint')) {
      console.log(chalk.gray('    Fixing linting errors...'));
      await this.functionCalling.callFunction({
        id: Date.now().toString(),
        name: 'run_command',
        arguments: { command: 'npm run lint -- --fix' }
      });
    }
  }

  private generateTaskId(): string {
    return `task-${Date.now()}-${Math.random().toString(36).substr(2, 9)}`;
  }
}

// Parallel task executor
class ParallelExecutor {
  async executeTasks<T extends AgentTask>(
    tasks: T[],
    executor: (task: T) => Promise<void>
  ): Promise<T[]> {
    const promises = tasks.map(task => 
      executor(task).catch(error => {
        task.status = 'failed';
        task.error = error instanceof Error ? error.message : String(error);
      })
    );
    
    await Promise.all(promises);
    return tasks;
  }
}

// Advanced file editing with AST understanding
export class SmartFileEditor extends FileEditor {
  private lspClient?: LSPClient;

  async editWithUnderstanding(
    filePath: string,
    intent: string
  ): Promise<void> {
    console.log(chalk.cyan(`\n🧠 Smart Edit: ${intent}\n`));
    
    // Read file
    const content = fs.readFileSync(filePath, 'utf-8');
    
    // Find relevant chunk
    const chunk = ChunkLocalizer.findRelevantChunk(content, intent);
    if (!chunk) {
      throw new Error('Could not locate relevant code section');
    }
    
    console.log(chalk.gray(`Found relevant code at lines ${chunk.startLine}-${chunk.endLine}`));
    
    // In real implementation, would use LLM to generate the edit
    // For now, just show what would happen
    console.log(chalk.yellow('Would apply intelligent edit based on intent'));
  }

  async applyBulkRefactoring(
    pattern: string,
    transformation: string
  ): Promise<void> {
    console.log(chalk.cyan(`\n🔄 Bulk Refactoring: ${pattern} → ${transformation}\n`));
    
    // Find all files matching pattern
    const files = await this.findFilesWithPattern(pattern);
    console.log(chalk.gray(`Found ${files.length} files to refactor`));
    
    // Apply transformation to each file
    for (const file of files) {
      console.log(chalk.gray(`  Refactoring ${file}...`));
      // Would apply transformation
    }
  }

  private async findFilesWithPattern(pattern: string): Promise<string[]> {
    // Simplified implementation
    const results: string[] = [];
    
    const walkDir = (dir: string) => {
      const files = fs.readdirSync(dir);
      for (const file of files) {
        const fullPath = path.join(dir, file);
        const stats = fs.statSync(fullPath);
        
        if (stats.isDirectory() && !file.startsWith('.') && file !== 'node_modules') {
          walkDir(fullPath);
        } else if (stats.isFile() && fullPath.endsWith('.ts')) {
          const content = fs.readFileSync(fullPath, 'utf-8');
          if (content.includes(pattern)) {
            results.push(fullPath);
          }
        }
      }
    };
    
    walkDir(process.cwd());
    return results;
  }
}

// Placeholder for LSP client (would integrate with real LSP)
class LSPClient {
  async initialize(rootPath: string): Promise<void> {
    console.log(chalk.gray(`Initializing LSP for ${rootPath}`));
  }
  
  async getDefinition(file: string, position: number): Promise<any> {
    // Would return actual definition location
    return null;
  }
  
  async getReferences(file: string, position: number): Promise<any[]> {
    // Would return all references
    return [];
  }
  
  async rename(file: string, position: number, newName: string): Promise<any> {
    // Would perform project-wide rename
    return null;
  }
}