feat(workflow): Expose updateId to update handlers

packages/test/src/test-integration-update.ts

@@ -461,6 +461,77 @@ test('Interruption of update by server long-poll timeout is invisible to client'
   });
 });
 
+export const currentInfoUpdate = wf.defineUpdate<string, []>('current-info-update');
+
+export async function workflowWithCurrentUpdateInfo(): Promise<string[]> {
+  const state: Promise<string>[] = [];
+  const getUpdateId = async (): Promise<string> => {
+    await wf.sleep(10);
+    const info = wf.currentUpdateInfo();
+    if (info === undefined) {
+      throw new Error('No current update info');
+    }
+    return info.id;
+  };
+  const updateHandler = async (): Promise<string> => {
+    const info = wf.currentUpdateInfo();
+    if (info === undefined || info.name !== 'current-info-update') {
+      throw new Error(`Invalid current update info in updateHandler: info ${info?.name}`);
+    }
+    const id = await getUpdateId();
+    if (info.id !== id) {
+      throw new Error(`Update id changed: before ${info.id} after ${id}`);
+    }
+
+    state.push(getUpdateId());
+    // Re-fetch and return
+    const infoAfter = wf.currentUpdateInfo();
+    if (infoAfter === undefined) {
+      throw new Error('Invalid current update info in updateHandler - after');
+    }
+    return infoAfter.id;
+  };
+
+  const validator = (): void => {
+    const info = wf.currentUpdateInfo();
+    if (info === undefined || info.name !== 'current-info-update') {
+      throw new Error(`Invalid current update info in validator: info ${info?.name}`);
+    }
+  };
+
+  wf.setHandler(currentInfoUpdate, updateHandler, { validator });
+
+  if (wf.currentUpdateInfo() !== undefined) {
+    throw new Error('Current update info not undefined outside handler');
+  }
+
+  await wf.condition(() => state.length === 5);
+
+  if (wf.currentUpdateInfo() !== undefined) {
+    throw new Error('Current update info not undefined outside handler - after');
+  }
+
+  return await Promise.all(state);
+}
+
+test('currentUpdateInfo returns the update id', async (t) => {
+  const { createWorker, startWorkflow } = helpers(t);
+  const worker = await createWorker();
+  await worker.runUntil(async () => {
+    const wfHandle = await startWorkflow(workflowWithCurrentUpdateInfo);
+    const updateIds = await Promise.all([
+      wfHandle.executeUpdate(currentInfoUpdate, { updateId: 'update1' }),
+      wfHandle.executeUpdate(currentInfoUpdate, { updateId: 'update2' }),
+      wfHandle.executeUpdate(currentInfoUpdate, { updateId: 'update3' }),
+      wfHandle.executeUpdate(currentInfoUpdate, { updateId: 'update4' }),
+      wfHandle.executeUpdate(currentInfoUpdate, { updateId: 'update5' }),
+    ]);
+    t.deepEqual(updateIds, ['update1', 'update2', 'update3', 'update4', 'update5']);
+    const wfResults = await wfHandle.result();
+    t.deepEqual(wfResults.sort(), ['update1', 'update2', 'update3', 'update4', 'update5']);
+  });
+});
+
 test('startUpdate throws WorkflowUpdateRPCTimeoutOrCancelledError with no worker', async (t) => {
   const { startWorkflow } = helpers(t);
   const wfHandle = await startWorkflow(workflowWithUpdates);

packages/workflow/src/interfaces.ts

@@ -190,6 +190,23 @@ export interface UnsafeWorkflowInfo {
   readonly isReplaying: boolean;
 }
 
+/**
+ * Information about a workflow update.
+ *
+ * @experimental
+ */
+export interface UpdateInfo {
+  /**
+   *  A workflow-unique identifier for this update.
+   */
+  readonly id: string;
+
+  /**
+   *  The update type name.
+   */
+  readonly name: string;
+}
+
 export interface ParentWorkflowInfo {
   workflowId: string;
   runId: string;

packages/workflow/src/internals.ts

@@ -21,6 +21,7 @@ import { checkExtends } from '@temporalio/common/lib/type-helpers';
 import type { coresdk, temporal } from '@temporalio/proto';
 import { alea, RNG } from './alea';
 import { RootCancellationScope } from './cancellation-scope';
+import { UpdateScope } from './update-scope';
 import { DeterminismViolationError, LocalActivityDoBackoff, isCancellation } from './errors';
 import { QueryInput, SignalInput, UpdateInput, WorkflowExecuteInput, WorkflowInterceptors } from './interceptors';
 import {
@@ -629,34 +630,37 @@ export class Activator implements ActivationHandler {
     //
     // Note that there is a deliberately unhandled promise rejection below.
     // These are caught elsewhere and fail the corresponding activation.
-    let input: UpdateInput;
-    try {
-      if (runValidator && this.updateHandlers.get(name)?.validator) {
-        const validate = composeInterceptors(
-          this.interceptors.inbound,
-          'validateUpdate',
-          this.validateUpdateNextHandler.bind(this)
-        );
-        validate(makeInput());
+    const doUpdateImpl = async () => {
+      let input: UpdateInput;
+      try {
+        if (runValidator && this.updateHandlers.get(name)?.validator) {
+          const validate = composeInterceptors(
+            this.interceptors.inbound,
+            'validateUpdate',
+            this.validateUpdateNextHandler.bind(this)
+          );
+          validate(makeInput());
+        }
+        input = makeInput();
+      } catch (error) {
+        this.rejectUpdate(protocolInstanceId, error);
+        return;
       }
-      input = makeInput();
-    } catch (error) {
-      this.rejectUpdate(protocolInstanceId, error);
-      return;
-    }
-    const execute = composeInterceptors(this.interceptors.inbound, 'handleUpdate', this.updateNextHandler.bind(this));
-    this.acceptUpdate(protocolInstanceId);
-    untrackPromise(
-      execute(input)
+      const execute = composeInterceptors(this.interceptors.inbound, 'handleUpdate', this.updateNextHandler.bind(this));
+      this.acceptUpdate(protocolInstanceId);
+      const res = execute(input)
         .then((result) => this.completeUpdate(protocolInstanceId, result))
         .catch((error) => {
           if (error instanceof TemporalFailure) {
             this.rejectUpdate(protocolInstanceId, error);
           } else {
             throw error;
           }
-        })
-    );
+        });
+      untrackPromise(res);
+      return res;
+    };
+    untrackPromise(UpdateScope.updateWithInfo(updateId, name, doUpdateImpl));
   }
 
   protected async updateNextHandler({ name, args }: UpdateInput): Promise<unknown> {

packages/workflow/src/update-scope.ts

@@ -0,0 +1,67 @@
+import type { AsyncLocalStorage as ALS } from 'node:async_hooks';
+
+/**
+ * Option for constructing a UpdateScope
+ */
+export interface UpdateScopeOptions {
+  /**
+   *  A workflow-unique identifier for this update.
+   */
+  id: string;
+
+  /**
+   *  The update type name.
+   */
+  name: string;
+}
+
+// AsyncLocalStorage is injected via vm module into global scope.
+// In case Workflow code is imported in Node.js context, replace with an empty class.
+export const AsyncLocalStorage: new <T>() => ALS<T> = (globalThis as any).AsyncLocalStorage ?? class {};
+
+export class UpdateScope {
+  /**
+   *  A workflow-unique identifier for this update.
+   */
+  public readonly id: string;
+
+  /**
+   *  The update type name.
+   */
+  public readonly name: string;
+
+  constructor(options: UpdateScopeOptions) {
+    this.id = options.id;
+    this.name = options.name;
+  }
+
+  /**
+   * Activate the scope as current and run the update handler `fn`.
+   *
+   * @return the result of `fn`
+   */
+  run<T>(fn: () => Promise<T>): Promise<T> {
+    return storage.run(this, fn);
+  }
+
+  /**
+   * Get the current "active" update scope.
+   */
+  static current(): UpdateScope | undefined {
+    return storage.getStore();
+  }
+
+  /** Alias to `new UpdateScope({ id, name }).run(fn)` */
+  static updateWithInfo<T>(id: string, name: string, fn: () => Promise<T>): Promise<T> {
+    return new this({ id, name }).run(fn);
+  }
+}
+
+const storage = new AsyncLocalStorage<UpdateScope>();
+
+/**
+ * Disable the async local storage for updates.
+ */
+export function disableUpdateStorage(): void {
+  storage.disable();
+}

packages/workflow/src/worker-interface.ts

@@ -8,6 +8,7 @@ import { msToTs, tsToMs } from '@temporalio/common/lib/time';
 import { composeInterceptors } from '@temporalio/common/lib/interceptors';
 import { coresdk } from '@temporalio/proto';
 import { disableStorage } from './cancellation-scope';
+import { disableUpdateStorage } from './update-scope';
 import { DeterminismViolationError } from './errors';
 import { WorkflowInterceptorsFactory } from './interceptors';
 import { WorkflowCreateOptionsInternal } from './interfaces';
@@ -295,6 +296,7 @@ export function shouldUnblockConditions(job: coresdk.workflow_activation.IWorkfl
 export function dispose(): void {
   const dispose = composeInterceptors(getActivator().interceptors.internals, 'dispose', async () => {
     disableStorage();
+    disableUpdateStorage();
   });
   dispose({});
 }

packages/workflow/src/workflow.ts

@@ -24,6 +24,7 @@ import { Duration, msOptionalToTs, msToNumber, msToTs, tsToMs } from '@temporali
 import { composeInterceptors } from '@temporalio/common/lib/interceptors';
 import { temporal } from '@temporalio/proto';
 import { CancellationScope, registerSleepImplementation } from './cancellation-scope';
+import { UpdateScope } from './update-scope';
 import {
   ActivityInput,
   LocalActivityInput,
@@ -44,6 +45,7 @@ import {
   SignalHandlerOptions,
   UpdateHandlerOptions,
   WorkflowInfo,
+  UpdateInfo,
 } from './interfaces';
 import { LocalActivityDoBackoff } from './errors';
 import { assertInWorkflowContext, getActivator, maybeGetActivator } from './global-attributes';
@@ -868,6 +870,19 @@ export function workflowInfo(): WorkflowInfo {
   return activator.info;
 }
 
+/**
+ * Get information about the current update if any.
+ *
+ * @return Info for the current update handler the code calling this is executing
+ * within if any.
+ *
+ * @experimental
+ */
+export function currentUpdateInfo(): UpdateInfo | undefined {
+  assertInWorkflowContext('Workflow.currentUpdateInfo(...) may only be used from a Workflow Execution.');
+  return UpdateScope.current();
+}
+
 /**
  * Returns whether or not code is executing in workflow context
  */