openclaw/src/infra/heartbeat-wake.test.ts

import { afterEach, describe, expect, it, vi } from "vitest";

async function loadWakeModule() {
  vi.resetModules();
  return import("./heartbeat-wake.js");
}

describe("heartbeat-wake", () => {
  afterEach(() => {
    vi.useRealTimers();
    vi.restoreAllMocks();
  });

  it("coalesces multiple wake requests into one run", async () => {
    vi.useFakeTimers();
    const wake = await loadWakeModule();
    const handler = vi.fn().mockResolvedValue({ status: "skipped", reason: "disabled" });
    wake.setHeartbeatWakeHandler(handler);

    wake.requestHeartbeatNow({ reason: "interval", coalesceMs: 200 });
    wake.requestHeartbeatNow({ reason: "exec-event", coalesceMs: 200 });
    wake.requestHeartbeatNow({ reason: "retry", coalesceMs: 200 });

    expect(wake.hasPendingHeartbeatWake()).toBe(true);

    await vi.advanceTimersByTimeAsync(199);
    expect(handler).not.toHaveBeenCalled();

    await vi.advanceTimersByTimeAsync(1);
    expect(handler).toHaveBeenCalledTimes(1);
    expect(handler).toHaveBeenCalledWith({ reason: "exec-event" });
    expect(wake.hasPendingHeartbeatWake()).toBe(false);
  });

  it("retries requests-in-flight after the default retry delay", 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);

    await vi.advanceTimersByTimeAsync(500);
    expect(handler).toHaveBeenCalledTimes(1);

    await vi.advanceTimersByTimeAsync(500);
    expect(handler).toHaveBeenCalledTimes(2);
    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();
    const handler = vi
      .fn()
      .mockRejectedValueOnce(new Error("boom"))
      .mockResolvedValueOnce({ status: "skipped", reason: "disabled" });
    wake.setHeartbeatWakeHandler(handler);

    wake.requestHeartbeatNow({ reason: "exec-event", coalesceMs: 0 });

    await vi.advanceTimersByTimeAsync(1);
    expect(handler).toHaveBeenCalledTimes(1);

    await vi.advanceTimersByTimeAsync(500);
    expect(handler).toHaveBeenCalledTimes(1);

    await vi.advanceTimersByTimeAsync(500);
    expect(handler).toHaveBeenCalledTimes(2);
    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();

    wake.requestHeartbeatNow({ reason: "manual", coalesceMs: 0 });
    await vi.advanceTimersByTimeAsync(1);
    expect(wake.hasPendingHeartbeatWake()).toBe(true);

    const handler = vi.fn().mockResolvedValue({ status: "skipped", reason: "disabled" });
    wake.setHeartbeatWakeHandler(handler);

    await vi.advanceTimersByTimeAsync(249);
    expect(handler).not.toHaveBeenCalled();

    await vi.advanceTimersByTimeAsync(1);
    expect(handler).toHaveBeenCalledTimes(1);
    expect(handler).toHaveBeenCalledWith({ reason: "manual" });
    expect(wake.hasPendingHeartbeatWake()).toBe(false);
  });
});