@multiplier-labs/stepflow 0.3.3 → 0.3.4

package/src/core/engine.ts

@@ -4,6 +4,10 @@
  * This is the primary entry point for the workflow engine.
  * It provides methods for registering workflows, starting runs,
  * and managing workflow execution.
+ *
+ * Internal concerns are decomposed into co-located submodules under
+ * `./engine/` (see `engine/concurrency.ts`, `engine/lifecycle.ts`).
+ * WorkflowEngine itself stays the thin public face.
  */
 
 import type {
@@ -25,6 +29,8 @@ import {
   WaitForRunTimeoutError,
 } from '../utils/errors';
 import { ConsoleLogger } from '../utils/logger';
+import { ConcurrencyManager, type QueuedRun } from './engine/concurrency';
+import { LifecycleManager } from './engine/lifecycle';
 
 // ============================================================================
 // Configuration Types
@@ -106,32 +112,35 @@ export interface StartRunOptions<TInput = Record<string, unknown>> {
  * const runId = await engine.startRun({ kind: 'my.workflow' });
  * ```
  */
-/**
- * Queued run waiting to be executed when capacity is available.
- */
-interface QueuedRun {
-  runId: string;
-  definition: WorkflowDefinition;
-  input: Record<string, unknown>;
-  metadata?: Record<string, unknown>;
-  priority: number;
-  queuedAt: Date;
-}
-
 export class WorkflowEngine {
   private registry = new Map<WorkflowKind, WorkflowDefinition>();
   private storage: StorageAdapter;
   private events: EventTransport;
   private logger: Logger;
-  private activeRuns = new Map<string, AbortController>();
   private settings: NonNullable<WorkflowEngineConfig['settings']>;
-  private runQueue: QueuedRun[] = [];
+  private concurrency: ConcurrencyManager;
+  private lifecycle: LifecycleManager;
 
   constructor(config: WorkflowEngineConfig = {}) {
     this.storage = config.storage ?? new MemoryStorageAdapter();
     this.events = config.events ?? new MemoryEventTransport();
     this.logger = config.logger ?? new ConsoleLogger();
     this.settings = config.settings ?? {};
+    this.concurrency = new ConcurrencyManager({
+      storage: this.storage,
+      events: this.events,
+      logger: this.logger,
+      maxConcurrency: this.settings.maxConcurrency,
+    });
+    this.lifecycle = new LifecycleManager({
+      storage: this.storage,
+      events: this.events,
+      logger: this.logger,
+      registry: this.registry,
+      concurrency: this.concurrency,
+      launchRun: (runId, def, input, metadata, delay, checkpoint) =>
+        this.launchRun(runId, def, input, metadata, delay, checkpoint),
+    });
   }
 
   /**
@@ -140,34 +149,21 @@ export class WorkflowEngine {
    * (e.g., PostgresStorageAdapter).
    */
   async initialize(): Promise<void> {
-    if (this.storage.initialize) {
-      await this.storage.initialize();
-    }
+    return this.lifecycle.initialize();
   }
 
   /**
    * Get the current number of active runs.
    */
   getActiveRunCount(): number {
-    return this.activeRuns.size;
+    return this.concurrency.getActiveCount();
   }
 
   /**
    * Get the number of queued runs waiting for capacity.
    */
   getQueuedRunCount(): number {
-    return this.runQueue.length;
-  }
-
-  /**
-   * Check if capacity is available for a new run.
-   */
-  private hasCapacity(): boolean {
-    const maxConcurrency = this.settings.maxConcurrency;
-    if (maxConcurrency === undefined || maxConcurrency <= 0) {
-      return true; // No limit
-    }
-    return this.activeRuns.size < maxConcurrency;
+    return this.concurrency.getQueuedCount();
   }
 
   // ============================================================================
@@ -265,7 +261,7 @@ export class WorkflowEngine {
     });
 
     // Check if we have capacity to start immediately
-    if (this.hasCapacity()) {
+    if (this.concurrency.hasCapacity()) {
       // Start execution immediately
       this.executeRun(
         runId,
@@ -276,7 +272,7 @@ export class WorkflowEngine {
       );
     } else {
       // Queue the run for later execution
-      this.queueRun({
+      this.concurrency.enqueue({
         runId,
         definition,
         input: (options.input ?? {}) as Record<string, unknown>,
@@ -285,7 +281,7 @@ export class WorkflowEngine {
         queuedAt: new Date(),
       });
 
-      this.logger.debug(`Run ${runId} queued (${this.runQueue.length} in queue)`);
+      this.logger.debug(`Run ${runId} queued (${this.concurrency.getQueuedCount()} in queue)`);
 
       // Emit queued event
       this.events.emit({
@@ -293,31 +289,13 @@ export class WorkflowEngine {
         kind: options.kind,
         eventType: 'run.queued',
         timestamp: new Date(),
-        payload: { queuePosition: this.runQueue.length },
+        payload: { queuePosition: this.concurrency.getQueuedCount() },
       });
     }
 
     return runId;
   }
 
-  /**
-   * Queue a run in priority order.
-   */
-  private queueRun(queuedRun: QueuedRun): void {
-    // Insert in priority order (higher priority first, then FIFO for same priority)
-    let inserted = false;
-    for (let i = 0; i < this.runQueue.length; i++) {
-      if (queuedRun.priority > this.runQueue[i].priority) {
-        this.runQueue.splice(i, 0, queuedRun);
-        inserted = true;
-        break;
-      }
-    }
-    if (!inserted) {
-      this.runQueue.push(queuedRun);
-    }
-  }
-
   /**
    * Execute a run (internal method).
    */
@@ -344,7 +322,7 @@ export class WorkflowEngine {
     checkpoint?: { completedStepKeys: Set<string>; results: Record<string, unknown> }
   ): void {
     const abortController = new AbortController();
-    this.activeRuns.set(runId, abortController);
+    this.concurrency.registerActive(runId, abortController);
 
     const execute = async () => {
       try {
@@ -361,41 +339,34 @@ export class WorkflowEngine {
           checkpoint,
         });
       } finally {
-        this.activeRuns.delete(runId);
+        this.concurrency.unregisterActive(runId);
         this.processQueue();
       }
     };
 
-    if (delay) {
-      setTimeout(execute, delay);
-    } else {
-      setImmediate(execute);
-    }
+    // executeWorkflow has its own try/catch that converts errors into a failed
+    // run record, but the surrounding `finally` (unregisterActive + processQueue)
+    // could still throw, and `setTimeout`/`setImmediate` discards the returned
+    // promise — leaving any rejection unhandled and capable of crashing Node.
+    // Wrap the launch so those rejections are logged instead.
+    const guarded = () => {
+      execute().catch((err) => {
+        this.concurrency.unregisterActive(runId);
+        this.logger.error(`Unhandled error launching run ${runId}`, err);
+      });
+    };
+
+    this.concurrency.scheduleLaunch(guarded, delay);
   }
 
   /**
-   * Process the queue and start runs if capacity is available.
+   * Drain the queue, dispatching each ready run via `executeRun`. Resilience
+   * against per-item dispatch failures lives in `ConcurrencyManager.processQueue`.
    */
   private processQueue(): void {
-    while (this.hasCapacity() && this.runQueue.length > 0) {
-      const next = this.runQueue.shift()!;
-      this.logger.debug(`Starting queued run ${next.runId}`);
-
-      // Emit dequeued event
-      this.events.emit({
-        runId: next.runId,
-        kind: next.definition.kind,
-        eventType: 'run.dequeued',
-        timestamp: new Date(),
-      });
-
-      this.executeRun(
-        next.runId,
-        next.definition,
-        next.input,
-        next.metadata
-      );
-    }
+    this.concurrency.processQueue((next) => {
+      this.executeRun(next.runId, next.definition, next.input, next.metadata);
+    });
   }
 
   /**
@@ -429,17 +400,9 @@ export class WorkflowEngine {
       return;
     }
 
-    // Remove from queue if queued
-    const queueIndex = this.runQueue.findIndex(q => q.runId === runId);
-    if (queueIndex !== -1) {
-      this.runQueue.splice(queueIndex, 1);
-    }
-
-    // Signal abort to the running workflow
-    const controller = this.activeRuns.get(runId);
-    if (controller) {
-      controller.abort();
-    }
+    // Remove from queue if queued, then signal abort if active
+    this.concurrency.removeFromQueue(runId);
+    this.concurrency.abortActive(runId);
 
     // Update status in storage (the orchestrator will also update on completion)
     await this.storage.updateRun(runId, {
@@ -483,46 +446,60 @@ export class WorkflowEngine {
   ): Promise<WorkflowRunRecord> {
     const timeout = options.timeout ?? 60000;
 
-    // Check storage first to handle already-completed runs (avoids race condition)
-    const existingRun = await this.storage.getRun(runId);
-    if (!existingRun) {
-      throw new RunNotFoundError(runId);
-    }
-    if (WorkflowEngine.TERMINAL_STATUSES.includes(existingRun.status)) {
-      return existingRun;
-    }
-
-    return new Promise<WorkflowRunRecord>((resolve, reject) => {
+    return new Promise<WorkflowRunRecord>(async (resolve, reject) => {
       let timeoutId: ReturnType<typeof setTimeout> | undefined;
       let unsubscribe: (() => void) | undefined;
+      let settled = false;
 
       const cleanup = () => {
         if (timeoutId) clearTimeout(timeoutId);
         if (unsubscribe) unsubscribe();
       };
 
+      const settle = (fn: () => void) => {
+        if (settled) return;
+        settled = true;
+        cleanup();
+        fn();
+      };
+
       // Set up timeout
       timeoutId = setTimeout(() => {
-        cleanup();
-        reject(new WaitForRunTimeoutError(runId, timeout));
+        settle(() => reject(new WaitForRunTimeoutError(runId, timeout)));
       }, timeout);
 
-      // Subscribe to run events
+      // Subscribe to run events FIRST to avoid missing terminal events
       unsubscribe = this.events.subscribe(runId, async (event) => {
         if (WorkflowEngine.TERMINAL_EVENT_TYPES.includes(event.eventType)) {
-          cleanup();
           try {
             const run = await this.storage.getRun(runId);
-            if (!run) {
-              reject(new RunNotFoundError(runId));
-            } else {
-              resolve(run);
-            }
+            settle(() => {
+              if (!run) {
+                reject(new RunNotFoundError(runId));
+              } else {
+                resolve(run);
+              }
+            });
           } catch (err) {
-            reject(err);
+            settle(() => reject(err));
           }
         }
       });
+
+      // THEN check storage for already-terminal runs
+      try {
+        const existingRun = await this.storage.getRun(runId);
+        if (!existingRun) {
+          settle(() => reject(new RunNotFoundError(runId)));
+          return;
+        }
+        if (WorkflowEngine.TERMINAL_STATUSES.includes(existingRun.status)) {
+          settle(() => resolve(existingRun));
+          return;
+        }
+      } catch (err) {
+        settle(() => reject(err));
+      }
     });
   }
 
@@ -541,39 +518,7 @@ export class WorkflowEngine {
    * @throws {WorkflowNotFoundError} If the workflow kind is not registered
    */
   async resumeRun(runId: string): Promise<string> {
-    const run = await this.storage.getRun(runId);
-    if (!run) {
-      throw new RunNotFoundError(runId);
-    }
-
-    // Check if run can be resumed
-    if (!['queued', 'running'].includes(run.status)) {
-      throw new Error(`Cannot resume run ${runId}: status is "${run.status}"`);
-    }
-
-    // Get the workflow definition
-    const definition = this.registry.get(run.kind);
-    if (!definition) {
-      throw new WorkflowNotFoundError(run.kind);
-    }
-
-    // Check if already being executed
-    if (this.activeRuns.has(runId)) {
-      this.logger.warn(`Run ${runId} is already active, skipping resume`);
-      return runId;
-    }
-
-    this.logger.info(`Resuming run ${runId} from checkpoint`);
-
-    // Get completed step keys from the dedicated completedSteps field
-    const completedStepKeys = new Set(run.completedSteps ?? []);
-
-    this.launchRun(runId, definition, run.input, run.metadata, undefined, {
-      completedStepKeys,
-      results: run.context,
-    });
-
-    return runId;
+    return this.lifecycle.resumeRun(runId);
   }
 
   /**
@@ -581,10 +526,7 @@ export class WorkflowEngine {
    * Returns runs with status 'queued' or 'running'.
    */
   async getResumableRuns(): Promise<WorkflowRunRecord[]> {
-    const result = await this.storage.listRuns({
-      status: ['queued', 'running'],
-    });
-    return result.items;
+    return this.lifecycle.getResumableRuns();
   }
 
   /**
@@ -594,24 +536,7 @@ export class WorkflowEngine {
    * @returns Array of resumed run IDs
    */
   async resumeAllInterrupted(): Promise<string[]> {
-    const resumableRuns = await this.getResumableRuns();
-    const resumedIds: string[] = [];
-
-    for (const run of resumableRuns) {
-      // Only resume if the workflow is registered
-      if (this.registry.has(run.kind)) {
-        try {
-          await this.resumeRun(run.id);
-          resumedIds.push(run.id);
-        } catch (error) {
-          this.logger.error(`Failed to resume run ${run.id}:`, error);
-        }
-      } else {
-        this.logger.warn(`Cannot resume run ${run.id}: workflow "${run.kind}" not registered`);
-      }
-    }
-
-    return resumedIds;
+    return this.lifecycle.resumeAllInterrupted();
   }
 
   // ============================================================================
@@ -668,25 +593,6 @@ export class WorkflowEngine {
    * Cancels all active runs and closes resources.
    */
   async shutdown(): Promise<void> {
-    this.logger.info('Shutting down workflow engine...');
-
-    // Cancel all active runs
-    for (const [runId, controller] of this.activeRuns) {
-      this.logger.debug(`Canceling active run: ${runId}`);
-      controller.abort();
-    }
-
-    // Close event transport
-    if (this.events.close) {
-      this.events.close();
-    }
-
-    // Close storage adapter
-    if (this.storage.close) {
-      await this.storage.close();
-    }
-
-    this.activeRuns.clear();
-    this.logger.info('Workflow engine shutdown complete');
+    return this.lifecycle.shutdown();
   }
 }

package/src/core/orchestrator/persistence.ts

@@ -0,0 +1,165 @@
+/**
+ * Orchestrator persistence layer.
+ *
+ * Wraps the storage adapter and event transport with intent-named helpers
+ * so the orchestrator and step-runner can stay focused on control flow
+ * instead of repeating field shapes for every status transition.
+ *
+ * All side-effect calls the orchestrator makes against storage or the event
+ * transport flow through this module — keeping the rest of the pipeline
+ * pure(-ish) and easy to read.
+ */
+import type { StorageAdapter, WorkflowRunStepRecord } from '../../storage/types';
+import type { EventTransport, WorkflowEvent } from '../../events/types';
+import type { Logger, WorkflowError } from '../types';
+import { WorkflowEngineError } from '../../utils/errors';
+
+/**
+ * Mark the run as running. On a fresh start, also stamp `startedAt`;
+ * on resume, leave the original `startedAt` alone.
+ */
+export async function markRunRunning(
+  storage: StorageAdapter,
+  runId: string,
+  isResume: boolean
+): Promise<void> {
+  await storage.updateRun(runId, {
+    status: 'running',
+    ...(isResume ? {} : { startedAt: new Date() }),
+  });
+}
+
+/**
+ * Persist the workflow checkpoint after a step completes:
+ * the accumulated step results plus the list of completed step keys.
+ */
+export async function saveCheckpoint(
+  storage: StorageAdapter,
+  runId: string,
+  results: Record<string, unknown>,
+  completedSteps: string[]
+): Promise<void> {
+  await storage.updateRun(runId, {
+    context: { ...results },
+    completedSteps: [...completedSteps],
+  });
+}
+
+/**
+ * Mark the run as succeeded, persisting results and finishedAt.
+ */
+export async function markRunSucceeded(
+  storage: StorageAdapter,
+  runId: string,
+  results: Record<string, unknown>
+): Promise<void> {
+  await storage.updateRun(runId, {
+    status: 'succeeded',
+    context: results,
+    finishedAt: new Date(),
+  });
+}
+
+/**
+ * Persist a non-success terminal state for the run (failed / canceled / timeout).
+ * `status` is the storage-level status — the orchestrator decides this from the
+ * raised error using `state-transitions.ts`.
+ */
+export async function markRunFinal(
+  storage: StorageAdapter,
+  runId: string,
+  status: 'failed' | 'canceled',
+  results: Record<string, unknown>,
+  error: WorkflowError
+): Promise<void> {
+  await storage.updateRun(runId, {
+    status,
+    context: results,
+    error,
+    finishedAt: new Date(),
+  });
+}
+
+/**
+ * Create a fresh step record in 'running' state for the given attempt.
+ * Returns the persisted record so the caller can read its generated id.
+ */
+export async function createStepRecord(
+  storage: StorageAdapter,
+  runId: string,
+  stepKey: string,
+  stepName: string,
+  attempt: number
+): Promise<WorkflowRunStepRecord> {
+  return storage.createStep({
+    runId,
+    stepKey,
+    stepName,
+    status: 'running',
+    attempt,
+    startedAt: new Date(),
+  });
+}
+
+/**
+ * Mark a step as succeeded, recording its result and finishedAt.
+ */
+export async function markStepSucceeded(
+  storage: StorageAdapter,
+  stepId: string,
+  result: unknown
+): Promise<void> {
+  await storage.updateStep(stepId, {
+    status: 'succeeded',
+    result,
+    finishedAt: new Date(),
+  });
+}
+
+/**
+ * Mark a step as failed (or canceled, if the abort signal already fired)
+ * with the normalized engine error attached.
+ */
+export async function markStepFailed(
+  storage: StorageAdapter,
+  stepId: string,
+  error: unknown,
+  isCanceled: boolean
+): Promise<void> {
+  await storage.updateStep(stepId, {
+    status: isCanceled ? 'canceled' : 'failed',
+    error: WorkflowEngineError.fromError(error),
+    finishedAt: new Date(),
+  });
+}
+
+/**
+ * Mark a step as canceled without an error payload — used when an in-flight
+ * retry sleep is interrupted by an abort signal.
+ */
+export async function markStepCanceled(
+  storage: StorageAdapter,
+  stepId: string
+): Promise<void> {
+  await storage.updateStep(stepId, {
+    status: 'canceled',
+    finishedAt: new Date(),
+  });
+}
+
+/**
+ * Emit a workflow event without ever throwing back into the caller.
+ * Event emission failures are logged but never abort the workflow,
+ * so a flaky transport can't take down a run.
+ */
+export function safeEmitEvent(
+  events: EventTransport,
+  logger: Logger,
+  event: WorkflowEvent
+): void {
+  try {
+    events.emit(event);
+  } catch (error) {
+    logger.error('Failed to emit event:', error);
+  }
+}