plan2code/plan2code-loop/src/controller.ts

import { agentRegistry, type Agent, type AgentExecutionResult } from './agents/index.js';
import { StateManager, type SessionConfig, type IterationLogEntry } from './state/index.js';
import { buildLoopPrompt } from './prompt/index.js';
import { checkForCompletion, logger, type CompletionCheckResult } from './utils/index.js';

export interface TaskCompleteInfo {
  marker: string;
  taskId?: string;
  taskName?: string;
}

export interface ControllerOptions {
  config: SessionConfig;
  stateManager: StateManager;
  onIteration?: (iteration: number, max: number) => void;
  onTaskComplete?: (info: TaskCompleteInfo) => void | Promise<void>;
  onLoopComplete?: () => void;
}

export interface LoopResult {
  completed: boolean;
  iterations: number;
  finalMarker?: string;
  exitReason: 'all_complete' | 'max_iterations' | 'interrupted' | 'error';
  tasksCompleted: number;
  error?: Error;
}

export class Controller {
  private readonly config: SessionConfig;
  private readonly stateManager: StateManager;
  private readonly agent: Agent;
  private readonly onIteration?: (iteration: number, max: number) => void;
  private readonly onTaskComplete?: (info: TaskCompleteInfo) => void | Promise<void>;
  private readonly onLoopComplete?: () => void;
  private interrupted = false;
  private abortController: AbortController | null = null;
  private tasksCompleted = 0;

  constructor(options: ControllerOptions) {
    this.config = options.config;
    this.stateManager = options.stateManager;
    this.onIteration = options.onIteration;
    this.onTaskComplete = options.onTaskComplete;
    this.onLoopComplete = options.onLoopComplete;

    const agent = agentRegistry.get(this.config.agent);
    if (!agent) {
      throw new Error(`Agent not found: ${this.config.agent}`);
    }
    this.agent = agent;
  }

  private async buildPrompt(): Promise<string> {
    return buildLoopPrompt({
      specPath: this.config.specPath,
      iteration: this.config.currentIteration + 1,
      maxIterations: this.config.maxIterations,
      stateManager: this.stateManager,
    });
  }

  private async executeIteration(prompt: string): Promise<AgentExecutionResult> {
    const timeoutMs = this.config.timeout * 60 * 1000;

    // Create new AbortController for this iteration
    this.abortController = new AbortController();

    const result = await this.agent.execute({
      prompt,
      model: this.config.model,
      timeout: timeoutMs,
      verbose: this.config.verbose,
      cwd: process.cwd(),
      signal: this.abortController.signal,
    });

    this.abortController = null;
    return result;
  }

  private createLogEntry(
    result: AgentExecutionResult,
    status: IterationLogEntry['status'],
    marker?: string
  ): IterationLogEntry {
    return {
      iteration: this.config.currentIteration + 1,
      timestamp: new Date().toISOString(),
      duration: result.duration,
      exitCode: result.exitCode,
      status,
      completionMarker: marker,
    };
  }

  private formatTaskDisplay(completion: CompletionCheckResult): string {
    if (completion.taskId && completion.taskName) {
      return `Task ${completion.taskId}: ${completion.taskName}`;
    } else if (completion.taskId) {
      return `Task ${completion.taskId}`;
    }
    return '';
  }

  private displayIterationResult(result: AgentExecutionResult, iterNum: number, completion: CompletionCheckResult): void {
    const duration = Math.round(result.duration / 1000);
    const taskDisplay = this.formatTaskDisplay(completion);

    // Show iteration completion with task info if available
    if (taskDisplay) {
      logger.iteration(iterNum, this.config.maxIterations, `${taskDisplay} (${duration}s)`);
    } else {
      logger.iteration(iterNum, this.config.maxIterations, `completed in ${duration}s`);
    }

    // Verbose mode: show full output
    if (this.config.verbose) {
      console.log();
      logger.dim('--- Agent Output ---');
      console.log(result.stdout);
      if (result.stderr) {
        logger.dim('--- Agent Stderr ---');
        console.log(result.stderr);
      }
      logger.dim('--- End Output ---');
      console.log();
    } else if (result.exitCode !== 0 && result.stderr.trim()) {
      logger.error(`  ${result.stderr.trim().split('\n')[0]}`);
    }
  }

  async run(): Promise<LoopResult> {
    logger.header('Starting Plan2Code Loop');
    logger.info(`Agent: ${this.agent.config.displayName}`);
    logger.info(`Model: ${this.config.model}`);
    logger.info(`Spec: ${this.config.specPath}`);
    logger.info(`Max iterations: ${this.config.maxIterations}`);
    console.log();

    while (this.config.currentIteration < this.config.maxIterations) {
      if (this.interrupted) {
        return {
          completed: false,
          iterations: this.config.currentIteration,
          exitReason: 'interrupted',
          tasksCompleted: this.tasksCompleted,
        };
      }

      const iterNum = this.config.currentIteration + 1;
      this.onIteration?.(iterNum, this.config.maxIterations);

      console.log();
      logger.info(`Iteration ${iterNum}/${this.config.maxIterations}`);

      // Build prompt - simple, just spec path and iteration info
      const prompt = await this.buildPrompt();

      const spinner = logger.spinner('Waiting for agent response');
      const startTime = Date.now();

      // Update spinner with elapsed time every second
      const elapsedInterval = setInterval(() => {
        const elapsed = Math.round((Date.now() - startTime) / 1000);
        spinner.text = `Waiting for agent response (${elapsed}s)`;
      }, 1000);

      try {
        const result = await this.executeIteration(prompt);
        clearInterval(elapsedInterval);
        spinner.stop();

        // Check if cancelled
        if (result.cancelled || this.interrupted) {
          logger.info('Agent process cancelled');

          const entry: IterationLogEntry = {
            iteration: iterNum,
            timestamp: new Date().toISOString(),
            duration: result.duration,
            exitCode: -1,
            status: 'interrupted',
          };
          await this.stateManager.appendIterationLog(entry);

          return {
            completed: false,
            iterations: this.config.currentIteration,
            exitReason: 'interrupted',
            tasksCompleted: this.tasksCompleted,
          };
        }

        // Check for completion markers in output (now includes task info)
        const completion = checkForCompletion(result.stdout + result.stderr);

        // Determine status
        let status: IterationLogEntry['status'] = 'running';
        if (completion.completed) {
          if (completion.marker === 'TASK_BLOCKED') {
            status = 'blocked';
          } else {
            status = 'completed';
          }
        } else if (result.timedOut) {
          status = 'timeout';
        } else if (result.exitCode !== 0) {
          status = 'error';
        }

        // Log iteration
        const logEntry = this.createLogEntry(result, status, completion.marker);
        await this.stateManager.appendIterationLog(logEntry);

        // Display iteration result with task info from completion marker
        this.displayIterationResult(result, iterNum, completion);

        // Note: Scratchpad updates are now handled by the LLM directly

        // Handle completion markers
        if (completion.completed) {
          const taskDisplay = this.formatTaskDisplay(completion);

          if (completion.marker === 'TASK_COMPLETE') {
            this.tasksCompleted++;
            await this.onTaskComplete?.({
              marker: completion.marker,
              taskId: completion.taskId,
              taskName: completion.taskName,
            });
            logger.success(taskDisplay ? `Completed: ${taskDisplay}` : 'Task completed!');
          } else if (completion.marker === 'TASK_BLOCKED') {
            const blockInfo = completion.taskId
              ? `Task ${completion.taskId} blocked: ${completion.reason || 'Unknown reason'}`
              : `Task blocked: ${completion.reason || 'Unknown reason'}`;
            logger.warning(blockInfo);
          } else if (completion.marker === 'LOOP_COMPLETE') {
            // Commit any final changes before completing
            this.tasksCompleted++;
            await this.onTaskComplete?.({
              marker: completion.marker,
              taskId: completion.taskId,
              taskName: completion.taskName || 'Final implementation complete',
            });
            this.onLoopComplete?.();
            logger.success('All tasks complete!');
            return {
              completed: true,
              iterations: iterNum,
              finalMarker: 'LOOP_COMPLETE',
              exitReason: 'all_complete',
              tasksCompleted: this.tasksCompleted,
            };
          }
        }

        // Increment iteration and update spec hash (so resume doesn't see false changes)
        await this.stateManager.incrementIteration();
        await this.stateManager.updateSpecHash(this.config.specPath);
        this.config.currentIteration++;

        // Handle timeout
        if (result.timedOut) {
          logger.warning(`Iteration ${iterNum} timed out, continuing...`);
        }

        // Handle error (but continue - LLM might recover)
        if (result.exitCode !== 0 && !result.timedOut && !completion.completed) {
          logger.warning(`Iteration ${iterNum} exited with code ${result.exitCode}, continuing...`);
        }

      } catch (error) {
        clearInterval(elapsedInterval);
        spinner.stop();
        logger.error(`Iteration ${iterNum} failed: ${error}`);

        // Log the error
        const entry: IterationLogEntry = {
          iteration: iterNum,
          timestamp: new Date().toISOString(),
          duration: 0,
          exitCode: -1,
          status: 'error',
        };
        await this.stateManager.appendIterationLog(entry);

        return {
          completed: false,
          iterations: this.config.currentIteration,
          exitReason: 'error',
          tasksCompleted: this.tasksCompleted,
          error: error instanceof Error ? error : new Error(String(error)),
        };
      }
    }

    // Max iterations reached
    logger.warning('Max iterations reached');
    return {
      completed: false,
      iterations: this.config.currentIteration,
      exitReason: 'max_iterations',
      tasksCompleted: this.tasksCompleted,
    };
  }

  interrupt(): void {
    this.interrupted = true;
    if (this.abortController) {
      this.abortController.abort();
    }
  }
}