mcp-workflow-server-enhanced/src/shared/workflow.ts

Enhanced MCP Workflow Server with smart problem routing, comprehensive validation, guide compliance, and robust error handling. Intelligently routes to appropriate AI functions based on problem type.

import {
  WorkflowContext,
  WorkflowConfig,
  WorkflowError,
  FunctionOutput,
  GuideCompliance,
} from './types.js';
import {
  createWorkflowContext,
  updateWorkflowContext,
  logWorkflowProgress,
  retryWithBackoff,
} from './utils.js';
import {
  validateGuideCompliance,
  validateStageTransition,
  ValidationError,
  rollbackToLastValidState,
  createValidationMetadata,
} from './validation.js';

/**
 * Workflow orchestrator that manages the execution of the 7-step process
 */
export class WorkflowOrchestrator {
  private config: WorkflowConfig;
  private functionMap: Map<string, (input: any, context: WorkflowContext) => Promise<FunctionOutput>>;

  constructor(config: WorkflowConfig) {
    this.config = config;
    this.functionMap = new Map();
  }

  /**
   * Register a workflow function
   */
  registerFunction(
    name: string,
    handler: (input: any, context: WorkflowContext) => Promise<FunctionOutput>
  ): void {
    this.functionMap.set(name, handler);
  }

  /**
   * Execute the complete workflow
   */
  async executeWorkflow(
    originalPrompt: string,
    startStep: string = 'improve-prompt'
  ): Promise<WorkflowContext> {
    let context = createWorkflowContext(originalPrompt, startStep as WorkflowContext['currentStep']);
    
    logWorkflowProgress(context, 'workflow', 'Starting workflow execution');

    const steps = this.config.enabledFunctions;
    const startIndex = steps.indexOf(startStep);
    
    if (startIndex === -1) {
      throw new WorkflowError(`Invalid start step: ${startStep}`, 'workflow', context);
    }

    for (let i = startIndex; i < steps.length; i++) {
      const stepName = steps[i];
      
      if (!this.functionMap.has(stepName)) {
        throw new WorkflowError(`Function not registered: ${stepName}`, stepName, context);
      }

      try {
        context = await this.executeStep(stepName, context);
        
        if (!this.config.autoAdvance) {
          // If auto-advance is disabled, return after each step
          break;
        }
      } catch (error) {
        const workflowError = error instanceof WorkflowError 
          ? error 
          : new WorkflowError(`Step failed: ${error.message}`, stepName, context, error as Error);

        if (this.config.errorHandling === 'stop') {
          throw workflowError;
        } else if (this.config.errorHandling === 'retry') {
          context = await this.retryStep(stepName, context, workflowError);
        } else {
          // Skip step
          logWorkflowProgress(context, stepName, `Skipping step due to error: ${error.message}`);
          continue;
        }
      }
    }

    logWorkflowProgress(context, 'workflow', 'Workflow execution completed');
    return context;
  }

  /**
   * Execute a single workflow step with validation checkpoints
   */
  private async executeStep(stepName: string, context: WorkflowContext): Promise<WorkflowContext> {
    logWorkflowProgress(context, stepName, 'Starting step execution');

    // Validate stage transition before execution
    if (context.currentStep !== stepName) {
      const transitionValidation = await validateStageTransition(
        context.currentStep,
        stepName,
        context
      );

      if (!transitionValidation.canProceed) {
        throw new ValidationError(
          `Stage transition validation failed: ${transitionValidation.issues.join(', ')}`,
          stepName,
          transitionValidation.issues
        );
      }
    }

    const handler = this.functionMap.get(stepName)!;
    const input = this.prepareStepInput(stepName, context);

    // Execute the step
    const result = await handler(input, context);

    if (!result.success) {
      throw new WorkflowError(`Step failed: ${JSON.stringify(result)}`, stepName, context);
    }

    // Validate guide compliance for the completed step
    const guideCompliance = await validateGuideCompliance(
      stepName,
      input,
      result,
      context
    );

    // Create validation metadata
    const validationMetadata = createValidationMetadata(
      stepName,
      true,
      undefined,
      guideCompliance
    );

    // Check if compliance score meets minimum threshold
    if (guideCompliance.complianceScore < 70) {
      logWorkflowProgress(context, stepName,
        `Warning: Low guide compliance score: ${guideCompliance.complianceScore}%`
      );

      // Optionally rollback if compliance is critically low
      if (guideCompliance.complianceScore < 50) {
        throw new ValidationError(
          `Critical guide compliance failure: ${guideCompliance.complianceScore}%`,
          stepName,
          [`Compliance score below minimum threshold`],
          guideCompliance
        );
      }
    }

    const updatedContext = updateWorkflowContext(
      context,
      stepName,
      result.result,
      result.nextStep as WorkflowContext['currentStep']
    );

    // Add validation metadata to context
    updatedContext.validationMetadata = validationMetadata;
    updatedContext.guideCompliance = guideCompliance;

    logWorkflowProgress(updatedContext, stepName,
      `Step completed successfully with ${guideCompliance.complianceScore}% guide compliance`
    );
    return updatedContext;
  }

  /**
   * Retry a failed step with exponential backoff and rollback support
   */
  private async retryStep(
    stepName: string,
    context: WorkflowContext,
    originalError: WorkflowError
  ): Promise<WorkflowContext> {
    logWorkflowProgress(context, stepName, `Retrying step (max retries: ${this.config.maxRetries})`);

    // If it's a validation error, consider rollback
    if (originalError instanceof ValidationError && originalError.guideCompliance?.complianceScore < 50) {
      logWorkflowProgress(context, stepName, 'Critical validation failure - attempting rollback');

      try {
        // Find previous successful step for rollback
        const steps = ['improve-prompt', 'research', 'cognitive', 'planner', 'task-generation', 'implementation'];
        const currentIndex = steps.indexOf(stepName);

        if (currentIndex > 0) {
          const previousStep = steps[currentIndex - 1];
          const rolledBackContext = await rollbackToLastValidState(context, previousStep);

          logWorkflowProgress(rolledBackContext, stepName,
            `Rolled back to ${previousStep} due to validation failure`
          );

          return rolledBackContext;
        }
      } catch (rollbackError) {
        logWorkflowProgress(context, stepName,
          `Rollback failed: ${rollbackError.message}, continuing with retry`
        );
      }
    }

    try {
      return await retryWithBackoff(
        () => this.executeStep(stepName, context),
        this.config.maxRetries
      );
    } catch (error) {
      throw new WorkflowError(
        `Step failed after ${this.config.maxRetries} retries: ${originalError.message}`,
        stepName,
        context,
        originalError
      );
    }
  }

  /**
   * Prepare input for a specific step based on previous results
   */
  private prepareStepInput(stepName: string, context: WorkflowContext): any {
    const stepResults = context.stepResults || {};
    
    switch (stepName) {
      case 'improve-prompt':
        return {
          userPrompt: context.originalPrompt,
          context,
        };
      
      case 'research':
        return {
          improvedPrompt: stepResults['improve-prompt']?.improvedPrompt || context.originalPrompt,
          researchTopics: stepResults['improve-prompt']?.researchTopics || [],
          codeAnalysis: true,
          context,
        };
      
      case 'cognitive':
        return {
          researchData: stepResults['research'],
          improvedPrompt: stepResults['improve-prompt']?.improvedPrompt || context.originalPrompt,
          analysisDepth: 'deep',
          context,
        };
      
      case 'planner':
        return {
          cognitiveAnalysis: stepResults['cognitive'],
          requirements: stepResults['cognitive']?.requirements || [],
          constraints: stepResults['cognitive']?.constraints || [],
          context,
        };
      
      case 'task-generation':
        return {
          plan: stepResults['planner'],
          granularity: 'medium',
          context,
        };
      
      case 'implementation':
        return {
          tasks: stepResults['task-generation']?.tasks || [],
          currentTaskIndex: 0,
          context,
        };
      
      case 'problem-solver':
        return {
          error: stepResults['implementation']?.error || 'Unknown error',
          context: stepResults['implementation'],
          previousAttempts: [],
          strictMode: true,
          workflowContext: context,
        };
      
      default:
        return { context };
    }
  }

  /**
   * Get the current workflow status
   */
  getWorkflowStatus(context: WorkflowContext): {
    currentStep: string;
    completedSteps: string[];
    remainingSteps: string[];
    progress: number;
  } {
    const completedSteps = Object.keys(context.stepResults || {});
    const allSteps = this.config.enabledFunctions;
    const currentIndex = allSteps.indexOf(context.currentStep);
    const remainingSteps = allSteps.slice(currentIndex + 1);
    
    return {
      currentStep: context.currentStep,
      completedSteps,
      remainingSteps,
      progress: (completedSteps.length / allSteps.length) * 100,
    };
  }
}