fix(sdk): enhance invoke handler to wait for running operations (#88)

ParidelPooya · ParidelPooya · commit 5e8ae84a69ac · 2025-11-01T23:12:49.000-07:00
fix(sdk): enhance invoke handler to wait for running operations before
terminating

- Wrap invoke logic in while loop to allow re-evaluation of step status
- Check for running operations before terminating STARTED invokes
- Use waitBeforeContinue utility when other operations are running
- Update unit tests to cover new waiting behavior and while loop logic
- Add test for scenario where invoke waits for other operations
- Remove outdated intercepted execution test that didn't fit new pattern
diff --git a/packages/aws-durable-execution-sdk-js/src/handlers/invoke-handler/invoke-handler.test.ts b/packages/aws-durable-execution-sdk-js/src/handlers/invoke-handler/invoke-handler.test.ts
@@ -14,13 +14,15 @@ jest.mock("../../utils/termination-helper");
 jest.mock("../../mocks/operation-interceptor");
 jest.mock("../../utils/logger/logger");
 jest.mock("../../errors/serdes-errors/serdes-errors");
+jest.mock("../../utils/wait-before-continue/wait-before-continue");
 
 import { terminate } from "../../utils/termination-helper";
 import { log } from "../../utils/logger/logger";
 import {
   safeSerialize,
   safeDeserialize,
 } from "../../errors/serdes-errors/serdes-errors";
+import { waitBeforeContinue } from "../../utils/wait-before-continue/wait-before-continue";
 
 const mockTerminate = terminate as jest.MockedFunction<typeof terminate>;
 const mockLog = log as jest.MockedFunction<typeof log>;
@@ -30,6 +32,9 @@ const mockSafeSerialize = safeSerialize as jest.MockedFunction<
 const mockSafeDeserialize = safeDeserialize as jest.MockedFunction<
   typeof safeDeserialize
 >;
+const mockWaitBeforeContinue = waitBeforeContinue as jest.MockedFunction<
+  typeof waitBeforeContinue
+>;
 
 describe("InvokeHandler", () => {
   let mockContext: ExecutionContext;
@@ -237,6 +242,7 @@ describe("InvokeHandler", () => {
       });
 
       mockContext.getStepData = mockGetStepData;
+      mockHasRunningOperations.mockReturnValue(false); // No other operations running
 
       const invokeHandler = createInvokeHandler(
         mockContext,
@@ -261,7 +267,49 @@ describe("InvokeHandler", () => {
       );
     });
 
+    it("should wait when operation is in progress and other operations are running", async () => {
+      const mockGetStepData = jest.fn()
+        .mockReturnValueOnce({ Status: OperationStatus.STARTED })
+        .mockReturnValueOnce({ Status: OperationStatus.SUCCEEDED, InvokeDetails: { Result: '{"result":"success"}' } });
+
+      mockContext.getStepData = mockGetStepData;
+      mockHasRunningOperations.mockReturnValue(true); // Other operations running
+      mockWaitBeforeContinue.mockResolvedValue({ reason: "status" });
+      mockSafeDeserialize.mockResolvedValue({ result: "success" });
+
+      const invokeHandler = createInvokeHandler(
+        mockContext,
+        mockCheckpointFn,
+        mockCreateStepId,
+        mockHasRunningOperations,
+      );
+
+      const result = await invokeHandler("test-function", { test: "data" });
+
+      expect(result).toEqual({ result: "success" });
+      expect(mockLog).toHaveBeenCalledWith(
+        true,
+        "⏳",
+        "Invoke test-function still in progress, waiting for other operations",
+      );
+      expect(mockWaitBeforeContinue).toHaveBeenCalledWith({
+        checkHasRunningOperations: true,
+        checkStepStatus: true,
+        checkTimer: false,
+        stepId: "test-step-1",
+        context: mockContext,
+        hasRunningOperations: mockHasRunningOperations,
+      });
+    });
+
     it("should create checkpoint and terminate for new invoke without name", async () => {
+      const mockGetStepData = jest.fn()
+        .mockReturnValueOnce(undefined) // First call - no step data
+        .mockReturnValueOnce({ Status: OperationStatus.STARTED }); // After checkpoint
+
+      mockContext.getStepData = mockGetStepData;
+      mockHasRunningOperations.mockReturnValue(false); // No other operations running
+
       mockOperationInterceptor.execute.mockImplementation(
         async (name: any, fn: any) => {
           return await fn();
@@ -305,11 +353,18 @@ describe("InvokeHandler", () => {
       expect(mockLog).toHaveBeenCalledWith(
         true,
         "🚀",
-        "Invoke test-function started, terminating for async execution",
+        "Invoke test-function started, re-checking status",
       );
     });
 
     it("should create checkpoint and terminate for new invoke with name", async () => {
+      const mockGetStepData = jest.fn()
+        .mockReturnValueOnce(undefined) // First call - no step data
+        .mockReturnValueOnce({ Status: OperationStatus.STARTED }); // After checkpoint
+
+      mockContext.getStepData = mockGetStepData;
+      mockHasRunningOperations.mockReturnValue(false); // No other operations running
+
       mockOperationInterceptor.execute.mockImplementation(
         async (name: any, fn: any) => {
           return await fn();
@@ -342,6 +397,13 @@ describe("InvokeHandler", () => {
     });
 
     it("should handle invoke with options", async () => {
+      const mockGetStepData = jest.fn()
+        .mockReturnValueOnce(undefined) // First call - no step data
+        .mockReturnValueOnce({ Status: OperationStatus.STARTED }); // After checkpoint
+
+      mockContext.getStepData = mockGetStepData;
+      mockHasRunningOperations.mockReturnValue(false); // No other operations running
+
       mockOperationInterceptor.execute.mockImplementation(
         async (name: any, fn: any) => {
           return await fn();
@@ -381,28 +443,5 @@ describe("InvokeHandler", () => {
         },
       });
     });
-
-    it("should handle intercepted execution", async () => {
-      const interceptedResult = { intercepted: true };
-      mockOperationInterceptor.execute.mockResolvedValue(interceptedResult);
-
-      const invokeHandler = createInvokeHandler(
-        mockContext,
-        mockCheckpointFn,
-        mockCreateStepId,
-        mockHasRunningOperations,
-      );
-
-      const result = await invokeHandler("test-function", { test: "data" });
-
-      expect(result).toBe(interceptedResult);
-      expect(OperationInterceptor.forExecution).toHaveBeenCalledWith(
-        "test-arn",
-      );
-      expect(mockOperationInterceptor.execute).toHaveBeenCalledWith(
-        undefined,
-        expect.any(Function),
-      );
-    });
   });
 });
diff --git a/packages/aws-durable-execution-sdk-js/src/handlers/invoke-handler/invoke-handler.ts b/packages/aws-durable-execution-sdk-js/src/handlers/invoke-handler/invoke-handler.ts
@@ -14,12 +14,13 @@ import {
   safeDeserialize,
 } from "../../errors/serdes-errors/serdes-errors";
 import { OperationInterceptor } from "../../mocks/operation-interceptor";
+import { waitBeforeContinue } from "../../utils/wait-before-continue/wait-before-continue";
 
 export const createInvokeHandler = (
   context: ExecutionContext,
   checkpoint: ReturnType<typeof createCheckpoint>,
   createStepId: () => string,
-  _hasRunningOperations: () => boolean,
+  hasRunningOperations: () => boolean,
 ): {
   <I, O>(funcId: string, input: I, config?: InvokeConfig<I, O>): Promise<O>;
   <I, O>(
@@ -58,88 +59,105 @@ export const createInvokeHandler = (
 
     log(context.isVerbose, "🔗", `Invoke ${name || funcId} (${stepId})`);
 
-    // Check if we have existing step data
-    const stepData = context.getStepData(stepId);
+    // Main invoke logic - can be re-executed if step status changes
+    while (true) {
+      // Check if we have existing step data
+      const stepData = context.getStepData(stepId);
 
-    if (stepData?.Status === OperationStatus.SUCCEEDED) {
-      // Return cached result - no need to check for errors in successful operations
-      const invokeDetails = stepData.InvokeDetails;
-      return await safeDeserialize(
-        config?.resultSerdes || defaultSerdes,
-        invokeDetails?.Result,
-        stepId,
-        name,
-        context.terminationManager,
-        context.isVerbose,
-        context.durableExecutionArn,
-      );
-    }
+      if (stepData?.Status === OperationStatus.SUCCEEDED) {
+        // Return cached result - no need to check for errors in successful operations
+        const invokeDetails = stepData.InvokeDetails;
+        return await safeDeserialize(
+          config?.resultSerdes || defaultSerdes,
+          invokeDetails?.Result,
+          stepId,
+          name,
+          context.terminationManager,
+          context.isVerbose,
+          context.durableExecutionArn,
+        );
+      }
 
-    if (stepData?.Status === OperationStatus.FAILED) {
-      // Operation failed, throw error
-      const invokeDetails = stepData.InvokeDetails;
-      const error = new Error(
-        invokeDetails?.Error?.ErrorMessage || "Invoke failed",
-      );
-      error.name = invokeDetails?.Error?.ErrorType || "InvokeError";
-      throw error;
-    }
+      if (stepData?.Status === OperationStatus.FAILED) {
+        // Operation failed, throw error
+        const invokeDetails = stepData.InvokeDetails;
+        const error = new Error(
+          invokeDetails?.Error?.ErrorMessage || "Invoke failed",
+        );
+        error.name = invokeDetails?.Error?.ErrorType || "InvokeError";
+        throw error;
+      }
 
-    if (stepData?.Status === OperationStatus.STARTED) {
-      // Operation is still running, terminate and wait for completion
-      // It's a temporary solution until we implement more sopesticated solution
-      log(
-        context.isVerbose,
-        "⏳",
-        `Invoke ${name || funcId} still in progress, terminating`,
-      );
-      return terminate(context, TerminationReason.OPERATION_TERMINATED, stepId);
-    }
+      if (stepData?.Status === OperationStatus.STARTED) {
+        // Operation is still running, check for other operations before terminating
+        if (hasRunningOperations()) {
+          log(
+            context.isVerbose,
+            "⏳",
+            `Invoke ${name || funcId} still in progress, waiting for other operations`,
+          );
+          await waitBeforeContinue({
+            checkHasRunningOperations: true,
+            checkStepStatus: true,
+            checkTimer: false,
+            stepId,
+            context,
+            hasRunningOperations,
+          });
+          continue; // Re-evaluate status after waiting
+        }
+        
+        // No other operations running, safe to terminate
+        log(
+          context.isVerbose,
+          "⏳",
+          `Invoke ${name || funcId} still in progress, terminating`,
+        );
+        return terminate(context, TerminationReason.OPERATION_TERMINATED, stepId);
+      }
 
-    // Execute with potential interception (testing)
-    const result = await OperationInterceptor.forExecution(
-      context.durableExecutionArn,
-    ).execute(name, async (): Promise<O> => {
-      // Serialize the input payload
-      const serializedPayload = await safeSerialize(
-        config?.payloadSerdes || defaultSerdes,
-        input,
-        stepId,
-        name,
-        context.terminationManager,
-        context.isVerbose,
+      // Execute with potential interception (testing)
+      await OperationInterceptor.forExecution(
         context.durableExecutionArn,
-      );
+      ).execute(name, async (): Promise<void> => {
+        // Serialize the input payload
+        const serializedPayload = await safeSerialize(
+          config?.payloadSerdes || defaultSerdes,
+          input,
+          stepId,
+          name,
+          context.terminationManager,
+          context.isVerbose,
+          context.durableExecutionArn,
+        );
 
-      // Create checkpoint for the invoke operation
-      await checkpoint(stepId, {
-        Id: stepId,
-        ParentId: context.parentId,
-        Action: OperationAction.START,
-        SubType: OperationSubType.INVOKE,
-        Type: OperationType.INVOKE,
-        Name: name,
-        Payload: serializedPayload,
-        InvokeOptions: {
-          FunctionName: funcId,
-          ...(config?.timeoutSeconds && {
-            TimeoutSeconds: config.timeoutSeconds,
-          }),
-        },
-      });
-
-      log(
-        context.isVerbose,
-        "🚀",
-        `Invoke ${name || funcId} started, terminating for async execution`,
-      );
+        // Create checkpoint for the invoke operation
+        await checkpoint(stepId, {
+          Id: stepId,
+          ParentId: context.parentId,
+          Action: OperationAction.START,
+          SubType: OperationSubType.INVOKE,
+          Type: OperationType.INVOKE,
+          Name: name,
+          Payload: serializedPayload,
+          InvokeOptions: {
+            FunctionName: funcId,
+            ...(config?.timeoutSeconds && {
+              TimeoutSeconds: config.timeoutSeconds,
+            }),
+          },
+        });
 
-      // Terminate to allow the invoke to execute asynchronously
-      // It's a temporary solution until we implement more sopesticated solution
-      return terminate(context, TerminationReason.OPERATION_TERMINATED, stepId);
-    });
+        log(
+          context.isVerbose,
+          "🚀",
+          `Invoke ${name || funcId} started, re-checking status`,
+        );
+      });
 
-    return result;
+      // Continue the loop to re-evaluate status (will hit STARTED case)
+      continue;
+    }
   }
 
   return invokeHandler;
