@pikku/mongodb 0.12.22 → 0.12.24

package/README.md

@@ -13,12 +13,20 @@ npm install @pikku/mongodb
 
 ```typescript
 import { PikkuMongoDB, MongoDBWorkflowService } from '@pikku/mongodb'
+import { RedisLeaseService } from '@pikku/redis'
 
 const { db } = new PikkuMongoDB(logger, 'mongodb://localhost:27017')
 
-const workflowService = new MongoDBWorkflowService(db)
+const leaseService = new RedisLeaseService('redis://localhost:6379')
+
+const workflowService = new MongoDBWorkflowService(db, { leaseService })
 ```
 
+The workflow service locks runs and steps on the `leaseService` you give it,
+and this package ships none: pair it with `RedisLeaseService`, a Kysely lease
+service, or `InMemoryLeaseService` from `@pikku/core/services` for a single
+process. Register the same instance as the app's `leaseService`.
+
 `PikkuMongoDB` also accepts an existing `MongoClient`, in which case it will not
 close the connection on shutdown.
 

package/dist/src/mongodb-deployment-service.js

@@ -40,7 +40,10 @@ export class MongoDBDeploymentService {
                 createdAt: now,
             },
         }, { upsert: true });
-        this.heartbeatTimer = setInterval(() => this.sendHeartbeat(), this.heartbeatInterval);
+        this.heartbeatTimer = setInterval(() => {
+            // sendHeartbeat catches its own errors; it retries on the next tick
+            void this.sendHeartbeat();
+        }, this.heartbeatInterval);
     }
     async stop() {
         if (this.heartbeatTimer) {

package/dist/src/mongodb-workflow-run-service.js

@@ -35,23 +35,29 @@ export class MongoDBWorkflowRunService {
             .find({ workflowRunId: runId })
             .sort({ createdAt: 1 })
             .toArray();
-        const stepIds = result.map((r) => r._id);
-        const historyCounts = await this.stepHistory
-            .aggregate([
-            { $match: { workflowStepId: { $in: stepIds } } },
-            { $group: { _id: '$workflowStepId', count: { $sum: 1 } } },
-        ])
+        const history = await this.stepHistory
+            .find({ workflowStepId: { $in: result.map((row) => row._id) } })
+            .sort({ createdAt: 1, attempt: 1 })
             .toArray();
-        const countMap = new Map(historyCounts.map((h) => [h._id, h.count]));
+        const latest = new Map();
+        for (const row of history)
+            latest.set(row.workflowStepId, row);
         return result.map((row) => ({
             stepId: row._id,
             stepName: row.stepName,
             rpcName: row.rpcName ?? undefined,
             data: row.data ?? undefined,
+            childRunId: row.childRunId ?? undefined,
+            succeededAt: latest.get(row._id)?.succeededAt
+                ? new Date(latest.get(row._id).succeededAt)
+                : undefined,
+            failedAt: latest.get(row._id)?.failedAt
+                ? new Date(latest.get(row._id).failedAt)
+                : undefined,
             status: row.status,
             result: row.result ?? undefined,
             error: row.error ?? undefined,
-            attemptCount: countMap.get(row._id) ?? 1,
+            attemptCount: row.attempt,
             retries: row.retries != null ? Number(row.retries) : undefined,
             retryDelay: row.retryDelay ?? undefined,
             createdAt: new Date(row.createdAt),
@@ -64,20 +70,18 @@ export class MongoDBWorkflowRunService {
         const stepIds = steps.map((s) => s._id);
         const history = await this.stepHistory
             .find({ workflowStepId: { $in: stepIds } })
-            .sort({ createdAt: 1 })
+            .sort({ createdAt: 1, attempt: 1 })
             .toArray();
-        let attemptCounters = {};
         return history.map((row) => {
             const stepId = row.workflowStepId;
             const step = stepMap.get(stepId);
-            attemptCounters[stepId] = (attemptCounters[stepId] ?? 0) + 1;
             return {
                 stepId,
                 stepName: step.stepName,
                 status: row.status,
                 result: row.result ?? undefined,
                 error: row.error ?? undefined,
-                attemptCount: attemptCounters[stepId],
+                attemptCount: row.attempt,
                 retries: step.retries != null ? Number(step.retries) : undefined,
                 retryDelay: step.retryDelay ?? undefined,
                 createdAt: new Date(row.createdAt),

package/dist/src/mongodb-workflow-service.d.ts

@@ -1,6 +1,6 @@
 import type { SerializedError } from '@pikku/core/errors';
 import { PikkuWorkflowService } from '@pikku/core/workflow';
-import type { WorkflowPlannedStep, WorkflowQueueOptions, WorkflowRun, WorkflowRunWire, StepState, StepStatus, WorkflowStatus, WorkflowVersionStatus } from '@pikku/core/workflow';
+import type { WorkflowPlannedStep, WorkflowServiceOptions, WorkflowRun, WorkflowRunWire, StepState, StepStatus, WorkflowStatus, WorkflowVersionStatus } from '@pikku/core/workflow';
 import type { Db } from 'mongodb';
 export declare class MongoDBWorkflowService extends PikkuWorkflowService {
     private initialized;
@@ -9,7 +9,7 @@ export declare class MongoDBWorkflowService extends PikkuWorkflowService {
     private steps;
     private stepHistory;
     private versions;
-    constructor(db: Db, options?: WorkflowQueueOptions);
+    constructor(db: Db, options: WorkflowServiceOptions);
     init(): Promise<void>;
     protected createRunImpl(workflowName: string, input: any, inline: boolean, graphHash: string, wire: WorkflowRunWire, options?: {
         deterministic?: boolean;
@@ -22,51 +22,71 @@ export declare class MongoDBWorkflowService extends PikkuWorkflowService {
         retryDelay?: string | number;
     }, fromStepName?: string): Promise<StepState>;
     getStepState(runId: string, stepName: string): Promise<StepState>;
+    private toStepState;
+    getRunSteps(runId: string): Promise<Array<StepState & {
+        stepName: string;
+        rpcName?: string;
+        data?: any;
+    }>>;
     getRunHistory(runId: string): Promise<Array<StepState & {
         stepName: string;
     }>>;
     protected setStepRunningImpl(stepId: string): Promise<void>;
     protected setStepScheduledImpl(stepId: string): Promise<void>;
     /**
-     * Move the current attempt of a step's history to a status.
+     * Move one attempt of a step's history to a status — `attempt`, or the
+     * newest when it is not given.
      *
      * Finding the newest attempt and then updating it by `_id` was two round
      * trips with a window in between; `findOneAndUpdate` applies the same sort
      * server-side and writes the row it selected, so nothing can move underneath
-     * it. The attempt is still identified by `createdAt` because Mongo keeps no
-     * attempt number — see `getStepState`, which counts history documents.
+     * it.
      */
     private writeLatestHistory;
     private insertHistoryRecord;
     private getTimestampFieldForStatus;
     protected setStepChildRunIdImpl(stepId: string, childRunId: string): Promise<void>;
-    protected setStepResultImpl(stepId: string, result: any): Promise<void>;
-    protected setStepErrorImpl(stepId: string, error: Error): Promise<void>;
-    protected createRetryAttemptImpl(stepId: string, status: 'pending' | 'running'): Promise<StepState>;
-    withRunLock<T>(id: string, fn: () => Promise<T>): Promise<T>;
+    refreshStepLease(stepId: string, leaseMs: number | null, attempt?: number): Promise<boolean>;
+    protected setStepResultImpl(stepId: string, result: any, attempt?: number): Promise<void>;
+    protected setStepErrorImpl(stepId: string, error: Error, attempt?: number): Promise<void>;
     /**
-     * A pass-through: MongoDB has no advisory-lock primitive to build one on.
-     * The one decision that must exclude — claiming a step to execute it — is
-     * made by `claimStepForExecution` below as a single conditional write, which
-     * needs no lock to be exclusive.
+     * Settle a step. With `attempt`, a step claimed again since matches nothing
+     * and throws `WorkflowStepSupersededError`, keeping the newer claim's state.
      */
-    withStepLock<T>(_runId: string, _stepName: string, fn: () => Promise<T>): Promise<T>;
+    private writeStepOutcome;
+    protected createRetryAttemptImpl(stepId: string, status: 'pending' | 'running'): Promise<StepState>;
     /**
-     * Claim the step with a status-guarded update and read the modified count:
-     * a single-document update is atomic in MongoDB, so two dispatches racing for
+     * Claim the step with a guarded update and read what it matched: a
+     * single-document update is atomic in MongoDB, so two dispatches racing for
      * the same step cannot both proceed.
      *
      * This replaces the read-then-write the base engine does under
-     * `withStepLock`, which excludes nothing here — the lock above is a
-     * pass-through. The winner then goes through the ordinary transition methods,
-     * so history rows see exactly what they saw before.
+     * `withStepLock`: one atomic write needs no lease round-trips around it. The
+     * winner then goes through the ordinary transition methods, so history rows
+     * see exactly what they saw before.
+     *
+     * A `running` step is taken back only once its lease has lapsed, and the
+     * lease is part of the same guarded update, so two dispatches that read one
+     * lapsed lease cannot both take the step.
      */
-    protected claimStepForExecution(runId: string, stepName: string, rpcName: string): Promise<StepState | null>;
+    protected claimStepForExecution(runId: string, stepName: string, rpcName: string, leaseMs: number): Promise<StepState | null>;
     /**
-     * Move a step to `running` only if it is still in one of `from`, reporting
-     * whether this caller is the one that moved it.
+     * Take back a failed or lapsed step as a new attempt. The attempt number is
+     * advanced by the same guarded update that grants the lease, so the worker
+     * it is taken from is fenced out of its step at once.
+     */
+    private reclaimStepAttempt;
+    /**
+     * Move a step to `running` under a fresh lease, only if it is still the
+     * attempt that was read and in one of `from`, reporting whether this caller
+     * is the one that moved it.
      */
     private claimStepStatus;
+    /**
+     * A `running` step matches only once its lease has lapsed; a null lease —
+     * a step parked on a child run — never does.
+     */
+    private claimFilter;
     getCompletedGraphState(runId: string): Promise<{
         completedNodeIds: string[];
         failedNodeIds: string[];
@@ -76,6 +96,7 @@ export declare class MongoDBWorkflowService extends PikkuWorkflowService {
         stepName: string;
         status: StepStatus;
         fromStepName?: string;
+        leaseExpiresAt?: Date;
     }>>;
     getNodeResults(runId: string, nodeIds: string[]): Promise<Record<string, any>>;
     protected setBranchTakenImpl(stepId: string, branchKey: string): Promise<void>;

package/dist/src/mongodb-workflow-service.js

@@ -1,5 +1,13 @@
-import { PikkuWorkflowService, WorkflowStepFunctionMismatchError, } from '@pikku/core/workflow';
+import { leaseAttemptsExhausted, PikkuWorkflowService, WorkflowStepFunctionMismatchError, WorkflowStepLeaseExpiredError, WorkflowStepSupersededError, } from '@pikku/core/workflow';
 import { MongoDBWorkflowRunService } from './mongodb-workflow-run-service.js';
+const ENCODED_STATE_KEY = '__enc_';
+/** Field names cannot hold `.` or `$`, and step names are free text. */
+const encodeStateKey = (name) => /^[a-zA-Z0-9_]+$/.test(name) && !name.startsWith(ENCODED_STATE_KEY)
+    ? name
+    : ENCODED_STATE_KEY + Buffer.from(name, 'utf8').toString('hex');
+const decodeStateKey = (key) => key.startsWith(ENCODED_STATE_KEY)
+    ? Buffer.from(key.slice(ENCODED_STATE_KEY.length), 'hex').toString('utf8')
+    : key;
 export class MongoDBWorkflowService extends PikkuWorkflowService {
     initialized = false;
     runService;
@@ -7,7 +15,7 @@ export class MongoDBWorkflowService extends PikkuWorkflowService {
     steps;
     stepHistory;
     versions;
-    constructor(db, options = {}) {
+    constructor(db, options) {
         super(options);
         this.runService = new MongoDBWorkflowRunService(db);
         this.runs = db.collection('workflow_runs');
@@ -78,10 +86,11 @@ export class MongoDBWorkflowService extends PikkuWorkflowService {
             retries: stepOptions?.retries ?? null,
             retryDelay: stepOptions?.retryDelay?.toString() ?? null,
             fromStepName: fromStepName ?? null,
+            attempt: 1,
             createdAt: now,
             updatedAt: now,
         });
-        await this.insertHistoryRecord(stepId, 'pending');
+        await this.insertHistoryRecord(stepId, 1, 'pending');
         return {
             stepId,
             status: 'pending',
@@ -104,23 +113,27 @@ export class MongoDBWorkflowService extends PikkuWorkflowService {
         if (!row) {
             throw new Error(`Step not found: runId=${runId}, stepName=${stepName}. Use insertStepState to create it.`);
         }
-        const attemptCount = await this.stepHistory.countDocuments({
-            workflowStepId: row._id,
-        });
+        return this.toStepState(row);
+    }
+    toStepState(row) {
         return {
             stepId: row._id,
             status: row.status,
             rpcName: row.rpcName ?? null,
             result: row.result ?? undefined,
             error: row.error ?? undefined,
-            attemptCount,
+            attemptCount: row.attempt,
             retries: row.retries != null ? Number(row.retries) : undefined,
             retryDelay: row.retryDelay ?? undefined,
             fromStepName: row.fromStepName ?? undefined,
+            leaseExpiresAt: row.leaseExpiresAt ?? undefined,
             createdAt: new Date(row.createdAt),
             updatedAt: new Date(row.updatedAt),
         };
     }
+    async getRunSteps(runId) {
+        return this.runService.getRunSteps(runId);
+    }
     async getRunHistory(runId) {
         return this.runService.getRunHistory(runId);
     }
@@ -133,28 +146,31 @@ export class MongoDBWorkflowService extends PikkuWorkflowService {
         await this.writeLatestHistory(stepId, 'scheduled');
     }
     /**
-     * Move the current attempt of a step's history to a status.
+     * Move one attempt of a step's history to a status — `attempt`, or the
+     * newest when it is not given.
      *
      * Finding the newest attempt and then updating it by `_id` was two round
      * trips with a window in between; `findOneAndUpdate` applies the same sort
      * server-side and writes the row it selected, so nothing can move underneath
-     * it. The attempt is still identified by `createdAt` because Mongo keeps no
-     * attempt number — see `getStepState`, which counts history documents.
+     * it.
      */
-    async writeLatestHistory(stepId, status, values = {}) {
+    async writeLatestHistory(stepId, status, values = {}, attempt) {
         const now = new Date();
         const update = { status, ...values };
         const timestampField = this.getTimestampFieldForStatus(status);
         if (timestampField !== 'createdAt') {
             update[timestampField] = now;
         }
-        await this.stepHistory.findOneAndUpdate({ workflowStepId: stepId }, { $set: update }, { sort: { createdAt: -1 } });
+        await this.stepHistory.findOneAndUpdate(attempt === undefined
+            ? { workflowStepId: stepId }
+            : { workflowStepId: stepId, attempt }, { $set: update }, { sort: { attempt: -1 } });
     }
-    async insertHistoryRecord(stepId, status, result, error) {
+    async insertHistoryRecord(stepId, attempt, status, result, error) {
         const now = new Date();
         const doc = {
             _id: crypto.randomUUID(),
             workflowStepId: stepId,
+            attempt,
             status,
             result: result ?? null,
             error: error ?? null,
@@ -192,108 +208,105 @@ export class MongoDBWorkflowService extends PikkuWorkflowService {
             },
         });
     }
-    async setStepResultImpl(stepId, result) {
-        await this.steps.updateOne({ _id: stepId }, {
+    async refreshStepLease(stepId, leaseMs, attempt) {
+        const renewed = await this.steps.updateOne(attempt === undefined ? { _id: stepId } : { _id: stepId, attempt }, {
             $set: {
-                status: 'succeeded',
-                result,
-                error: null,
-                updatedAt: new Date(),
+                leaseExpiresAt: leaseMs === null ? null : new Date(Date.now() + leaseMs),
             },
         });
-        await this.writeLatestHistory(stepId, 'succeeded', { result });
+        return renewed.matchedCount > 0;
     }
-    async setStepErrorImpl(stepId, error) {
+    async setStepResultImpl(stepId, result, attempt) {
+        await this.writeStepOutcome(stepId, 'succeeded', { result, error: null }, attempt);
+        await this.writeLatestHistory(stepId, 'succeeded', { result }, attempt);
+    }
+    async setStepErrorImpl(stepId, error, attempt) {
         const serializedError = {
             message: error.message,
             stack: error.stack,
             code: error.code,
         };
-        await this.steps.updateOne({ _id: stepId }, {
-            $set: {
-                status: 'failed',
-                error: serializedError,
-                result: null,
-                updatedAt: new Date(),
-            },
-        });
-        await this.writeLatestHistory(stepId, 'failed', {
-            error: serializedError,
-        });
+        await this.writeStepOutcome(stepId, 'failed', { error: serializedError, result: null }, attempt);
+        await this.writeLatestHistory(stepId, 'failed', { error: serializedError }, attempt);
+    }
+    /**
+     * Settle a step. With `attempt`, a step claimed again since matches nothing
+     * and throws `WorkflowStepSupersededError`, keeping the newer claim's state.
+     */
+    async writeStepOutcome(stepId, status, values, attempt) {
+        const written = await this.steps.updateOne(attempt === undefined ? { _id: stepId } : { _id: stepId, attempt }, { $set: { ...values, status, updatedAt: new Date() } });
+        if (attempt !== undefined && written.matchedCount === 0) {
+            throw new WorkflowStepSupersededError(stepId, attempt);
+        }
     }
     async createRetryAttemptImpl(stepId, status) {
-        await this.steps.updateOne({ _id: stepId }, {
+        const row = await this.steps.findOneAndUpdate({ _id: stepId }, {
             $set: {
                 status,
                 result: null,
                 error: null,
                 updatedAt: new Date(),
             },
-        });
-        await this.insertHistoryRecord(stepId, status);
-        const row = await this.steps.findOne({ _id: stepId });
+            $inc: { attempt: 1 },
+        }, { returnDocument: 'after' });
         if (!row)
             throw new Error(`Step not found: ${stepId}`);
-        const attemptCount = await this.stepHistory.countDocuments({
-            workflowStepId: stepId,
-        });
-        return {
-            stepId: row._id,
-            status: row.status,
-            rpcName: row.rpcName ?? null,
-            result: row.result ?? undefined,
-            error: row.error ?? undefined,
-            attemptCount,
-            retries: row.retries != null ? Number(row.retries) : undefined,
-            retryDelay: row.retryDelay ?? undefined,
-            fromStepName: row.fromStepName ?? undefined,
-            createdAt: new Date(row.createdAt),
-            updatedAt: new Date(row.updatedAt),
-        };
-    }
-    async withRunLock(id, fn) {
-        return fn();
-    }
-    /**
-     * A pass-through: MongoDB has no advisory-lock primitive to build one on.
-     * The one decision that must exclude — claiming a step to execute it — is
-     * made by `claimStepForExecution` below as a single conditional write, which
-     * needs no lock to be exclusive.
-     */
-    async withStepLock(_runId, _stepName, fn) {
-        return fn();
+        await this.insertHistoryRecord(stepId, row.attempt, status);
+        return this.toStepState(row);
     }
     /**
-     * Claim the step with a status-guarded update and read the modified count:
-     * a single-document update is atomic in MongoDB, so two dispatches racing for
+     * Claim the step with a guarded update and read what it matched: a
+     * single-document update is atomic in MongoDB, so two dispatches racing for
      * the same step cannot both proceed.
      *
      * This replaces the read-then-write the base engine does under
-     * `withStepLock`, which excludes nothing here — the lock above is a
-     * pass-through. The winner then goes through the ordinary transition methods,
-     * so history rows see exactly what they saw before.
+     * `withStepLock`: one atomic write needs no lease round-trips around it. The
+     * winner then goes through the ordinary transition methods, so history rows
+     * see exactly what they saw before.
+     *
+     * A `running` step is taken back only once its lease has lapsed, and the
+     * lease is part of the same guarded update, so two dispatches that read one
+     * lapsed lease cannot both take the step.
      */
-    async claimStepForExecution(runId, stepName, rpcName) {
+    async claimStepForExecution(runId, stepName, rpcName, leaseMs) {
         const stepState = await this.getStepState(runId, stepName);
         // knowledge: decisions/security/a-step-runs-the-function-the-workflow-dispatched-it-with.md
         if (stepState.rpcName !== undefined &&
             stepState.rpcName !== (rpcName ?? null)) {
             throw new WorkflowStepFunctionMismatchError(runId, stepName);
         }
-        if (stepState.status === 'succeeded' || stepState.status === 'running') {
+        if (stepState.status === 'succeeded') {
             return null;
         }
-        if (stepState.status === 'failed') {
-            if (!(await this.claimStepStatus(stepState.stepId, ['failed']))) {
+        if (stepState.status === 'running' && leaseAttemptsExhausted(stepState)) {
+            // Fail it only if its lease is still lapsed at the moment of writing:
+            // a worker that renewed since the read above still owns the step.
+            if (!(await this.claimStepStatus(stepState, ['running'], leaseMs))) {
                 return null;
             }
-            return this.createRetryAttempt(stepState.stepId, 'running');
+            await this.setStepError(stepState.stepId, new WorkflowStepLeaseExpiredError(runId, stepName, stepState.attemptCount), stepState.attemptCount);
+            // A null claim reads as "another dispatch owns it", so nothing else
+            // would requeue the orchestrator to see this failure.
+            await this.resumeWorkflow(runId);
+            return null;
+        }
+        // A step that has spent every attempt is settled, however it failed; a
+        // redelivered message must not buy it another one.
+        if (stepState.status === 'failed' && leaseAttemptsExhausted(stepState)) {
+            return null;
+        }
+        if (stepState.status === 'failed' || stepState.status === 'running') {
+            return this.mirrored(() => this.reclaimStepAttempt(stepState, leaseMs), async (mirror, newStep) => {
+                if (newStep) {
+                    await mirror.createRetryAttempt(stepState.stepId, {
+                        ...newStep,
+                        stepName,
+                    });
+                }
+            });
         }
         if (stepState.status === 'pending' || stepState.status === 'scheduled') {
-            if (!(await this.claimStepStatus(stepState.stepId, [
-                'pending',
-                'scheduled',
-            ]))) {
+            if (!(await this.claimStepStatus(stepState, ['pending', 'scheduled'], leaseMs))) {
                 return null;
             }
             await this.setStepRunning(stepState.stepId);
@@ -302,12 +315,54 @@ export class MongoDBWorkflowService extends PikkuWorkflowService {
         return stepState;
     }
     /**
-     * Move a step to `running` only if it is still in one of `from`, reporting
-     * whether this caller is the one that moved it.
+     * Take back a failed or lapsed step as a new attempt. The attempt number is
+     * advanced by the same guarded update that grants the lease, so the worker
+     * it is taken from is fenced out of its step at once.
+     */
+    async reclaimStepAttempt(stepState, leaseMs) {
+        const now = new Date();
+        const row = await this.steps.findOneAndUpdate(this.claimFilter(stepState, [stepState.status], now), {
+            $set: {
+                status: 'running',
+                result: null,
+                error: null,
+                leaseExpiresAt: new Date(now.getTime() + leaseMs),
+                updatedAt: now,
+            },
+            $inc: { attempt: 1 },
+        }, { returnDocument: 'after' });
+        if (!row)
+            return null;
+        await this.insertHistoryRecord(row._id, row.attempt, 'running');
+        return this.toStepState(row);
+    }
+    /**
+     * Move a step to `running` under a fresh lease, only if it is still the
+     * attempt that was read and in one of `from`, reporting whether this caller
+     * is the one that moved it.
+     */
+    async claimStepStatus(stepState, from, leaseMs) {
+        const now = new Date();
+        const claimed = await this.steps.updateOne(this.claimFilter(stepState, from, now), {
+            $set: {
+                status: 'running',
+                leaseExpiresAt: new Date(now.getTime() + leaseMs),
+                updatedAt: now,
+            },
+        });
+        return claimed.matchedCount > 0;
+    }
+    /**
+     * A `running` step matches only once its lease has lapsed; a null lease —
+     * a step parked on a child run — never does.
      */
-    async claimStepStatus(stepId, from) {
-        const claimed = await this.steps.updateOne({ _id: stepId, status: { $in: from } }, { $set: { status: 'running', updatedAt: new Date() } });
-        return claimed.modifiedCount > 0;
+    claimFilter(stepState, from, now) {
+        return {
+            _id: stepState.stepId,
+            attempt: stepState.attemptCount,
+            status: { $in: from },
+            ...(from.includes('running') ? { leaseExpiresAt: { $lt: now } } : {}),
+        };
     }
     async getCompletedGraphState(runId) {
         const results = await this.steps
@@ -329,10 +384,7 @@ export class MongoDBWorkflowService extends PikkuWorkflowService {
             }
             else if (row.status === 'failed') {
                 const maxAttempts = (row.retries ?? 0) + 1;
-                const attemptCount = await this.stepHistory.countDocuments({
-                    workflowStepId: row._id,
-                });
-                if (attemptCount >= maxAttempts) {
+                if (row.attempt >= maxAttempts) {
                     failedNodeIds.push(nodeId);
                 }
             }
@@ -342,12 +394,13 @@ export class MongoDBWorkflowService extends PikkuWorkflowService {
     async getStepInstances(runId) {
         const rows = await this.steps
             .find({ workflowRunId: runId })
-            .project({ stepName: 1, status: 1, fromStepName: 1 })
+            .project({ stepName: 1, status: 1, fromStepName: 1, leaseExpiresAt: 1 })
             .toArray();
         return rows.map((r) => ({
             stepName: r.stepName,
             status: r.status,
             fromStepName: r.fromStepName ?? undefined,
+            leaseExpiresAt: r.leaseExpiresAt ?? undefined,
         }));
     }
     async getNodeResults(runId, nodeIds) {
@@ -370,12 +423,9 @@ export class MongoDBWorkflowService extends PikkuWorkflowService {
         await this.steps.updateOne({ _id: stepId }, { $set: { branchTaken: branchKey, updatedAt: new Date() } });
     }
     async updateRunStateImpl(runId, name, value) {
-        if (!/^[a-zA-Z0-9_]+$/.test(name)) {
-            throw new Error('Invalid state key name');
-        }
         await this.runs.updateOne({ _id: runId }, {
             $set: {
-                [`state.${name}`]: value,
+                [`state.${encodeStateKey(name)}`]: value,
                 updatedAt: new Date(),
             },
         });
@@ -384,7 +434,10 @@ export class MongoDBWorkflowService extends PikkuWorkflowService {
         const row = await this.runs.findOne({ _id: runId }, { projection: { state: 1 } });
         if (!row)
             return {};
-        return row.state ?? {};
+        return Object.fromEntries(Object.entries(row.state ?? {}).map(([key, value]) => [
+            decodeStateKey(key),
+            value,
+        ]));
     }
     async upsertWorkflowVersionImpl(name, graphHash, graph, source, status) {
         await this.versions.updateOne({ workflowName: name, graphHash }, {

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@pikku/mongodb",
-  "version": "0.12.22",
+  "version": "0.12.24",
   "repository": {
     "type": "git",
     "url": "git+https://github.com/pikkujs/pikku.git",
@@ -15,24 +15,30 @@
   "scripts": {
     "tsc": "tsc",
     "ncu": "npx npm-check-updates",
-    "build": "tsc -b",
+    "build": "bun run clean && tsc -b",
     "test": "bash run-tests.sh",
     "test:watch": "bash run-tests.sh --watch",
     "test:coverage": "bash run-tests.sh --coverage",
     "release": "npm run build && npm test",
-    "prepublishOnly": "bun run build"
+    "prepublishOnly": "bun run build",
+    "clean": "rm -rf dist tsconfig.tsbuildinfo"
   },
   "peerDependencies": {
-    "@pikku/core": "^0.12.108"
+    "@pikku/core": "^0.12.144"
   },
   "dependencies": {
-    "mongodb": "^7.5.0"
+    "mongodb": "^7.7.0"
   },
   "devDependencies": {
-    "mongodb": "^7.5.0",
-    "mongodb-memory-server": "^11.2.0"
+    "mongodb": "^7.7.0",
+    "mongodb-memory-server": "^11.3.0"
   },
   "exports": {
-    ".": "./dist/src/index.js"
-  }
+    ".": "./dist/src/index.js",
+    "./package.json": "./package.json"
+  },
+  "files": [
+    "dist",
+    "!dist/**/*.tsbuildinfo"
+  ]
 }