boocode/packages/ion/src/schema/__tests__/dag-node.test.ts

import { describe, it, expect } from 'vitest';
import {
  dagNodeSchema,
  promptNodeSchema,
  commandNodeSchema,
  bashNodeSchema,
  scriptNodeSchema,
  loopNodeSchema,
  approvalNodeSchema,
  cancelNodeSchema,
  loopNodeConfigSchema,
  stepRetryConfigSchema,
  triggerRuleSchema,
  isBashNode,
  isLoopNode,
  isApprovalNode,
  isCancelNode,
  isScriptNode,
  isPromptNode,
  isCommandNode,
  effortLevelSchema,
  thinkingConfigSchema,
  approvalOnRejectSchema,
  dagNodeBaseSchema,
} from '../index.js';
import type {
  DagNode,
  PromptNode,
  CommandNode,
  BashNode,
  ScriptNode,
  LoopNode,
  ApprovalNode,
  CancelNode,
} from '../index.js';

const validBase = {
  id: 'test-node',
  depends_on: [],
};

describe('dagNodeSchema', () => {
  it('validates a prompt node', () => {
    const node = { ...validBase, kind: 'prompt' as const, prompt: 'Do the thing' };
    const result = dagNodeSchema.safeParse(node);
    expect(result.success).toBe(true);
  });

  it('validates a command node', () => {
    const node = { ...validBase, kind: 'command' as const, command: 'echo hello' };
    const result = dagNodeSchema.safeParse(node);
    expect(result.success).toBe(true);
  });

  it('validates a bash node', () => {
    const node = { ...validBase, kind: 'bash' as const, bash: 'echo hello' };
    const result = dagNodeSchema.safeParse(node);
    expect(result.success).toBe(true);
  });

  it('validates a script node', () => {
    const node = { ...validBase, kind: 'script' as const, script: 'print(1)', runtime: 'bun' as const };
    const result = dagNodeSchema.safeParse(node);
    expect(result.success).toBe(true);
  });

  it('validates a loop node', () => {
    const node = {
      ...validBase,
      kind: 'loop' as const,
      config: { prompt: 'iterate', until: 'done', max_iterations: 5 },
    };
    const result = dagNodeSchema.safeParse(node);
    expect(result.success).toBe(true);
  });

  it('validates an approval node', () => {
    const node = { ...validBase, kind: 'approval' as const, message: 'Approve this?' };
    const result = dagNodeSchema.safeParse(node);
    expect(result.success).toBe(true);
  });

  it('validates a cancel node', () => {
    const node = { ...validBase, kind: 'cancel' as const };
    const result = dagNodeSchema.safeParse(node);
    expect(result.success).toBe(true);
  });

  it('rejects a node with no kind field', () => {
    const result = dagNodeSchema.safeParse({ id: 'x', prompt: 'y' });
    expect(result.success).toBe(false);
  });

  it('rejects a node with an invalid kind', () => {
    const result = dagNodeSchema.safeParse({ ...validBase, kind: 'unknown' });
    expect(result.success).toBe(false);
  });
});

describe('node type unique fields', () => {
  it('prompt node requires prompt or command_file', () => {
    const result = promptNodeSchema.safeParse({ ...validBase, kind: 'prompt' });
    expect(result.success).toBe(true);
  });

  it('command node requires command', () => {
    const result = commandNodeSchema.safeParse({ ...validBase, kind: 'command' });
    expect(result.success).toBe(false);
  });

  it('bash node requires bash', () => {
    const result = bashNodeSchema.safeParse({ ...validBase, kind: 'bash' });
    expect(result.success).toBe(false);
  });

  it('script node requires script and runtime', () => {
    const noScript = scriptNodeSchema.safeParse({ ...validBase, kind: 'script', runtime: 'bun' });
    expect(noScript.success).toBe(false);

    const noRuntime = scriptNodeSchema.safeParse({ ...validBase, kind: 'script', script: 'print(1)' });
    expect(noRuntime.success).toBe(false);

    const valid = scriptNodeSchema.safeParse({
      ...validBase,
      kind: 'script',
      script: 'print(1)',
      runtime: 'bun',
    });
    expect(valid.success).toBe(true);
  });

  it('loop node requires config', () => {
    const result = loopNodeSchema.safeParse({ ...validBase, kind: 'loop' });
    expect(result.success).toBe(false);
  });

  it('approval node requires message', () => {
    const result = approvalNodeSchema.safeParse({ ...validBase, kind: 'approval' });
    expect(result.success).toBe(false);
  });

  it('cancel node does not require reason', () => {
    const result = cancelNodeSchema.safeParse({ ...validBase, kind: 'cancel' });
    expect(result.success).toBe(true);
  });
});

describe('loopNodeConfigSchema', () => {
  it('validates a minimal config', () => {
    const result = loopNodeConfigSchema.safeParse({
      prompt: 'iterate',
      until: 'done',
      max_iterations: 5,
    });
    expect(result.success).toBe(true);
  });

  it('rejects empty prompt', () => {
    const result = loopNodeConfigSchema.safeParse({
      prompt: '',
      until: 'done',
      max_iterations: 5,
    });
    expect(result.success).toBe(false);
  });

  it('rejects empty until', () => {
    const result = loopNodeConfigSchema.safeParse({
      prompt: 'iterate',
      until: '',
      max_iterations: 5,
    });
    expect(result.success).toBe(false);
  });

  it('rejects non-positive max_iterations', () => {
    const zeroResult = loopNodeConfigSchema.safeParse({
      prompt: 'iterate',
      until: 'done',
      max_iterations: 0,
    });
    expect(zeroResult.success).toBe(false);

    const negResult = loopNodeConfigSchema.safeParse({
      prompt: 'iterate',
      until: 'done',
      max_iterations: -1,
    });
    expect(negResult.success).toBe(false);
  });

  it('requires gate_message when interactive is true', () => {
    const noGate = loopNodeConfigSchema.safeParse({
      prompt: 'iterate',
      until: 'done',
      max_iterations: 5,
      interactive: true,
    });
    expect(noGate.success).toBe(false);

    const withGate = loopNodeConfigSchema.safeParse({
      prompt: 'iterate',
      until: 'done',
      max_iterations: 5,
      interactive: true,
      gate_message: 'Approve this iteration?',
    });
    expect(withGate.success).toBe(true);
  });

  it('allows interactive=false without gate_message', () => {
    const result = loopNodeConfigSchema.safeParse({
      prompt: 'iterate',
      until: 'done',
      max_iterations: 5,
      interactive: false,
    });
    expect(result.success).toBe(true);
  });
});

describe('stepRetryConfigSchema', () => {
  it('validates a minimal retry config', () => {
    const result = stepRetryConfigSchema.safeParse({ max_attempts: 3 });
    expect(result.success).toBe(true);
  });

  it('rejects max_attempts below 1', () => {
    const result = stepRetryConfigSchema.safeParse({ max_attempts: 0 });
    expect(result.success).toBe(false);
  });

  it('rejects max_attempts above 5', () => {
    const result = stepRetryConfigSchema.safeParse({ max_attempts: 6 });
    expect(result.success).toBe(false);
  });

  it('accepts max_attempts at boundaries (1 and 5)', () => {
    expect(stepRetryConfigSchema.safeParse({ max_attempts: 1 }).success).toBe(true);
    expect(stepRetryConfigSchema.safeParse({ max_attempts: 5 }).success).toBe(true);
  });

  it('rejects delay_ms below 1000', () => {
    const result = stepRetryConfigSchema.safeParse({ max_attempts: 2, delay_ms: 500 });
    expect(result.success).toBe(false);
  });

  it('rejects delay_ms above 60000', () => {
    const result = stepRetryConfigSchema.safeParse({ max_attempts: 2, delay_ms: 70000 });
    expect(result.success).toBe(false);
  });

  it('accepts delay_ms at boundaries (1000 and 60000)', () => {
    expect(stepRetryConfigSchema.safeParse({ max_attempts: 2, delay_ms: 1000 }).success).toBe(true);
    expect(stepRetryConfigSchema.safeParse({ max_attempts: 2, delay_ms: 60000 }).success).toBe(true);
  });

  it('defaults on_error to transient', () => {
    const result = stepRetryConfigSchema.parse({ max_attempts: 2 });
    expect(result.on_error).toBe('transient');
  });
});

describe('triggerRuleSchema', () => {
  it('validates all trigger rules', () => {
    const rules = ['all_success', 'one_success', 'all_done', 'none_failed_min_one_success'];
    for (const rule of rules) {
      expect(triggerRuleSchema.safeParse(rule).success).toBe(true);
    }
  });

  it('rejects invalid trigger rules', () => {
    const result = triggerRuleSchema.safeParse('invalid_rule');
    expect(result.success).toBe(false);
  });
});

describe('type guards', () => {
  const nodes: DagNode[] = [
    { ...validBase, kind: 'prompt', prompt: 'test' } as PromptNode,
    { ...validBase, kind: 'command', command: 'echo' } as CommandNode,
    { ...validBase, kind: 'bash', bash: 'echo' } as BashNode,
    { ...validBase, kind: 'script', script: 'print(1)', runtime: 'bun' } as ScriptNode,
    { ...validBase, kind: 'loop', config: { prompt: 'x', until: 'y', max_iterations: 3 } } as LoopNode,
    { ...validBase, kind: 'approval', message: 'ok?' } as ApprovalNode,
    { ...validBase, kind: 'cancel' } as CancelNode,
  ];

  it('isPromptNode identifies prompt nodes', () => {
    expect(isPromptNode(nodes[0])).toBe(true);
    expect(isPromptNode(nodes[1])).toBe(false);
  });

  it('isCommandNode identifies command nodes', () => {
    expect(isCommandNode(nodes[1])).toBe(true);
    expect(isCommandNode(nodes[0])).toBe(false);
  });

  it('isBashNode identifies bash nodes', () => {
    expect(isBashNode(nodes[2])).toBe(true);
    expect(isBashNode(nodes[0])).toBe(false);
  });

  it('isScriptNode identifies script nodes', () => {
    expect(isScriptNode(nodes[3])).toBe(true);
    expect(isScriptNode(nodes[0])).toBe(false);
  });

  it('isLoopNode identifies loop nodes', () => {
    expect(isLoopNode(nodes[4])).toBe(true);
    expect(isLoopNode(nodes[0])).toBe(false);
  });

  it('isApprovalNode identifies approval nodes', () => {
    expect(isApprovalNode(nodes[5])).toBe(true);
    expect(isApprovalNode(nodes[0])).toBe(false);
  });

  it('isCancelNode identifies cancel nodes', () => {
    expect(isCancelNode(nodes[6])).toBe(true);
    expect(isCancelNode(nodes[0])).toBe(false);
  });
});

describe('effortLevelSchema', () => {
  it('validates all effort levels', () => {
    expect(effortLevelSchema.safeParse('low').success).toBe(true);
    expect(effortLevelSchema.safeParse('medium').success).toBe(true);
    expect(effortLevelSchema.safeParse('high').success).toBe(true);
  });

  it('rejects invalid effort levels', () => {
    expect(effortLevelSchema.safeParse('extreme').success).toBe(false);
  });
});

describe('thinkingConfigSchema', () => {
  it('validates a minimal thinking config', () => {
    const result = thinkingConfigSchema.safeParse({});
    expect(result.success).toBe(true);
    if (result.success) {
      expect(result.data.enabled).toBe(false);
    }
  });

  it('validates a full thinking config', () => {
    const result = thinkingConfigSchema.safeParse({ enabled: true, max_tokens: 4096 });
    expect(result.success).toBe(true);
  });
});

describe('approvalOnRejectSchema', () => {
  it('validates all on-reject actions', () => {
    expect(approvalOnRejectSchema.safeParse('retry').success).toBe(true);
    expect(approvalOnRejectSchema.safeParse('fail').success).toBe(true);
    expect(approvalOnRejectSchema.safeParse('skip').success).toBe(true);
  });

  it('rejects invalid on-reject actions', () => {
    expect(approvalOnRejectSchema.safeParse('restart').success).toBe(false);
  });
});

describe('dagNodeBaseSchema', () => {
  it('validates a minimal base node', () => {
    const result = dagNodeBaseSchema.safeParse({
      id: 'node-1',
      kind: 'prompt',
    });
    expect(result.success).toBe(true);
  });

  it('defaults depends_on to empty array', () => {
    const result = dagNodeBaseSchema.parse({ id: 'node-1', kind: 'prompt' });
    expect(result.depends_on).toEqual([]);
  });
});