fix(testing-sdk): support PENDING status reinvocation and execution failure (#397)

*Issue #, if available:*

#384

*Description of changes:*

When the client (language SDK) has stale data, it should be reinvoked if
it exited without knowing the latest state of the checkpoint API. This
is how the actual service works, and so it should be implemented in the
testing library as well. This also fixes an issue where the language SDK
exits at the same time that a callback completes, and hangs
forever/throws an error. It will now re-invoke in this case.

Additionally, the testing SDK will now throw an error when the language
SDK returns PENDING with no outstanding pending operations, timers, or
stale data. This will fix the issue where it would hang forever when
this happened.


By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice.

Author: anthonyting

packages/aws-durable-execution-sdk-js-examples/src/examples/wait-for-callback/multiple-invocations/wait-for-callback-multiple-invocations.test.ts

@@ -8,6 +8,9 @@ import { createTests } from "../../../utils/test-helper";
 createTests({
   handler,
   invocationType: InvocationType.Event,
+  localRunnerConfig: {
+    skipTime: false,
+  },
   tests: (runner) => {
     it("should handle multiple invocations tracking with waitForCallback operations", async () => {
       // Get operations for verification

packages/aws-durable-execution-sdk-js-testing/src/checkpoint-server/handlers/__tests__/execution-handlers.test.ts

@@ -141,7 +141,10 @@ describe("execution handlers", () => {
           Error: { Payload: undefined },
         },
       };
-      completeInvocationSpy.mockReturnValue(mockEvent);
+      completeInvocationSpy.mockReturnValue({
+        event: mockEvent,
+        hasDirtyOperations: false,
+      });
 
       const errorObject: ErrorObject = {
         ErrorType: "TestError",
@@ -160,7 +163,10 @@ describe("execution handlers", () => {
         invocationId,
         errorObject,
       );
-      expect(result).toEqual(mockEvent);
+      expect(result).toEqual({
+        event: mockEvent,
+        hasDirtyOperations: false,
+      });
     });
   });
 });

packages/aws-durable-execution-sdk-js-testing/src/checkpoint-server/handlers/execution-handlers.ts

@@ -1,4 +1,4 @@
-import { ErrorObject, Event } from "@aws-sdk/client-lambda";
+import { ErrorObject } from "@aws-sdk/client-lambda";
 import {
   createExecutionId,
   ExecutionId,
@@ -12,6 +12,7 @@ import {
   StartDurableExecutionRequest,
   StartInvocationRequest,
 } from "../worker-api/worker-api-request";
+import { CompleteInvocationResponse } from "../worker-api/worker-api-response";
 
 /**
  * Starts a durable execution. Returns the data needed for the handler invocation event.
@@ -42,6 +43,6 @@ export function processCompleteInvocation(
   invocationId: InvocationId,
   error: ErrorObject | undefined,
   executionManager: ExecutionManager,
-): Event {
+): CompleteInvocationResponse {
   return executionManager.completeInvocation(executionId, invocationId, error);
 }

packages/aws-durable-execution-sdk-js-testing/src/checkpoint-server/storage/__tests__/checkpoint-manager-dirty.test.ts

@@ -36,6 +36,7 @@ describe("CheckpointManager dirty operation tracking", () => {
         const dirtyOperations = storage.getDirtyOperations();
 
         expect(dirtyOperations).toEqual([]);
+        expect(storage.hasDirtyOperations()).toBe(false);
       });
 
       it("should return dirty operations and clear the dirty set", () => {
@@ -51,8 +52,11 @@ describe("CheckpointManager dirty operation tracking", () => {
         expect(firstCall).toHaveLength(1);
         expect(firstCall[0].Id).toBe("dirty-step");
 
+        expect(storage.hasDirtyOperations()).toBe(false);
+
         const secondCall = storage.getDirtyOperations();
         expect(secondCall).toEqual([]);
+        expect(storage.hasDirtyOperations()).toBe(false);
       });
 
       it("should return multiple dirty operations in correct order", () => {
@@ -86,6 +90,7 @@ describe("CheckpointManager dirty operation tracking", () => {
           "second-op",
           "third-op",
         ]);
+        expect(storage.hasDirtyOperations()).toBe(false);
       });
     });
 
@@ -100,16 +105,20 @@ describe("CheckpointManager dirty operation tracking", () => {
         storage.registerUpdate(stepUpdate);
 
         // Verify operation is dirty before invocation
+        expect(storage.hasDirtyOperations()).toBe(true);
         const dirtyBefore = storage.getDirtyOperations();
         expect(dirtyBefore).toHaveLength(1);
+        expect(storage.hasDirtyOperations()).toBe(false);
 
         // Start new invocation should clear dirty set
         const invocationId = createInvocationId("test-invocation");
         storage.startInvocation(invocationId);
+        expect(storage.hasDirtyOperations()).toBe(false);
 
         // Verify dirty set is cleared after invocation
         const dirtyAfter = storage.getDirtyOperations();
         expect(dirtyAfter).toEqual([]);
+        expect(storage.hasDirtyOperations()).toBe(false);
       });
 
       it("should return all operations from startInvocation", () => {

packages/aws-durable-execution-sdk-js-testing/src/checkpoint-server/storage/__tests__/execution-manager.test.ts

@@ -4,6 +4,7 @@ import {
   OperationType,
   EventType,
   ErrorObject,
+  OperationAction,
 } from "@aws-sdk/client-lambda";
 import { ExecutionManager, StartExecutionParams } from "../execution-manager";
 import { CheckpointManager } from "../checkpoint-manager";
@@ -520,7 +521,10 @@ describe("execution-manager", () => {
         );
 
         // Verify the result
-        expect(result).toBe(mockHistoryEvent);
+        expect(result).toEqual({
+          event: mockHistoryEvent,
+          hasDirtyOperations: false,
+        });
 
         // Verify CheckpointManager.completeInvocation was called
         expect(completeInvocationSpy).toHaveBeenCalledWith(invocationId);
@@ -593,7 +597,10 @@ describe("execution-manager", () => {
         );
 
         // Verify the result
-        expect(result).toBe(mockHistoryEvent);
+        expect(result).toEqual({
+          event: mockHistoryEvent,
+          hasDirtyOperations: false,
+        });
 
         // Verify CheckpointManager.completeInvocation was called
         expect(completeInvocationSpy).toHaveBeenCalledWith(invocationId);
@@ -642,6 +649,91 @@ describe("execution-manager", () => {
         // Verify CheckpointManager.completeInvocation was called
         expect(completeInvocationSpy).toHaveBeenCalledWith(invocationId);
       });
+
+      it("should return hasDirtyOperations as true when checkpoint storage has dirty operations", () => {
+        // Create an execution first
+        const executionId = createExecutionId("test-execution-id");
+        const invocationId = createInvocationId("test-invocation-id");
+        executionManager.startExecution({ executionId, invocationId });
+
+        const storage = executionManager.getCheckpointsByExecution(executionId);
+        expect(storage).toBeDefined();
+
+        // Register an operation to make the storage dirty
+        const stepUpdate = {
+          Id: "test-step-operation",
+          Type: OperationType.STEP,
+          Action: OperationAction.START,
+        };
+        storage!.registerUpdate(stepUpdate);
+
+        // Mock the completeInvocation method on CheckpointManager
+        const mockTimestamps = {
+          startTimestamp: new Date("2023-01-01T00:00:00.000Z"),
+          endTimestamp: new Date("2023-01-01T00:01:00.000Z"),
+        };
+
+        const completeInvocationSpy = jest
+          .spyOn(storage!, "completeInvocation")
+          .mockReturnValue(mockTimestamps);
+
+        // Mock the eventProcessor.createHistoryEvent method
+        const mockHistoryEvent = {
+          EventType: EventType.InvocationCompleted,
+          Timestamp: new Date(),
+          InvocationCompletedDetails: {
+            StartTimestamp: mockTimestamps.startTimestamp,
+            EndTimestamp: mockTimestamps.endTimestamp,
+            Error: {
+              Payload: undefined,
+            },
+            RequestId: invocationId,
+          },
+        };
+
+        const createHistoryEventSpy = jest
+          .spyOn(storage!.eventProcessor, "createHistoryEvent")
+          .mockReturnValue(mockHistoryEvent);
+
+        // Mock hasDirtyOperations to return true
+        const hasDirtyOperationsSpy = jest
+          .spyOn(storage!, "hasDirtyOperations")
+          .mockReturnValue(true);
+
+        // Call the method
+        const result = executionManager.completeInvocation(
+          executionId,
+          invocationId,
+          undefined,
+        );
+
+        // Verify the result shows dirty operations exist
+        expect(result).toEqual({
+          event: mockHistoryEvent,
+          hasDirtyOperations: true,
+        });
+
+        // Verify CheckpointManager.completeInvocation was called
+        expect(completeInvocationSpy).toHaveBeenCalledWith(invocationId);
+
+        // Verify hasDirtyOperations was called
+        expect(hasDirtyOperationsSpy).toHaveBeenCalled();
+
+        // Verify eventProcessor.createHistoryEvent was called with correct parameters
+        expect(createHistoryEventSpy).toHaveBeenCalledWith(
+          EventType.InvocationCompleted,
+          undefined,
+          "InvocationCompletedDetails",
+          {
+            StartTimestamp: mockTimestamps.startTimestamp,
+            EndTimestamp: mockTimestamps.endTimestamp,
+            Error: {
+              Payload: undefined,
+            },
+            RequestId: invocationId,
+          },
+        );
+      });
     });
   });
 });

packages/aws-durable-execution-sdk-js-testing/src/checkpoint-server/storage/checkpoint-manager.ts

@@ -50,6 +50,10 @@ export class CheckpointManager {
     return this.operationDataMap;
   }
 
+  hasDirtyOperations(): boolean {
+    return !!this.dirtyOperationIds.size;
+  }
+
   getDirtyOperations(): Operation[] {
     const dirtyOperations = Array.from(this.dirtyOperationIds).map((id) => {
       const operationEvents = this.operationDataMap.get(id);

packages/aws-durable-execution-sdk-js-testing/src/checkpoint-server/storage/execution-manager.ts

@@ -9,11 +9,12 @@ import { CallbackId, CheckpointToken } from "../utils/tagged-strings";
 import { ExecutionId, InvocationId } from "../utils/tagged-strings";
 import { decodeCallbackId } from "../utils/callback-id";
 import { OperationEvents } from "../../test-runner/common/operations/operation-with-data";
-import { ErrorObject, EventType, Event } from "@aws-sdk/client-lambda";
+import { ErrorObject, EventType } from "@aws-sdk/client-lambda";
 import {
   StartDurableExecutionRequest,
   StartInvocationRequest,
 } from "../worker-api/worker-api-request";
+import { CompleteInvocationResponse } from "../worker-api/worker-api-response";
 
 export interface InvocationResult {
   checkpointToken: CheckpointToken;
@@ -105,7 +106,7 @@ export class ExecutionManager {
     executionId: ExecutionId,
     invocationId: InvocationId,
     error: ErrorObject | undefined,
-  ): Event {
+  ): CompleteInvocationResponse {
     const checkpointStorage = this.executions.get(executionId);
 
     if (!checkpointStorage) {
@@ -115,7 +116,7 @@ export class ExecutionManager {
     const { startTimestamp, endTimestamp } =
       checkpointStorage.completeInvocation(invocationId);
 
-    return checkpointStorage.eventProcessor.createHistoryEvent(
+    const event = checkpointStorage.eventProcessor.createHistoryEvent(
       EventType.InvocationCompleted,
       undefined,
       "InvocationCompletedDetails",
@@ -128,6 +129,11 @@ export class ExecutionManager {
         RequestId: invocationId,
       },
     );
+
+    return {
+      event,
+      hasDirtyOperations: checkpointStorage.hasDirtyOperations(),
+    };
   }
 
   /**

packages/aws-durable-execution-sdk-js-testing/src/checkpoint-server/worker-api/worker-api-response.ts

@@ -16,10 +16,15 @@ export interface PollCheckpointDataResponse {
   operations: CheckpointOperation[];
 }
 
+export interface CompleteInvocationResponse {
+  event: Event;
+  hasDirtyOperations: boolean;
+}
+
 export interface WorkerApiResponseMapping {
   [ApiType.StartDurableExecution]: InvocationResult;
   [ApiType.StartInvocation]: InvocationResult;
-  [ApiType.CompleteInvocation]: Event;
+  [ApiType.CompleteInvocation]: CompleteInvocationResponse;
   [ApiType.UpdateCheckpointData]: Record<string, never>;
   [ApiType.PollCheckpointData]: Promise<PollCheckpointDataResponse>;
   [ApiType.GetDurableExecutionState]: GetDurableExecutionStateResponse;

packages/aws-durable-execution-sdk-js-testing/src/checkpoint-server/worker/__tests__/checkpoint-worker.test.ts

@@ -306,7 +306,10 @@ describe("CheckpointWorker", () => {
         Timestamp: new Date(),
         Type: "InvocationComplete",
       };
-      mockApiHandlerInstance.performApiCall.mockReturnValue(mockEvent);
+      mockApiHandlerInstance.performApiCall.mockReturnValue({
+        event: mockEvent,
+        hasDirtyOperations: false,
+      });
 
       messageHandler(command);
 
@@ -318,7 +321,10 @@ describe("CheckpointWorker", () => {
         data: {
           type: ApiType.CompleteInvocation,
           requestId: "complete-request-456",
-          response: mockEvent,
+          response: {
+            event: mockEvent,
+            hasDirtyOperations: false,
+          },
         },
       });
     });

packages/aws-durable-execution-sdk-js-testing/src/test-runner/local/__tests__/test-execution-orchestrator-invocation-order.test.ts

@@ -25,6 +25,7 @@ import { ILocalDurableTestRunnerFactory } from "../interfaces/durable-test-runne
 import { DurableApiClient } from "../../common/create-durable-api-client";
 import { CheckpointApiClient } from "../api-client/checkpoint-api-client";
 import { InvocationResult } from "../../../checkpoint-server/storage/execution-manager";
+import { CompleteInvocationResponse } from "../../../checkpoint-server/worker-api/worker-api-response";
 
 // Mock dependencies
 jest.mock("../operations/local-operation-storage");
@@ -75,15 +76,20 @@ describe("TestExecutionOrchestrator - Invocation History Ordering", () => {
   // Tracking arrays for call order verification
   let callOrder: string[];
   let addHistoryEventSpy: jest.SpyInstance;
-  let completeInvocationSpy: jest.SpyInstance;
-
-  const mockInvocationCompletedEvent: Event = {
-    EventType: EventType.InvocationCompleted,
-    InvocationCompletedDetails: {
-      RequestId: "invocation-request-id",
-      StartTimestamp: new Date(),
-      EndTimestamp: new Date(),
+  let completeInvocationSpy: jest.SpyInstance<
+    Promise<CompleteInvocationResponse>
+  >;
+
+  const mockInvocationCompletedEvent: CompleteInvocationResponse = {
+    event: {
+      EventType: EventType.InvocationCompleted,
+      InvocationCompletedDetails: {
+        RequestId: "invocation-request-id",
+        StartTimestamp: new Date(),
+        EndTimestamp: new Date(),
+      },
     },
+    hasDirtyOperations: false,
   };
 
   const nonResolvingPromise = new Promise<never>(() => {
@@ -135,7 +141,10 @@ describe("TestExecutionOrchestrator - Invocation History Ordering", () => {
       }),
       completeInvocation: jest.fn().mockImplementation(() => {
         callOrder.push("completeInvocation");
-        return Promise.resolve(mockInvocationCompletedEvent);
+        return Promise.resolve({
+          event: mockInvocationCompletedEvent,
+          hasDirtyOperations: false,
+        });
       }),
     };
 
@@ -509,8 +518,11 @@ describe("TestExecutionOrchestrator - Invocation History Ordering", () => {
         invocationCount++;
         callOrder.push(`completeInvocation${invocationCount}`);
         return Promise.resolve({
-          ...mockInvocationCompletedEvent,
-          EventId: invocationCount,
+          event: {
+            ...mockInvocationCompletedEvent,
+            EventId: invocationCount,
+          },
+          hasDirtyOperations: false,
         });
       });