fix(sdk): fix race condition when checkpoint completion happens before waitForStatusChange is called

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

@@ -0,0 +1,37 @@
+import { InvocationType } from "@aws-sdk/client-lambda";
+import { handler } from "./wait-for-callback-quick-completion";
+import { createTests } from "../../../utils/test-helper";
+
+createTests({
+  handler,
+  invocationType: InvocationType.Event,
+  localRunnerConfig: {
+    skipTime: false,
+  },
+  tests: (runner, { isCloud }) => {
+    it("should handle waitForCallback when callback completes before ", async () => {
+      const callbackOp = runner.getOperationByIndex(0);
+
+      const executionPromise = runner.run({
+        payload: isCloud
+          ? {
+              submitterDelay: 5,
+            }
+          : {},
+      });
+
+      // only wait for started status
+      await callbackOp.waitForData();
+
+      await callbackOp.sendCallbackSuccess("{}");
+
+      const result = await executionPromise;
+
+      expect(result.getResult()).toEqual({
+        callbackResult: "{}",
+        success: true,
+      });
+      expect(result.getInvocations().length).toBe(1);
+    });
+  },
+});

packages/aws-durable-execution-sdk-js-examples/src/examples/wait-for-callback/quick-completion/wait-for-callback-quick-completion.ts

@@ -0,0 +1,29 @@
+import {
+  DurableContext,
+  withDurableExecution,
+} from "@aws/durable-execution-sdk-js";
+import { ExampleConfig } from "../../../types";
+
+export const config: ExampleConfig = {
+  name: "Wait for Callback - Quick Completion",
+  description:
+    "Demonstrates waitForCallback invocation-level completion scenario",
+};
+
+export const handler = withDurableExecution(
+  async (event: { submitterDelay: number }, context: DurableContext) => {
+    const result = await context.waitForCallback(async () => {
+      if (event.submitterDelay) {
+        await new Promise((resolve) =>
+          setTimeout(resolve, event.submitterDelay * 1000),
+        );
+      }
+      return Promise.resolve();
+    });
+
+    return {
+      callbackResult: result,
+      success: true,
+    };
+  },
+);