@plumbus/core 0.2.9 → 0.3.1

package/dist/flows/engine.js

@@ -1,15 +1,64 @@
-import { and, eq, inArray, isNull, lte, or } from 'drizzle-orm';
+import { and, eq, inArray, isNull, lte, or, sql } from 'drizzle-orm';
 import { randomUUID } from 'node:crypto';
+import { hostname } from 'node:os';
+import { FlowCancelledError, LeaseLostError } from '../errors/index.js';
 import { BackoffStrategy } from '../types/enums.js';
 import { flowExecutionsTable } from './schema.js';
 import { assertTransition, FlowStatus, isTerminal, StepStatus, } from './state-machine.js';
-import { buildHistoryEntry, executeStep } from './step-executor.js';
+import { buildHistoryEntry, executeStep, } from './step-executor.js';
+/** Default lease duration: 5 minutes */
+const DEFAULT_LEASE_DURATION_MS = 300_000;
+/** Default claim batch size */
+const DEFAULT_CLAIM_BATCH_SIZE = 50;
+/**
+ * Read the affected-row count from a Drizzle UPDATE/INSERT/DELETE result in a
+ * driver-agnostic way. postgres-js exposes this as `.count` (the array-like
+ * `Result` has `length === 0` when no RETURNING clause is used), while
+ * node-postgres exposes it as `.rowCount`. Returning `length` is only a
+ * correct fallback for queries that include `.returning(...)`.
+ *
+ * Historical bug: reading only `rowCount ?? length` silently returned 0 for
+ * every postgres-js UPDATE without RETURNING, which caused `extendLease` to
+ * always report "lease lost" even when the row was successfully updated.
+ */
+function getRowsAffected(result) {
+    const r = result;
+    if (typeof r?.rowCount === 'number')
+        return r.rowCount;
+    if (typeof r?.count === 'number')
+        return r.count;
+    if (typeof r?.length === 'number')
+        return r.length;
+    return 0;
+}
+/**
+ * Generate a deterministic-ish worker id for the current process.
+ */
+export function generateWorkerId() {
+    return `${hostname()}:${process.pid}:${randomUUID().slice(0, 8)}`;
+}
 /**
  * Creates the flow execution engine that manages flow lifecycle:
  * start, run steps, handle wait/resume, and persist state.
+ *
+ * Supports lease-based claiming to prevent duplicate step execution
+ * across multiple workers sharing the same database.
  */
 export function createFlowEngine(config) {
     const { db, registry, stepDeps, audit, onFlowError, createEventService } = config;
+    const workerId = config.workerId ?? generateWorkerId();
+    const leaseDurationMs = config.flowLeaseDurationMs ?? DEFAULT_LEASE_DURATION_MS;
+    const heartbeatIntervalMs = config.flowHeartbeatIntervalMs ?? Math.floor(leaseDurationMs / 3);
+    const claimBatchSize = config.flowClaimBatchSize ?? DEFAULT_CLAIM_BATCH_SIZE;
+    const logger = config.logger;
+    /**
+     * AbortControllers for currently-executing steps in THIS engine instance,
+     * keyed by executionId. `cancel()` uses this to abort the running step
+     * synchronously when the cancel request arrives in the same process as
+     * the worker. Cross-process cancels still converge via the next heartbeat
+     * (which now also fails when status != 'running').
+     */
+    const activeStepAborts = new Map();
     /**
      * Start a new flow execution.
      */
@@ -59,9 +108,125 @@ export function createFlowEngine(config) {
             status: FlowStatus.Created,
         };
     }
+    /**
+     * Atomically claim up to `batchSize` runnable flow executions.
+     * Uses FOR UPDATE SKIP LOCKED so concurrent workers each get different rows.
+     * Also reclaims rows whose lease has expired (crash recovery).
+     */
+    async function claimNext(batchSize) {
+        const limit = batchSize ?? claimBatchSize;
+        const leaseDurationInterval = `${leaseDurationMs} milliseconds`;
+        const result = await db.execute(sql `
+      UPDATE flow_executions
+      SET status = ${FlowStatus.Running},
+          lease_owner = ${workerId},
+          lease_expires_at = now() + ${leaseDurationInterval}::interval,
+          updated_at = now()
+      WHERE id IN (
+        SELECT id
+        FROM flow_executions
+        WHERE
+             status = ${FlowStatus.Created}
+          OR (status = ${FlowStatus.Running} AND lease_expires_at < now())
+          OR (status = ${FlowStatus.Waiting} AND waiting_for_event IS NULL AND wake_at <= now())
+        ORDER BY created_at
+        FOR UPDATE SKIP LOCKED
+        LIMIT ${limit}
+      )
+      RETURNING *
+    `);
+        const rows = result;
+        for (const row of rows) {
+            if (audit) {
+                const wasExpired = row.status === FlowStatus.Running;
+                if (wasExpired) {
+                    await audit.record('flow.lease.expired', {
+                        executionId: row.id,
+                        flowName: row.flowName,
+                        workerId,
+                        previousOwner: row.leaseOwner,
+                        step: row.currentStep,
+                    });
+                }
+                await audit.record('flow.lease.claimed', {
+                    executionId: row.id,
+                    flowName: row.flowName,
+                    workerId,
+                    step: row.currentStep,
+                });
+            }
+        }
+        return rows;
+    }
+    /**
+     * Extend the lease for a flow execution.
+     * Returns true if the lease was extended, false if the lease was lost.
+     *
+     * The WHERE clause also requires `status = 'running'`. If another actor
+     * (e.g. `cancel()`) has flipped the status to `cancelled`/`failed`, the
+     * UPDATE matches 0 rows and we take the same "lease lost" path — which
+     * the heartbeat uses to abort the per-step AbortController so the running
+     * capability can cooperatively stop its in-flight AI / HTTP work.
+     */
+    async function extendLease(executionId) {
+        const leaseDurationInterval = `${leaseDurationMs} milliseconds`;
+        const result = await db
+            .update(flowExecutionsTable)
+            .set({
+            leaseExpiresAt: sql `now() + ${leaseDurationInterval}::interval`,
+            updatedAt: new Date(),
+        })
+            .where(and(eq(flowExecutionsTable.id, executionId), eq(flowExecutionsTable.leaseOwner, workerId), eq(flowExecutionsTable.status, FlowStatus.Running)));
+        const rowsAffected = getRowsAffected(result);
+        if (rowsAffected === 0 && logger) {
+            // D1 diagnostic: the heartbeat's status+owner guard didn't match, which
+            // means something mutated the row out from under us (status flipped by
+            // cancel/complete/failure, or a different worker stole the lease). Read
+            // the row back and log its actual shape plus our own identity so we can
+            // tell at a glance which dimension diverged. Best-effort: a failure
+            // here must not break the heartbeat.
+            try {
+                const rows = await db
+                    .select({
+                    status: flowExecutionsTable.status,
+                    leaseOwner: flowExecutionsTable.leaseOwner,
+                    leaseExpiresAt: flowExecutionsTable.leaseExpiresAt,
+                    updatedAt: flowExecutionsTable.updatedAt,
+                    currentStep: flowExecutionsTable.currentStep,
+                })
+                    .from(flowExecutionsTable)
+                    .where(eq(flowExecutionsTable.id, executionId))
+                    .limit(1);
+                const row = rows[0];
+                logger.warn('extendLease matched 0 rows', {
+                    executionId,
+                    expectedWorkerId: workerId,
+                    expectedStatus: FlowStatus.Running,
+                    actual: row
+                        ? {
+                            status: row.status,
+                            leaseOwner: row.leaseOwner,
+                            leaseExpiresAt: row.leaseExpiresAt?.toISOString() ?? null,
+                            updatedAt: row.updatedAt?.toISOString() ?? null,
+                            currentStep: row.currentStep,
+                        }
+                        : null,
+                    nowMs: Date.now(),
+                });
+            }
+            catch (diagErr) {
+                logger.warn('extendLease diagnostic SELECT failed', {
+                    executionId,
+                    error: diagErr instanceof Error ? diagErr.message : String(diagErr),
+                });
+            }
+        }
+        return rowsAffected > 0;
+    }
     /**
      * Run the next step(s) for a flow execution.
      * This is designed to be called by a worker process.
+     * The execution must have been claimed by this worker (via claimNext) first.
      */
     async function runNext(executionId, ctx) {
         const rows = await db
@@ -87,13 +252,15 @@ export function createFlowEngine(config) {
         if (row.status === FlowStatus.Waiting && row.wakeAt && row.wakeAt.getTime() > Date.now()) {
             return { id: row.id, flowName: row.flowName, status: FlowStatus.Waiting };
         }
-        // Transition to running
+        // Transition to running (only if not already claimed via claimNext)
         if (row.status === FlowStatus.Created || row.status === FlowStatus.Waiting) {
             assertTransition(row.status, FlowStatus.Running);
-            await updateExecution(executionId, {
+            await guardedUpdate(executionId, {
                 status: FlowStatus.Running,
                 waitingForEvent: null,
                 wakeAt: null,
+                leaseOwner: workerId,
+                leaseExpiresAt: new Date(Date.now() + leaseDurationMs),
             });
         }
         // Find current step
@@ -117,19 +284,121 @@ export function createFlowEngine(config) {
         };
         const flowData = config.createDataService ? config.createDataService(flowAuth) : ctx.data;
         const flowEvents = createEventService ? createEventService(flowAuth) : ctx.events;
+        // Create a heartbeat function for manual use by step handlers
+        const heartbeatFn = async () => {
+            const extended = await extendLease(executionId);
+            if (!extended) {
+                throw new LeaseLostError(`Lease lost for flow execution "${executionId}" (worker: ${workerId})`, { executionId, workerId, step: currentStepName });
+            }
+        };
+        // Per-step AbortController: aborted when the lease is lost (either
+        // because another worker claimed the execution OR because cancel() flipped
+        // the status out of 'running'). Capability handlers can subscribe via
+        // `ctx.signal` and pass it to AI/HTTP calls to stop cooperatively rather
+        // than continuing to burn budget as a zombie.
+        const stepAc = new AbortController();
+        activeStepAborts.set(executionId, stepAc);
+        // Auto-thread `signal` into AI calls so capability authors don't have to
+        // pass it explicitly. An explicit `signal` in the call overrides the
+        // step default. See AIService type for the per-method signal parameter.
+        const flowAi = wrapAiWithDefaultSignal(ctx.ai, stepAc.signal);
         const flowCtx = {
             ...ctx,
             auth: flowAuth,
             data: flowData,
             events: flowEvents,
+            ai: flowAi,
             state: row.state,
             step: currentStepName,
             flowId: executionId,
+            workerId,
+            signal: stepAc.signal,
+            flows: {
+                ...ctx.flows,
+                heartbeat: heartbeatFn,
+            },
         };
+        // Start automatic heartbeat timer
+        let leaseLost = false;
+        let heartbeatTickCount = 0;
+        const heartbeatStartedAtMs = Date.now();
+        const heartbeatTimer = setInterval(async () => {
+            heartbeatTickCount += 1;
+            const tick = heartbeatTickCount;
+            const firedAtMs = Date.now();
+            const expectedFireAtMs = heartbeatStartedAtMs + tick * heartbeatIntervalMs;
+            const driftMs = firedAtMs - expectedFireAtMs;
+            // D2 diagnostic: emit one log line per heartbeat tick so we can see
+            // whether ticks are firing on schedule (event-loop starvation shows up
+            // here as large positive `driftMs`). First tick is logged at `info`
+            // because that's typically enough for a post-mortem; subsequent ticks
+            // at `debug` to avoid log spam for long-running flows.
+            if (logger) {
+                const payload = {
+                    executionId,
+                    workerId,
+                    step: currentStepName,
+                    tick,
+                    firedAtMs,
+                    expectedFireAtMs,
+                    driftMs,
+                    heartbeatIntervalMs,
+                };
+                if (tick === 1) {
+                    logger.info('flow.heartbeat.tick', payload);
+                }
+                else {
+                    logger.debug('flow.heartbeat.tick', payload);
+                }
+            }
+            const extended = await extendLease(executionId);
+            if (extended) {
+                if (audit) {
+                    await audit.record('flow.lease.extended', {
+                        executionId,
+                        flowName: row.flowName,
+                        workerId,
+                        step: currentStepName,
+                    });
+                }
+            }
+            else {
+                leaseLost = true;
+                if (!stepAc.signal.aborted) {
+                    stepAc.abort(new LeaseLostError(`Lease lost for flow execution "${executionId}" (worker: ${workerId})`, { executionId, workerId, step: currentStepName }));
+                }
+                if (audit) {
+                    await audit.record('flow.lease.lost', {
+                        executionId,
+                        flowName: row.flowName,
+                        workerId,
+                        step: currentStepName,
+                    });
+                }
+            }
+        }, heartbeatIntervalMs);
         // Execute the step
+        let result;
         const startedAt = new Date();
-        const result = await executeStep(step, flowCtx, row.input, row.state, stepDeps);
+        try {
+            result = await executeStep(step, flowCtx, row.input, row.state, stepDeps);
+        }
+        finally {
+            clearInterval(heartbeatTimer);
+            // Abort any stray detached tasks once the step settles so they observe
+            // the abort promptly. If already aborted (lease loss), this is a no-op.
+            if (!stepAc.signal.aborted) {
+                stepAc.abort();
+            }
+            if (activeStepAborts.get(executionId) === stepAc) {
+                activeStepAborts.delete(executionId);
+            }
+        }
         const completedAt = new Date();
+        // Check if lease was lost during step execution
+        if (leaseLost) {
+            throw new LeaseLostError(`Lease lost during step "${currentStepName}" of flow execution "${executionId}"`, { executionId, workerId, step: currentStepName });
+        }
         const historyEntry = buildHistoryEntry(currentStepName, result, startedAt, completedAt);
         const history = Array.isArray(row.stepHistory) ? row.stepHistory : [];
         history.push(historyEntry);
@@ -149,14 +418,23 @@ export function createFlowEngine(config) {
         }
         // Wait step — pause the flow
         if (result.waitEvent) {
-            await updateExecution(executionId, {
+            const ok = await guardedUpdate(executionId, {
                 status: FlowStatus.Waiting,
                 state: nextState,
                 stepHistory: history,
                 currentStep: currentStepName, // stay on current step until resumed
                 waitingForEvent: result.waitEvent,
                 wakeAt: null,
+                leaseOwner: null,
+                leaseExpiresAt: null,
             });
+            if (!ok) {
+                throw new LeaseLostError(`Lease lost committing wait for flow execution "${executionId}"`, {
+                    executionId,
+                    workerId,
+                    step: currentStepName,
+                });
+            }
             return { id: executionId, flowName: row.flowName, status: FlowStatus.Waiting };
         }
         // Delay step — schedule next step after duration
@@ -169,13 +447,15 @@ export function createFlowEngine(config) {
             catch (err) {
                 return handleStepFailure(executionId, row, flow, history, err instanceof Error ? err.message : String(err));
             }
-            await updateExecution(executionId, {
+            await guardedUpdate(executionId, {
                 status: FlowStatus.Waiting,
                 state: nextState,
                 stepHistory: history,
                 currentStep: nextStep,
                 waitingForEvent: null,
                 wakeAt: new Date(Date.now() + delayMs),
+                leaseOwner: null,
+                leaseExpiresAt: null,
             });
             return { id: executionId, flowName: row.flowName, status: FlowStatus.Waiting };
         }
@@ -205,7 +485,7 @@ export function createFlowEngine(config) {
                 return handleStepFailure(executionId, row, flow, history, 'One or more parallel branches failed');
             }
             const nextStep = getNextStepName(flow.steps, currentStepName);
-            await updateExecution(executionId, {
+            await guardedUpdate(executionId, {
                 state: parallelState,
                 stepHistory: history,
                 currentStep: nextStep,
@@ -218,7 +498,7 @@ export function createFlowEngine(config) {
         }
         // Conditional step — jump to the chosen branch
         if (result.nextStep) {
-            await updateExecution(executionId, {
+            await guardedUpdate(executionId, {
                 state: nextState,
                 stepHistory: history,
                 currentStep: result.nextStep,
@@ -228,11 +508,11 @@ export function createFlowEngine(config) {
         // Normal completion — advance to next step
         const nextStep = getNextStepName(flow.steps, currentStepName);
         if (!nextStep) {
-            await updateExecution(executionId, { state: nextState, stepHistory: history });
+            await guardedUpdate(executionId, { state: nextState, stepHistory: history });
             await completeFlow(executionId, row.flowName);
             return { id: executionId, flowName: row.flowName, status: FlowStatus.Completed };
         }
-        await updateExecution(executionId, {
+        await guardedUpdate(executionId, {
             state: nextState,
             stepHistory: history,
             currentStep: nextStep,
@@ -257,27 +537,47 @@ export function createFlowEngine(config) {
         const flow = registry.get(row.flowName);
         if (!flow)
             throw new Error(`Flow "${row.flowName}" not registered`);
-        // Advance past the wait/delay step
+        // Advance past the wait/delay step — set to created so claimNext picks it up
         const nextStep = getNextStepName(flow.steps, row.currentStep ?? '');
         await updateExecution(executionId, {
-            status: FlowStatus.Running,
+            status: FlowStatus.Created,
             currentStep: nextStep,
             state: signal !== undefined ? signal : row.state,
             waitingForEvent: null,
             wakeAt: null,
+            leaseOwner: null,
+            leaseExpiresAt: null,
         });
     }
     /**
      * Resume all flows waiting on a specific event type.
+     * Uses FOR UPDATE SKIP LOCKED to prevent duplicate resumes across workers.
      * Returns the number of executions resumed.
      */
     async function resumeWaitingByEvent(eventType, signal) {
-        const rows = await db
-            .select({ id: flowExecutionsTable.id })
-            .from(flowExecutionsTable)
-            .where(and(eq(flowExecutionsTable.status, FlowStatus.Waiting), eq(flowExecutionsTable.waitingForEvent, eventType)));
-        for (const row of rows) {
-            await resume(row.id, signal);
+        const result = await db.execute(sql `
+      UPDATE flow_executions
+      SET status = ${FlowStatus.Created},
+          waiting_for_event = NULL,
+          wake_at = NULL,
+          lease_owner = NULL,
+          lease_expires_at = NULL,
+          updated_at = now()
+      WHERE id IN (
+        SELECT id
+        FROM flow_executions
+        WHERE status = ${FlowStatus.Waiting}
+          AND waiting_for_event = ${eventType}
+        FOR UPDATE SKIP LOCKED
+      )
+      RETURNING id
+    `);
+        const rows = result;
+        // If a signal was provided, update state for each resumed execution
+        if (signal !== undefined) {
+            for (const row of rows) {
+                await updateExecution(row.id, { state: signal });
+            }
         }
         return rows.length;
     }
@@ -313,6 +613,13 @@ export function createFlowEngine(config) {
             status: FlowStatus.Cancelled,
             completedAt: new Date(),
         });
+        // Synchronously abort the step's AbortController if it is running in THIS
+        // engine instance. Cross-process cancellations still converge via the next
+        // heartbeat (which fails because status is no longer 'running').
+        const activeAc = activeStepAborts.get(executionId);
+        if (activeAc && !activeAc.signal.aborted) {
+            activeAc.abort(new FlowCancelledError(executionId));
+        }
         if (audit) {
             await audit.record(`flow.cancelled.${row.flowName}`, {
                 executionId,
@@ -340,7 +647,7 @@ export function createFlowEngine(config) {
         const maxRetries = flow.retry?.attempts ?? 0;
         if (retryCount <= maxRetries) {
             // Retry: keep current step, increment counter
-            await updateExecution(executionId, {
+            await guardedUpdate(executionId, {
                 stepHistory: history,
                 retryCount,
                 lastError: error ?? null,
@@ -352,21 +659,29 @@ export function createFlowEngine(config) {
         return { id: executionId, flowName: row.flowName, status: FlowStatus.Failed };
     }
     async function completeFlow(executionId, flowName) {
-        await updateExecution(executionId, {
+        const ok = await guardedUpdate(executionId, {
             status: FlowStatus.Completed,
             currentStep: null,
             completedAt: new Date(),
+            leaseOwner: null,
+            leaseExpiresAt: null,
         });
+        if (!ok)
+            return; // Lease lost — another worker owns this execution
         if (audit) {
             await audit.record(`flow.completed.${flowName}`, { executionId, flowName });
         }
     }
     async function failFlow(executionId, flowName, error, row) {
-        await updateExecution(executionId, {
+        const ok = await guardedUpdate(executionId, {
             status: FlowStatus.Failed,
             lastError: error ?? 'Unknown error',
             completedAt: new Date(),
+            leaseOwner: null,
+            leaseExpiresAt: null,
         });
+        if (!ok)
+            return; // Lease lost — another worker owns this execution
         if (audit) {
             await audit.record(`flow.failed.${flowName}`, {
                 executionId,
@@ -397,7 +712,96 @@ export function createFlowEngine(config) {
             .set({ ...updates, updatedAt: new Date() })
             .where(eq(flowExecutionsTable.id, id));
     }
-    return { start, runNext, resume, resumeWaitingByEvent, listRunnable, cancel, status };
+    /**
+     * Finalize an execution as `failed` from the flow runner's perspective when
+     * `runNext` rejects with an unhandled error (e.g. `LeaseLostError` from a
+     * heartbeat). Without this, a zombie row left by a previously-crashed worker
+     * would be re-claimed by `claimNext` on every poll cycle (via the
+     * `status='running' AND lease_expires_at < now()` crash-recovery branch),
+     * retry forever, and pile up log noise.
+     *
+     * The guard is `status IN ('running','created','waiting')` (a set-guard, NOT
+     * a `lease_owner = workerId` check): on `LeaseLostError` the lease is by
+     * definition no longer ours, and we still want to stop the loop. A
+     * concurrent successful completion or an explicit `cancel()` will have
+     * already moved the row to a terminal status, and the `IN (...)` clause
+     * turns this call into a no-op in those races.
+     *
+     * Idempotent: second call on an already-finalized row is a no-op.
+     */
+    async function markFailedFromRunner(executionId, error) {
+        const message = error instanceof Error ? error.message : typeof error === 'string' ? error : 'Unknown error';
+        const result = await db
+            .update(flowExecutionsTable)
+            .set({
+            status: FlowStatus.Failed,
+            lastError: message,
+            completedAt: new Date(),
+            leaseOwner: null,
+            leaseExpiresAt: null,
+            updatedAt: new Date(),
+        })
+            .where(and(eq(flowExecutionsTable.id, executionId), inArray(flowExecutionsTable.status, [
+            FlowStatus.Running,
+            FlowStatus.Created,
+            FlowStatus.Waiting,
+        ])));
+        const rowsAffected = getRowsAffected(result);
+        if (rowsAffected > 0 && audit) {
+            await audit.record('flow.runner.finalized', {
+                executionId,
+                workerId,
+                error: message,
+            });
+        }
+        return rowsAffected > 0;
+    }
+    /**
+     * Update execution with lease ownership guard.
+     * Returns true if the update succeeded (we still own the lease),
+     * false if the lease was lost (0 rows affected).
+     */
+    async function guardedUpdate(id, updates) {
+        const result = await db
+            .update(flowExecutionsTable)
+            .set({ ...updates, updatedAt: new Date() })
+            .where(and(eq(flowExecutionsTable.id, id), eq(flowExecutionsTable.leaseOwner, workerId)));
+        const rowsAffected = getRowsAffected(result);
+        return rowsAffected > 0;
+    }
+    return {
+        start,
+        runNext,
+        resume,
+        resumeWaitingByEvent,
+        listRunnable,
+        claimNext,
+        cancel,
+        status,
+        markFailedFromRunner,
+        get workerId() {
+            return workerId;
+        },
+    };
+}
+/**
+ * Wrap an AIService so that every method defaults `signal` to the flow step's
+ * AbortSignal when the caller doesn't pass one explicitly. `streamGenerate`
+ * is a sync function returning an AsyncIterable and is wrapped accordingly.
+ */
+function wrapAiWithDefaultSignal(ai, defaultSignal) {
+    const withSignal = (params) => ({
+        ...params,
+        signal: params.signal ?? defaultSignal,
+    });
+    return {
+        generate: (params) => ai.generate(withSignal(params)),
+        generateWithUsage: (params) => ai.generateWithUsage(withSignal(params)),
+        streamGenerate: (params) => ai.streamGenerate(withSignal(params)),
+        extract: (params) => ai.extract(withSignal(params)),
+        classify: (params) => ai.classify(withSignal(params)),
+        retrieve: (params) => ai.retrieve(withSignal(params)),
+    };
 }
 function mergeExecutionState(currentState, stepData) {
     if (stepData === undefined) {