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

import { z } from 'zod';
import { stepRetryConfigSchema } from './retry.js';
import { loopNodeConfigSchema } from './loop.js';
import { triggerRuleSchema } from './trigger-rule.js';

// ---------------------------------------------------------------------------
// Effort level
// ---------------------------------------------------------------------------

/** Effort level for AI model calls. */
export const effortLevelSchema = z.enum(['low', 'medium', 'high']);

export type EffortLevel = z.infer<typeof effortLevelSchema>;

// ---------------------------------------------------------------------------
// Thinking configuration
// ---------------------------------------------------------------------------

/** Configuration for extended thinking / chain-of-thought. */
export const thinkingConfigSchema = z.object({
  /** Whether thinking is enabled. */
  enabled: z.boolean().default(false),
  /** Maximum thinking tokens. */
  max_tokens: z.number().int().positive().optional(),
});

export type ThinkingConfig = z.infer<typeof thinkingConfigSchema>;

// ---------------------------------------------------------------------------
// Approval on-reject action
// ---------------------------------------------------------------------------

/** What to do when an approval node is rejected. */
export const approvalOnRejectSchema = z.enum(['retry', 'fail', 'skip']);

export type ApprovalOnReject = z.infer<typeof approvalOnRejectSchema>;

// ---------------------------------------------------------------------------
// Base DAG node
// ---------------------------------------------------------------------------

/** The kind of a DAG node determines how it executes. */
export const dagNodeKindSchema = z.enum([
  'prompt',
  'command',
  'bash',
  'script',
  'approval',
  'loop',
  'cancel',
]);

/** Base fields shared by all DAG node types. */
export const dagNodeBaseSchema = z.object({
  id: z.string(),
  kind: dagNodeKindSchema,
  name: z.string().optional(),
  when: z.string().optional(),
  depends_on: z.array(z.string()).default([]),
  trigger_rule: triggerRuleSchema.optional(),
  retry: stepRetryConfigSchema.optional(),
  env: z.record(z.string()).optional(),
});

export type DagNodeBase = z.infer<typeof dagNodeBaseSchema>;

export type DagNodeKind = z.infer<typeof dagNodeKindSchema>;

// ---------------------------------------------------------------------------
// Prompt node — sends a prompt to an AI provider
// ---------------------------------------------------------------------------

export const promptNodeSchema = z.object({
  id: z.string(),
  kind: z.literal('prompt'),
  /** Human-readable name for display. */
  name: z.string().optional(),
  /** The prompt text (inline). Mutually exclusive with command_file. */
  prompt: z.string().optional(),
  /** Path to a command file containing the prompt. */
  command_file: z.string().optional(),
  /** Provider id to use (overrides workflow default). */
  provider: z.string().optional(),
  /** Model override for the provider. */
  model: z.string().optional(),
  /** Structured output format definition. */
  output_format: z
    .record(z.unknown())
    .optional(),
  /** Condition expression; node runs only when truthy. */
  when: z.string().optional(),
  /** Node ids this node depends on. */
  depends_on: z.array(z.string()).default([]),
  /** Trigger rule for evaluating dependency states. */
  trigger_rule: triggerRuleSchema.optional(),
  /** Retry configuration. */
  retry: stepRetryConfigSchema.optional(),
  /** Idle timeout in milliseconds. */
  idle_timeout_ms: z.number().positive().optional(),
  /** Environment variable overrides for this node. */
  env: z.record(z.string()).optional(),
});

// ---------------------------------------------------------------------------
// Command node — runs a shell command
// ---------------------------------------------------------------------------

export const commandNodeSchema = z.object({
  id: z.string(),
  kind: z.literal('command'),
  name: z.string().optional(),
  /** The command string to execute. */
  command: z.string(),
  /** Working directory override. */
  cwd: z.string().optional(),
  when: z.string().optional(),
  depends_on: z.array(z.string()).default([]),
  trigger_rule: triggerRuleSchema.optional(),
  retry: stepRetryConfigSchema.optional(),
  env: z.record(z.string()).optional(),
});

// ---------------------------------------------------------------------------
// Bash node — runs a bash script
// ---------------------------------------------------------------------------

export const bashNodeSchema = z.object({
  id: z.string(),
  kind: z.literal('bash'),
  name: z.string().optional(),
  /** Bash script content to execute. */
  bash: z.string(),
  /** Timeout in milliseconds. */
  timeout_ms: z.number().positive().optional(),
  when: z.string().optional(),
  depends_on: z.array(z.string()).default([]),
  trigger_rule: triggerRuleSchema.optional(),
  retry: stepRetryConfigSchema.optional(),
  env: z.record(z.string()).optional(),
});

// ---------------------------------------------------------------------------
// Script node — runs a script with a specific runtime
// ---------------------------------------------------------------------------

export const scriptNodeSchema = z.object({
  id: z.string(),
  kind: z.literal('script'),
  name: z.string().optional(),
  /** Script content to execute. */
  script: z.string(),
  /** Runtime: 'bun' or 'uv'. */
  runtime: z.enum(['bun', 'uv']),
  /** Dependencies to install before running. */
  deps: z.array(z.string()).default([]),
  /** Timeout in milliseconds. */
  timeout_ms: z.number().positive().optional(),
  when: z.string().optional(),
  depends_on: z.array(z.string()).default([]),
  trigger_rule: triggerRuleSchema.optional(),
  retry: stepRetryConfigSchema.optional(),
  env: z.record(z.string()).optional(),
});

// ---------------------------------------------------------------------------
// Approval node — pauses for human approval
// ---------------------------------------------------------------------------

export const approvalNodeSchema = z.object({
  id: z.string(),
  kind: z.literal('approval'),
  name: z.string().optional(),
  /** Message shown to the approver. */
  message: z.string(),
  /** Prompt to execute if the approval is rejected. */
  on_reject: z.string().optional(),
  when: z.string().optional(),
  depends_on: z.array(z.string()).default([]),
  trigger_rule: triggerRuleSchema.optional(),
  env: z.record(z.string()).optional(),
});

// ---------------------------------------------------------------------------
// Loop node — iterates until a condition is met
// ---------------------------------------------------------------------------

export const loopNodeSchema = z.object({
  id: z.string(),
  kind: z.literal('loop'),
  name: z.string().optional(),
  /** Loop configuration (prompt, until, max_iterations, etc.). */
  config: loopNodeConfigSchema,
  /** Provider id to use (overrides workflow default). */
  provider: z.string().optional(),
  /** Model override for the provider. */
  model: z.string().optional(),
  when: z.string().optional(),
  depends_on: z.array(z.string()).default([]),
  trigger_rule: triggerRuleSchema.optional(),
  retry: stepRetryConfigSchema.optional(),
  env: z.record(z.string()).optional(),
});

// ---------------------------------------------------------------------------
// Cancel node — cancels the workflow
// ---------------------------------------------------------------------------

export const cancelNodeSchema = z.object({
  id: z.string(),
  kind: z.literal('cancel'),
  name: z.string().optional(),
  /** Reason for cancellation. */
  reason: z.string().optional(),
  when: z.string().optional(),
  depends_on: z.array(z.string()).default([]),
  trigger_rule: triggerRuleSchema.optional(),
  env: z.record(z.string()).optional(),
});

// ---------------------------------------------------------------------------
// Union type — any DAG node
// ---------------------------------------------------------------------------

export const dagNodeSchema = z.discriminatedUnion('kind', [
  promptNodeSchema,
  commandNodeSchema,
  bashNodeSchema,
  scriptNodeSchema,
  approvalNodeSchema,
  loopNodeSchema,
  cancelNodeSchema,
]);

export type DagNode = z.infer<typeof dagNodeSchema>;
export type PromptNode = z.infer<typeof promptNodeSchema>;
export type CommandNode = z.infer<typeof commandNodeSchema>;
export type BashNode = z.infer<typeof bashNodeSchema>;
export type ScriptNode = z.infer<typeof scriptNodeSchema>;
export type ApprovalNode = z.infer<typeof approvalNodeSchema>;
export type LoopNode = z.infer<typeof loopNodeSchema>;
export type CancelNode = z.infer<typeof cancelNodeSchema>;

// ---------------------------------------------------------------------------
// Type guards
// ---------------------------------------------------------------------------

export function isBashNode(node: DagNode): node is BashNode {
  return node.kind === 'bash';
}

export function isScriptNode(node: DagNode): node is ScriptNode {
  return node.kind === 'script';
}

export function isLoopNode(node: DagNode): node is LoopNode {
  return node.kind === 'loop';
}

export function isApprovalNode(node: DagNode): node is ApprovalNode {
  return node.kind === 'approval';
}

export function isCancelNode(node: DagNode): node is CancelNode {
  return node.kind === 'cancel';
}

export function isPromptNode(node: DagNode): node is PromptNode {
  return node.kind === 'prompt';
}

export function isCommandNode(node: DagNode): node is CommandNode {
  return node.kind === 'command';
}