Add support for orchestration tags in TaskHubGrpcClient and related components

Author: torosent

packages/durabletask-js/src/client/client.ts

@@ -22,10 +22,26 @@ import { OrchestrationQuery, ListInstanceIdsOptions, DEFAULT_PAGE_SIZE } from ".
 import { Page, AsyncPageable, createAsyncPageable } from "../orchestration/page";
 import { FailureDetails } from "../task/failure-details";
 import { Logger, ConsoleLogger } from "../types/logger.type";
+import { StartOrchestrationOptions } from "../task/options";
 
 // Re-export MetadataGenerator for backward compatibility
 export { MetadataGenerator } from "../utils/grpc-helper.util";
 
+function mapToRecord(tagsMap?: { forEach: (cb: (value: string, key: string) => void) => void }):
+  | Record<string, string>
+  | undefined {
+  if (!tagsMap) {
+    return;
+  }
+
+  const tags: Record<string, string> = {};
+  tagsMap.forEach((value, key) => {
+    tags[key] = value;
+  });
+
+  return Object.keys(tags).length > 0 ? tags : undefined;
+}
+
 /**
  * Options for creating a TaskHubGrpcClient.
  */
@@ -133,13 +149,46 @@ export class TaskHubGrpcClient {
     input?: TInput,
     instanceId?: string,
     startAt?: Date,
+  ): Promise<string>;
+  /**
+   * Schedules a new orchestrator using the DurableTask client.
+   *
+   * @param {TOrchestrator | string} orchestrator - The orchestrator or the name of the orchestrator to be scheduled.
+   * @param {TInput} input - Optional input for the orchestrator.
+   * @param {StartOrchestrationOptions} options - Options for instance ID, start time, and tags.
+   * @return {Promise<string>} A Promise resolving to the unique ID of the scheduled orchestrator instance.
+   */
+  async scheduleNewOrchestration(
+    orchestrator: TOrchestrator | string,
+    input?: TInput,
+    options?: StartOrchestrationOptions,
+  ): Promise<string>;
+  async scheduleNewOrchestration(
+    orchestrator: TOrchestrator | string,
+    input?: TInput,
+    instanceIdOrOptions?: string | StartOrchestrationOptions,
+    startAt?: Date,
   ): Promise<string> {
     let name;
     if (typeof orchestrator === "string") {
       name = orchestrator;
     } else {
       name = getName(orchestrator);
     }
+
+    const instanceId =
+      typeof instanceIdOrOptions === "string" || instanceIdOrOptions === undefined
+        ? instanceIdOrOptions
+        : instanceIdOrOptions.instanceId;
+    const scheduledStartAt =
+      typeof instanceIdOrOptions === "string" || instanceIdOrOptions === undefined
+        ? startAt
+        : instanceIdOrOptions.startAt;
+    const tags =
+      typeof instanceIdOrOptions === "string" || instanceIdOrOptions === undefined
+        ? undefined
+        : instanceIdOrOptions.tags;
+
     const req = new pb.CreateInstanceRequest();
     req.setName(name);
     req.setInstanceid(instanceId ?? randomUUID());
@@ -148,11 +197,18 @@ export class TaskHubGrpcClient {
     i.setValue(JSON.stringify(input));
 
     const ts = new Timestamp();
-    ts.fromDate(new Date(startAt?.getTime() ?? 0));
+    ts.fromDate(new Date(scheduledStartAt?.getTime() ?? 0));
 
     req.setInput(i);
     req.setScheduledstarttimestamp(ts);
 
+    if (tags) {
+      const tagsMap = req.getTagsMap();
+      for (const [key, value] of Object.entries(tags)) {
+        tagsMap.set(key, value);
+      }
+    }
+
     this._logger.info(`Starting new ${name} instance with ID = ${req.getInstanceid()}`);
 
     const res = await callWithMetadata<pb.CreateInstanceRequest, pb.CreateInstanceResponse>(
@@ -779,6 +835,8 @@ export class TaskHubGrpcClient {
       );
     }
 
+    const tags = mapToRecord(protoState.getTagsMap());
+
     return new OrchestrationState(
       instanceId,
       name ?? "",
@@ -789,6 +847,7 @@ export class TaskHubGrpcClient {
       serializedOutput,
       serializedCustomStatus,
       failureDetails,
+      tags,
     );
   }
 }

packages/durabletask-js/src/index.ts

@@ -27,7 +27,13 @@ export { Task } from "./task/task";
 
 // Retry policies and task options
 export { RetryPolicy, RetryPolicyOptions } from "./task/retry";
-export { TaskOptions, SubOrchestrationOptions, taskOptionsFromRetryPolicy, subOrchestrationOptionsFromRetryPolicy } from "./task/options";
+export {
+	TaskOptions,
+	SubOrchestrationOptions,
+	StartOrchestrationOptions,
+	taskOptionsFromRetryPolicy,
+	subOrchestrationOptionsFromRetryPolicy,
+} from "./task/options";
 
 // Types
 export { TOrchestrator } from "./types/orchestrator.type";

packages/durabletask-js/src/orchestration/index.ts

@@ -6,6 +6,21 @@ import { FailureDetails } from "../task/failure-details";
 import { fromProtobuf } from "./enum/orchestration-status.enum";
 import { OrchestrationState } from "./orchestration-state";
 
+function mapToRecord(tagsMap?: { forEach: (cb: (value: string, key: string) => void) => void }):
+  | Record<string, string>
+  | undefined {
+  if (!tagsMap) {
+    return;
+  }
+
+  const tags: Record<string, string> = {};
+  tagsMap.forEach((value, key) => {
+    tags[key] = value;
+  });
+
+  return Object.keys(tags).length > 0 ? tags : undefined;
+}
+
 export function newOrchestrationState(
   instanceId: string,
   res?: pb.GetInstanceResponse,
@@ -43,6 +58,8 @@ export function newOrchestrationState(
     tsUpdatedParsed = new Date(tsUpdated.getSeconds() * 1000 + tsUpdated.getNanos() / 1000000);
   }
 
+  const tags = mapToRecord(state?.getTagsMap());
+
   return new OrchestrationState(
     instanceId,
     state?.getName() ?? "",
@@ -53,5 +70,6 @@ export function newOrchestrationState(
     state?.getOutput()?.toString(),
     state?.getCustomstatus()?.toString(),
     failureDetails,
+    tags,
   );
 }

packages/durabletask-js/src/orchestration/orchestration-state.ts

@@ -15,6 +15,7 @@ export class OrchestrationState {
   serializedOutput?: string;
   serializedCustomStatus?: string;
   failureDetails?: FailureDetails;
+  tags?: Record<string, string>;
 
   constructor(
     instanceId: string,
@@ -26,6 +27,7 @@ export class OrchestrationState {
     serializedOutput?: string,
     serializedCustomStatus?: string,
     failureDetails?: FailureDetails,
+    tags?: Record<string, string>,
   ) {
     this.instanceId = instanceId;
     this.name = name;
@@ -36,6 +38,7 @@ export class OrchestrationState {
     this.serializedOutput = serializedOutput;
     this.serializedCustomStatus = serializedCustomStatus;
     this.failureDetails = failureDetails;
+    this.tags = tags;
   }
 
   raiseIfFailed(): void {

packages/durabletask-js/src/task/options/index.ts

@@ -4,6 +4,7 @@
 export {
   TaskOptions,
   SubOrchestrationOptions,
+  StartOrchestrationOptions,
   taskOptionsFromRetryPolicy,
   subOrchestrationOptionsFromRetryPolicy,
 } from "./task-options";

packages/durabletask-js/src/task/options/task-options.ts

@@ -12,6 +12,10 @@ export interface TaskOptions {
    * Controls how many times a task is retried and the delay between retries.
    */
   retry?: RetryPolicy;
+  /**
+   * The tags to associate with the task.
+   */
+  tags?: Record<string, string>;
 }
 
 /**
@@ -26,6 +30,26 @@ export interface SubOrchestrationOptions extends TaskOptions {
   instanceId?: string;
 }
 
+/**
+ * Options for scheduling new orchestrations via the client.
+ */
+export interface StartOrchestrationOptions {
+  /**
+   * The unique ID to use for the orchestration instance.
+   * If not specified, a random UUID will be generated.
+   */
+  instanceId?: string;
+  /**
+   * The time when the orchestration should start executing.
+   * If not specified or in the past, it will start immediately.
+   */
+  startAt?: Date;
+  /**
+   * The tags to associate with the orchestration instance.
+   */
+  tags?: Record<string, string>;
+}
+
 /**
  * Creates a TaskOptions instance from a RetryPolicy.
  *

packages/durabletask-js/src/utils/pb-helper.util.ts

@@ -243,6 +243,19 @@ export function getStringValue(val?: string): StringValue | undefined {
   return stringValue;
 }
 
+function setTagsMap(
+  tagsMap: { set: (key: string, value: string) => void },
+  tags?: Record<string, string>,
+): void {
+  if (!tags) {
+    return;
+  }
+
+  for (const [key, value] of Object.entries(tags)) {
+    tagsMap.set(key, value);
+  }
+}
+
 export function newCompleteOrchestrationAction(
   id: number,
   status: pb.OrchestrationStatus,
@@ -277,10 +290,16 @@ export function newCreateTimerAction(id: number, fireAt: Date): pb.OrchestratorA
   return action;
 }
 
-export function newScheduleTaskAction(id: number, name: string, encodedInput?: string): pb.OrchestratorAction {
+export function newScheduleTaskAction(
+  id: number,
+  name: string,
+  encodedInput?: string,
+  tags?: Record<string, string>,
+): pb.OrchestratorAction {
   const scheduleTaskAction = new pb.ScheduleTaskAction();
   scheduleTaskAction.setName(name);
   scheduleTaskAction.setInput(getStringValue(encodedInput));
+  setTagsMap(scheduleTaskAction.getTagsMap(), tags);
 
   const action = new pb.OrchestratorAction();
   action.setId(id);
@@ -300,11 +319,13 @@ export function newCreateSubOrchestrationAction(
   name: string,
   instanceId?: string | null,
   encodedInput?: string,
+  tags?: Record<string, string>,
 ): pb.OrchestratorAction {
   const createSubOrchestrationAction = new pb.CreateSubOrchestrationAction();
   createSubOrchestrationAction.setName(name);
   createSubOrchestrationAction.setInstanceid(instanceId || "");
   createSubOrchestrationAction.setInput(getStringValue(encodedInput));
+  setTagsMap(createSubOrchestrationAction.getTagsMap(), tags);
 
   const action = new pb.OrchestratorAction();
   action.setId(id);

packages/durabletask-js/src/worker/runtime-orchestration-context.ts

@@ -15,6 +15,21 @@ import { TOrchestrator } from "../types/orchestrator.type";
 import { Task } from "../task/task";
 import { StopIterationError } from "./exception/stop-iteration-error";
 
+function mapToRecord(tagsMap?: { forEach: (cb: (value: string, key: string) => void) => void }):
+  | Record<string, string>
+  | undefined {
+  if (!tagsMap) {
+    return;
+  }
+
+  const tags: Record<string, string> = {};
+  tagsMap.forEach((value, key) => {
+    tags[key] = value;
+  });
+
+  return Object.keys(tags).length > 0 ? tags : undefined;
+}
+
 export class RuntimeOrchestrationContext extends OrchestrationContext {
   _generator?: Generator<Task<any>, any, any>;
   _previousTask?: Task<any>;
@@ -262,7 +277,7 @@ export class RuntimeOrchestrationContext extends OrchestrationContext {
     const id = this.nextSequenceNumber();
     const name = typeof activity === "string" ? activity : getName(activity);
     const encodedInput = input ? JSON.stringify(input) : undefined;
-    const action = ph.newScheduleTaskAction(id, name, encodedInput);
+    const action = ph.newScheduleTaskAction(id, name, encodedInput, options?.tags);
     this._pendingActions[action.getId()] = action;
 
     // If a retry policy is provided, create a RetryableTask
@@ -306,7 +321,7 @@ export class RuntimeOrchestrationContext extends OrchestrationContext {
     }
 
     const encodedInput = input ? JSON.stringify(input) : undefined;
-    const action = ph.newCreateSubOrchestrationAction(id, name, instanceId, encodedInput);
+    const action = ph.newCreateSubOrchestrationAction(id, name, instanceId, encodedInput, options?.tags);
     this._pendingActions[action.getId()] = action;
 
     // If a retry policy is provided, create a RetryableTask
@@ -520,7 +535,8 @@ export class RuntimeOrchestrationContext extends OrchestrationContext {
       }
       const name = scheduleTask.getName();
       const input = scheduleTask.getInput()?.getValue();
-      newAction = ph.newScheduleTaskAction(newId, name, input);
+      const tags = mapToRecord(scheduleTask.getTagsMap());
+      newAction = ph.newScheduleTaskAction(newId, name, input, tags);
     } else {
       // Reschedule a sub-orchestration task
       const subOrch = originalAction.getCreatesuborchestration();
@@ -530,7 +546,8 @@ export class RuntimeOrchestrationContext extends OrchestrationContext {
       const name = subOrch.getName();
       const instanceId = subOrch.getInstanceid();
       const input = subOrch.getInput()?.getValue();
-      newAction = ph.newCreateSubOrchestrationAction(newId, name, instanceId, input);
+      const tags = mapToRecord(subOrch.getTagsMap());
+      newAction = ph.newCreateSubOrchestrationAction(newId, name, instanceId, input, tags);
     }
 
     // Register the new action

packages/durabletask-js/test/orchestration_executor.spec.ts

@@ -149,6 +149,25 @@ describe("Orchestration Executor", () => {
     expect(result.actions[0]?.getId()).toEqual(1);
     expect(result.actions[0]?.getScheduletask()?.getName()).toEqual("dummyActivity");
   });
+
+  it("should include tags on scheduled activity actions", async () => {
+    const dummyActivity = async (_: ActivityContext) => {
+      // do nothing
+    };
+    const orchestrator: TOrchestrator = async function* (ctx: OrchestrationContext, _: any) {
+      yield ctx.callActivity(dummyActivity, undefined, { tags: { env: "test", owner: "durable" } });
+      return "done";
+    };
+    const registry = new Registry();
+    const name = registry.addOrchestrator(orchestrator);
+    const newEvents = [newExecutionStartedEvent(name, TEST_INSTANCE_ID, undefined)];
+    const executor = new OrchestrationExecutor(registry, testLogger);
+    const result = await executor.execute(TEST_INSTANCE_ID, [], newEvents);
+    const scheduleTask = result.actions[0]?.getScheduletask();
+
+    expect(scheduleTask?.getTagsMap().get("env")).toEqual("test");
+    expect(scheduleTask?.getTagsMap().get("owner")).toEqual("durable");
+  });
   it("should test the successful completion of an activity task", async () => {
     const dummyActivity = async (_: ActivityContext) => {
       // do nothing
@@ -348,6 +367,26 @@ describe("Orchestration Executor", () => {
     expect(completeAction?.getOrchestrationstatus()).toEqual(pb.OrchestrationStatus.ORCHESTRATION_STATUS_COMPLETED);
     expect(completeAction?.getResult()?.getValue()).toEqual("42");
   });
+
+  it("should include tags on scheduled sub-orchestration actions", async () => {
+    const subOrchestrator = async (_: OrchestrationContext) => {
+      // do nothing
+    };
+    const orchestrator: TOrchestrator = async function* (ctx: OrchestrationContext, _: any): any {
+      yield ctx.callSubOrchestrator(subOrchestrator, undefined, { tags: { env: "test" } });
+      return "done";
+    };
+    const registry = new Registry();
+    const subOrchestratorName = registry.addOrchestrator(subOrchestrator);
+    const orchestratorName = registry.addOrchestrator(orchestrator);
+    const newEvents = [newExecutionStartedEvent(orchestratorName, TEST_INSTANCE_ID, undefined)];
+    const executor = new OrchestrationExecutor(registry, testLogger);
+    const result = await executor.execute(TEST_INSTANCE_ID, [], newEvents);
+    const createSubOrch = result.actions[0]?.getCreatesuborchestration();
+
+    expect(createSubOrch?.getName()).toEqual(subOrchestratorName);
+    expect(createSubOrch?.getTagsMap().get("env")).toEqual("test");
+  });
   it("should test that a sub-orchestration task is completed when the sub-orchestration fails", async () => {
     const subOrchestrator = async (_: OrchestrationContext) => {
       // do nothing