keystone-cli 0.1.0

package/src/runner/mcp-server.ts

@@ -0,0 +1,436 @@
+import * as readline from 'node:readline';
+import { WorkflowDb } from '../db/workflow-db';
+import { WorkflowParser } from '../parser/workflow-parser';
+import { generateMermaidGraph } from '../utils/mermaid';
+import { WorkflowRegistry } from '../utils/workflow-registry';
+import { WorkflowSuspendedError } from './step-executor';
+import { WorkflowRunner } from './workflow-runner';
+
+interface MCPMessage {
+  jsonrpc: '2.0';
+  method: string;
+  params?: unknown;
+  id?: string | number;
+}
+
+export class MCPServer {
+  private db: WorkflowDb;
+
+  constructor(db?: WorkflowDb) {
+    this.db = db || new WorkflowDb();
+  }
+
+  async start() {
+    const rl = readline.createInterface({
+      input: process.stdin,
+      terminal: false,
+    });
+
+    rl.on('line', async (line) => {
+      if (!line.trim()) return;
+
+      try {
+        const message = JSON.parse(line) as MCPMessage;
+        const response = await this.handleMessage(message);
+        if (response) {
+          process.stdout.write(`${JSON.stringify(response)}\n`);
+        }
+      } catch (error) {
+        console.error('Error handling MCP message:', error);
+      }
+    });
+  }
+
+  private async handleMessage(message: MCPMessage) {
+    const { method, params, id } = message;
+
+    switch (method) {
+      case 'initialize':
+        return {
+          jsonrpc: '2.0',
+          id,
+          result: {
+            protocolVersion: '2024-11-05',
+            capabilities: {
+              tools: {},
+            },
+            serverInfo: {
+              name: 'keystone-mcp',
+              version: '0.1.0',
+            },
+          },
+        };
+
+      case 'tools/list':
+        return {
+          jsonrpc: '2.0',
+          id,
+          result: {
+            tools: [
+              {
+                name: 'list_workflows',
+                description: 'List all available workflows and their required inputs.',
+                inputSchema: {
+                  type: 'object',
+                  properties: {},
+                },
+              },
+              {
+                name: 'run_workflow',
+                description: 'Execute a workflow by name.',
+                inputSchema: {
+                  type: 'object',
+                  properties: {
+                    workflow_name: {
+                      type: 'string',
+                      description: 'The name of the workflow to run (e.g., "deploy", "cleanup")',
+                    },
+                    inputs: {
+                      type: 'object',
+                      description: 'Key-value pairs for workflow inputs',
+                    },
+                  },
+                  required: ['workflow_name'],
+                },
+              },
+              {
+                name: 'get_run_logs',
+                description: 'Get the logs and status of a specific workflow run.',
+                inputSchema: {
+                  type: 'object',
+                  properties: {
+                    run_id: { type: 'string' },
+                  },
+                  required: ['run_id'],
+                },
+              },
+              {
+                name: 'get_workflow_graph',
+                description: 'Get a visual diagram (Mermaid.js) of the workflow structure.',
+                inputSchema: {
+                  type: 'object',
+                  properties: {
+                    workflow_name: { type: 'string' },
+                  },
+                  required: ['workflow_name'],
+                },
+              },
+              {
+                name: 'answer_human_input',
+                description:
+                  'Provide input to a workflow that is paused waiting for human interaction.',
+                inputSchema: {
+                  type: 'object',
+                  properties: {
+                    run_id: { type: 'string', description: 'The ID of the paused run' },
+                    input: {
+                      type: 'string',
+                      description: 'The text input or "confirm" for confirmation steps',
+                    },
+                  },
+                  required: ['run_id', 'input'],
+                },
+              },
+            ],
+          },
+        };
+
+      case 'tools/call': {
+        const toolParams = params as { name: string; arguments: Record<string, unknown> };
+
+        try {
+          // --- Tool: list_workflows ---
+          if (toolParams.name === 'list_workflows') {
+            const workflows = WorkflowRegistry.listWorkflows();
+            return {
+              jsonrpc: '2.0',
+              id,
+              result: {
+                content: [{ type: 'text', text: JSON.stringify(workflows, null, 2) }],
+              },
+            };
+          }
+
+          // --- Tool: run_workflow ---
+          if (toolParams.name === 'run_workflow') {
+            const { workflow_name, inputs } = toolParams.arguments as {
+              workflow_name: string;
+              inputs: Record<string, unknown>;
+            };
+
+            const path = WorkflowRegistry.resolvePath(workflow_name);
+            const workflow = WorkflowParser.loadWorkflow(path);
+
+            // Use a custom logger that captures logs for the MCP response
+            const logs: string[] = [];
+            const logger = {
+              log: (msg: string) => logs.push(msg),
+              error: (msg: string) => logs.push(`ERROR: ${msg}`),
+              warn: (msg: string) => logs.push(`WARN: ${msg}`),
+            };
+
+            const runner = new WorkflowRunner(workflow, {
+              inputs,
+              logger,
+            });
+
+            // Note: This waits for completion. For long workflows, we might want to
+            // return the run_id immediately and let the agent poll via get_run_logs.
+            // For now, synchronous is easier for the agent to reason about.
+            let outputs: Record<string, unknown> | undefined;
+            try {
+              outputs = await runner.run();
+            } catch (error) {
+              if (error instanceof WorkflowSuspendedError) {
+                return {
+                  jsonrpc: '2.0',
+                  id,
+                  result: {
+                    content: [
+                      {
+                        type: 'text',
+                        text: JSON.stringify(
+                          {
+                            status: 'paused',
+                            run_id: runner.getRunId(),
+                            message: error.message,
+                            step_id: error.stepId,
+                            input_type: error.inputType,
+                            instructions:
+                              error.inputType === 'confirm'
+                                ? 'Use answer_human_input with input="confirm" to proceed.'
+                                : 'Use answer_human_input with the required text input.',
+                          },
+                          null,
+                          2
+                        ),
+                      },
+                    ],
+                  },
+                };
+              }
+              // Even if it fails, we return the logs so the agent knows why
+              return {
+                jsonrpc: '2.0',
+                id,
+                result: {
+                  isError: true,
+                  content: [
+                    {
+                      type: 'text',
+                      text: `Workflow failed.\n\nLogs:\n${logs.join('\n')}`,
+                    },
+                  ],
+                },
+              };
+            }
+
+            return {
+              jsonrpc: '2.0',
+              id,
+              result: {
+                content: [
+                  {
+                    type: 'text',
+                    text: JSON.stringify(
+                      {
+                        status: 'success',
+                        outputs,
+                        logs: logs.slice(-20), // Return last 20 lines to avoid token limits
+                      },
+                      null,
+                      2
+                    ),
+                  },
+                ],
+              },
+            };
+          }
+
+          // --- Tool: get_run_logs ---
+          if (toolParams.name === 'get_run_logs') {
+            const { run_id } = toolParams.arguments as { run_id: string };
+            const run = this.db.getRun(run_id);
+
+            if (!run) {
+              throw new Error(`Run ID ${run_id} not found`);
+            }
+
+            const steps = this.db.getStepsByRun(run_id);
+            const summary = {
+              workflow: run.workflow_name,
+              status: run.status,
+              error: run.error,
+              steps: steps.map((s) => ({
+                step: s.step_id,
+                status: s.status,
+                error: s.error,
+                output: s.output ? JSON.parse(s.output) : null,
+              })),
+            };
+
+            return {
+              jsonrpc: '2.0',
+              id,
+              result: {
+                content: [{ type: 'text', text: JSON.stringify(summary, null, 2) }],
+              },
+            };
+          }
+
+          // --- Tool: get_workflow_graph ---
+          if (toolParams.name === 'get_workflow_graph') {
+            const { workflow_name } = toolParams.arguments as { workflow_name: string };
+            const path = WorkflowRegistry.resolvePath(workflow_name);
+            const workflow = WorkflowParser.loadWorkflow(path);
+
+            const mermaid = generateMermaidGraph(workflow);
+
+            return {
+              jsonrpc: '2.0',
+              id,
+              result: {
+                content: [
+                  {
+                    type: 'text',
+                    text: `Here is the graph for **${workflow_name}**:\n\n\`\`\`mermaid\n${mermaid}\n\`\`\``,
+                  },
+                ],
+              },
+            };
+          }
+
+          // --- Tool: answer_human_input ---
+          if (toolParams.name === 'answer_human_input') {
+            const { run_id, input } = toolParams.arguments as { run_id: string; input: string };
+            const run = this.db.getRun(run_id);
+            if (!run) {
+              throw new Error(`Run ID ${run_id} not found`);
+            }
+
+            if (run.status !== 'paused') {
+              throw new Error(`Run ${run_id} is not paused (status: ${run.status})`);
+            }
+
+            // Find the pending human step
+            const steps = this.db.getStepsByRun(run_id);
+            const pendingStep = steps.find((s) => s.status === 'pending');
+            if (!pendingStep) {
+              throw new Error(`No pending step found for run ${run_id}`);
+            }
+
+            // Fulfill the step in the DB
+            const output = input === 'confirm' ? true : input;
+            await this.db.completeStep(pendingStep.id, 'success', output);
+
+            // Resume the workflow
+            const path = WorkflowRegistry.resolvePath(run.workflow_name);
+            const workflow = WorkflowParser.loadWorkflow(path);
+
+            const logs: string[] = [];
+            const logger = {
+              log: (msg: string) => logs.push(msg),
+              error: (msg: string) => logs.push(`ERROR: ${msg}`),
+              warn: (msg: string) => logs.push(`WARN: ${msg}`),
+            };
+
+            const runner = new WorkflowRunner(workflow, {
+              resumeRunId: run_id,
+              logger,
+            });
+
+            let outputs: Record<string, unknown> | undefined;
+            try {
+              outputs = await runner.run();
+            } catch (error) {
+              if (error instanceof WorkflowSuspendedError) {
+                return {
+                  jsonrpc: '2.0',
+                  id,
+                  result: {
+                    content: [
+                      {
+                        type: 'text',
+                        text: JSON.stringify(
+                          {
+                            status: 'paused',
+                            run_id: runner.getRunId(),
+                            message: error.message,
+                            step_id: error.stepId,
+                            input_type: error.inputType,
+                            instructions:
+                              error.inputType === 'confirm'
+                                ? 'Use answer_human_input with input="confirm" to proceed.'
+                                : 'Use answer_human_input with the required text input.',
+                          },
+                          null,
+                          2
+                        ),
+                      },
+                    ],
+                  },
+                };
+              }
+
+              return {
+                jsonrpc: '2.0',
+                id,
+                result: {
+                  isError: true,
+                  content: [
+                    {
+                      type: 'text',
+                      text: `Workflow failed after resume.\n\nLogs:\n${logs.join('\n')}`,
+                    },
+                  ],
+                },
+              };
+            }
+
+            return {
+              jsonrpc: '2.0',
+              id,
+              result: {
+                content: [
+                  {
+                    type: 'text',
+                    text: JSON.stringify(
+                      {
+                        status: 'success',
+                        outputs,
+                        logs: logs.slice(-20),
+                      },
+                      null,
+                      2
+                    ),
+                  },
+                ],
+              },
+            };
+          }
+
+          throw new Error(`Unknown tool: ${toolParams.name}`);
+        } catch (error) {
+          return {
+            jsonrpc: '2.0',
+            id,
+            error: {
+              code: -32000,
+              message: error instanceof Error ? error.message : String(error),
+            },
+          };
+        }
+      }
+
+      default:
+        return {
+          jsonrpc: '2.0',
+          id,
+          error: {
+            code: -32601,
+            message: `Method not found: ${method}`,
+          },
+        };
+    }
+  }
+}

package/src/runner/retry.test.ts

@@ -0,0 +1,52 @@
+import { describe, expect, spyOn, test } from 'bun:test';
+import { withRetry } from './retry';
+
+describe('withRetry', () => {
+  test('should return result if fn succeeds on first try', async () => {
+    const fn = async () => 'success';
+    const result = await withRetry(fn, { count: 3, backoff: 'linear' });
+    expect(result).toBe('success');
+  });
+
+  test('should retry and succeed', async () => {
+    let attempts = 0;
+    const fn = async () => {
+      attempts++;
+      if (attempts < 3) throw new Error('fail');
+      return 'success';
+    };
+
+    // Use a small delay for tests if possible, but retry.ts has hardcoded 1000ms base delay.
+    // This test might take a few seconds.
+    const result = await withRetry(fn, { count: 3, backoff: 'linear' });
+    expect(result).toBe('success');
+    expect(attempts).toBe(3);
+  }, 10000); // 10s timeout
+
+  test('should throw after exhausting retries', async () => {
+    let attempts = 0;
+    const fn = async () => {
+      attempts++;
+      throw new Error('fail');
+    };
+
+    await expect(withRetry(fn, { count: 2, backoff: 'linear' })).rejects.toThrow('fail');
+    expect(attempts).toBe(3); // 1 original + 2 retries
+  }, 10000);
+
+  test('should call onRetry callback', async () => {
+    let attempts = 0;
+    const onRetry = (attempt: number, error: Error) => {
+      expect(attempt).toBeGreaterThan(0);
+      expect(error.message).toBe('fail');
+    };
+
+    const fn = async () => {
+      attempts++;
+      if (attempts < 2) throw new Error('fail');
+      return 'success';
+    };
+
+    await withRetry(fn, { count: 1, backoff: 'linear' }, onRetry);
+  }, 5000);
+});

package/src/runner/retry.ts

@@ -0,0 +1,58 @@
+import type { RetryConfig } from '../parser/schema.ts';
+
+/**
+ * Calculate backoff delay in milliseconds
+ */
+function calculateBackoff(
+  attempt: number,
+  backoff: 'linear' | 'exponential',
+  baseDelay = 1000
+): number {
+  if (backoff === 'exponential') {
+    return baseDelay * 2 ** attempt;
+  }
+  // Linear backoff
+  return baseDelay * (attempt + 1);
+}
+
+/**
+ * Sleep for a given duration
+ */
+function sleep(ms: number): Promise<void> {
+  return new Promise((resolve) => setTimeout(resolve, ms));
+}
+
+/**
+ * Execute a function with retry logic
+ */
+export async function withRetry<T>(
+  fn: () => Promise<T>,
+  retry?: RetryConfig,
+  onRetry?: (attempt: number, error: Error) => void
+): Promise<T> {
+  const maxRetries = retry?.count || 0;
+  const backoffType = retry?.backoff || 'linear';
+  const baseDelay = retry?.baseDelay ?? 1000;
+
+  let lastError: Error | undefined;
+
+  for (let attempt = 0; attempt <= maxRetries; attempt++) {
+    try {
+      return await fn();
+    } catch (error) {
+      lastError = error instanceof Error ? error : new Error(String(error));
+
+      // Don't retry if we've exhausted attempts
+      if (attempt >= maxRetries) {
+        break;
+      }
+
+      // Calculate delay and wait before retry
+      const delay = calculateBackoff(attempt, backoffType, baseDelay);
+      onRetry?.(attempt + 1, lastError);
+      await sleep(delay);
+    }
+  }
+
+  throw lastError || new Error('Operation failed with no error details');
+}

package/src/runner/shell-executor.test.ts

@@ -0,0 +1,123 @@
+import { describe, expect, it } from 'bun:test';
+import type { ExpressionContext } from '../expression/evaluator';
+import type { ShellStep } from '../parser/schema';
+import { escapeShellArg, executeShell } from './shell-executor';
+
+describe('shell-executor', () => {
+  describe('escapeShellArg', () => {
+    it('should wrap in single quotes', () => {
+      expect(escapeShellArg('hello')).toBe("'hello'");
+    });
+
+    it('should escape single quotes', () => {
+      expect(escapeShellArg("don't")).toBe("'don'\\''t'");
+    });
+  });
+
+  describe('executeShell', () => {
+    const context: ExpressionContext = {
+      inputs: {},
+      steps: {},
+      env: {},
+    };
+
+    it('should execute a simple command', async () => {
+      const step: ShellStep = {
+        id: 'test',
+        type: 'shell',
+        run: 'echo "hello world"',
+      };
+
+      const result = await executeShell(step, context);
+      expect(result.stdout.trim()).toBe('hello world');
+      expect(result.exitCode).toBe(0);
+    });
+
+    it('should evaluate expressions in the command', async () => {
+      const step: ShellStep = {
+        id: 'test',
+        type: 'shell',
+        run: 'echo "${{ inputs.name }}"',
+      };
+      const customContext: ExpressionContext = {
+        ...context,
+        inputs: { name: 'world' },
+      };
+
+      const result = await executeShell(step, customContext);
+      expect(result.stdout.trim()).toBe('world');
+    });
+
+    it('should handle environment variables', async () => {
+      const step: ShellStep = {
+        id: 'test',
+        type: 'shell',
+        run: 'echo $TEST_VAR',
+        env: {
+          TEST_VAR: 'env-value',
+        },
+      };
+
+      const result = await executeShell(step, context);
+      expect(result.stdout.trim()).toBe('env-value');
+    });
+
+    it('should handle working directory', async () => {
+      const step: ShellStep = {
+        id: 'test',
+        type: 'shell',
+        run: 'pwd',
+        dir: '/tmp',
+      };
+
+      const result = await executeShell(step, context);
+      expect(result.stdout.trim()).toMatch(/\/tmp$/);
+    });
+
+    it('should capture stderr', async () => {
+      const step: ShellStep = {
+        id: 'test',
+        type: 'shell',
+        run: 'echo "error" >&2',
+      };
+
+      const result = await executeShell(step, context);
+      expect(result.stderr.trim()).toBe('error');
+    });
+
+    it('should handle non-zero exit codes', async () => {
+      const step: ShellStep = {
+        id: 'test',
+        type: 'shell',
+        run: 'exit 1',
+      };
+
+      const result = await executeShell(step, context);
+      expect(result.exitCode).toBe(1);
+    });
+
+    it('should warn about shell injection risk', async () => {
+      const spy = console.warn;
+      let warned = false;
+      console.warn = (...args: unknown[]) => {
+        const msg = args[0];
+        if (
+          typeof msg === 'string' &&
+          msg.includes('WARNING: Command contains shell metacharacters')
+        ) {
+          warned = true;
+        }
+      };
+
+      const step: ShellStep = {
+        id: 'test',
+        type: 'shell',
+        run: 'echo "hello" ; rm -rf /tmp/foo',
+      };
+
+      await executeShell(step, context);
+      expect(warned).toBe(true);
+      console.warn = spy; // Restore
+    });
+  });
+});