@donkeylabs/server 2.0.27 → 2.0.29

package/src/core/workflows.ts

@@ -59,7 +59,7 @@ type InferZodOutput<T extends ZodSchema> = z.infer<T>;
 // Step Types
 // ============================================
 
-export type StepType = "task" | "parallel" | "choice" | "pass";
+export type StepType = "task" | "parallel" | "choice" | "pass" | "poll" | "loop";
 
 export interface BaseStepDefinition {
   name: string;
@@ -146,11 +146,54 @@ export interface PassStepDefinition extends BaseStepDefinition {
   result?: any;
 }
 
+export interface PollStepResult<T = any> {
+  done: boolean;
+  result?: T;
+}
+
+export interface PollStepDefinition<
+  TInput extends ZodSchema = ZodSchema,
+  TOutput extends ZodSchema = ZodSchema,
+> extends BaseStepDefinition {
+  type: "poll";
+  /** Wait duration between checks in ms */
+  interval: number;
+  /** Max total time before failing this step (ms) */
+  timeout?: number;
+  /** Max number of check cycles before failing */
+  maxAttempts?: number;
+  /** Input schema or mapper */
+  inputSchema?: TInput | ((prev: any, workflowInput: any) => InferZodOutput<TInput>);
+  /** Output schema for the final result */
+  outputSchema?: TOutput;
+  /** Check handler: return done:true to proceed */
+  check: (
+    input: InferZodOutput<TInput>,
+    ctx: WorkflowContext
+  ) => Promise<PollStepResult<InferZodOutput<TOutput>>> | PollStepResult<InferZodOutput<TOutput>>;
+}
+
+export interface LoopStepDefinition extends BaseStepDefinition {
+  type: "loop";
+  /** Condition to continue looping */
+  condition: (ctx: WorkflowContext) => boolean;
+  /** Step name to jump back to when condition is true */
+  target: string;
+  /** Optional delay before looping (ms) */
+  interval?: number;
+  /** Max total time before failing this loop (ms) */
+  timeout?: number;
+  /** Max number of loop iterations before failing */
+  maxIterations?: number;
+}
+
 export type StepDefinition =
   | TaskStepDefinition
   | ParallelStepDefinition
   | ChoiceStepDefinition
-  | PassStepDefinition;
+  | PassStepDefinition
+  | PollStepDefinition
+  | LoopStepDefinition;
 
 // ============================================
 // Workflow Definition
@@ -203,6 +246,11 @@ export interface StepResult {
   startedAt?: Date;
   completedAt?: Date;
   attempts: number;
+  pollCount?: number;
+  lastPolledAt?: Date;
+  loopCount?: number;
+  lastLoopedAt?: Date;
+  loopStartedAt?: Date;
 }
 
 export interface WorkflowInstance {
@@ -569,6 +617,69 @@ export class WorkflowBuilder {
     return this;
   }
 
+  loop(
+    name: string,
+    config: {
+      condition: (ctx: WorkflowContext) => boolean;
+      target: string;
+      interval?: number;
+      timeout?: number;
+      maxIterations?: number;
+      next?: string;
+      end?: boolean;
+    }
+  ): this {
+    const step: LoopStepDefinition = {
+      name,
+      type: "loop",
+      condition: config.condition,
+      target: config.target,
+      interval: config.interval,
+      timeout: config.timeout,
+      maxIterations: config.maxIterations,
+      next: config.next,
+      end: config.end,
+    };
+
+    this.addStep(step);
+    return this;
+  }
+
+  poll<TInput extends ZodSchema = ZodSchema, TOutput extends ZodSchema = ZodSchema>(
+    name: string,
+    config: {
+      check: (
+        input: InferZodOutput<TInput>,
+        ctx: WorkflowContext
+      ) => Promise<PollStepResult<InferZodOutput<TOutput>>> | PollStepResult<InferZodOutput<TOutput>>;
+      interval: number;
+      timeout?: number;
+      maxAttempts?: number;
+      inputSchema?: TInput | ((prev: any, workflowInput: any) => InferZodOutput<TInput>);
+      outputSchema?: TOutput;
+      retry?: RetryConfig;
+      next?: string;
+      end?: boolean;
+    }
+  ): this {
+    const step: PollStepDefinition<TInput, TOutput> = {
+      name,
+      type: "poll",
+      check: config.check,
+      interval: config.interval,
+      timeout: config.timeout,
+      maxAttempts: config.maxAttempts,
+      inputSchema: config.inputSchema,
+      outputSchema: config.outputSchema,
+      retry: config.retry,
+      next: config.next,
+      end: config.end,
+    };
+
+    this.addStep(step);
+    return this;
+  }
+
   /** Add an end step (shorthand for pass with end: true) */
   end(name: string = "end"): this {
     return this.pass(name, { end: true });
@@ -649,12 +760,22 @@ export interface WorkflowsConfig {
   heartbeatTimeout?: number;
   /** Timeout waiting for isolated subprocess readiness (ms, default: 10000) */
   readyTimeout?: number;
+  /** Grace period before SIGKILL when terminating isolated subprocesses (ms, default: 5000) */
+  killGraceMs?: number;
+  /** SQLite pragmas for isolated subprocess connections */
+  sqlitePragmas?: SqlitePragmaConfig;
   /** Resume strategy for orphaned workflows (default: "blocking") */
   resumeStrategy?: WorkflowResumeStrategy;
 }
 
 export type WorkflowResumeStrategy = "blocking" | "background" | "skip";
 
+export interface SqlitePragmaConfig {
+  busyTimeout?: number;
+  synchronous?: "OFF" | "NORMAL" | "FULL" | "EXTRA";
+  journalMode?: "DELETE" | "TRUNCATE" | "PERSIST" | "MEMORY" | "WAL" | "OFF";
+}
+
 /** Options for registering a workflow */
 export interface WorkflowRegisterOptions {
   /**
@@ -743,6 +864,8 @@ class WorkflowsImpl implements Workflows {
   private dbPath?: string;
   private heartbeatTimeoutMs: number;
   private readyTimeoutMs: number;
+  private killGraceMs: number;
+  private sqlitePragmas?: SqlitePragmaConfig;
   private resumeStrategy!: WorkflowResumeStrategy;
   private workflowModulePaths = new Map<string, string>();
   private isolatedProcesses = new Map<string, IsolatedProcessInfo>();
@@ -777,6 +900,8 @@ class WorkflowsImpl implements Workflows {
     this.dbPath = config.dbPath;
     this.heartbeatTimeoutMs = config.heartbeatTimeout ?? 60000;
     this.readyTimeoutMs = config.readyTimeout ?? 10000;
+    this.killGraceMs = config.killGraceMs ?? 5000;
+    this.sqlitePragmas = config.sqlitePragmas;
     this.resumeStrategy = config.resumeStrategy ?? "blocking";
   }
 
@@ -938,11 +1063,7 @@ class WorkflowsImpl implements Workflows {
     // Kill isolated process if running
     const isolatedInfo = this.isolatedProcesses.get(instanceId);
     if (isolatedInfo) {
-      try {
-        process.kill(isolatedInfo.pid, "SIGTERM");
-      } catch {
-        // Process might already be dead
-      }
+      await killProcessWithGrace(isolatedInfo.pid, this.killGraceMs);
       if (isolatedInfo.timeout) clearTimeout(isolatedInfo.timeout);
       if (isolatedInfo.heartbeatTimeout) clearTimeout(isolatedInfo.heartbeatTimeout);
       this.isolatedProcesses.delete(instanceId);
@@ -1198,6 +1319,51 @@ class WorkflowsImpl implements Workflows {
           });
         }
       },
+      onStepPoll: (id, stepName, pollCount, done, result) => {
+        this.emitEvent("workflow.step.poll", {
+          instanceId: id,
+          stepName,
+          pollCount,
+          done,
+          result,
+        });
+        if (this.sse) {
+          this.sse.broadcast(`workflow:${id}`, "step.poll", {
+            stepName,
+            pollCount,
+            done,
+            result,
+          });
+          this.sse.broadcast("workflows:all", "workflow.step.poll", {
+            instanceId: id,
+            stepName,
+            pollCount,
+            done,
+            result,
+          });
+        }
+      },
+      onStepLoop: (id, stepName, loopCount, target) => {
+        this.emitEvent("workflow.step.loop", {
+          instanceId: id,
+          stepName,
+          loopCount,
+          target,
+        });
+        if (this.sse) {
+          this.sse.broadcast(`workflow:${id}`, "step.loop", {
+            stepName,
+            loopCount,
+            target,
+          });
+          this.sse.broadcast("workflows:all", "workflow.step.loop", {
+            instanceId: id,
+            stepName,
+            loopCount,
+            target,
+          });
+        }
+      },
       onStepRetry: (id, stepName, attempt, max, delayMs) => {
         this.emitEvent("workflow.step.retry", {
           instanceId: id,
@@ -1314,6 +1480,7 @@ class WorkflowsImpl implements Workflows {
       pluginModulePaths: this.pluginModulePaths,
       pluginConfigs,
       coreConfig,
+      sqlitePragmas: this.sqlitePragmas,
     };
 
     // Spawn the subprocess
@@ -1484,6 +1651,55 @@ class WorkflowsImpl implements Workflows {
         break;
       }
 
+      case "step.poll": {
+        await this.emitEvent("workflow.step.poll", {
+          instanceId,
+          stepName: event.stepName,
+          pollCount: event.pollCount,
+          done: event.done,
+          result: event.result,
+        });
+        if (this.sse) {
+          this.sse.broadcast(`workflow:${instanceId}`, "step.poll", {
+            stepName: event.stepName,
+            pollCount: event.pollCount,
+            done: event.done,
+            result: event.result,
+          });
+          this.sse.broadcast("workflows:all", "workflow.step.poll", {
+            instanceId,
+            stepName: event.stepName,
+            pollCount: event.pollCount,
+            done: event.done,
+            result: event.result,
+          });
+        }
+        break;
+      }
+
+      case "step.loop": {
+        await this.emitEvent("workflow.step.loop", {
+          instanceId,
+          stepName: event.stepName,
+          loopCount: event.loopCount,
+          target: event.target,
+        });
+        if (this.sse) {
+          this.sse.broadcast(`workflow:${instanceId}`, "step.loop", {
+            stepName: event.stepName,
+            loopCount: event.loopCount,
+            target: event.target,
+          });
+          this.sse.broadcast("workflows:all", "workflow.step.loop", {
+            instanceId,
+            stepName: event.stepName,
+            loopCount: event.loopCount,
+            target: event.target,
+          });
+        }
+        break;
+      }
+
       case "progress": {
         await this.emitEvent("workflow.progress", {
           instanceId,
@@ -1790,6 +2006,11 @@ class WorkflowsImpl implements Workflows {
       }
 
       console.error(`[Workflows] No heartbeat from isolated workflow ${instanceId} for ${this.heartbeatTimeoutMs}ms`);
+      await this.emitEvent("workflow.watchdog.stale", {
+        instanceId,
+        reason: "heartbeat",
+        timeoutMs: this.heartbeatTimeoutMs,
+      });
       await this.handleIsolatedTimeout(instanceId, pid);
     }, this.heartbeatTimeoutMs);
   }
@@ -1801,12 +2022,12 @@ class WorkflowsImpl implements Workflows {
     const info = this.isolatedProcesses.get(instanceId);
     if (!info) return;
 
-    // Kill the process
-    try {
-      process.kill(pid, "SIGKILL");
-    } catch {
-      // Process might already be dead
-    }
+    await killProcessWithGrace(pid, this.killGraceMs);
+    await this.emitEvent("workflow.watchdog.killed", {
+      instanceId,
+      reason: "timeout",
+      timeoutMs: this.heartbeatTimeoutMs,
+    });
 
     // Clean up
     if (info.timeout) clearTimeout(info.timeout);
@@ -1943,3 +2164,29 @@ function isPlainObject(value: Record<string, any>): boolean {
 export function createWorkflows(config?: WorkflowsConfig): Workflows {
   return new WorkflowsImpl(config);
 }
+
+async function killProcessWithGrace(pid: number, graceMs: number): Promise<void> {
+  try {
+    process.kill(pid, "SIGTERM");
+  } catch {
+    return;
+  }
+
+  if (graceMs <= 0) {
+    try {
+      process.kill(pid, "SIGKILL");
+    } catch {
+      return;
+    }
+    return;
+  }
+
+  await new Promise((resolve) => setTimeout(resolve, graceMs));
+
+  try {
+    process.kill(pid, 0);
+    process.kill(pid, "SIGKILL");
+  } catch {
+    // Process already exited
+  }
+}