From 4fca8e6c5491ef338760df7ac617999e483dfc82 Mon Sep 17 00:00:00 2001
From: Cristian <seeve.c@gmail.com>
Date: Mon, 7 Sep 2026 02:08:14 +0000
Subject: [PATCH] fix: release waitFor listeners on cancellation

---
 .changeset/wait-for-cleanup.md |  5 +++
 AGENTS.md                      |  2 ++
 src/actor.ts                   | 23 +++++++------
 test/wait-for-cleanup.test.ts  | 59 ++++++++++++++++++++++++++++++++++
 4 files changed, 77 insertions(+), 12 deletions(-)
 create mode 100644 .changeset/wait-for-cleanup.md
 create mode 100644 test/wait-for-cleanup.test.ts

diff --git a/.changeset/wait-for-cleanup.md b/.changeset/wait-for-cleanup.md
new file mode 100644
index 0000000..0722fba
--- /dev/null
+++ b/.changeset/wait-for-cleanup.md
@@ -0,0 +1,5 @@
+---
+"effect-machine": patch
+---
+
+Release waitFor listeners when the caller cancels or the subscription recheck fails. This also applies to timed waits, awaitFinal, and sendAndWait.
diff --git a/AGENTS.md b/AGENTS.md
index 6947806..0534093 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -279,6 +279,8 @@ const count = yield* actor.ask(Event.GetCount);  // number
 
 ## Gotchas
 
+- `waitFor` owns its listener with `Effect.acquireUseRelease`. Keep the post-subscription recheck inside that lifetime. Cancellation and predicate defects must remove the listener.
+
 - `actor.stop` cancels pending startup and waits for recovery cleanup. Stop callers can cancel their own wait without cancelling shutdown. A stop from recovery marks that startup interrupted. Shutdown waits for its protected regions and finalizers to finish. Recovery cleanup defects reach every stop caller.
 - Protect both shutdown owner creation and its cache publication from interruption. Protecting only the cached body can cache cancellation when its protected region ends.
 - Recovery self-stop must also identify the supervisor fiber. A protected supervisor must not join an owner that waits for that same supervisor. Preserve supervised recovery cleanup defects on stop.
diff --git a/src/actor.ts b/src/actor.ts
index 20f0a79..b8a74f7 100644
--- a/src/actor.ts
+++ b/src/actor.ts
@@ -555,18 +555,17 @@ const buildActorRefCore = <
       }
     };
     // @effect-diagnostics runEffectInsideEffect:on
-    listeners.add(listener);
-
-    // Re-check after subscribing to close the race window
-    const afterSubscribe = yield* SubscriptionRef.get(stateRef);
-    if (predicate(afterSubscribe)) {
-      listeners.delete(listener);
-      return afterSubscribe;
-    }
-
-    const result = yield* Deferred.await(done);
-    listeners.delete(listener);
-    return result;
+    return yield* Effect.acquireUseRelease(
+      Effect.sync(() => listeners.add(listener)),
+      () =>
+        Effect.gen(function* () {
+          // Re-check after subscribing to close the race window.
+          const afterSubscribe = yield* SubscriptionRef.get(stateRef);
+          if (predicate(afterSubscribe)) return afterSubscribe;
+          return yield* Deferred.await(done);
+        }),
+      () => Effect.sync(() => listeners.delete(listener)),
+    );
   });
 
   const awaitFinal = waitFor((state) => machine._isFinal(state._tag)).pipe(
diff --git a/test/wait-for-cleanup.test.ts b/test/wait-for-cleanup.test.ts
new file mode 100644
index 0000000..3ed0231
--- /dev/null
+++ b/test/wait-for-cleanup.test.ts
@@ -0,0 +1,59 @@
+import { Effect, Exit, Fiber, Queue } from "effect";
+import { describe, expect, it } from "effect-bun-test";
+
+import { ActorSystemDefault, Event, Machine, State } from "../src/index.js";
+
+const TestState = State({ Idle: {}, Done: {} });
+const TestEvent = Event({ Finish: {} });
+const machine = Machine.make({ state: TestState, event: TestEvent, initial: TestState.Idle })
+  .on(TestState.Idle, TestEvent.Finish, () => TestState.Done)
+  .final(TestState.Done);
+
+describe("waitFor listener lifetime", () => {
+  it.scopedLive("stops calling a cancelled wait predicate", () =>
+    Effect.gen(function* () {
+      const actor = yield* Machine.spawn(machine);
+      yield* actor.start;
+      yield* Effect.addFinalizer(() => actor.stop);
+      const observations = yield* Queue.unbounded<string>();
+      const observed: string[] = [];
+      const waiting = yield* actor
+        .waitFor((state) => {
+          observed.push(state._tag);
+          Queue.offerUnsafe(observations, state._tag);
+          return false;
+        })
+        .pipe(Effect.forkScoped);
+      yield* Queue.take(observations);
+      yield* Queue.take(observations);
+      yield* Fiber.interrupt(waiting);
+      expect(Exit.hasInterrupts(yield* Fiber.await(waiting))).toBe(true);
+      const before = observed.length;
+      yield* actor.call(TestEvent.Finish);
+      expect((yield* actor.snapshot)._tag).toBe("Done");
+      expect(observed).toHaveLength(before);
+    }).pipe(Effect.provide(ActorSystemDefault)),
+  );
+
+  it.scopedLive("removes a listener when the subscription recheck predicate defects", () =>
+    Effect.gen(function* () {
+      const actor = yield* Machine.spawn(machine);
+      yield* actor.start;
+      yield* Effect.addFinalizer(() => actor.stop);
+      let observations = 0;
+      const failure = yield* actor
+        .waitFor(() => {
+          observations += 1;
+          // oxlint-disable-next-line effect/noThrowStatement, effect/noNewError -- Exercise a defect from the synchronous public predicate.
+          if (observations === 2) throw new Error("predicate failed");
+          return false;
+        })
+        .pipe(Effect.exit);
+      expect(Exit.hasDies(failure)).toBe(true);
+      const before = observations;
+      yield* actor.call(TestEvent.Finish);
+      expect((yield* actor.snapshot)._tag).toBe("Done");
+      expect(observations).toBe(before);
+    }).pipe(Effect.provide(ActorSystemDefault)),
+  );
+});