feat: add entity registration methods to Azure managed worker builder

packages/durabletask-js-azuremanaged/src/worker-builder.ts

@@ -10,6 +10,7 @@ import {
   TActivity,
   TInput,
   TOutput,
+  EntityFactory,
   Logger,
   ConsoleLogger,
   VersioningOptions,
@@ -25,6 +26,7 @@ export class DurableTaskAzureManagedWorkerBuilder {
   private _grpcChannelOptions: grpc.ChannelOptions = {};
   private _orchestrators: { name?: string; fn: TOrchestrator }[] = [];
   private _activities: { name?: string; fn: TActivity<TInput, TOutput> }[] = [];
+  private _entities: { name?: string; factory: EntityFactory }[] = [];
   private _logger: Logger = new ConsoleLogger();
   private _shutdownTimeoutMs?: number;
   private _versioning?: VersioningOptions;
@@ -185,6 +187,30 @@ export class DurableTaskAzureManagedWorkerBuilder {
     return this;
   }
 
+  /**
+   * Registers an entity factory with the worker.
+   * The entity name is derived from the factory function name.
+   *
+   * @param factory The entity factory function.
+   * @returns This builder instance.
+   */
+  addEntity(factory: EntityFactory): DurableTaskAzureManagedWorkerBuilder {
+    this._entities.push({ factory });
+    return this;
+  }
+
+  /**
+   * Registers a named entity factory with the worker.
+   *
+   * @param name The name of the entity.
+   * @param factory The entity factory function.
+   * @returns This builder instance.
+   */
+  addNamedEntity(name: string, factory: EntityFactory): DurableTaskAzureManagedWorkerBuilder {
+    this._entities.push({ name, factory });
+    return this;
+  }
+
   /**
    * Sets the logger to use for logging.
    * Defaults to ConsoleLogger.
@@ -289,6 +315,15 @@ export class DurableTaskAzureManagedWorkerBuilder {
       }
     }
 
+    // Register all entities
+    for (const { name, factory } of this._entities) {
+      if (name) {
+        worker.addNamedEntity(name, factory);
+      } else {
+        worker.addEntity(factory);
+      }
+    }
+
     return worker;
   }
 }

packages/durabletask-js-azuremanaged/test/unit/worker-builder.spec.ts

@@ -0,0 +1,93 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import { DurableTaskAzureManagedWorkerBuilder, createAzureManagedWorkerBuilder } from "../../src/worker-builder";
+import { TaskEntity, ITaskEntity, TaskEntityOperation } from "@microsoft/durabletask-js";
+
+// Simple test entity for registration testing
+class CounterEntity extends TaskEntity<number> {
+  add(operation: TaskEntityOperation): void {
+    const amount = operation.getInput<number>() ?? 1;
+    this.state = (this.state ?? 0) + amount;
+  }
+}
+
+// Factory functions for testing
+function createCounterEntity(): ITaskEntity {
+  return new CounterEntity();
+}
+
+describe("DurableTaskAzureManagedWorkerBuilder", () => {
+  const ENDPOINT = "localhost:8080";
+  const TASKHUB = "test";
+
+  describe("addEntity", () => {
+    it("should register an entity factory and return the builder for chaining", () => {
+      const builder = new DurableTaskAzureManagedWorkerBuilder();
+
+      const result = builder.endpoint(ENDPOINT, TASKHUB, null).addEntity(createCounterEntity);
+
+      expect(result).toBe(builder);
+    });
+
+    it("should register an entity that gets added to the worker on build", () => {
+      const builder = new DurableTaskAzureManagedWorkerBuilder();
+
+      const worker = builder.endpoint(ENDPOINT, TASKHUB, null).addEntity(createCounterEntity).build();
+
+      // The worker should have the entity registered. We verify by checking that
+      // attempting to register it again with the same name throws a duplicate error.
+      expect(() => worker.addNamedEntity("createCounterEntity", createCounterEntity)).toThrow();
+    });
+  });
+
+  describe("addNamedEntity", () => {
+    it("should register a named entity factory and return the builder for chaining", () => {
+      const builder = new DurableTaskAzureManagedWorkerBuilder();
+
+      const result = builder.endpoint(ENDPOINT, TASKHUB, null).addNamedEntity("MyCounter", createCounterEntity);
+
+      expect(result).toBe(builder);
+    });
+
+    it("should register a named entity that gets added to the worker on build", () => {
+      const builder = new DurableTaskAzureManagedWorkerBuilder();
+
+      const worker = builder.endpoint(ENDPOINT, TASKHUB, null).addNamedEntity("MyCounter", createCounterEntity).build();
+
+      // Entity names are lowercased. Registering the same lowercased name should throw.
+      expect(() => worker.addNamedEntity("mycounter", createCounterEntity)).toThrow();
+    });
+  });
+
+  describe("fluent chaining with entities", () => {
+    it("should support registering orchestrators, activities, and entities together", () => {
+      const orchestrator = async function* testOrchestrator() {
+        yield;
+        return "done";
+      };
+      const activity = async () => "result";
+
+      const builder = new DurableTaskAzureManagedWorkerBuilder();
+
+      const worker = builder
+        .endpoint(ENDPOINT, TASKHUB, null)
+        .addNamedOrchestrator("myOrch", orchestrator)
+        .addNamedActivity("myActivity", activity)
+        .addNamedEntity("MyCounter", createCounterEntity)
+        .build();
+
+      expect(worker).toBeDefined();
+    });
+  });
+
+  describe("createAzureManagedWorkerBuilder", () => {
+    it("should create a builder that supports entity registration", () => {
+      const builder = createAzureManagedWorkerBuilder(ENDPOINT, TASKHUB, null);
+
+      const result = builder.addNamedEntity("MyCounter", createCounterEntity);
+
+      expect(result).toBe(builder);
+    });
+  });
+});