keystone-cli 0.8.0 → 1.0.0

package/src/runner/workflow-runner.test.ts

@@ -1,13 +1,19 @@
-import { afterAll, afterEach, beforeEach, describe, expect, it, mock, spyOn } from 'bun:test';
+import { afterAll, afterEach, describe, expect, it, mock, spyOn } from 'bun:test';
 import { existsSync, rmSync } from 'node:fs';
 import { WorkflowDb } from '../db/workflow-db';
 import type { Workflow } from '../parser/schema';
 import { WorkflowParser } from '../parser/workflow-parser';
+import { ConfigLoader } from '../utils/config-loader';
 import { WorkflowRegistry } from '../utils/workflow-registry';
 import { WorkflowRunner } from './workflow-runner';
 
 describe('WorkflowRunner', () => {
   const dbPath = ':memory:';
+  const activeSpies: Array<{ mockRestore: () => void }> = [];
+  const trackSpy = <T extends { mockRestore: () => void }>(spy: T): T => {
+    activeSpies.push(spy);
+    return spy;
+  };
 
   afterAll(() => {
     if (existsSync('test-resume.db')) {
@@ -18,8 +24,12 @@ describe('WorkflowRunner', () => {
     }
   });
 
-  beforeEach(() => {
-    mock.restore();
+  afterEach(() => {
+    for (const spy of activeSpies) {
+      spy.mockRestore();
+    }
+    activeSpies.length = 0;
+    ConfigLoader.clear();
   });
 
   const workflow: Workflow = {
@@ -49,6 +59,63 @@ describe('WorkflowRunner', () => {
     expect(outputs.final).toBe('hello world');
   });
 
+  it('should expose workflow env to shell steps', async () => {
+    const envWorkflow: Workflow = {
+      name: 'env-workflow',
+      inputs: {
+        token: { type: 'string', default: 'env-token' },
+      },
+      env: {
+        TOKEN: '${{ inputs.token }}',
+      },
+      steps: [
+        {
+          id: 'print',
+          type: 'shell',
+          run: 'echo $TOKEN',
+          needs: [],
+        },
+      ],
+      outputs: {
+        token: '${{ steps.print.output.stdout.trim() }}',
+      },
+    } as unknown as Workflow;
+
+    const runner = new WorkflowRunner(envWorkflow, { dbPath });
+    const outputs = await runner.run();
+    expect(outputs.token).toBe('env-token');
+  });
+
+  it('should skip workflow env entries that fail to evaluate', async () => {
+    const envWorkflow: Workflow = {
+      name: 'env-skip-workflow',
+      env: {
+        LATER: '${{ steps.after.output.stdout.trim() }}',
+      },
+      steps: [
+        {
+          id: 'before',
+          type: 'shell',
+          run: 'echo "start"',
+          needs: [],
+        },
+        {
+          id: 'after',
+          type: 'shell',
+          run: 'echo "later"',
+          needs: ['before'],
+        },
+      ],
+      outputs: {
+        first: '${{ steps.before.output.stdout.trim() }}',
+      },
+    } as unknown as Workflow;
+
+    const runner = new WorkflowRunner(envWorkflow, { dbPath });
+    const outputs = await runner.run();
+    expect(outputs.first).toBe('start');
+  });
+
   it('should handle foreach steps', async () => {
     const foreachWorkflow: Workflow = {
       name: 'foreach-workflow',
@@ -155,6 +222,85 @@ describe('WorkflowRunner', () => {
     expect(outputs).toBeDefined();
   });
 
+  it('should validate step input schema', async () => {
+    const schemaDbPath = 'test-step-input-schema.db';
+    const workflowWithInputSchema: Workflow = {
+      name: 'step-input-schema-wf',
+      steps: [
+        {
+          id: 's1',
+          type: 'shell',
+          run: 'echo "hello"',
+          needs: [],
+          inputSchema: {
+            type: 'object',
+            properties: { run: { type: 'number' } },
+            required: ['run'],
+            additionalProperties: false,
+          },
+        },
+      ],
+    } as unknown as Workflow;
+
+    const runner = new WorkflowRunner(workflowWithInputSchema, { dbPath: schemaDbPath });
+    await expect(runner.run()).rejects.toThrow(/Step s1 failed/);
+
+    const db = new WorkflowDb(schemaDbPath);
+    const steps = await db.getStepsByRun(runner.runId);
+    db.close();
+
+    expect(steps[0]?.error || '').toMatch(/Input schema validation failed/);
+    if (existsSync(schemaDbPath)) rmSync(schemaDbPath);
+  });
+
+  it('should validate step output schema', async () => {
+    const schemaDbPath = 'test-step-output-schema.db';
+    const workflowWithOutputSchema: Workflow = {
+      name: 'step-output-schema-wf',
+      steps: [
+        {
+          id: 's1',
+          type: 'shell',
+          run: 'echo "hello"',
+          needs: [],
+          outputSchema: {
+            type: 'object',
+            properties: { ok: { const: true } },
+            required: ['ok'],
+            additionalProperties: false,
+          },
+        },
+      ],
+    } as unknown as Workflow;
+
+    const runner = new WorkflowRunner(workflowWithOutputSchema, { dbPath: schemaDbPath });
+    await expect(runner.run()).rejects.toThrow(/Step s1 failed/);
+
+    const db = new WorkflowDb(schemaDbPath);
+    const steps = await db.getStepsByRun(runner.runId);
+    db.close();
+
+    expect(steps[0]?.error || '').toMatch(/Output schema validation failed/);
+    if (existsSync(schemaDbPath)) rmSync(schemaDbPath);
+  });
+
+  it('should enforce input enums', async () => {
+    const workflowWithEnums: Workflow = {
+      name: 'enum-wf',
+      inputs: {
+        mode: { type: 'string', values: ['fast', 'slow'] },
+      },
+      steps: [{ id: 's1', type: 'shell', run: 'echo "${{ inputs.mode }}"', needs: [] }],
+    } as unknown as Workflow;
+
+    const runner = new WorkflowRunner(workflowWithEnums, {
+      dbPath,
+      inputs: { mode: 'invalid' },
+    });
+
+    await expect(runner.run()).rejects.toThrow(/must be one of/);
+  });
+
   it('should handle step failure and workflow failure', async () => {
     const failWorkflow: Workflow = {
       name: 'fail-wf',
@@ -164,6 +310,185 @@ describe('WorkflowRunner', () => {
     await expect(runner.run()).rejects.toThrow(/Step fail failed/);
   });
 
+  it('should execute errors block when a step fails', async () => {
+    let errorsBlockExecuted = false;
+    let lastFailedStepId: string | undefined;
+    const runnerLogger = {
+      log: (msg: string) => {
+        if (msg.includes('Executing errors step: err1')) {
+          errorsBlockExecuted = true;
+        }
+      },
+      error: () => {},
+      warn: () => {},
+      info: () => {},
+      debug: () => {},
+    };
+
+    const errorsWorkflow: Workflow = {
+      name: 'errors-wf',
+      steps: [{ id: 's1', type: 'shell', run: 'exit 1', needs: [] }],
+      errors: [
+        {
+          id: 'err1',
+          type: 'shell',
+          run: 'echo "Handling failure of ${{ last_failed_step.id }}"',
+          needs: [],
+        },
+      ],
+    } as unknown as Workflow;
+
+    const runner = new WorkflowRunner(errorsWorkflow, { dbPath, logger: runnerLogger });
+    try {
+      await runner.run();
+    } catch (e) {
+      // Expected to fail
+    }
+
+    expect(errorsBlockExecuted).toBe(true);
+  });
+
+  it('should continue when allowFailure is true', async () => {
+    const allowFailureWorkflow: Workflow = {
+      name: 'allow-failure-wf',
+      steps: [
+        { id: 'fail', type: 'shell', run: 'exit 1', needs: [], allowFailure: true },
+        { id: 'next', type: 'shell', run: 'echo ok', needs: ['fail'] },
+      ],
+      outputs: {
+        status: '${{ steps.fail.status }}',
+        error: '${{ steps.fail.error }}',
+      },
+    } as unknown as Workflow;
+
+    const runner = new WorkflowRunner(allowFailureWorkflow, { dbPath });
+    const outputs = await runner.run();
+    expect(outputs.status).toBe('success');
+    expect(String(outputs.error || '')).toMatch(/exit 1|Step failed|Shell command exited/);
+  });
+
+  it('should deduplicate steps using idempotencyKey within a run', async () => {
+    const idempotencyDbPath = 'test-idempotency.db';
+    if (existsSync(idempotencyDbPath)) rmSync(idempotencyDbPath);
+
+    let idempotencyHitCount = 0;
+    const runnerLogger = {
+      log: (msg: string) => {
+        if (msg.includes('idempotency hit')) {
+          idempotencyHitCount++;
+        }
+      },
+      error: () => {},
+      warn: () => {},
+      info: () => {},
+      debug: () => {},
+    };
+
+    const idempotencyWorkflow: Workflow = {
+      name: 'idempotency-wf',
+      steps: [
+        {
+          id: 's1',
+          type: 'shell',
+          run: 'echo "executed"',
+          needs: [],
+          idempotencyKey: '"fixed-key-123"',
+        },
+        {
+          id: 's2',
+          type: 'shell',
+          run: 'echo "second"',
+          needs: ['s1'],
+          idempotencyKey: '"fixed-key-123"',
+        },
+      ],
+      outputs: {
+        out1: '${{ steps.s1.output.stdout.trim() }}',
+        out2: '${{ steps.s2.output.stdout.trim() }}',
+      },
+    } as unknown as Workflow;
+
+    const runner = new WorkflowRunner(idempotencyWorkflow, {
+      dbPath: idempotencyDbPath,
+      logger: runnerLogger,
+    });
+    const outputs = await runner.run();
+    expect(outputs.out1).toBe('executed');
+    expect(outputs.out2).toBe('executed');
+    expect(idempotencyHitCount).toBe(1);
+
+    if (existsSync(idempotencyDbPath)) rmSync(idempotencyDbPath);
+  });
+
+  it('should allow disabling idempotency deduplication', async () => {
+    const idempotencyDbPath = 'test-idempotency-disabled.db';
+    if (existsSync(idempotencyDbPath)) rmSync(idempotencyDbPath);
+
+    const idempotencyWorkflow: Workflow = {
+      name: 'idempotency-disabled-wf',
+      steps: [
+        {
+          id: 's1',
+          type: 'shell',
+          run: 'echo "first"',
+          needs: [],
+          idempotencyKey: '"fixed-key-123"',
+        },
+        {
+          id: 's2',
+          type: 'shell',
+          run: 'echo "second"',
+          needs: ['s1'],
+          idempotencyKey: '"fixed-key-123"',
+        },
+      ],
+      outputs: {
+        out1: '${{ steps.s1.output.stdout.trim() }}',
+        out2: '${{ steps.s2.output.stdout.trim() }}',
+      },
+    } as unknown as Workflow;
+
+    const runner = new WorkflowRunner(idempotencyWorkflow, {
+      dbPath: idempotencyDbPath,
+      dedup: false,
+    });
+    const outputs = await runner.run();
+    expect(outputs.out1).toBe('first');
+    expect(outputs.out2).toBe('second');
+
+    if (existsSync(idempotencyDbPath)) rmSync(idempotencyDbPath);
+  });
+
+  it('should detect in-flight idempotency keys', async () => {
+    const idempotencyDbPath = 'test-idempotency-inflight.db';
+    if (existsSync(idempotencyDbPath)) rmSync(idempotencyDbPath);
+
+    const idempotencyWorkflow: Workflow = {
+      name: 'idempotency-inflight-wf',
+      steps: [
+        {
+          id: 's1',
+          type: 'sleep',
+          duration: 50,
+          needs: [],
+          idempotencyKey: '"same-key"',
+        },
+        {
+          id: 's2',
+          type: 'sleep',
+          duration: 50,
+          needs: [],
+          idempotencyKey: '"same-key"',
+        },
+      ],
+    } as unknown as Workflow;
+
+    const runner = new WorkflowRunner(idempotencyWorkflow, { dbPath: idempotencyDbPath });
+    await expect(runner.run()).rejects.toThrow(/Idempotency key already in-flight/);
+
+    if (existsSync(idempotencyDbPath)) rmSync(idempotencyDbPath);
+  });
+
   it('should execute steps in parallel', async () => {
     const parallelWorkflow: Workflow = {
       name: 'parallel-wf',
@@ -211,12 +536,12 @@ describe('WorkflowRunner', () => {
         },
       ],
       outputs: {
-        final: '${{ steps.sub.output.out }}',
+        final: '${{ steps.sub.output.outputs.out }}',
       },
     } as unknown as Workflow;
 
-    spyOn(WorkflowRegistry, 'resolvePath').mockReturnValue('child.yaml');
-    spyOn(WorkflowParser, 'loadWorkflow').mockReturnValue(childWorkflow);
+    trackSpy(spyOn(WorkflowRegistry, 'resolvePath')).mockReturnValue('child.yaml');
+    trackSpy(spyOn(WorkflowParser, 'loadWorkflow')).mockReturnValue(childWorkflow);
 
     const runner = new WorkflowRunner(parentWorkflow, { dbPath });
     const outputs = await runner.run();
@@ -304,7 +629,7 @@ describe('WorkflowRunner', () => {
     try {
       await runner1.run();
     } catch (e) {
-      runId = runner1.getRunId();
+      runId = runner1.runId;
     }
 
     const fixedWorkflow: Workflow = {
@@ -333,21 +658,60 @@ describe('WorkflowRunner', () => {
       },
     } as unknown as Workflow;
 
-    // @ts-ignore
-    spyOn(WorkflowRunner.prototype, 'loadSecrets').mockReturnValue({
-      MY_SECRET: 'my-super-secret',
+    const secretValue = 'my-super-secret';
+    const runner = new WorkflowRunner(workflow, {
+      dbPath,
+      secrets: { KEYSTONE_TEST_REDACTION_SECRET: secretValue },
     });
+    await runner.run();
 
-    const runner = new WorkflowRunner(workflow, { dbPath });
+    expect(runner.redact(secretValue)).toBe('***REDACTED***');
+  });
+
+  it('should redact secret inputs at rest', async () => {
+    const dbFile = 'test-secret-at-rest.db';
+    const workflow: Workflow = {
+      name: 'secret-input-wf',
+      inputs: {
+        token: { type: 'string', secret: true },
+      },
+      steps: [{ id: 's1', type: 'shell', run: 'echo "ok"', needs: [] }],
+    } as unknown as Workflow;
+
+    ConfigLoader.setConfig({
+      default_provider: 'openai',
+      providers: {
+        openai: { type: 'openai' },
+      },
+      model_mappings: {},
+      storage: { retention_days: 30, redact_secrets_at_rest: true },
+      mcp_servers: {},
+      engines: { allowlist: {}, denylist: [] },
+      concurrency: { default: 10, pools: { llm: 2, shell: 5, http: 10, engine: 2 } },
+    });
+
+    const runner = new WorkflowRunner(workflow, {
+      dbPath: dbFile,
+      inputs: { token: 'super-secret' },
+    });
     await runner.run();
 
-    expect(runner.redact('my-super-secret')).toBe('***REDACTED***');
+    const db = new WorkflowDb(dbFile);
+    const run = await db.getRun(runner.runId);
+    db.close();
+
+    expect(run).toBeTruthy();
+    const persistedInputs = run ? JSON.parse(run.inputs) : {};
+    expect(persistedInputs.token).toBe('***REDACTED***');
+
+    ConfigLoader.clear();
+    if (existsSync(dbFile)) rmSync(dbFile);
   });
 
   it('should return run ID', () => {
     const runner = new WorkflowRunner(workflow, { dbPath });
-    expect(runner.getRunId()).toBeDefined();
-    expect(typeof runner.getRunId()).toBe('string');
+    expect(runner.runId).toBeDefined();
+    expect(typeof runner.runId).toBe('string');
   });
 
   it('should continue even if finally step fails', async () => {
@@ -455,7 +819,7 @@ describe('WorkflowRunner', () => {
         : undefined
     ).toBe('WorkflowSuspendedError');
 
-    const runId = runner1.getRunId();
+    const runId = runner1.runId;
 
     // Check DB status - parent should be 'paused' and step should be 'suspended'
     const db = new WorkflowDb(resumeDbPath);
@@ -544,4 +908,117 @@ describe('WorkflowRunner', () => {
 
     if (existsSync(resumeDbPath)) rmSync(resumeDbPath);
   });
+
+  it('should parse outputRetries and repairStrategy in step schema', async () => {
+    // This test verifies that the schema accepts outputRetries and repairStrategy
+    // Actual LLM repair would require mocking, so we just verify the config is parsed
+    const workflowWithRepair: Workflow = {
+      name: 'output-repair-wf',
+      steps: [
+        {
+          id: 's1',
+          type: 'shell',
+          run: 'echo "hello"',
+          needs: [],
+          outputRetries: 2,
+          repairStrategy: 'reask',
+        },
+      ],
+    } as unknown as Workflow;
+
+    // Verify the workflow can be created with these fields
+    expect(workflowWithRepair.steps[0].outputRetries).toBe(2);
+    expect(workflowWithRepair.steps[0].repairStrategy).toBe('reask');
+
+    // The runner should accept these fields without error
+    const runner = new WorkflowRunner(workflowWithRepair, { dbPath });
+    const outputs = await runner.run();
+    expect(outputs).toBeDefined();
+  });
+
+  it('should allow cancellation via abortSignal', async () => {
+    const cancelDbPath = 'test-cancel.db';
+    if (existsSync(cancelDbPath)) rmSync(cancelDbPath);
+
+    const workflow: Workflow = {
+      name: 'cancel-wf',
+      steps: [
+        { id: 's1', type: 'sleep', duration: 10, needs: [] },
+        { id: 's2', type: 'sleep', duration: 10, needs: ['s1'] },
+      ],
+    } as unknown as Workflow;
+
+    const runner = new WorkflowRunner(workflow, { dbPath: cancelDbPath, preventExit: true });
+
+    // Verify the abort signal is exposed and not yet aborted
+    expect(runner.abortSignal).toBeDefined();
+    expect(runner.abortSignal instanceof AbortSignal).toBe(true);
+    expect(runner.abortSignal.aborted).toBe(false);
+
+    // Run the workflow (short duration so it completes quickly)
+    await runner.run();
+
+    if (existsSync(cancelDbPath)) rmSync(cancelDbPath);
+  });
+
+  it('should resume from canceled state', async () => {
+    const resumeDbPath = 'test-cancel-resume.db';
+    if (existsSync(resumeDbPath)) rmSync(resumeDbPath);
+
+    const workflow: Workflow = {
+      name: 'cancel-resume-wf',
+      steps: [
+        { id: 's1', type: 'shell', run: 'echo "one"', needs: [] },
+        { id: 's2', type: 'shell', run: 'echo "two"', needs: ['s1'] },
+      ],
+      outputs: {
+        out: '${{ steps.s1.output.stdout.trim() }}-${{ steps.s2.output.stdout.trim() }}',
+      },
+    } as unknown as Workflow;
+
+    // Manually create a "canceled" state in the DB
+    const db = new WorkflowDb(resumeDbPath);
+    const runId = crypto.randomUUID();
+    await db.createRun(runId, workflow.name, {});
+    await db.updateRunStatus(runId, 'canceled', undefined, 'Canceled by user');
+
+    // Create a completed step 1
+    const step1Id = crypto.randomUUID();
+    await db.createStep(step1Id, runId, 's1');
+    await db.startStep(step1Id);
+    await db.completeStep(step1Id, 'success', { stdout: 'one\n', exitCode: 0 });
+    db.close();
+
+    // Resume from canceled state
+    let loggedResume = false;
+    const logger = {
+      log: (msg: string) => {
+        if (msg.includes('Resuming a previously canceled run')) {
+          loggedResume = true;
+        }
+      },
+      error: () => {},
+      warn: () => {},
+      info: () => {},
+    };
+
+    const runner = new WorkflowRunner(workflow, {
+      dbPath: resumeDbPath,
+      resumeRunId: runId,
+      // @ts-ignore
+      logger: logger,
+    });
+
+    const outputs = await runner.run();
+    expect(outputs.out).toBe('one-two');
+    expect(loggedResume).toBe(true);
+
+    // Verify final status is success
+    const finalDb = new WorkflowDb(resumeDbPath);
+    const finalRun = await finalDb.getRun(runId);
+    expect(finalRun?.status).toBe('success');
+    finalDb.close();
+
+    if (existsSync(resumeDbPath)) rmSync(resumeDbPath);
+  });
 });