Refactor entity locking and orchestration execution logic; improve test coverage for entity operations

Author: YunchuWang

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

@@ -1031,7 +1031,7 @@ export class TaskHubGrpcClient {
   getEntities<T = unknown>(query?: EntityQuery): AsyncPageable<EntityMetadata<T>> {
     const includeState = query?.includeState ?? true;
 
-    return AsyncPageable.create(async (continuationToken?: string): Promise<Page<EntityMetadata<T>>> => {
+    return createAsyncPageable(async (continuationToken?: string): Promise<Page<EntityMetadata<T>>> => {
       const req = new pb.QueryEntitiesRequest();
       const protoQuery = new pb.EntityQuery();
 
@@ -1081,10 +1081,7 @@ export class TaskHubGrpcClient {
       const entities = res.getEntitiesList();
       const values = entities.map((protoMetadata) => this.convertEntityMetadata<T>(protoMetadata, includeState));
 
-      return {
-        values,
-        continuationToken: res.getContinuationtoken()?.getValue(),
-      };
+      return new Page(values, res.getContinuationtoken()?.getValue());
     });
   }
 
@@ -1152,7 +1149,7 @@ export class TaskHubGrpcClient {
       return createEntityMetadataWithoutState(entityId, lastModifiedTime, backlogQueueSize, lockedBy) as EntityMetadata<T>;
     }
   }
-  
+
   /**
    * Helper method to create an OrchestrationState from a protobuf OrchestrationState.
    */

packages/durabletask-js/src/index.ts

@@ -105,7 +105,7 @@ export {
 // Entity types - Worker-side operation types (Step 3)
 export { SignalEntityOptions, CallEntityOptions } from "./entities/signal-entity-options";
 export { TaskEntityState } from "./entities/task-entity-state";
-export { TaskEntityContext, StartOrchestrationOptions } from "./entities/task-entity-context";
+export { TaskEntityContext } from "./entities/task-entity-context";
 export { TaskEntityOperation } from "./entities/task-entity-operation";
 
 // Entity interface and base class (Step 4)

packages/durabletask-js/src/worker/orchestration-executor.ts

@@ -583,6 +583,26 @@ export class OrchestrationExecutor {
             }
           }
           break;
+        // This history event confirms that the lock request was successfully scheduled.
+        // Remove the action from the pending action list so we don't schedule it again.
+        // The pending lock request in _entityFeature.pendingLockRequests remains to receive the granted event.
+        case pb.HistoryEvent.EventtypeCase.ENTITYLOCKREQUESTED:
+          {
+            const eventId = event.getEventid();
+            const action = ctx._pendingActions[eventId];
+            delete ctx._pendingActions[eventId];
+
+            const isSendEntityMessageAction = action?.hasSendentitymessage();
+
+            if (!action) {
+              throw getNonDeterminismError(eventId, "lockEntities");
+            } else if (!isSendEntityMessageAction) {
+              throw getWrongActionTypeError(eventId, "lockEntities", action);
+            } else if (!action.getSendentitymessage()?.hasEntitylockrequested()) {
+              throw getWrongActionTypeError(eventId, "lockEntities (EntityLockRequested)", action);
+            }
+          }
+          break;
         case pb.HistoryEvent.EventtypeCase.ENTITYOPERATIONCOMPLETED:
           {
             const completedEvent = event.getEntityoperationcompleted();

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

@@ -37,7 +37,7 @@ export class RuntimeOrchestrationContext extends OrchestrationContext {
   _newGuidCounter: number;
   _currentUtcDatetime: Date;
   _instanceId: string;
-  _executionId: string;
+  _executionId: string = "";
   _version: string;
   _parent?: ParentOrchestrationInstance;
   _completionStatus?: pb.OrchestrationStatus;
@@ -69,7 +69,7 @@ export class RuntimeOrchestrationContext extends OrchestrationContext {
     this._newInput = undefined;
     this._saveEvents = false;
     this._customStatus = undefined;
-    this._entityFeature = new RuntimeOrchestrationEntityFeature();
+    this._entityFeature = new RuntimeOrchestrationEntityFeature(this);
   }
 
   get instanceId(): string {