@hatchet-dev/typescript-sdk 1.13.0 → 1.14.0

package/v1/client/worker/worker-internal.js

@@ -19,22 +19,25 @@ var __importDefault = (this && this.__importDefault) || function (mod) {
     return (mod && mod.__esModule) ? mod : { "default": mod };
 };
 Object.defineProperty(exports, "__esModule", { value: true });
-exports.V1Worker = void 0;
+exports.InternalWorker = void 0;
+exports.mapRateLimitPb = mapRateLimitPb;
 /* eslint-disable no-underscore-dangle */
 /* eslint-disable no-nested-ternary */
 const hatchet_error_1 = __importDefault(require("../../../util/errors/hatchet-error"));
 const dispatcher_1 = require("../../../protoc/dispatcher");
 const hatchet_promise_1 = __importDefault(require("../../../util/hatchet-promise/hatchet-promise"));
 const workflows_1 = require("../../../protoc/workflows");
+const logger_1 = require("../../../util/logger");
 const task_1 = require("../../task");
 const transformer_1 = require("../../conditions/transformer");
 const zod_to_json_schema_1 = require("zod-to-json-schema");
-const step_1 = require("../../../step");
 const apply_namespace_1 = require("../../../util/apply-namespace");
+const sleep_1 = __importDefault(require("../../../util/sleep"));
+const abort_error_1 = require("../../../util/abort-error");
 const context_1 = require("./context");
 const parent_run_context_vars_1 = require("../../parent-run-context-vars");
 const health_server_1 = require("./health-server");
-class V1Worker {
+class InternalWorker {
     constructor(client, options) {
         var _a, _b, _c, _d;
         this.workflow_registry = [];
@@ -90,48 +93,7 @@ class V1Worker {
         this.status = status;
         this.logger.debug(`Worker status changed to: ${status}`);
     }
-    registerActions(workflow) {
-        var _a;
-        const newActions = workflow.steps.reduce((acc, step) => {
-            acc[`${workflow.id}:${step.name.toLowerCase()}`] = step.run;
-            return acc;
-        }, {});
-        const onFailureAction = workflow.onFailure
-            ? {
-                [`${workflow.id}-on-failure:${workflow.onFailure.name}`]: workflow.onFailure.run,
-            }
-            : {};
-        this.action_registry = Object.assign(Object.assign(Object.assign({}, this.action_registry), newActions), onFailureAction);
-        this.action_registry =
-            ((_a = workflow.concurrency) === null || _a === void 0 ? void 0 : _a.name) && workflow.concurrency.key
-                ? Object.assign(Object.assign({}, this.action_registry), { [`${workflow.id}:${workflow.concurrency.name.toLowerCase()}`]: workflow.concurrency.key }) : Object.assign({}, this.action_registry);
-    }
-    getHandler(workflows) {
-        throw new Error('Not implemented');
-        // TODO v1
-        // for (const workflow of workflows) {
-        //   const wf: Workflow = {
-        //     ...workflow,
-        //     id: this.client.config.namespace + workflow.id,
-        //   };
-        //   this.registerActions(wf);
-        // }
-        // return new WebhookHandler(this, workflows);
-    }
-    registerWebhook(webhook) {
-        return __awaiter(this, void 0, void 0, function* () {
-            return this.client._v0.admin.registerWebhook(Object.assign({}, webhook));
-        });
-    }
-    /**
-     * @deprecated use registerWorkflow instead
-     */
-    register_workflow(initWorkflow) {
-        return __awaiter(this, void 0, void 0, function* () {
-            return this.registerWorkflow(initWorkflow);
-        });
-    }
-    registerDurableActionsV1(workflow) {
+    registerDurableActions(workflow) {
         const newActions = workflow._durableTasks
             .filter((task) => !!task.fn)
             .reduce((acc, task) => {
@@ -140,7 +102,7 @@ class V1Worker {
         }, {});
         this.action_registry = Object.assign(Object.assign({}, this.action_registry), newActions);
     }
-    registerActionsV1(workflow) {
+    registerActions(workflow) {
         const newActions = workflow._tasks
             .filter((task) => !!task.fn)
             .reduce((acc, task) => {
@@ -159,9 +121,9 @@ class V1Worker {
             : {};
         this.action_registry = Object.assign(Object.assign(Object.assign({}, this.action_registry), newActions), onFailureAction);
     }
-    registerWorkflowV1(initWorkflow_1) {
+    registerWorkflow(initWorkflow_1) {
         return __awaiter(this, arguments, void 0, function* (initWorkflow, durable = false) {
-            var _a, _b, _c, _d, _e, _f, _g, _h, _j, _k, _l, _m, _o, _p, _q, _r, _s, _t, _u, _v, _w, _x;
+            var _a, _b, _c, _d, _e, _f, _g, _h, _j, _k, _l, _m, _o, _p, _q, _r, _s, _t, _u, _v;
             // patch the namespace
             const workflow = Object.assign(Object.assign({}, initWorkflow.definition), { name: (0, apply_namespace_1.applyNamespace)(initWorkflow.definition.name, this.client.config.namespace).toLowerCase() });
             try {
@@ -192,8 +154,8 @@ class V1Worker {
                         inputs: '{}',
                         parents: [],
                         retries: onFailure.retries || ((_c = workflow.taskDefaults) === null || _c === void 0 ? void 0 : _c.retries) || 0,
-                        rateLimits: (0, step_1.mapRateLimit)(onFailure.rateLimits || ((_d = workflow.taskDefaults) === null || _d === void 0 ? void 0 : _d.rateLimits)),
-                        workerLabels: toPbWorkerLabel(onFailure.desiredWorkerLabels || ((_e = workflow.taskDefaults) === null || _e === void 0 ? void 0 : _e.workerLabels)),
+                        rateLimits: mapRateLimitPb(onFailure.rateLimits || ((_d = workflow.taskDefaults) === null || _d === void 0 ? void 0 : _d.rateLimits)),
+                        workerLabels: mapWorkerLabelPb(onFailure.desiredWorkerLabels || ((_e = workflow.taskDefaults) === null || _e === void 0 ? void 0 : _e.workerLabels)),
                         concurrency: [],
                         backoffFactor: ((_f = onFailure.backoff) === null || _f === void 0 ? void 0 : _f.factor) || ((_h = (_g = workflow.taskDefaults) === null || _g === void 0 ? void 0 : _g.backoff) === null || _h === void 0 ? void 0 : _h.factor),
                         backoffMaxSeconds: ((_j = onFailure.backoff) === null || _j === void 0 ? void 0 : _j.maxSeconds) || ((_l = (_k = workflow.taskDefaults) === null || _k === void 0 ? void 0 : _k.backoff) === null || _l === void 0 ? void 0 : _l.maxSeconds),
@@ -207,11 +169,11 @@ class V1Worker {
                     onSuccessTask = {
                         name: 'on-success-task',
                         fn: workflow.onSuccess,
-                        timeout: '60s',
+                        executionTimeout: '60s',
                         parents,
                         retries: 0,
                         rateLimits: [],
-                        desiredWorkerLabels: {},
+                        desiredWorkerLabels: undefined,
                         concurrency: [],
                     };
                 }
@@ -221,7 +183,7 @@ class V1Worker {
                     onSuccessTask = {
                         name: 'on-success-task',
                         fn: onSuccess.fn,
-                        timeout: onSuccess.executionTimeout || ((_m = workflow.taskDefaults) === null || _m === void 0 ? void 0 : _m.executionTimeout) || '60s',
+                        executionTimeout: onSuccess.executionTimeout || ((_m = workflow.taskDefaults) === null || _m === void 0 ? void 0 : _m.executionTimeout) || '60s',
                         scheduleTimeout: onSuccess.scheduleTimeout || ((_o = workflow.taskDefaults) === null || _o === void 0 ? void 0 : _o.scheduleTimeout),
                         parents,
                         retries: onSuccess.retries || ((_p = workflow.taskDefaults) === null || _p === void 0 ? void 0 : _p.retries) || 0,
@@ -234,19 +196,21 @@ class V1Worker {
                 if (onSuccessTask) {
                     workflow._tasks.push(onSuccessTask);
                 }
-                // cron and event triggers
-                if (workflow.on) {
-                    this.logger.warn(`\`on\` for event and cron triggers is deprecated and will be removed soon, use \`onEvents\` and \`onCrons\` instead for ${workflow.name}`);
-                }
                 const eventTriggers = [
                     ...(workflow.onEvents || []).map((event) => (0, apply_namespace_1.applyNamespace)(event, this.client.config.namespace)),
-                    ...(((_u = workflow.on) === null || _u === void 0 ? void 0 : _u.event)
-                        ? [(0, apply_namespace_1.applyNamespace)(workflow.on.event, this.client.config.namespace)]
+                    ...(workflow.on && 'event' in workflow.on && workflow.on.event
+                        ? Array.isArray(workflow.on.event)
+                            ? workflow.on.event.map((event) => (0, apply_namespace_1.applyNamespace)(event, this.client.config.namespace))
+                            : [(0, apply_namespace_1.applyNamespace)(workflow.on.event, this.client.config.namespace)]
                         : []),
                 ];
                 const cronTriggers = [
                     ...(workflow.onCrons || []),
-                    ...(((_v = workflow.on) === null || _v === void 0 ? void 0 : _v.cron) ? [workflow.on.cron] : []),
+                    ...(workflow.on && 'cron' in workflow.on && workflow.on.cron
+                        ? Array.isArray(workflow.on.cron)
+                            ? workflow.on.cron
+                            : [workflow.on.cron]
+                        : []),
                 ];
                 const concurrencyArr = Array.isArray(concurrency) ? concurrency : [];
                 const concurrencySolo = !Array.isArray(concurrency) ? concurrency : undefined;
@@ -258,13 +222,34 @@ class V1Worker {
                     inputJsonSchema = new TextEncoder().encode(JSON.stringify(jsonSchema));
                 }
                 const durableTaskSet = new Set(workflow._durableTasks);
-                const registeredWorkflow = this.client._v0.admin.putWorkflowV1({
+                let stickyStrategy;
+                // `workflow.sticky` is optional. When omitted, we don't set any sticky strategy.
+                //
+                // When provided, `workflow.sticky` is a v1 (non-protobuf) config which may also include
+                // legacy protobuf enum values for backwards compatibility.
+                if (workflow.sticky != null) {
+                    switch (workflow.sticky) {
+                        case 'soft':
+                        case 'SOFT':
+                        case 0:
+                            stickyStrategy = workflows_1.StickyStrategy.SOFT;
+                            break;
+                        case 'hard':
+                        case 'HARD':
+                        case 1:
+                            stickyStrategy = workflows_1.StickyStrategy.HARD;
+                            break;
+                        default:
+                            throw new hatchet_error_1.default(`Invalid sticky strategy: ${workflow.sticky}`);
+                    }
+                }
+                const registeredWorkflow = this.client.admin.putWorkflow({
                     name: workflow.name,
                     description: workflow.description || '',
                     version: workflow.version || '',
                     eventTriggers,
                     cronTriggers,
-                    sticky: workflow.sticky,
+                    sticky: stickyStrategy,
                     concurrencyArr,
                     onFailureTask,
                     defaultPriority: workflow.defaultPriority,
@@ -283,8 +268,8 @@ class V1Worker {
                             parents: (_d = (_c = task.parents) === null || _c === void 0 ? void 0 : _c.map((p) => p.name)) !== null && _d !== void 0 ? _d : [],
                             userData: '{}',
                             retries: task.retries || ((_e = workflow.taskDefaults) === null || _e === void 0 ? void 0 : _e.retries) || 0,
-                            rateLimits: (0, step_1.mapRateLimit)(task.rateLimits || ((_f = workflow.taskDefaults) === null || _f === void 0 ? void 0 : _f.rateLimits)),
-                            workerLabels: toPbWorkerLabel(task.desiredWorkerLabels || ((_g = workflow.taskDefaults) === null || _g === void 0 ? void 0 : _g.workerLabels)),
+                            rateLimits: mapRateLimitPb(task.rateLimits || ((_f = workflow.taskDefaults) === null || _f === void 0 ? void 0 : _f.rateLimits)),
+                            workerLabels: mapWorkerLabelPb(task.desiredWorkerLabels || ((_g = workflow.taskDefaults) === null || _g === void 0 ? void 0 : _g.workerLabels)),
                             backoffFactor: ((_h = task.backoff) === null || _h === void 0 ? void 0 : _h.factor) || ((_k = (_j = workflow.taskDefaults) === null || _j === void 0 ? void 0 : _j.backoff) === null || _k === void 0 ? void 0 : _k.factor),
                             backoffMaxSeconds: ((_l = task.backoff) === null || _l === void 0 ? void 0 : _l.maxSeconds) || ((_o = (_m = workflow.taskDefaults) === null || _m === void 0 ? void 0 : _m.backoff) === null || _o === void 0 ? void 0 : _o.maxSeconds),
                             conditions: (0, transformer_1.taskConditionsToPb)(task),
@@ -302,94 +287,11 @@ class V1Worker {
                         });
                     }),
                     concurrency: concurrencySolo,
-                    defaultFilters: (_x = (_w = workflow.defaultFilters) === null || _w === void 0 ? void 0 : _w.map((f) => ({
+                    defaultFilters: (_v = (_u = workflow.defaultFilters) === null || _u === void 0 ? void 0 : _u.map((f) => ({
                         scope: f.scope,
                         expression: f.expression,
                         payload: f.payload ? new TextEncoder().encode(JSON.stringify(f.payload)) : undefined,
-                    }))) !== null && _x !== void 0 ? _x : [],
-                });
-                this.registeredWorkflowPromises.push(registeredWorkflow);
-                yield registeredWorkflow;
-                this.workflow_registry.push(workflow);
-            }
-            catch (e) {
-                throw new hatchet_error_1.default(`Could not register workflow: ${e.message}`);
-            }
-            this.registerActionsV1(workflow);
-        });
-    }
-    registerWorkflow(initWorkflow) {
-        return __awaiter(this, void 0, void 0, function* () {
-            var _a, _b, _c;
-            const workflow = Object.assign(Object.assign({}, initWorkflow), { id: (0, apply_namespace_1.applyNamespace)(initWorkflow.id, this.client.config.namespace).toLowerCase() });
-            try {
-                if (((_a = workflow.concurrency) === null || _a === void 0 ? void 0 : _a.key) && workflow.concurrency.expression) {
-                    throw new hatchet_error_1.default('Cannot have both key function and expression in workflow concurrency configuration');
-                }
-                const concurrency = ((_b = workflow.concurrency) === null || _b === void 0 ? void 0 : _b.name) || ((_c = workflow.concurrency) === null || _c === void 0 ? void 0 : _c.expression)
-                    ? {
-                        action: !workflow.concurrency.expression
-                            ? `${workflow.id}:${workflow.concurrency.name}`
-                            : undefined,
-                        maxRuns: workflow.concurrency.maxRuns || 1,
-                        expression: workflow.concurrency.expression,
-                        limitStrategy: workflow.concurrency.limitStrategy || workflows_1.ConcurrencyLimitStrategy.CANCEL_IN_PROGRESS,
-                    }
-                    : undefined;
-                const onFailureJob = workflow.onFailure
-                    ? {
-                        name: `${workflow.id}-on-failure`,
-                        description: workflow.description,
-                        steps: [
-                            {
-                                readableId: workflow.onFailure.name,
-                                action: `${workflow.id}-on-failure:${workflow.onFailure.name}`,
-                                timeout: workflow.onFailure.timeout || '60s',
-                                inputs: '{}',
-                                parents: [],
-                                userData: '{}',
-                                retries: workflow.onFailure.retries || 0,
-                                rateLimits: (0, step_1.mapRateLimit)(workflow.onFailure.rate_limits),
-                                workerLabels: {}, // no worker labels for on failure steps
-                            },
-                        ],
-                    }
-                    : undefined;
-                const registeredWorkflow = this.client._v0.admin.putWorkflow({
-                    name: workflow.id,
-                    description: workflow.description,
-                    version: workflow.version || '',
-                    eventTriggers: workflow.on && workflow.on.event
-                        ? [(0, apply_namespace_1.applyNamespace)(workflow.on.event, this.client.config.namespace)]
-                        : [],
-                    cronTriggers: workflow.on && workflow.on.cron ? [workflow.on.cron] : [],
-                    scheduledTriggers: [],
-                    concurrency,
-                    scheduleTimeout: workflow.scheduleTimeout,
-                    onFailureJob,
-                    sticky: workflow.sticky,
-                    jobs: [
-                        {
-                            name: workflow.id,
-                            description: workflow.description,
-                            steps: workflow.steps.map((step) => {
-                                var _a, _b, _c;
-                                return ({
-                                    readableId: step.name,
-                                    action: `${workflow.id}:${step.name}`,
-                                    timeout: step.timeout || '60s',
-                                    inputs: '{}',
-                                    parents: (_a = step.parents) !== null && _a !== void 0 ? _a : [],
-                                    userData: '{}',
-                                    retries: step.retries || 0,
-                                    rateLimits: (0, step_1.mapRateLimit)(step.rate_limits),
-                                    workerLabels: toPbWorkerLabel(step.worker_labels),
-                                    backoffFactor: (_b = step.backoff) === null || _b === void 0 ? void 0 : _b.factor,
-                                    backoffMaxSeconds: (_c = step.backoff) === null || _c === void 0 ? void 0 : _c.maxSeconds,
-                                });
-                            }),
-                        },
-                    ],
+                    }))) !== null && _v !== void 0 ? _v : [],
                 });
                 this.registeredWorkflowPromises.push(registeredWorkflow);
                 yield registeredWorkflow;
@@ -401,12 +303,9 @@ class V1Worker {
             this.registerActions(workflow);
         });
     }
-    registerAction(actionId, action) {
-        this.action_registry[actionId.toLowerCase()] = action;
-    }
     handleStartStepRun(action) {
         return __awaiter(this, void 0, void 0, function* () {
-            const { actionId, taskRunExternalId } = action;
+            const { actionId, taskRunExternalId, taskName } = action;
             try {
                 // Note: we always use a DurableContext since its a superset of the Context class
                 const context = new context_1.DurableContext(action, this.client, this);
@@ -418,32 +317,37 @@ class V1Worker {
                     return;
                 }
                 const run = () => __awaiter(this, void 0, void 0, function* () {
-                    parent_run_context_vars_1.parentRunContextManager.setContext({
-                        parentId: action.workflowRunId,
-                        parentTaskRunExternalId: taskRunExternalId,
-                        childIndex: 0,
-                        desiredWorkerId: this.workerId || '',
-                    });
                     const { middleware } = this.client.config;
                     if (middleware === null || middleware === void 0 ? void 0 : middleware.before) {
                         const hooks = Array.isArray(middleware.before) ? middleware.before : [middleware.before];
                         for (const hook of hooks) {
-                            const returned = yield hook(context.input, context);
-                            if (returned !== undefined) {
-                                context.input = returned;
+                            const extra = yield hook(context.input, context);
+                            if (extra !== undefined) {
+                                const merged = Object.assign(Object.assign({}, context.input), extra);
+                                context.input = merged;
                                 if (context.data && typeof context.data === 'object') {
-                                    context.data.input = returned;
+                                    context.data.input = merged;
                                 }
                             }
                         }
                     }
-                    let result = yield step(context);
+                    let result = yield parent_run_context_vars_1.parentRunContextManager.runWithContext({
+                        parentId: action.workflowRunId,
+                        parentTaskRunExternalId: taskRunExternalId,
+                        childIndex: 0,
+                        desiredWorkerId: this.workerId || '',
+                        signal: context.abortController.signal,
+                    }, () => {
+                        // Precheck: if cancellation already happened, don't execute user code.
+                        (0, abort_error_1.throwIfAborted)(context.abortController.signal);
+                        return step(context);
+                    });
                     if (middleware === null || middleware === void 0 ? void 0 : middleware.after) {
                         const hooks = Array.isArray(middleware.after) ? middleware.after : [middleware.after];
                         for (const hook of hooks) {
-                            const returned = yield hook(result, context, context.input);
-                            if (returned !== undefined) {
-                                result = returned;
+                            const extra = yield hook(result, context, context.input);
+                            if (extra !== undefined) {
+                                result = extra;
                             }
                         }
                     }
@@ -454,17 +358,17 @@ class V1Worker {
                         if (context.cancelled) {
                             return;
                         }
-                        this.logger.info(`Task run ${taskRunExternalId} succeeded`);
+                        this.logger.info((0, logger_1.taskRunLog)(taskName, taskRunExternalId, 'completed'));
                         // Send the action event to the dispatcher
                         const event = this.getStepActionEvent(action, dispatcher_1.StepActionEventType.STEP_EVENT_TYPE_COMPLETED, false, result || null, action.retryCount);
-                        yield this.client._v0.dispatcher.sendStepActionEvent(event);
+                        yield this.client.dispatcher.sendStepActionEvent(event);
                     }
                     catch (actionEventError) {
                         this.logger.error(`Could not send completed action event: ${actionEventError.message || actionEventError}`);
                         // send a failure event
                         const failureEvent = this.getStepActionEvent(action, dispatcher_1.StepActionEventType.STEP_EVENT_TYPE_FAILED, false, actionEventError.message, action.retryCount);
                         try {
-                            yield this.client._v0.dispatcher.sendStepActionEvent(failureEvent);
+                            yield this.client.dispatcher.sendStepActionEvent(failureEvent);
                         }
                         catch (failureEventError) {
                             this.logger.error(`Could not send failed action event: ${failureEventError.message || failureEventError}`);
@@ -483,7 +387,7 @@ class V1Worker {
                         if (context.cancelled) {
                             return;
                         }
-                        this.logger.error(`Task run ${taskRunExternalId} failed: ${error.message}`);
+                        this.logger.error((0, logger_1.taskRunLog)(taskName, taskRunExternalId, `failed: ${error.message}`));
                         if (error.stack) {
                             this.logger.error(error.stack);
                         }
@@ -492,7 +396,7 @@ class V1Worker {
                             message: error === null || error === void 0 ? void 0 : error.message,
                             stack: error === null || error === void 0 ? void 0 : error.stack,
                         }, action.retryCount);
-                        yield this.client._v0.dispatcher.sendStepActionEvent(event);
+                        yield this.client.dispatcher.sendStepActionEvent(event);
                     }
                     catch (e) {
                         this.logger.error(`Could not send action event: ${e.message}`);
@@ -512,12 +416,21 @@ class V1Worker {
                         yield failure(e);
                         return;
                     }
+                    // Postcheck: user code may swallow AbortError; don't report completion after cancellation.
+                    // If we reached this point and the signal is aborted, the task likely caught/ignored cancellation.
+                    if (context.abortController.signal.aborted) {
+                        this.logger.warn(`Cancellation: task run ${taskRunExternalId} returned after cancellation was signaled. ` +
+                            `This usually means an AbortError was caught and not propagated. ` +
+                            `See https://docs.hatchet.run/home/cancellation`);
+                        return;
+                    }
+                    (0, abort_error_1.throwIfAborted)(context.abortController.signal);
                     yield success(result);
                 }))());
                 this.futures[taskRunExternalId] = future;
                 // Send the action event to the dispatcher
                 const event = this.getStepActionEvent(action, dispatcher_1.StepActionEventType.STEP_EVENT_TYPE_STARTED, false, undefined, action.retryCount);
-                this.client._v0.dispatcher.sendStepActionEvent(event).catch((e) => {
+                this.client.dispatcher.sendStepActionEvent(event).catch((e) => {
                     this.logger.error(`Could not send action event: ${e.message}`);
                 });
                 try {
@@ -525,10 +438,8 @@ class V1Worker {
                 }
                 catch (e) {
                     const message = (e === null || e === void 0 ? void 0 : e.message) || String(e);
-                    if (message.includes('Cancelled')) {
-                        this.logger.debug(`Task run ${taskRunExternalId} was cancelled`);
-                    }
-                    else {
+                    // TODO is this cased correctly...
+                    if (!message.includes('Cancelled')) {
                         this.logger.error(`Could not wait for task run ${taskRunExternalId} to finish. ` +
                             `See https://docs.hatchet.run/home/cancellation for best practices on handling cancellation: `, e);
                     }
@@ -541,7 +452,7 @@ class V1Worker {
     }
     handleStartGroupKeyRun(action) {
         return __awaiter(this, void 0, void 0, function* () {
-            const { actionId, getGroupKeyRunId, taskRunExternalId } = action;
+            const { actionId, getGroupKeyRunId, taskRunExternalId, taskName } = action;
             this.logger.error('Concurrency Key Functions have been deprecated and will be removed in a future release. Use Concurrency Expressions instead.');
             try {
                 const context = new context_1.Context(action, this.client, this);
@@ -561,11 +472,11 @@ class V1Worker {
                     return step(context);
                 });
                 const success = (result) => {
-                    this.logger.info(`Task run ${taskRunExternalId} succeeded`);
+                    this.logger.info((0, logger_1.taskRunLog)(taskName, taskRunExternalId, 'completed'));
                     try {
                         // Send the action event to the dispatcher
                         const event = this.getGroupKeyActionEvent(action, dispatcher_1.GroupKeyActionEventType.GROUP_KEY_EVENT_TYPE_COMPLETED, result);
-                        this.client._v0.dispatcher.sendGroupKeyActionEvent(event).catch((e) => {
+                        this.client.dispatcher.sendGroupKeyActionEvent(event).catch((e) => {
                             this.logger.error(`Could not send action event: ${e.message}`);
                         });
                     }
@@ -579,11 +490,11 @@ class V1Worker {
                     }
                 };
                 const failure = (error) => {
-                    this.logger.error(`Task run ${key} failed: ${error.message}`);
+                    this.logger.error((0, logger_1.taskRunLog)(taskName, taskRunExternalId, `failed: ${error.message}`));
                     try {
                         // Send the action event to the dispatcher
                         const event = this.getGroupKeyActionEvent(action, dispatcher_1.GroupKeyActionEventType.GROUP_KEY_EVENT_TYPE_FAILED, error);
-                        this.client._v0.dispatcher.sendGroupKeyActionEvent(event).catch((e) => {
+                        this.client.dispatcher.sendGroupKeyActionEvent(event).catch((e) => {
                             this.logger.error(`Could not send action event: ${e.message}`);
                         });
                     }
@@ -600,7 +511,7 @@ class V1Worker {
                 this.futures[key] = future;
                 // Send the action event to the dispatcher
                 const event = this.getGroupKeyActionEvent(action, dispatcher_1.GroupKeyActionEventType.GROUP_KEY_EVENT_TYPE_STARTED);
-                this.client._v0.dispatcher.sendGroupKeyActionEvent(event).catch((e) => {
+                this.client.dispatcher.sendGroupKeyActionEvent(event).catch((e) => {
                     this.logger.error(`Could not send action event: ${e.message}`);
                 });
                 yield future.promise;
@@ -641,25 +552,60 @@ class V1Worker {
     }
     handleCancelStepRun(action) {
         return __awaiter(this, void 0, void 0, function* () {
-            const { taskRunExternalId } = action;
+            var _a, _b, _c;
+            const { taskRunExternalId, taskName } = action;
             try {
-                this.logger.info(`Cancelling task run ${taskRunExternalId}`);
                 const future = this.futures[taskRunExternalId];
                 const context = this.contexts[taskRunExternalId];
                 if (context && context.abortController) {
-                    context.abortController.abort('Cancelled by worker');
+                    context.abortController.abort('Cancelled by worker'); // TODO this reason is nonsensical
                 }
                 if (future) {
-                    future.promise.catch(() => {
-                        this.logger.info(`Cancelled task run ${taskRunExternalId}`);
-                    });
-                    future.cancel('Cancelled by worker');
-                    yield future.promise;
+                    const start = Date.now();
+                    const warningThresholdMs = (_a = this.client.config.cancellation_warning_threshold) !== null && _a !== void 0 ? _a : 300;
+                    const gracePeriodMs = (_b = this.client.config.cancellation_grace_period) !== null && _b !== void 0 ? _b : 1000;
+                    const warningMs = Math.max(0, warningThresholdMs);
+                    const graceMs = Math.max(0, gracePeriodMs);
+                    // Ensure cancelling this future doesn't create an unhandled rejection in cases
+                    // where the main action handler isn't currently awaiting `future.promise`.
+                    future.promise.catch(() => undefined);
+                    // Cancel the future (rejects the wrapper); user code must still cooperate with AbortSignal.
+                    future.cancel('Cancelled by worker'); // TODO this reason is nonsensical
+                    // Track completion of the underlying work (not the cancelable wrapper).
+                    // Ensure this promise never throws into our supervision flow.
+                    const completion = ((_c = future.inner) !== null && _c !== void 0 ? _c : future.promise).catch(() => undefined);
+                    // Wait until warning threshold, then log if still running.
+                    if (warningMs > 0) {
+                        const winner = yield Promise.race([
+                            completion.then(() => 'done'),
+                            (0, sleep_1.default)(warningMs).then(() => 'warn'),
+                        ]);
+                        if (winner === 'warn') {
+                            const milliseconds = Date.now() - start;
+                            this.logger.warn(`Cancellation: task run ${taskRunExternalId} has not cancelled after ${milliseconds}ms. Consider checking for blocking operations. ` +
+                                `See https://docs.hatchet.run/home/cancellation`);
+                        }
+                    }
+                    // Wait until grace period (total), then log if still running.
+                    const elapsedMs = Date.now() - start;
+                    const remainingMs = graceMs - elapsedMs;
+                    const winner = yield Promise.race([
+                        completion.then(() => 'done'),
+                        (0, sleep_1.default)(Math.max(0, remainingMs)).then(() => 'grace'),
+                    ]);
+                    if (winner === 'done') {
+                        this.logger.info((0, logger_1.taskRunLog)(taskName, taskRunExternalId, 'cancelled'));
+                    }
+                    else {
+                        const totalElapsedMs = Date.now() - start;
+                        this.logger.error(`Cancellation: task run ${taskRunExternalId} still running after cancellation grace period ` +
+                            `${totalElapsedMs}ms.\n` +
+                            `JavaScript cannot force-kill user code; see: https://docs.hatchet.run/home/cancellation`);
+                    }
                 }
             }
             catch (e) {
-                // Expected: the promise rejects when cancelled
-                this.logger.debug(`Task run ${taskRunExternalId} cancellation completed`);
+                this.logger.error(`Cancellation: error while supervising cancellation for task run ${taskRunExternalId}: ${(e === null || e === void 0 ? void 0 : e.message) || e}`);
             }
             finally {
                 delete this.futures[taskRunExternalId];
@@ -708,7 +654,7 @@ class V1Worker {
      */
     createListener() {
         return __awaiter(this, void 0, void 0, function* () {
-            return this.client._v0.dispatcher.getActionListener({
+            return this.client.dispatcher.getActionListener({
                 workerName: this.name,
                 services: ['default'],
                 actions: Object.keys(this.action_registry),
@@ -747,7 +693,8 @@ class V1Worker {
                         _c = generator_1_1.value;
                         _d = false;
                         const action = _c;
-                        this.logger.info(`Worker ${this.name} received action ${action.actionId}:${action.actionType}`);
+                        const receivedType = (0, logger_1.actionMap)(action.actionType);
+                        this.logger.info((0, logger_1.taskRunLog)(action.taskName, action.taskRunExternalId, `${receivedType}`));
                         void this.handleAction(action);
                     }
                 }
@@ -796,41 +743,41 @@ class V1Worker {
                 this.logger.warn('Worker not registered.');
                 return this.labels;
             }
-            this.client._v0.dispatcher.upsertWorkerLabels(this.workerId, labels);
+            this.client.dispatcher.upsertWorkerLabels(this.workerId, labels);
             return this.labels;
         });
     }
 }
-exports.V1Worker = V1Worker;
-function toPbWorkerLabel(in_) {
+exports.InternalWorker = InternalWorker;
+function mapWorkerLabelPb(in_) {
     if (!in_) {
         return {};
     }
-    return Object.entries(in_).reduce((acc, [key, value]) => {
-        if (!value) {
+    return Object.entries(in_).reduce((acc, [key, label]) => {
+        if (!label) {
             return Object.assign(Object.assign({}, acc), { [key]: {
                     strValue: undefined,
                     intValue: undefined,
                 } });
         }
-        if (typeof value === 'string') {
+        if (typeof label === 'string') {
             return Object.assign(Object.assign({}, acc), { [key]: {
-                    strValue: value,
+                    strValue: label,
                     intValue: undefined,
                 } });
         }
-        if (typeof value === 'number') {
+        if (typeof label === 'number') {
             return Object.assign(Object.assign({}, acc), { [key]: {
                     strValue: undefined,
-                    intValue: value,
+                    intValue: label,
                 } });
         }
         return Object.assign(Object.assign({}, acc), { [key]: {
-                strValue: typeof value.value === 'string' ? value.value : undefined,
-                intValue: typeof value.value === 'number' ? value.value : undefined,
-                required: value.required,
-                weight: value.weight,
-                comparator: value.comparator,
+                strValue: typeof label.value === 'string' ? label.value : undefined,
+                intValue: typeof label.value === 'number' ? label.value : undefined,
+                required: label.required,
+                weight: label.weight,
+                comparator: label.comparator,
             } });
     }, {});
 }
@@ -843,3 +790,72 @@ function getLeaves(tasks) {
 function isLeafTask(task, allTasks) {
     return !allTasks.some((t) => { var _a; return (_a = t.parents) === null || _a === void 0 ? void 0 : _a.some((p) => p.name === task.name); });
 }
+function mapRateLimitPb(limits) {
+    if (!limits)
+        return [];
+    return limits.map((l) => {
+        let key = l.staticKey;
+        const keyExpression = l.dynamicKey;
+        if (l.key !== undefined) {
+            // eslint-disable-next-line no-console
+            console.warn('key is deprecated and will be removed in a future release, please use staticKey instead');
+            key = l.key;
+        }
+        if (keyExpression !== undefined) {
+            if (key !== undefined) {
+                throw new Error('Cannot have both static key and dynamic key set');
+            }
+            key = keyExpression;
+            if (!validateCelExpression(keyExpression)) {
+                throw new Error(`Invalid CEL expression: ${keyExpression}`);
+            }
+        }
+        if (key === undefined) {
+            throw new Error(`Invalid key`);
+        }
+        let units;
+        let unitsExpression;
+        if (typeof l.units === 'number') {
+            units = l.units;
+        }
+        else {
+            if (!validateCelExpression(l.units)) {
+                throw new Error(`Invalid CEL expression: ${l.units}`);
+            }
+            unitsExpression = l.units;
+        }
+        let limitExpression;
+        if (l.limit !== undefined) {
+            if (typeof l.limit === 'number') {
+                limitExpression = `${l.limit}`;
+            }
+            else {
+                if (!validateCelExpression(l.limit)) {
+                    throw new Error(`Invalid CEL expression: ${l.limit}`);
+                }
+                limitExpression = l.limit;
+            }
+        }
+        if (keyExpression !== undefined && limitExpression === undefined) {
+            throw new Error('CEL based keys requires limit to be set');
+        }
+        if (limitExpression === undefined) {
+            limitExpression = `-1`;
+        }
+        return {
+            key,
+            keyExpr: keyExpression,
+            units,
+            unitsExpr: unitsExpression,
+            limitValuesExpr: limitExpression,
+            duration: l.duration,
+        };
+    });
+}
+// Helper function to validate CEL expressions
+// eslint-disable-next-line @typescript-eslint/no-unused-vars
+function validateCelExpression(_expr) {
+    // FIXME: this is a placeholder. In a real implementation, you'd need to use a CEL parser or validator.
+    // For now, we'll just return true to mimic the behavior.
+    return true;
+}