fix: prevent heartbeat scheduler silent death from wake handler race (#15108)

Merged via /review-pr -> /prepare-pr -> /merge-pr.

Prepared head SHA: fd7165b935
Co-authored-by: joeykrug <5925937+joeykrug@users.noreply.github.com>
Co-authored-by: gumadeiras <5599352+gumadeiras@users.noreply.github.com>
Reviewed-by: @gumadeiras

src/infra/heartbeat-wake.test.ts

@@ -28,7 +28,7 @@ describe("heartbeat-wake", () => {
 
     await vi.advanceTimersByTimeAsync(1);
     expect(handler).toHaveBeenCalledTimes(1);
-    expect(handler).toHaveBeenCalledWith({ reason: "retry" });
+    expect(handler).toHaveBeenCalledWith({ reason: "exec-event" });
     expect(wake.hasPendingHeartbeatWake()).toBe(false);
   });
 
@@ -54,6 +54,29 @@ describe("heartbeat-wake", () => {
     expect(handler.mock.calls[1]?.[0]).toEqual({ reason: "interval" });
   });
 
+  it("keeps retry cooldown even when a sooner request arrives", async () => {
+    vi.useFakeTimers();
+    const wake = await loadWakeModule();
+    const handler = vi
+      .fn()
+      .mockResolvedValueOnce({ status: "skipped", reason: "requests-in-flight" })
+      .mockResolvedValueOnce({ status: "ran", durationMs: 1 });
+    wake.setHeartbeatWakeHandler(handler);
+
+    wake.requestHeartbeatNow({ reason: "interval", coalesceMs: 0 });
+    await vi.advanceTimersByTimeAsync(1);
+    expect(handler).toHaveBeenCalledTimes(1);
+
+    // Retry is now waiting for 1000ms. This should not preempt cooldown.
+    wake.requestHeartbeatNow({ reason: "hook:wake", coalesceMs: 0 });
+    await vi.advanceTimersByTimeAsync(998);
+    expect(handler).toHaveBeenCalledTimes(1);
+
+    await vi.advanceTimersByTimeAsync(1);
+    expect(handler).toHaveBeenCalledTimes(2);
+    expect(handler.mock.calls[1]?.[0]).toEqual({ reason: "hook:wake" });
+  });
+
   it("retries thrown handler errors after the default retry delay", async () => {
     vi.useFakeTimers();
     const wake = await loadWakeModule();
@@ -76,6 +99,81 @@ describe("heartbeat-wake", () => {
     expect(handler.mock.calls[1]?.[0]).toEqual({ reason: "exec-event" });
   });
 
+  it("stale disposer does not clear a newer handler", async () => {
+    vi.useFakeTimers();
+    const wake = await loadWakeModule();
+    const handlerA = vi.fn().mockResolvedValue({ status: "ran", durationMs: 1 });
+    const handlerB = vi.fn().mockResolvedValue({ status: "ran", durationMs: 1 });
+
+    // Runner A registers its handler
+    const disposeA = wake.setHeartbeatWakeHandler(handlerA);
+
+    // Runner B registers its handler (replaces A)
+    const disposeB = wake.setHeartbeatWakeHandler(handlerB);
+
+    // Runner A's stale cleanup runs — should NOT clear handlerB
+    disposeA();
+    expect(wake.hasHeartbeatWakeHandler()).toBe(true);
+
+    // handlerB should still work
+    wake.requestHeartbeatNow({ reason: "interval", coalesceMs: 0 });
+    await vi.advanceTimersByTimeAsync(1);
+    expect(handlerB).toHaveBeenCalledTimes(1);
+    expect(handlerA).not.toHaveBeenCalled();
+
+    // Runner B's dispose should work
+    disposeB();
+    expect(wake.hasHeartbeatWakeHandler()).toBe(false);
+  });
+
+  it("preempts existing timer when a sooner schedule is requested", async () => {
+    vi.useFakeTimers();
+    const wake = await loadWakeModule();
+    const handler = vi.fn().mockResolvedValue({ status: "ran", durationMs: 1 });
+    wake.setHeartbeatWakeHandler(handler);
+
+    // Schedule for 5 seconds from now
+    wake.requestHeartbeatNow({ reason: "slow", coalesceMs: 5000 });
+
+    // Schedule for 100ms from now — should preempt the 5s timer
+    wake.requestHeartbeatNow({ reason: "fast", coalesceMs: 100 });
+
+    await vi.advanceTimersByTimeAsync(100);
+    expect(handler).toHaveBeenCalledTimes(1);
+    // The reason should be "fast" since it was set last
+    expect(handler).toHaveBeenCalledWith({ reason: "fast" });
+  });
+
+  it("keeps existing timer when later schedule is requested", async () => {
+    vi.useFakeTimers();
+    const wake = await loadWakeModule();
+    const handler = vi.fn().mockResolvedValue({ status: "ran", durationMs: 1 });
+    wake.setHeartbeatWakeHandler(handler);
+
+    // Schedule for 100ms from now
+    wake.requestHeartbeatNow({ reason: "fast", coalesceMs: 100 });
+
+    // Schedule for 5 seconds from now — should NOT preempt
+    wake.requestHeartbeatNow({ reason: "slow", coalesceMs: 5000 });
+
+    await vi.advanceTimersByTimeAsync(100);
+    expect(handler).toHaveBeenCalledTimes(1);
+  });
+
+  it("does not downgrade a higher-priority pending reason", async () => {
+    vi.useFakeTimers();
+    const wake = await loadWakeModule();
+    const handler = vi.fn().mockResolvedValue({ status: "ran", durationMs: 1 });
+    wake.setHeartbeatWakeHandler(handler);
+
+    wake.requestHeartbeatNow({ reason: "exec-event", coalesceMs: 100 });
+    wake.requestHeartbeatNow({ reason: "retry", coalesceMs: 100 });
+
+    await vi.advanceTimersByTimeAsync(100);
+    expect(handler).toHaveBeenCalledTimes(1);
+    expect(handler).toHaveBeenCalledWith({ reason: "exec-event" });
+  });
+
   it("drains pending wake once a handler is registered", async () => {
     vi.useFakeTimers();
     const wake = await loadWakeModule();