jay-code/src/models/planning-agent.ts

Streamlined AI CLI orchestration engine with mathematical rigor and enterprise-grade reliability

/**
 * Jay-Code Planning Agent: Graph-based Task Decomposition with Mathematical Analysis
 * Implements formal complexity analysis and dependency resolution using graph theory
 */
import { EventEmitter } from 'events';

export interface Task {
  id: string;
  description: string;
  complexity: number; // 0-1 scale based on mathematical analysis
  dependencies: string[];
  requirements: string[];
  estimatedTokens: number;
  priority: number;
}

export interface SubTask extends Task {
  parentId: string;
  sequence: number;
  completionCriteria: string[];
}

export interface ExecutionPlan {
  id: string;
  originalTask: Task;
  subTasks: SubTask[];
  dependencyGraph: Map<string, string[]>;
  criticalPath: string[];
  estimatedDuration: number;
  prompt: string;
  context: any;
}

export interface ProgressUpdate {
  taskId: string;
  status: 'pending' | 'in-progress' | 'completed' | 'failed';
  completedSubTasks: number;
  totalSubTasks: number;
  estimatedRemaining: number;
  currentSubTask?: string;
}

export class PlanningAgent extends EventEmitter {
  private complexityThreshold: number = 0.7;
  private maxSubTasks: number = 20;
  private activePlans: Map<string, ExecutionPlan> = new Map();

  constructor(options: { complexityThreshold?: number; maxSubTasks?: number } = {}) {
    super();
    this.complexityThreshold = options.complexityThreshold || 0.7;
    this.maxSubTasks = options.maxSubTasks || 20;
  }

  /**
   * Mathematical task breakdown using graph theory and complexity analysis
   */
  async breakdownTask(task: Task): Promise<SubTask[]> {
    const complexity = this.calculateComplexity(task);
    
    if (complexity < this.complexityThreshold) {
      // Simple task - no breakdown needed
      return [{
        ...task,
        parentId: task.id,
        sequence: 1,
        completionCriteria: [task.description]
      }];
    }

    // Complex task - apply mathematical decomposition
    const subTasks = await this.decomposeComplexTask(task);
    this.emit('task-decomposed', { taskId: task.id, subTaskCount: subTasks.length, complexity });
    
    return subTasks;
  }

  private calculateComplexity(task: Task): number {
    // Mathematical complexity analysis using multiple factors
    let complexity = 0;

    // Description complexity (entropy-based analysis)
    const descriptionComplexity = this.calculateTextComplexity(task.description);
    complexity += descriptionComplexity * 0.4;

    // Dependency complexity (graph theory)
    const dependencyComplexity = task.dependencies.length / 10; // Normalized
    complexity += Math.min(dependencyComplexity, 0.3);

    // Requirements complexity
    const requirementsComplexity = task.requirements.length / 15; // Normalized
    complexity += Math.min(requirementsComplexity, 0.2);

    // Token estimation complexity
    const tokenComplexity = Math.min(task.estimatedTokens / 10000, 0.1);
    complexity += tokenComplexity;

    return Math.min(complexity, 1.0);
  }

  private calculateTextComplexity(text: string): number {
    // Information theory approach - calculate entropy
    const charFreq = new Map<string, number>();
    const textLength = text.length;
    
    for (const char of text.toLowerCase()) {
      charFreq.set(char, (charFreq.get(char) || 0) + 1);
    }

    let entropy = 0;
    for (const freq of charFreq.values()) {
      const probability = freq / textLength;
      entropy -= probability * Math.log2(probability);
    }

    // Normalize entropy (max theoretical entropy for English is ~4.7 bits)
    return Math.min(entropy / 4.7, 1.0);
  }

  private async decomposeComplexTask(task: Task): Promise<SubTask[]> {
    const subTasks: SubTask[] = [];
    const keywords = this.extractKeywords(task.description);
    
    // Mathematical pattern matching for common code generation patterns
    const patterns = this.identifyPatterns(task.description, keywords);
    
    for (let i = 0; i < patterns.length && i < this.maxSubTasks; i++) {
      const pattern = patterns[i];
      subTasks.push({
        id: `${task.id}_sub_${i + 1}`,
        description: pattern.description,
        complexity: pattern.complexity,
        dependencies: i > 0 ? [`${task.id}_sub_${i}`] : task.dependencies,
        requirements: pattern.requirements,
        estimatedTokens: Math.ceil(task.estimatedTokens / patterns.length),
        priority: task.priority + (patterns.length - i) * 0.1, // Higher priority for earlier tasks
        parentId: task.id,
        sequence: i + 1,
        completionCriteria: pattern.criteria
      });
    }

    return subTasks;
  }

  private extractKeywords(description: string): string[] {
    // Simple keyword extraction - could be enhanced with NLP
    const commonWords = ['the', 'and', 'or', 'but', 'in', 'on', 'at', 'to', 'for', 'of', 'with', 'by', 'a', 'an'];
    return description
      .toLowerCase()
      .split(/\W+/)
      .filter(word => word.length > 2 && !commonWords.includes(word))
      .slice(0, 10); // Limit keywords for complexity control
  }

  private identifyPatterns(description: string, keywords: string[]): Array<{
    description: string;
    complexity: number;
    requirements: string[];
    criteria: string[];
  }> {
    const patterns = [];
    const desc = description.toLowerCase();

    // Function/method patterns
    if (desc.includes('function') || desc.includes('method') || desc.includes('class')) {
      patterns.push({
        description: 'Define function signatures and interfaces',
        complexity: 0.3,
        requirements: ['TypeScript', 'proper typing'],
        criteria: ['Function signature defined', 'Parameters typed', 'Return type specified']
      });
      
      patterns.push({
        description: 'Implement core logic',
        complexity: 0.6,
        requirements: ['algorithmic thinking', 'error handling'],
        criteria: ['Core algorithm implemented', 'Edge cases handled', 'Error handling added']
      });
    }

    // API patterns
    if (desc.includes('api') || desc.includes('endpoint') || desc.includes('route')) {
      patterns.push({
        description: 'Design API structure and validation',
        complexity: 0.4,
        requirements: ['REST principles', 'input validation'],
        criteria: ['API structure defined', 'Validation rules set', 'Response format specified']
      });
    }

    // Database patterns
    if (desc.includes('database') || desc.includes('sql') || desc.includes('store')) {
      patterns.push({
        description: 'Design data schema and queries',
        complexity: 0.5,
        requirements: ['database design', 'query optimization'],
        criteria: ['Schema designed', 'Queries optimized', 'Indexes planned']
      });
    }

    // Testing patterns
    if (desc.includes('test') || desc.includes('testing')) {
      patterns.push({
        description: 'Create comprehensive test suite',
        complexity: 0.4,
        requirements: ['test framework', 'coverage analysis'],
        criteria: ['Unit tests written', 'Integration tests added', 'Coverage >90%']
      });
    }

    // Fallback pattern if no specific patterns found
    if (patterns.length === 0) {
      patterns.push({
        description: 'Implement core functionality',
        complexity: 0.5,
        requirements: ['clean code', 'documentation'],
        criteria: ['Functionality implemented', 'Code documented', 'Requirements met']
      });
    }

    return patterns;
  }

  /**
   * Create execution plan with dependency graph and critical path analysis
   */
  async coordinateExecution(plan: ExecutionPlan): Promise<void> {
    this.activePlans.set(plan.id, plan);
    
    // Build dependency graph using mathematical graph theory
    const dependencyGraph = this.buildDependencyGraph(plan.subTasks);
    plan.dependencyGraph = dependencyGraph;
    
    // Calculate critical path using longest path algorithm
    plan.criticalPath = this.calculateCriticalPath(plan.subTasks, dependencyGraph);
    
    this.emit('execution-coordinated', {
      planId: plan.id,
      totalSubTasks: plan.subTasks.length,
      criticalPathLength: plan.criticalPath.length
    });
  }

  private buildDependencyGraph(subTasks: SubTask[]): Map<string, string[]> {
    const graph = new Map<string, string[]>();
    
    for (const task of subTasks) {
      graph.set(task.id, [...task.dependencies]);
    }
    
    return graph;
  }

  private calculateCriticalPath(subTasks: SubTask[], dependencyGraph: Map<string, string[]>): string[] {
    // Topological sort with longest path calculation
    const inDegree = new Map<string, number>();
    const distances = new Map<string, number>();
    const predecessors = new Map<string, string>();
    
    // Initialize
    for (const task of subTasks) {
      inDegree.set(task.id, task.dependencies.length);
      distances.set(task.id, 0);
    }
    
    // Find longest path (critical path)
    const queue: string[] = [];
    for (const [taskId, degree] of inDegree) {
      if (degree === 0) queue.push(taskId);
    }
    
    while (queue.length > 0) {
      const current = queue.shift()!;
      const currentTask = subTasks.find(t => t.id === current)!;
      
      // Update distances to dependent tasks
      for (const task of subTasks) {
        if (task.dependencies.includes(current)) {
          const newDistance = distances.get(current)! + currentTask.estimatedTokens;
          if (newDistance > distances.get(task.id)!) {
            distances.set(task.id, newDistance);
            predecessors.set(task.id, current);
          }
          
          const newInDegree = inDegree.get(task.id)! - 1;
          inDegree.set(task.id, newInDegree);
          if (newInDegree === 0) queue.push(task.id);
        }
      }
    }
    
    // Reconstruct critical path
    let maxDistance = 0;
    let endTask = '';
    for (const [taskId, distance] of distances) {
      if (distance > maxDistance) {
        maxDistance = distance;
        endTask = taskId;
      }
    }
    
    const criticalPath: string[] = [];
    let current = endTask;
    while (current) {
      criticalPath.unshift(current);
      current = predecessors.get(current) || '';
    }
    
    return criticalPath;
  }

  /**
   * Monitor progress with mathematical performance tracking
   */
  async monitorProgress(taskId: string): Promise<ProgressUpdate> {
    const plan = this.activePlans.get(taskId);
    if (!plan) {
      throw new Error(`No active plan found for task: ${taskId}`);
    }
    
    const completedSubTasks = plan.subTasks.filter(st => st.id.includes('completed')).length;
    const totalSubTasks = plan.subTasks.length;
    const progress = completedSubTasks / totalSubTasks;
    
    // Mathematical estimation of remaining time
    const averageTaskTime = plan.estimatedDuration / totalSubTasks;
    const estimatedRemaining = (totalSubTasks - completedSubTasks) * averageTaskTime;
    
    const progressUpdate: ProgressUpdate = {
      taskId,
      status: progress === 1 ? 'completed' : progress > 0 ? 'in-progress' : 'pending',
      completedSubTasks,
      totalSubTasks,
      estimatedRemaining,
      currentSubTask: plan.subTasks.find(st => !st.id.includes('completed'))?.id
    };
    
    this.emit('progress-updated', progressUpdate);
    return progressUpdate;
  }
}