fix: waitForState timeout handler fails to remove stale waiter (#330)

Author: YunchuWang

packages/durabletask-js/src/testing/in-memory-backend.ts

@@ -429,32 +429,49 @@ export class InMemoryOrchestrationBackend {
     return new Promise((resolve, reject) => {
       // When timeoutMs is 0, no timeout is applied — the waiter will only be
       // resolved by a matching state change or rejected by reset().
+      // `waiter` is declared before the timer so the timeout callback can find it by
+      // object identity; the waiter reads `timer` only when invoked, by which point
+      // it has been assigned.
       let timer: ReturnType<typeof setTimeout> | undefined;
+
+      const waiter: StateWaiter = {
+        resolve: (result) => {
+          if (timer !== undefined) {
+            clearTimeout(timer);
+            this.pendingTimers.delete(timer);
+          }
+          resolve(result);
+        },
+        reject: (error) => {
+          if (timer !== undefined) {
+            clearTimeout(timer);
+            this.pendingTimers.delete(timer);
+          }
+          reject(error);
+        },
+        predicate,
+      };
+
       if (timeoutMs > 0) {
         timer = setTimeout(() => {
+          if (timer !== undefined) {
+            this.pendingTimers.delete(timer);
+          }
           const waiters = this.stateWaiters.get(instanceId);
           if (waiters) {
-            const index = waiters.findIndex((w) => w.resolve === resolve);
+            const index = waiters.indexOf(waiter);
             if (index >= 0) {
               waiters.splice(index, 1);
             }
+            if (waiters.length === 0) {
+              this.stateWaiters.delete(instanceId);
+            }
           }
           reject(new Error(`Timeout waiting for orchestration '${instanceId}'`));
         }, timeoutMs);
+        this.pendingTimers.add(timer);
       }
 
-      const waiter: StateWaiter = {
-        resolve: (result) => {
-          if (timer !== undefined) clearTimeout(timer);
-          resolve(result);
-        },
-        reject: (error) => {
-          if (timer !== undefined) clearTimeout(timer);
-          reject(error);
-        },
-        predicate,
-      };
-
       let waiters = this.stateWaiters.get(instanceId);
       if (!waiters) {
         waiters = [];

packages/durabletask-js/test/in-memory-backend.spec.ts

@@ -933,4 +933,58 @@ describe("In-Memory Backend", () => {
       expect(backend.toClientStatus(suspendedInstance!.status)).toEqual(OrchestrationStatus.SUSPENDED);
     });
   });
+
+  it("should clean up stale waiters after waitForState timeout", async () => {
+    const instanceId = "timeout-cleanup-test";
+    backend.createInstance(instanceId, "testOrch");
+
+    // Call waitForState with a predicate that will never match and a short timeout
+    await expect(
+      backend.waitForState(
+        instanceId,
+        () => false, // Will never match
+        50, // 50ms timeout
+      ),
+    ).rejects.toThrow("Timeout waiting for orchestration");
+
+    // After the timeout, the stale waiter should be cleaned up.
+    // Access internal stateWaiters to verify the waiter was removed.
+    const stateWaitersMap = (backend as any).stateWaiters as Map<string, any[]>;
+
+    // The timed-out waiter should have been removed, and since it was the only
+    // waiter, the instance entry should be removed from the map entirely.
+    expect(stateWaitersMap.has(instanceId)).toBe(false);
+  });
+
+  it("should remove only the timed-out waiter when multiple waiters exist", async () => {
+    const instanceId = "multi-waiter-timeout-test";
+    backend.createInstance(instanceId, "testOrch");
+
+    // Start a waiter with a long timeout (won't time out during the test)
+    const longWaitPromise = backend.waitForState(
+      instanceId,
+      () => false, // Never matches
+      60000, // 60 second timeout — won't fire
+    );
+
+    // Start a waiter with a very short timeout
+    await expect(
+      backend.waitForState(
+        instanceId,
+        () => false, // Never matches
+        50, // 50ms timeout
+      ),
+    ).rejects.toThrow("Timeout waiting for orchestration");
+
+    // After the short timeout, only the long-lived waiter should remain
+    const stateWaitersMap = (backend as any).stateWaiters as Map<string, any[]>;
+    const waiters = stateWaitersMap.get(instanceId);
+
+    expect(waiters).toBeDefined();
+    expect(waiters!.length).toBe(1);
+
+    // Clean up: reset to clear the remaining waiter and its timer
+    backend.reset();
+    await longWaitPromise.catch(() => {}); // Ignore the reset rejection
+  });
 });