From 08e59e14c71a3289946bc168b6dd6fae3c6a7d3f Mon Sep 17 00:00:00 2001
From: Copilot <copilot@github.com>
Date: Tue, 24 Mar 2026 08:38:47 +0000
Subject: [PATCH] fix: waitForState timeout handler fails to remove stale
 waiter

The timeout handler in InMemoryOrchestrationBackend.waitForState() used
findIndex with `w.resolve === resolve` to locate and remove the timed-out
waiter. However, the waiter's resolve property is a wrapper function that
calls clearTimeout before delegating to the original resolve, so the
identity check always failed and the stale waiter was never removed.

Fix:
- Move waiter declaration before timer so the timeout callback can use
  indexOf(waiter) for correct object-identity lookup.
- Track waitForState timers in pendingTimers so reset() cleans them up.
- Remove timer from pendingTimers on resolve, reject, and timeout.
- Delete the stateWaiters map entry when the last waiter is removed.

Fixes #201

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 .../src/testing/in-memory-backend.ts          | 33 ++++++++----
 .../test/in-memory-backend.spec.ts            | 54 +++++++++++++++++++
 2 files changed, 77 insertions(+), 10 deletions(-)

diff --git a/packages/durabletask-js/src/testing/in-memory-backend.ts b/packages/durabletask-js/src/testing/in-memory-backend.ts
index 2c96acb1..38eac32a 100644
--- a/packages/durabletask-js/src/testing/in-memory-backend.ts
+++ b/packages/durabletask-js/src/testing/in-memory-backend.ts
@@ -337,29 +337,42 @@ export class InMemoryOrchestrationBackend {
     }
 
     return new Promise((resolve, reject) => {
-      const timer = setTimeout(() => {
-        const waiters = this.stateWaiters.get(instanceId);
-        if (waiters) {
-          const index = waiters.findIndex((w) => w.resolve === resolve);
-          if (index >= 0) {
-            waiters.splice(index, 1);
-          }
-        }
-        reject(new Error(`Timeout waiting for orchestration '${instanceId}'`));
-      }, timeoutMs);
+      // Declare timer with let so we can reference it inside the waiter
+      // and reference the waiter inside the timer callback.
+      // eslint-disable-next-line prefer-const
+      let timer: ReturnType<typeof setTimeout>;
 
       const waiter: StateWaiter = {
         resolve: (result) => {
           clearTimeout(timer);
+          this.pendingTimers.delete(timer);
           resolve(result);
         },
         reject: (error) => {
           clearTimeout(timer);
+          this.pendingTimers.delete(timer);
           reject(error);
         },
         predicate,
       };
 
+      timer = setTimeout(() => {
+        this.pendingTimers.delete(timer);
+        const waiters = this.stateWaiters.get(instanceId);
+        if (waiters) {
+          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);
+
       let waiters = this.stateWaiters.get(instanceId);
       if (!waiters) {
         waiters = [];
diff --git a/packages/durabletask-js/test/in-memory-backend.spec.ts b/packages/durabletask-js/test/in-memory-backend.spec.ts
index 66fe6a96..220ce6fc 100644
--- a/packages/durabletask-js/test/in-memory-backend.spec.ts
+++ b/packages/durabletask-js/test/in-memory-backend.spec.ts
@@ -377,4 +377,58 @@ describe("In-Memory Backend", () => {
     expect(state?.runtimeStatus).toEqual(OrchestrationStatus.COMPLETED);
     expect(state?.serializedOutput).toEqual(JSON.stringify(42));
   });
+
+  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
+  });
 });