@dbos-inc/dbos-sdk 5.4.10-preview → 5.4.11-preview

package/dist/src/dbos.js

@@ -48,6 +48,7 @@ const crontab_1 = require("./scheduler/crontab");
 const wfqueue_1 = require("./wfqueue");
 const enqueue_workflow_1 = require("./enqueue_workflow");
 const node_assert_1 = __importDefault(require("node:assert"));
+const shard_keys_1 = require("./shard_keys");
 function resolveTimeoutSeconds(options) {
     if (options === undefined)
         return undefined;
@@ -259,8 +260,29 @@ class DBOS {
      * @param systemDatabaseUrl - The system database to create or migrate
      * @param options.schemaName - The schema holding the DBOS system tables (default: `dbos`)
      * @param options.applicationRole - A role to grant access to the system schema once migrated
+     * @param options.sharding - Create or migrate a sharded system database (requires @dbos-inc/dbos-enterprise)
      */
     static async migrate(systemDatabaseUrl, options = {}) {
+        if (options.sharding) {
+            // Sharded migrations come from Transact Enterprise; the grant is the same.
+            const sysdb = (0, system_database_1.createSystemDatabase)({
+                systemDatabaseUrl,
+                logger: new logs_1.GlobalLogger(),
+                serializer: serialization_1.DBOSJSON,
+                schemaName: options.schemaName ?? 'dbos',
+                sharding: true,
+            });
+            try {
+                await sysdb.runMigrations();
+            }
+            finally {
+                await sysdb.destroy();
+            }
+            if (options.applicationRole) {
+                await (0, system_database_1.grantDbosSchemaPermissions)(systemDatabaseUrl, options.applicationRole, new logs_1.GlobalLogger(), options.schemaName ?? 'dbos');
+            }
+            return;
+        }
         await (0, system_database_1.migrateSystemDatabase)(systemDatabaseUrl, new logs_1.GlobalLogger(), options.schemaName ?? 'dbos', options.applicationRole);
     }
     /**
@@ -289,7 +311,11 @@ class DBOS {
         }
         (0, decorators_1.finalizeClassRegistrations)();
         // Every enterprise feature requires Conductor to report the org entitled to it, and the package.
-        const features = enterprise.enabledFeatures(config);
+        let features = enterprise.enabledFeatures(config);
+        if (internalConfig.sharding) {
+            // Auto-upgrade is disabled on a sharded system database, so it needs no entitlement.
+            features = features.filter((f) => f !== enterprise.AUTO_UPGRADE);
+        }
         const entitlementOptions = {
             enterpriseKey: config.transactEnterpriseKey,
             conductorKey: options?.conductorKey,
@@ -304,7 +330,7 @@ class DBOS {
                 : undefined;
         // Workflows can also enroll for auto-upgrade individually, which the config flag's check above does not cover.
         const autoUpgradableWorkflowNames = (0, auto_upgrade_1.autoUpgradableWorkflows)(internalConfig.autoUpgrade);
-        if (autoUpgradableWorkflowNames.length > 0 && !internalConfig.autoUpgrade) {
+        if (autoUpgradableWorkflowNames.length > 0 && !internalConfig.autoUpgrade && !internalConfig.sharding) {
             await enterprise.requireEnterprise([enterprise.AUTO_UPGRADE], entitlementOptions);
             loadedEnterprise ??= enterprise.load(enterprise.AUTO_UPGRADE.description);
         }
@@ -395,7 +421,11 @@ class DBOS {
             executor.conductor.start();
         }
         // Start the auto-upgrade loop if any workflow takes part.
-        if (autoUpgradableWorkflowNames.length > 0) {
+        if (autoUpgradableWorkflowNames.length > 0 && executor.systemDatabase.sharded) {
+            // Its queries aren't keyed by shard, so it doesn't run on a sharded database.
+            executor.logger.warn('Auto-upgrade is disabled: it is not supported on a sharded system database.');
+        }
+        else if (autoUpgradableWorkflowNames.length > 0) {
             (0, node_assert_1.default)(loadedEnterprise, 'the enterprise package is loaded whenever a workflow enrolls in auto-upgrade');
             const upgradePass = new loadedEnterprise.AutoUpgradePass(executor, {
                 graceMs: internalConfig.autoUpgradeGraceMs ?? auto_upgrade_1.DEFAULT_AUTO_UPGRADE_GRACE_MS,
@@ -459,7 +489,7 @@ class DBOS {
                     executor.autoUpgrader = undefined;
                 }
                 await executor.awaitRunningWorkflows(drainTimeoutMS);
-                await executor.stopWorkflowTimeoutSweep();
+                await executor.stopBackgroundJobs();
             }
             // Stop the conductor
             if (executor?.conductor) {
@@ -541,6 +571,11 @@ class DBOS {
     static get workflowID() {
         return (0, context_1.getCurrentContextStore)()?.workflowId;
     }
+    /** The shard key of the currently executing workflow; undefined outside a workflow or on an unsharded database. */
+    static get shardKey() {
+        const ctx = (0, context_1.getCurrentContextStore)();
+        return ctx?.workflowId !== undefined ? ctx.shardKey : undefined;
+    }
     /** Use portable serialization by default? */
     static get #defaultSerializationType() {
         return (0, context_1.getCurrentContextStore)()?.serializationType;
@@ -603,21 +638,23 @@ class DBOS {
      * @param workflowID - ID of the workflow
      * @returns status of the workflow as `WorkflowStatus`, or `null` if there is no workflow with `workflowID`
      */
-    static getWorkflowStatus(workflowID) {
+    static async getWorkflowStatus(workflowID, options) {
         (0, decorators_1.ensureDBOSIsLaunched)('getWorkflowStatus');
-        if (DBOS.isWithinWorkflow()) {
-            if (DBOS.isInStep()) {
-                // OK to use directly
-                return DBOS.#executor.getWorkflowStatus(workflowID);
-            }
-            else if (DBOS.isInWorkflow()) {
-                return DBOS.#executor.getWorkflowStatus(workflowID, DBOS.workflowID, (0, context_1.functionIDGetIncrement)());
-            }
-            else {
-                throw new error_1.DBOSInvalidWorkflowTransitionError('Invalid call to `getWorkflowStatus` inside a `transaction`');
+        return await (0, shard_keys_1.bindShardKeys)([[workflowID, options?.shardKey]], () => {
+            if (DBOS.isWithinWorkflow()) {
+                if (DBOS.isInStep()) {
+                    // OK to use directly
+                    return DBOS.#executor.getWorkflowStatus(workflowID);
+                }
+                else if (DBOS.isInWorkflow()) {
+                    return DBOS.#executor.getWorkflowStatus(workflowID, DBOS.workflowID, (0, context_1.functionIDGetIncrement)());
+                }
+                else {
+                    throw new error_1.DBOSInvalidWorkflowTransitionError('Invalid call to `getWorkflowStatus` inside a `transaction`');
+                }
             }
-        }
-        return DBOS.#executor.getWorkflowStatus(workflowID);
+            return DBOS.#executor.getWorkflowStatus(workflowID);
+        });
     }
     /**
      * Get the workflow result, given a workflow ID
@@ -634,7 +671,8 @@ class DBOS {
             // Reserve the function ID synchronously, before any await.
             timerFuncID = (0, context_1.functionIDGetIncrement)();
         }
-        return await getResultInternal(workflowID, timeoutSeconds, timerFuncID, undefined, pollingIntervalMs);
+        const shardKey = typeof options === 'object' ? options.shardKey : undefined;
+        return await (0, shard_keys_1.bindShardKeys)([[workflowID, shardKey]], () => getResultInternal(workflowID, timeoutSeconds, timerFuncID, undefined, pollingIntervalMs));
     }
     /**
      * Wait for any one of the given workflow handles to complete and return it.
@@ -658,9 +696,9 @@ class DBOS {
             handleMap.set(handle.workflowID, handle);
         }
         const workflowIds = [...handleMap.keys()];
-        const completedId = await runInternalStep(async () => {
+        const completedId = await (0, shard_keys_1.bindShardKeys)(handles.map((h) => [h.workflowID, h.shardKey]), () => runInternalStep(async () => {
             return await dbos_executor_1.DBOSExecutor.globalInstance.systemDatabase.awaitFirstWorkflowId(workflowIds, DBOS.workflowID, pollingIntervalMs);
-        }, 'DBOS.waitFirst');
+        }, 'DBOS.waitFirst'));
         return handleMap.get(completedId);
     }
     /**
@@ -677,9 +715,9 @@ class DBOS {
         }
         const pollingIntervalMs = resolvePollingIntervalMs(options);
         const workflowIds = [...new Set(handles.map((handle) => handle.workflowID))];
-        await runInternalStep(async () => {
+        await (0, shard_keys_1.bindShardKeys)(handles.map((h) => [h.workflowID, h.shardKey]), () => runInternalStep(async () => {
             await dbos_executor_1.DBOSExecutor.globalInstance.systemDatabase.awaitWorkflowIds(workflowIds, DBOS.workflowID, pollingIntervalMs);
-        }, 'DBOS.waitAll');
+        }, 'DBOS.waitAll'));
         return handles;
     }
     /**
@@ -689,15 +727,14 @@ class DBOS {
      * @param workflowID - ID of the workflow
      * @returns `WorkflowHandle` that can be used to poll for the status or result of any workflow with `workflowID`
      */
-    static retrieveWorkflow(workflowID) {
+    static retrieveWorkflow(workflowID, options) {
         (0, decorators_1.ensureDBOSIsLaunched)('retrieveWorkflow');
-        if (DBOS.isWithinWorkflow()) {
-            if (!DBOS.isInWorkflow()) {
-                throw new error_1.DBOSInvalidWorkflowTransitionError('Invalid call to `retrieveWorkflow` inside a `transaction` or `step`');
-            }
-            return new workflow_1.RetrievedHandle(workflowID);
+        if (DBOS.isWithinWorkflow() && !DBOS.isInWorkflow()) {
+            throw new error_1.DBOSInvalidWorkflowTransitionError('Invalid call to `retrieveWorkflow` inside a `transaction` or `step`');
         }
-        return DBOS.#executor.retrieveWorkflow(workflowID);
+        // The resolved key, so the handle also works outside this binding.
+        const shardKey = (0, shard_keys_1.bindShardKeys)([[workflowID, options?.shardKey]], () => DBOS.#executor.systemDatabase.keyHint(workflowID));
+        return new workflow_1.RetrievedHandle(workflowID, shardKey);
     }
     /**
      * Query the system database for all workflows matching the provided predicate
@@ -727,14 +764,14 @@ class DBOS {
      * @param workflowID - ID of the workflow
      * @param priority - Priority value (0 to 2,147,483,647). Lower values are dequeued first.
      */
-    static async setWorkflowPriority(workflowID, priority) {
+    static async setWorkflowPriority(workflowID, priority, options) {
         (0, decorators_1.ensureDBOSIsLaunched)('setWorkflowPriority');
         if (priority < dbos_executor_1.DBOS_QUEUE_MIN_PRIORITY || priority > dbos_executor_1.DBOS_QUEUE_MAX_PRIORITY) {
             throw new error_1.DBOSInvalidQueuePriorityError(priority, dbos_executor_1.DBOS_QUEUE_MIN_PRIORITY, dbos_executor_1.DBOS_QUEUE_MAX_PRIORITY);
         }
-        return runInternalStep(async () => {
+        return (0, shard_keys_1.bindShardKeys)([[workflowID, options?.shardKey]], () => runInternalStep(async () => {
             return dbos_executor_1.DBOSExecutor.globalInstance.systemDatabase.setWorkflowPriority(workflowID, priority);
-        }, 'DBOS.setWorkflowPriority');
+        }, 'DBOS.setWorkflowPriority'));
     }
     /**
      * Update the delay on a DELAYED workflow.
@@ -746,9 +783,10 @@ class DBOS {
     static async setWorkflowDelay(workflowID, options) {
         (0, decorators_1.ensureDBOSIsLaunched)('setWorkflowDelay');
         const delayUntilEpochMS = resolveDelayEpochMS(options);
-        return runInternalStep(async () => {
+        const shardKey = typeof options === 'object' ? options.shardKey : undefined;
+        return (0, shard_keys_1.bindShardKeys)([[workflowID, shardKey]], () => runInternalStep(async () => {
             return dbos_executor_1.DBOSExecutor.globalInstance.systemDatabase.setWorkflowDelay(workflowID, delayUntilEpochMS);
-        }, 'DBOS.setWorkflowDelay');
+        }, 'DBOS.setWorkflowDelay'));
     }
     /**
      * Retrieve the steps of a workflow
@@ -758,9 +796,9 @@ class DBOS {
      */
     static async listWorkflowSteps(workflowID, options) {
         (0, decorators_1.ensureDBOSIsLaunched)('listWorkflowSteps');
-        return await runInternalStep(async () => {
+        return await (0, shard_keys_1.bindShardKeys)([[workflowID, options?.shardKey]], () => runInternalStep(async () => {
             return await DBOS.#executor.listWorkflowSteps(workflowID, true, options);
-        }, 'DBOS.listWorkflowSteps');
+        }, 'DBOS.listWorkflowSteps'));
     }
     /**
      * Cancel a workflow given its ID.
@@ -779,9 +817,9 @@ class DBOS {
      */
     static async cancelWorkflows(workflowIDs, options) {
         (0, decorators_1.ensureDBOSIsLaunched)('cancelWorkflows');
-        return runInternalStep(async () => {
+        return (0, shard_keys_1.bindShardKeys)(workflowIDs.map((wid) => [wid, options?.shardKey]), () => runInternalStep(async () => {
             return dbos_executor_1.DBOSExecutor.globalInstance.systemDatabase.cancelWorkflows(workflowIDs, options?.cancelChildren);
-        }, 'DBOS.cancelWorkflow');
+        }, 'DBOS.cancelWorkflow'));
     }
     /**
      * Resume a workflow given its ID.
@@ -789,8 +827,11 @@ class DBOS {
      * @throws DBOSNonExistentWorkflowError if the workflow does not exist
      */
     static async resumeWorkflow(workflowID, options) {
-        const handles = await this.resumeWorkflows([workflowID], options);
-        return handles[0];
+        return await (0, shard_keys_1.bindShardKeys)([[workflowID, options?.shardKey]], async () => {
+            await this.resumeWorkflows([workflowID], options);
+            // The resolved key, so the handle also works outside this binding.
+            return new workflow_1.RetrievedHandle(workflowID, DBOS.#executor.systemDatabase.keyHint(workflowID));
+        });
     }
     /**
      * Resume multiple workflows given their IDs.
@@ -800,10 +841,12 @@ class DBOS {
      */
     static async resumeWorkflows(workflowIDs, options) {
         (0, decorators_1.ensureDBOSIsLaunched)('resumeWorkflows');
-        await runInternalStep(async () => {
-            return dbos_executor_1.DBOSExecutor.globalInstance.systemDatabase.resumeWorkflows(workflowIDs, options?.queueName);
-        }, 'DBOS.resumeWorkflow');
-        return workflowIDs.map((wfid) => this.retrieveWorkflow(wfid));
+        return await (0, shard_keys_1.bindShardKeys)(workflowIDs.map((wid) => [wid, options?.shardKey]), async () => {
+            await runInternalStep(async () => {
+                return dbos_executor_1.DBOSExecutor.globalInstance.systemDatabase.resumeWorkflows(workflowIDs, options?.queueName);
+            }, 'DBOS.resumeWorkflow');
+            return workflowIDs.map((wfid) => new workflow_1.RetrievedHandle(wfid, options?.shardKey));
+        });
     }
     /**
      * Delete a workflow and optionally all its child workflows.
@@ -814,8 +857,8 @@ class DBOS {
      * @param workflowID - ID of the workflow to delete
      * @param deleteChildren - If true, also delete all child workflows recursively (default: false)
      */
-    static async deleteWorkflow(workflowID, deleteChildren = false) {
-        return this.deleteWorkflows([workflowID], deleteChildren);
+    static async deleteWorkflow(workflowID, deleteChildren = false, options) {
+        return this.deleteWorkflows([workflowID], deleteChildren, options);
     }
     /**
      * Delete multiple workflows and optionally all their child workflows.
@@ -825,11 +868,11 @@ class DBOS {
      * @param workflowIDs - IDs of the workflows to delete
      * @param deleteChildren - If true, also delete all child workflows recursively (default: false)
      */
-    static async deleteWorkflows(workflowIDs, deleteChildren = false) {
+    static async deleteWorkflows(workflowIDs, deleteChildren = false, options) {
         (0, decorators_1.ensureDBOSIsLaunched)('deleteWorkflows');
-        return runInternalStep(async () => {
+        return (0, shard_keys_1.bindShardKeys)(workflowIDs.map((wid) => [wid, options?.shardKey]), () => runInternalStep(async () => {
             return dbos_executor_1.DBOSExecutor.globalInstance.systemDatabase.deleteWorkflows(workflowIDs, deleteChildren);
-        }, 'DBOS.deleteWorkflow');
+        }, 'DBOS.deleteWorkflow'));
     }
     /**
      * Fork a workflow given its ID.
@@ -841,10 +884,14 @@ class DBOS {
      */
     static async forkWorkflow(workflowID, startStep, options) {
         (0, decorators_1.ensureDBOSIsLaunched)('forkWorkflow');
-        const forkedID = await runInternalStep(async () => {
-            return await DBOS.#executor.forkWorkflow(workflowID, startStep, options);
-        }, 'DBOS.forkWorkflow');
-        return this.retrieveWorkflow(forkedID);
+        return await (0, shard_keys_1.bindShardKeys)([[workflowID, options?.shardKey]], async () => {
+            const forkedID = await runInternalStep(async () => {
+                return await DBOS.#executor.forkWorkflow(workflowID, startStep, options);
+            }, 'DBOS.forkWorkflow');
+            // A fork lives under its original's shard key, however the caller supplied it.
+            const forkKey = DBOS.#executor.systemDatabase.keyHint(workflowID);
+            return new workflow_1.RetrievedHandle(forkedID, forkKey);
+        });
     }
     /**
      * Rewind a workflow given its ID: drop its recorded history from `startStep` onwards
@@ -868,10 +915,13 @@ class DBOS {
     static async rewindWorkflow(workflowID, options) {
         (0, decorators_1.ensureDBOSIsLaunched)('rewindWorkflow');
         const startStep = options?.startStep ?? 0;
-        await runInternalStep(async () => {
-            return await DBOS.#executor.rewindWorkflow(workflowID, startStep, options ?? {});
-        }, 'DBOS.rewindWorkflow');
-        return this.retrieveWorkflow(workflowID);
+        return await (0, shard_keys_1.bindShardKeys)([[workflowID, options?.shardKey]], async () => {
+            await runInternalStep(async () => {
+                return await DBOS.#executor.rewindWorkflow(workflowID, startStep, options ?? {});
+            }, 'DBOS.rewindWorkflow');
+            // The resolved key, so the handle also works outside this binding.
+            return new workflow_1.RetrievedHandle(workflowID, DBOS.#executor.systemDatabase.keyHint(workflowID));
+        });
     }
     /**
      * Sleep for the specified amount of time.
@@ -986,6 +1036,13 @@ class DBOS {
         (0, decorators_1.ensureDBOSIsLaunched)('workflows');
         return DBOS.#withTopContext({ workflowTimeoutMS: timeoutMS }, callback);
     }
+    /** Set the shard key of each workflow started, enqueued, or called in `callback`; their children don't inherit it. */
+    static async withShardKey(shardKey, callback) {
+        if (typeof shardKey !== 'string' || shardKey === '') {
+            throw new error_1.DBOSError('shardKey must be a non-empty string');
+        }
+        return DBOS.#withTopContext({ shardKeyForNextWorkflow: shardKey }, callback);
+    }
     static async #withTopContext(options, callback) {
         const pctx = (0, context_1.getCurrentContextStore)();
         if (pctx) {
@@ -1073,16 +1130,16 @@ class DBOS {
             // calls in the same workflow don't race for IDs.
             const functionID = (0, context_1.functionIDGetIncrement)();
             const sermsg = await (0, serialization_1.serializeValue)(message, DBOS.#executor.serializer, options?.serializationType ?? DBOS.#defaultSerializationType);
-            return await dbos_executor_1.DBOSExecutor.globalInstance.systemDatabase.send(DBOS.workflowID, functionID, destinationID, sermsg.serializedValue, topic, sermsg.serialization, idempotencyKey);
+            return await dbos_executor_1.DBOSExecutor.globalInstance.systemDatabase.send(DBOS.workflowID, functionID, destinationID, sermsg.serializedValue, topic, sermsg.serialization, idempotencyKey, options?.shardKey);
         }
         else if (DBOS.isInStep()) {
             // Inside a step: not recorded, but fenced on the enclosing workflow's ownership.
             const sermsg = await (0, serialization_1.serializeValue)(message, DBOS.#executor.serializer, options?.serializationType ?? DBOS.#defaultSerializationType);
-            return await dbos_executor_1.DBOSExecutor.globalInstance.systemDatabase.sendFromStep(DBOS.workflowID, destinationID, sermsg.serializedValue, topic, sermsg.serialization, idempotencyKey);
+            return await dbos_executor_1.DBOSExecutor.globalInstance.systemDatabase.sendFromStep(DBOS.workflowID, destinationID, sermsg.serializedValue, topic, sermsg.serialization, idempotencyKey, options?.shardKey);
         }
         else {
             const sermsg = await (0, serialization_1.serializeValue)(message, DBOS.#executor.serializer, options?.serializationType ?? DBOS.#defaultSerializationType);
-            return dbos_executor_1.DBOSExecutor.globalInstance.systemDatabase.sendDirect(destinationID, sermsg.serializedValue, topic, sermsg.serialization, idempotencyKey);
+            return dbos_executor_1.DBOSExecutor.globalInstance.systemDatabase.sendDirect(destinationID, sermsg.serializedValue, topic, sermsg.serialization, idempotencyKey, undefined, options?.shardKey);
         }
     }
     /**
@@ -1150,6 +1207,10 @@ class DBOS {
      */
     static async getEvent(workflowID, key, options) {
         (0, decorators_1.ensureDBOSIsLaunched)('getEvent');
+        const shardKey = typeof options === 'object' ? options.shardKey : undefined;
+        return await (0, shard_keys_1.bindShardKeys)([[workflowID, shardKey]], () => DBOS.#getEvent(workflowID, key, options));
+    }
+    static async #getEvent(workflowID, key, options) {
         const timeoutSeconds = resolveTimeoutSeconds(options);
         const pollingIntervalMs = resolvePollingIntervalMs(options);
         if (DBOS.isWithinWorkflow()) {
@@ -1227,6 +1288,7 @@ class DBOS {
             timeoutMS: resolveStreamTimeoutMS(options),
             functionName: system_database_1.DBOS_FUNCNAME_READSTREAM,
             checkpoint: true,
+            shardKey: options.shardKey,
         });
     }
     /**
@@ -1250,6 +1312,7 @@ class DBOS {
             timeoutMS: resolveStreamTimeoutMS(options),
             functionName: system_database_1.DBOS_FUNCNAME_READSTREAMOFFSET,
             checkpoint: true,
+            shardKey: options.shardKey,
         });
     }
     //////
@@ -1306,6 +1369,8 @@ class DBOS {
         const ppctx = (0, context_1.getCurrentContextStore)();
         const queueName = params.queueName ?? ppctx?.queueAssignedForWorkflows;
         const timeoutMS = params.timeoutMS ?? ppctx?.workflowTimeoutMS;
+        // Set by DBOS.withShardKey; a workflow never inherits its caller's own key.
+        const shardKey = (0, context_1.nextWorkflowShardKey)(ppctx);
         const instance = $this === undefined || typeof $this === 'function' ? undefined : $this;
         if (instance && !(instance instanceof decorators_1.ConfiguredInstance)) {
             throw new error_1.DBOSInvalidWorkflowTransitionError('Attempt to call a `workflow` function on an object that is not a `ConfiguredInstance`');
@@ -1343,7 +1408,7 @@ class DBOS {
         if (isChild) {
             const pwfid = ppctx.workflowId;
             const wfParams = {
-                workflowUUID: wfId || pwfid + '-' + funcId,
+                workflowUUID: wfId || (0, context_1.childWorkflowID)(pwfid, funcId, ppctx?.shardKey, shardKey),
                 configuredInstance: instance,
                 queueName,
                 timeoutMS,
@@ -1355,6 +1420,7 @@ class DBOS {
                 workflowAttributes: params.workflowAttributes,
                 authenticatedUser: params.authenticatedUser,
                 authenticatedRoles: params.authenticatedRoles,
+                shardKey,
             };
             return await invokeRegOp(wfParams, pwfid, funcId);
         }
@@ -1370,6 +1436,7 @@ class DBOS {
                 workflowAttributes: params.workflowAttributes,
                 authenticatedUser: params.authenticatedUser,
                 authenticatedRoles: params.authenticatedRoles,
+                shardKey,
             };
             return await invokeRegOp(wfParams, undefined, undefined);
         }
@@ -1396,6 +1463,8 @@ class DBOS {
         // Reserve the parent's child-funcID once, before any await. Each loop iteration passes it
         // through so the return-existing retry loop consumes a single funcID instead of one per try.
         const childFuncId = DBOS.isInWorkflow() ? (0, context_1.functionIDGetIncrement)() : undefined;
+        // The key the attached-to workflow lives under, as the start would have used.
+        const shardKey = (0, context_1.nextWorkflowShardKey)((0, context_1.getCurrentContextStore)());
         // Resolve the workflow ID once: if the caller supplied one explicitly or via
         // `withNextWorkflowID`, reuse it on every iteration.
         const resolvedWorkflowID = (0, context_1.getNextWFID)(params.workflowID);
@@ -1416,7 +1485,7 @@ class DBOS {
                         const now = Date.now();
                         await dbos_executor_1.DBOSExecutor.globalInstance.systemDatabase.recordOperationResult(DBOS.workflowID, childFuncId, regOP.name, true, now, now, { childWorkflowID: existingID });
                     }
-                    return new workflow_1.RetrievedHandle(existingID);
+                    return new workflow_1.RetrievedHandle(existingID, shardKey);
                 }
                 // The prior workflow's deduplication_id was cleared between our INSERT and
                 // the lookup (it completed or was cancelled). Loop and try to claim the slot.
@@ -1700,6 +1769,7 @@ class DBOS {
             cronTimezone: options.options?.cronTimezone ?? null,
             queueName: options.options?.queueName ?? null,
             applicationName: utils_1.globalParams.appName,
+            shardKey: options.options?.shardKey,
         };
         await runTransactionalInternalStep((client) => dbos_executor_1.DBOSExecutor.globalInstance.systemDatabase.createSchedule(schedInternal, client), 'DBOS.createSchedule');
     }
@@ -1785,6 +1855,7 @@ class DBOS {
                 automaticBackfill: sched.automaticBackfill ?? false,
                 cronTimezone: sched.cronTimezone ?? null,
                 queueName: sched.queueName ?? null,
+                shardKey: sched.shardKey,
             });
         }
         await dbos_executor_1.DBOSExecutor.globalInstance.systemDatabase.applySchedules(internals);
@@ -1796,7 +1867,7 @@ class DBOS {
         }
         const executor = dbos_executor_1.DBOSExecutor.globalInstance;
         const workflowID = await (0, scheduler_1.triggerSchedule)(executor.systemDatabase, executor.serializer, name);
-        return new workflow_1.RetrievedHandle(workflowID);
+        return new workflow_1.RetrievedHandle(workflowID, await (0, scheduler_1.scheduleShardKey)(executor.systemDatabase, name));
     }
     static async backfillSchedule(name, start, end) {
         (0, decorators_1.ensureDBOSIsLaunched)('backfillSchedule');
@@ -1805,7 +1876,8 @@ class DBOS {
         }
         const executor = dbos_executor_1.DBOSExecutor.globalInstance;
         const workflowIDs = await (0, scheduler_1.backfillSchedule)(executor.systemDatabase, executor.serializer, name, start, end);
-        return workflowIDs.map((id) => new workflow_1.RetrievedHandle(id));
+        const shardKey = await (0, scheduler_1.scheduleShardKey)(executor.systemDatabase, name);
+        return workflowIDs.map((id) => new workflow_1.RetrievedHandle(id, shardKey));
     }
     // ==================== Application Versions ====================
     /**