@multiplier-labs/stepflow 0.3.1 → 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 {
@@ -12,24 +16,21 @@ import type {
   WorkflowContext,
   Logger,
   SpawnChildOptions,
-} from "./types";
-import type { StorageAdapter, WorkflowRunRecord } from "../storage/types";
-import type {
-  EventTransport,
-  EventCallback,
-  Unsubscribe,
-  WorkflowEvent,
-} from "../events/types";
-import { MemoryStorageAdapter } from "../storage/memory";
-import { MemoryEventTransport } from "../events/memory";
-import { executeWorkflow } from "./orchestrator";
+} from './types';
+import type { StorageAdapter, WorkflowRunRecord } from '../storage/types';
+import type { EventTransport, EventCallback, Unsubscribe, WorkflowEvent } from '../events/types';
+import { MemoryStorageAdapter } from '../storage/memory';
+import { MemoryEventTransport } from '../events/memory';
+import { executeWorkflow } from './orchestrator';
 import {
   WorkflowNotFoundError,
   WorkflowAlreadyRegisteredError,
   RunNotFoundError,
   WaitForRunTimeoutError,
-} from "../utils/errors";
-import { ConsoleLogger } from "../utils/logger";
+} from '../utils/errors';
+import { ConsoleLogger } from '../utils/logger';
+import { ConcurrencyManager, type QueuedRun } from './engine/concurrency';
+import { LifecycleManager } from './engine/lifecycle';
 
 // ============================================================================
 // Configuration Types
@@ -111,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 settings: NonNullable<WorkflowEngineConfig['settings']>;
+  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),
+    });
   }
 
   /**
@@ -145,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();
   }
 
   // ============================================================================
@@ -187,7 +178,7 @@ export class WorkflowEngine {
    * @throws WorkflowAlreadyRegisteredError if already registered
    */
   registerWorkflow<TInput = Record<string, unknown>>(
-    definition: WorkflowDefinition<TInput>,
+    definition: WorkflowDefinition<TInput>
   ): void {
     if (this.registry.has(definition.kind)) {
       throw new WorkflowAlreadyRegisteredError(definition.kind);
@@ -241,7 +232,7 @@ export class WorkflowEngine {
    * @throws WorkflowNotFoundError if the workflow kind is not registered
    */
   async startRun<TInput = Record<string, unknown>>(
-    options: StartRunOptions<TInput>,
+    options: StartRunOptions<TInput>
   ): Promise<string> {
     const definition = this.registry.get(options.kind);
     if (!definition) {
@@ -251,7 +242,7 @@ export class WorkflowEngine {
     // Create run record (status: queued)
     const run = await this.storage.createRun({
       kind: options.kind,
-      status: "queued",
+      status: 'queued',
       parentRunId: options.parentRunId,
       input: (options.input ?? {}) as Record<string, unknown>,
       metadata: options.metadata ?? {},
@@ -264,28 +255,24 @@ export class WorkflowEngine {
     this.events.emit({
       runId,
       kind: options.kind,
-      eventType: "run.created",
+      eventType: 'run.created',
       timestamp: new Date(),
-      payload: {
-        input: options.input,
-        metadata: options.metadata,
-        priority: options.priority ?? 0,
-      },
+      payload: { input: options.input, metadata: options.metadata, priority: options.priority ?? 0 },
     });
 
     // Check if we have capacity to start immediately
-    if (this.hasCapacity()) {
+    if (this.concurrency.hasCapacity()) {
       // Start execution immediately
       this.executeRun(
         runId,
         definition,
         (options.input ?? {}) as Record<string, unknown>,
         options.metadata,
-        options.delay,
+        options.delay
       );
     } else {
       // Queue the run for later execution
-      this.queueRun({
+      this.concurrency.enqueue({
         runId,
         definition,
         input: (options.input ?? {}) as Record<string, unknown>,
@@ -294,41 +281,21 @@ 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({
         runId,
         kind: options.kind,
-        eventType: "run.queued",
+        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).
    */
@@ -337,7 +304,7 @@ export class WorkflowEngine {
     definition: WorkflowDefinition,
     input: Record<string, unknown>,
     metadata?: Record<string, unknown>,
-    delay?: number,
+    delay?: number
   ): void {
     this.launchRun(runId, definition, input, metadata, delay);
   }
@@ -352,13 +319,10 @@ export class WorkflowEngine {
     input: Record<string, unknown>,
     metadata?: Record<string, unknown>,
     delay?: number,
-    checkpoint?: {
-      completedStepKeys: Set<string>;
-      results: Record<string, unknown>;
-    },
+    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 {
@@ -371,51 +335,45 @@ export class WorkflowEngine {
           events: this.events,
           logger: this.logger,
           abortController,
-          spawnChild: (childOptions: SpawnChildOptions) =>
-            this.spawnChild(runId, childOptions),
+          spawnChild: (childOptions: SpawnChildOptions) => this.spawnChild(runId, childOptions),
           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.concurrency.processQueue((next) => {
       this.executeRun(next.runId, next.definition, next.input, next.metadata);
-    }
+    });
   }
 
   /**
    * Start a child workflow from within a parent workflow.
    * Called internally by the context.spawnChild helper.
    */
-  private async spawnChild(
-    parentRunId: string,
-    options: SpawnChildOptions,
-  ): Promise<string> {
+  private async spawnChild(parentRunId: string, options: SpawnChildOptions): Promise<string> {
     return this.startRun({
       kind: options.kind,
       input: options.input,
@@ -438,32 +396,24 @@ export class WorkflowEngine {
     }
 
     // Don't overwrite terminal statuses
-    if (["succeeded", "failed", "canceled", "timeout"].includes(run.status)) {
+    if (['succeeded', 'failed', 'canceled', 'timeout'].includes(run.status)) {
       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, {
-      status: "canceled",
+      status: 'canceled',
       finishedAt: new Date(),
     });
 
     this.events.emit({
       runId,
       kind: run.kind,
-      eventType: "run.canceled",
+      eventType: 'run.canceled',
       timestamp: new Date(),
     });
   }
@@ -478,18 +428,8 @@ export class WorkflowEngine {
     return this.storage.getRun(runId);
   }
 
-  private static readonly TERMINAL_STATUSES = [
-    "succeeded",
-    "failed",
-    "canceled",
-    "timeout",
-  ];
-  private static readonly TERMINAL_EVENT_TYPES = [
-    "run.completed",
-    "run.failed",
-    "run.canceled",
-    "run.timeout",
-  ];
+  private static readonly TERMINAL_STATUSES = ['succeeded', 'failed', 'canceled', 'timeout'];
+  private static readonly TERMINAL_EVENT_TYPES = ['run.completed', 'run.failed', 'run.canceled', 'run.timeout'];
 
   /**
    * Wait for a run to complete.
@@ -502,50 +442,64 @@ export class WorkflowEngine {
    */
   async waitForRun(
     runId: string,
-    options: { timeout?: number } = {},
+    options: { timeout?: number } = {}
   ): 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));
+      }
     });
   }
 
@@ -564,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 context
-    const completedStepKeys = new Set(Object.keys(run.context));
-
-    this.launchRun(runId, definition, run.input, run.metadata, undefined, {
-      completedStepKeys,
-      results: run.context,
-    });
-
-    return runId;
+    return this.lifecycle.resumeRun(runId);
   }
 
   /**
@@ -604,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();
   }
 
   /**
@@ -617,26 +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();
   }
 
   // ============================================================================
@@ -693,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/index.ts

@@ -2,6 +2,6 @@
  * Core module exports.
  */
 
-export * from "./types";
-export * from "./engine";
-export * from "./orchestrator";
+export * from './types';
+export * from './engine';
+export * from './orchestrator';