mcp-workflow-server-enhanced
Version:
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.
552 lines (506 loc) • 17.9 kB
text/typescript
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
import { z } from 'zod';
import { WorkflowOrchestrator } from './shared/workflow.js';
import { WorkflowConfigSchema, WorkflowContextSchema } from './shared/types.js';
import { createWorkflowContext, logWorkflowProgress } from './shared/utils.js';
import { validateStageTransition, validateGuideCompliance, ValidationError } from './shared/validation.js';
// Import individual function implementations
import { createImprovePromptFunction } from './functions/improve-prompt.js';
import { createResearchFunction } from './functions/research.js';
import { createCognitiveFunction } from './functions/cognitive.js';
import { createPlannerFunction } from './functions/planner.js';
import { createTaskGenerationFunction } from './functions/task-generation.js';
import { createImplementationFunction } from './functions/implementation.js';
import { createProblemSolverFunction } from './functions/problem-solver.js';
/**
* Main MCP Workflow Server
*/
export class MCPWorkflowServer {
private server: McpServer;
private orchestrator: WorkflowOrchestrator;
private activeWorkflows: Map<string, any> = new Map();
constructor() {
this.server = new McpServer({
name: 'mcp-workflow-server',
version: '1.0.0',
});
// Initialize workflow orchestrator with default config
this.orchestrator = new WorkflowOrchestrator({
enabledFunctions: [
'improve-prompt',
'research',
'cognitive',
'planner',
'task-generation',
'implementation'
],
autoAdvance: true,
errorHandling: 'retry',
maxRetries: 3,
});
this.setupFunctions();
this.setupTools();
}
/**
* Setup individual workflow functions
*/
private setupFunctions(): void {
// Register all workflow functions with the orchestrator
this.orchestrator.registerFunction('improve-prompt', createImprovePromptFunction());
this.orchestrator.registerFunction('research', createResearchFunction());
this.orchestrator.registerFunction('cognitive', createCognitiveFunction());
this.orchestrator.registerFunction('planner', createPlannerFunction());
this.orchestrator.registerFunction('task-generation', createTaskGenerationFunction());
this.orchestrator.registerFunction('implementation', createImplementationFunction());
this.orchestrator.registerFunction('problem-solver', createProblemSolverFunction());
}
/**
* Setup MCP tools for external interaction
*/
private setupTools(): void {
// Main workflow execution tool
this.server.registerTool(
'execute-workflow',
{
title: 'Execute Smart AI Workflow',
description: 'Intelligently routes to the appropriate AI function based on the problem type',
inputSchema: {
userPrompt: z.string().describe('The user prompt to process through the workflow'),
config: WorkflowConfigSchema.optional().describe('Optional workflow configuration'),
startStep: z.string().optional().describe('Step to start from (default: improve-prompt)'),
},
},
async ({ userPrompt, config, startStep }) => {
try {
// Update orchestrator config if provided
if (config) {
this.orchestrator = new WorkflowOrchestrator(config);
this.setupFunctions();
}
// Smart problem routing
const routingResult = await this.smartProblemRouting(userPrompt);
if (routingResult.directFunction) {
// Execute specific function directly for targeted problems
const context = createWorkflowContext(userPrompt, routingResult.directFunction);
const functionHandler = this.functionMap.get(routingResult.directFunction);
if (functionHandler) {
const result = await functionHandler(routingResult.input, context);
return {
content: [
{
type: 'text',
text: JSON.stringify({
approach: 'direct-function',
function: routingResult.directFunction,
reasoning: routingResult.reasoning,
result: result,
sessionId: context.sessionId,
}, null, 2),
},
],
};
}
}
// Fallback to full workflow for complex problems
const context = await this.orchestrator.executeWorkflow(
userPrompt,
startStep || routingResult.startStep || 'improve-prompt'
);
// Store the workflow for potential continuation
this.activeWorkflows.set(context.sessionId, context);
return {
content: [
{
type: 'text',
text: JSON.stringify({
approach: 'full-workflow',
sessionId: context.sessionId,
status: 'completed',
results: context.stepResults,
finalStep: context.currentStep,
}, null, 2),
},
],
};
} catch (error) {
return {
content: [
{
type: 'text',
text: `Workflow execution failed: ${error.message}`,
},
],
isError: true,
};
}
}
);
// Individual step execution tools
this.server.registerTool(
'execute-step',
{
title: 'Execute Single Workflow Step',
description: 'Execute a single step in the workflow process',
inputSchema: {
stepName: z.enum([
'improve-prompt',
'research',
'cognitive',
'planner',
'task-generation',
'implementation',
'problem-solver'
]).describe('The workflow step to execute'),
input: z.any().describe('Input data for the step'),
context: WorkflowContextSchema.optional().describe('Existing workflow context'),
},
},
async ({ stepName, input, context }) => {
try {
const workflowContext = context || createWorkflowContext(
input.userPrompt || 'Single step execution',
stepName
);
// Get the function handler
const handler = this.orchestrator['functionMap'].get(stepName);
if (!handler) {
throw new Error(`Function not found: ${stepName}`);
}
const result = await handler(input, workflowContext);
return {
content: [
{
type: 'text',
text: JSON.stringify(result, null, 2),
},
],
};
} catch (error) {
return {
content: [
{
type: 'text',
text: `Step execution failed: ${error.message}`,
},
],
isError: true,
};
}
}
);
// Workflow status tool
this.server.registerTool(
'get-workflow-status',
{
title: 'Get Workflow Status',
description: 'Get the current status of a workflow execution',
inputSchema: {
sessionId: z.string().describe('The workflow session ID'),
},
},
async ({ sessionId }) => {
const workflow = this.activeWorkflows.get(sessionId);
if (!workflow) {
return {
content: [
{
type: 'text',
text: `Workflow not found: ${sessionId}`,
},
],
isError: true,
};
}
const status = this.orchestrator.getWorkflowStatus(workflow);
return {
content: [
{
type: 'text',
text: JSON.stringify(status, null, 2),
},
],
};
}
);
// Configuration tool
this.server.registerTool(
'configure-workflow',
{
title: 'Configure Workflow',
description: 'Update workflow configuration settings',
inputSchema: {
config: WorkflowConfigSchema.describe('New workflow configuration'),
},
},
async ({ config }) => {
try {
this.orchestrator = new WorkflowOrchestrator(config);
this.setupFunctions();
return {
content: [
{
type: 'text',
text: 'Workflow configuration updated successfully',
},
],
};
} catch (error) {
return {
content: [
{
type: 'text',
text: `Configuration update failed: ${error.message}`,
},
],
isError: true,
};
}
}
);
// Validation checkpoint tool
this.server.registerTool(
'validate-and-continue',
{
title: 'Validate Current Workflow Stage and Continue',
description: 'Validate current workflow stage and continue to next stage',
inputSchema: {
sessionId: z.string().describe('Workflow session ID'),
validation: z.enum(['proceed', 'revise', 'clarify']).describe('Validation decision'),
currentStage: z.string().describe('Current workflow stage'),
feedback: z.string().optional().describe('Optional feedback for revision'),
},
},
async ({ sessionId, validation, currentStage, feedback }) => {
try {
const workflow = this.activeWorkflows.get(sessionId);
if (!workflow) {
return {
content: [
{
type: 'text',
text: `No active workflow found for session: ${sessionId}`,
},
],
isError: true,
};
}
let result;
switch (validation) {
case 'proceed':
// Validate stage transition and proceed
const transitionValidation = await validateStageTransition(
currentStage,
workflow.nextStep || 'research',
workflow
);
if (!transitionValidation.canProceed) {
result = {
status: 'validation_failed',
issues: transitionValidation.issues,
recommendation: 'Address validation issues before proceeding',
};
} else {
result = {
status: 'validated',
message: 'Stage validation passed, proceeding to next stage',
nextStage: workflow.nextStep,
};
}
break;
case 'revise':
result = {
status: 'revision_requested',
message: 'Revision requested for current stage',
feedback: feedback || 'No specific feedback provided',
action: 'Workflow will be revised based on feedback',
};
break;
case 'clarify':
result = {
status: 'clarification_needed',
message: 'Clarification requested',
action: 'Please provide additional information or requirements',
};
break;
}
return {
content: [
{
type: 'text',
text: JSON.stringify(result, null, 2),
},
],
};
} catch (error) {
return {
content: [
{
type: 'text',
text: `Validation error: ${error.message}`,
},
],
isError: true,
};
}
}
);
// Guide compliance check tool
this.server.registerTool(
'check-guide-compliance',
{
title: 'Check Guide Compliance',
description: 'Check guide compliance for current workflow state',
inputSchema: {
sessionId: z.string().describe('Workflow session ID'),
stepName: z.string().describe('Step name to check compliance for'),
},
},
async ({ sessionId, stepName }) => {
try {
const workflow = this.activeWorkflows.get(sessionId);
if (!workflow) {
return {
content: [
{
type: 'text',
text: `No active workflow found for session: ${sessionId}`,
},
],
isError: true,
};
}
const guideCompliance = await validateGuideCompliance(
stepName,
workflow.stepResults?.[stepName]?.input || {},
workflow.stepResults?.[stepName] || {},
workflow
);
return {
content: [
{
type: 'text',
text: JSON.stringify({
stepName,
complianceScore: guideCompliance.complianceScore,
maximsApplied: guideCompliance.maximsApplied,
heuristicsApplied: guideCompliance.heuristicsApplied,
validationCheckpoints: guideCompliance.validationCheckpoints,
recommendation: guideCompliance.complianceScore >= 70
? 'Guide compliance is satisfactory'
: 'Guide compliance needs improvement',
}, null, 2),
},
],
};
} catch (error) {
return {
content: [
{
type: 'text',
text: `Guide compliance check error: ${error.message}`,
},
],
isError: true,
};
}
}
);
}
/**
* Smart problem routing based on user prompt analysis
*/
private async smartProblemRouting(userPrompt: string): Promise<{
directFunction?: string;
input: any;
reasoning: string;
startStep?: string;
}> {
const promptLower = userPrompt.toLowerCase();
// Problem-specific routing patterns
// 1. Direct problem-solving scenarios
if (promptLower.includes('error') || promptLower.includes('fix') || promptLower.includes('debug') ||
promptLower.includes('issue') || promptLower.includes('problem') || promptLower.includes('bug') ||
promptLower.includes('does not send') || promptLower.includes('opens folder settings') ||
promptLower.includes('command detection') || promptLower.includes('wrong command')) {
// Check for specific error types
if (promptLower.includes('postmessage') || promptLower.includes('webview') ||
promptLower.includes('vscode') || promptLower.includes('extension')) {
return {
directFunction: 'problem-solver',
input: {
error: userPrompt,
context: {
domain: 'vscode-extension',
component: 'webview',
action: 'postMessage'
},
strictMode: true,
workflowContext: createWorkflowContext(userPrompt, 'problem-solver')
},
reasoning: 'Detected VSCode extension error - routing directly to problem solver'
};
}
// Generic error/problem
return {
directFunction: 'problem-solver',
input: {
error: userPrompt,
context: { userRequest: true },
strictMode: true,
workflowContext: createWorkflowContext(userPrompt, 'problem-solver')
},
reasoning: 'Detected error/problem - routing directly to problem solver'
};
}
// 2. Analysis-only requests
if (promptLower.includes('analyze') || promptLower.includes('understand') ||
promptLower.includes('explain') || promptLower.includes('what is') ||
promptLower.includes('how does')) {
return {
directFunction: 'cognitive',
input: {
researchData: { findings: ['User analysis request'] },
improvedPrompt: userPrompt,
analysisDepth: 'comprehensive',
context: createWorkflowContext(userPrompt, 'cognitive')
},
reasoning: 'Detected analysis request - routing directly to cognitive analysis'
};
}
// 3. Research requests
if (promptLower.includes('research') || promptLower.includes('find') ||
promptLower.includes('learn about') || promptLower.includes('information')) {
return {
directFunction: 'research',
input: {
improvedPrompt: userPrompt,
researchTopics: [userPrompt],
codeAnalysis: true,
context: createWorkflowContext(userPrompt, 'research')
},
reasoning: 'Detected research request - routing directly to research function'
};
}
// Default: Use full workflow for complex requests
return {
startStep: 'improve-prompt',
input: {},
reasoning: 'Complex request - using full workflow for comprehensive handling'
};
}
/**
* Start the MCP server
*/
async start(): Promise<void> {
const transport = new StdioServerTransport();
await this.server.connect(transport);
console.log('MCP Workflow Server started successfully');
}
}
// Start the server if this file is run directly
if (import.meta.url === `file://${process.argv[1]}`) {
const server = new MCPWorkflowServer();
server.start().catch(console.error);
}