fix(daemon): persist rate-limit resumeAfter so restart doesn't strand sessions

saltbo · saltbo · commit 2288bcd75015 · 2026-04-20T14:26:15.000-04:00
Before: RateLimiter held the reset window only in memory (setTimeout). On
daemon restart the timer was lost and session files had no resumeAfter, so
resumeBackoffSessions skipped them — recovery only happened if a new task
re-hit the same runtime's limit.

Now routeRateLimit writes resumeAfter to the session file before notifying
the in-memory sink, so a fresh DaemonLoop picks up stranded sessions on its
first tick.
diff --git a/packages/cli/src/daemon/runtimePool.ts b/packages/cli/src/daemon/runtimePool.ts
@@ -244,7 +244,7 @@ async function routeEvent(
 
   switch (event.type) {
     case "turn.rate_limit":
-      routeRateLimit(agent, event, rateLimitSink);
+      await routeRateLimit(agent, event, rateLimitSink);
       return;
     case "turn.error":
       logger.warn(`Agent error on task ${agent.taskId} (${agent.providerName}): ${event.detail}`);
@@ -263,20 +263,28 @@ async function routeEvent(
   }
 }
 
-function routeRateLimit(agent: AgentFlags, event: Extract<AgentEvent, { type: "turn.rate_limit" }>, sink: RateLimitSink): void {
+async function routeRateLimit(agent: AgentFlags, event: Extract<AgentEvent, { type: "turn.rate_limit" }>, sink: RateLimitSink): Promise<void> {
   const runtime = agent.providerName;
   if (event.status === "rejected") {
     const mainReset = event.resetAt;
     const overageReset = event.overage?.status === "rejected" ? event.overage.resetAt : undefined;
-    const candidates = [mainReset, overageReset].filter((x): x is string => !!x);
+    const resetTimestamps = [mainReset, overageReset].filter((x): x is string => !!x);
     const pauseUntil =
-      candidates.length > 0
-        ? candidates.reduce((a, b) => (new Date(a).getTime() >= new Date(b).getTime() ? a : b))
+      resetTimestamps.length > 0
+        ? resetTimestamps.reduce((a, b) => (new Date(a).getTime() >= new Date(b).getTime() ? a : b))
         : new Date(Date.now() + 60 * 60 * 1000).toISOString();
     logger.warn(
       `Rate limited on task ${agent.taskId} (${runtime}), pausing until ${pauseUntil}${event.isUsingOverage ? " — agent continues via overage" : ""}`,
     );
     agent.rateLimited = true;
+    // Persist first, then notify the in-memory sink. Reverse order would leave
+    // a crash window: if the daemon died between sink update and disk write,
+    // the session would come back up without resumeAfter and be stranded —
+    // exactly the bug this fix exists to prevent.
+    const resumeAfter = new Date(pauseUntil).getTime();
+    await getSessionManager()
+      .patch(agent.sessionId, { resumeAfter })
+      .catch((e) => logger.warn(`Failed to persist resumeAfter for ${agent.sessionId.slice(0, 8)}: ${errMessage(e)}`));
     sink.onRateLimited(runtime, pauseUntil);
     return;
   }
diff --git a/tests/loop-rate-limit-prompt.test.ts b/tests/loop-rate-limit-prompt.test.ts
@@ -63,6 +63,8 @@ vi.mock("../packages/cli/src/logger.js", () => ({
 // ---- Imports after mocks are set up -----------------------------------------
 
 import { DaemonLoop } from "../packages/cli/src/daemon/loop.js";
+import { _setSessionManagerForTest, SessionManager } from "../packages/cli/src/session/manager.js";
+import { clearAllSessions, writeSession } from "../packages/cli/src/session/store.js";
 import type { SessionFile } from "../packages/cli/src/session/types.js";
 
 // ---- Minimal fakes ----------------------------------------------------------
@@ -230,6 +232,111 @@ describe("DaemonLoop — RATE_LIMIT_RESUME_PROMPT is non-empty", () => {
     });
   });
 
+  describe("post-crash restart: resumeAfter on disk → fresh DaemonLoop resumes (full chain)", () => {
+    /**
+     * Integration test for the daemon-restart scenario described in the bug fix:
+     *
+     * 1. routeRateLimit writes a session file to disk with status="rate_limited"
+     *    and resumeAfter=<past epoch ms> — exactly what happens when an agent hits
+     *    the rate limit and then the daemon process is killed.
+     * 2. A fresh DaemonLoop is constructed with a brand-new SessionManager (no
+     *    in-memory rate-limit state) — simulating a daemon restart.
+     * 3. resumeBackoffSessions() scans disk, finds the expired session, and
+     *    calls resumeOneSession.
+     *
+     * This verifies the FULL CHAIN, not just either half in isolation.
+     * Note: DaemonLoop is constructed AFTER _setSessionManagerForTest so the
+     * class field `private sessions = getSessionManager()` picks up the real
+     * SessionManager from the singleton. No internal stomping is used.
+     */
+
+    afterEach(() => {
+      clearAllSessions();
+      _setSessionManagerForTest(null);
+    });
+
+    it("resumes a rate_limited session whose resumeAfter was written to disk before daemon restart", async () => {
+      // Step 1: write a rate_limited session to disk (simulating post-crash state)
+      const sessionId = "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee";
+      const onDiskSession: SessionFile = {
+        type: "worker",
+        agentId: "agent-restart-test",
+        sessionId,
+        runtime: "claude" as any,
+        startedAt: Date.now() - 120_000,
+        apiUrl: "http://localhost",
+        privateKeyJwk: {} as any,
+        taskId: "task-restart-test",
+        status: "rate_limited",
+        resumeAfter: Date.now() - 5_000, // expired 5 seconds ago
+      };
+      writeSession(onDiskSession);
+
+      // Step 2: seed the singleton with a fresh SessionManager, then construct
+      // DaemonLoop so its class field `sessions = getSessionManager()` binds to it.
+      _setSessionManagerForTest(new SessionManager());
+
+      const pool = {
+        activeCount: 0,
+        hasTask: (_id: string) => false,
+        getActiveTaskIds: () => [],
+      } as any;
+
+      const freshLoop = new DaemonLoop(makeClient(), pool, makeRateLimiter(), makePrMonitor(), {
+        maxConcurrent: 4,
+        pollInterval: 1000,
+      });
+      (freshLoop as any).running = true;
+
+      // Step 3: call resumeBackoffSessions — must find the on-disk session
+      await (freshLoop as any).resumeBackoffSessions();
+
+      // Assert: resumeOneSession was called for the on-disk session
+      expect(resumeOneSessionMock).toHaveBeenCalledTimes(1);
+      const sessionArg: SessionFile = resumeOneSessionMock.mock.calls[0][0];
+      expect(sessionArg.sessionId).toBe(sessionId);
+      expect(sessionArg.taskId).toBe("task-restart-test");
+      expect(sessionArg.status).toBe("rate_limited");
+    });
+
+    it("does NOT resume a rate_limited session whose resumeAfter is still in the future after restart", async () => {
+      // Write a session whose backoff has NOT expired yet
+      const sessionId = "ffffffff-aaaa-bbbb-cccc-dddddddddddd";
+      const onDiskSession: SessionFile = {
+        type: "worker",
+        agentId: "agent-future-test",
+        sessionId,
+        runtime: "claude" as any,
+        startedAt: Date.now() - 60_000,
+        apiUrl: "http://localhost",
+        privateKeyJwk: {} as any,
+        taskId: "task-future-test",
+        status: "rate_limited",
+        resumeAfter: Date.now() + 600_000, // still 10 minutes away
+      };
+      writeSession(onDiskSession);
+
+      // Seed singleton before constructing DaemonLoop so the class field binds correctly.
+      _setSessionManagerForTest(new SessionManager());
+
+      const pool = {
+        activeCount: 0,
+        hasTask: (_id: string) => false,
+        getActiveTaskIds: () => [],
+      } as any;
+
+      const freshLoop = new DaemonLoop(makeClient(), pool, makeRateLimiter(), makePrMonitor(), {
+        maxConcurrent: 4,
+        pollInterval: 1000,
+      });
+      (freshLoop as any).running = true;
+
+      await (freshLoop as any).resumeBackoffSessions();
+
+      expect(resumeOneSessionMock).not.toHaveBeenCalled();
+    });
+  });
+
   describe("both callers pass the same non-empty prompt string", () => {
     it("resumeRateLimitedSessions and resumeBackoffSessions pass identical message text", async () => {
       // Call each function in isolation so call indices are unambiguous.
diff --git a/tests/runtimePool-coverage.test.ts b/tests/runtimePool-coverage.test.ts
@@ -700,6 +700,74 @@ describe("routeRateLimit — rejected status", () => {
   });
 });
 
+describe("routeRateLimit — persists resumeAfter to session file", () => {
+  it("persists resumeAfter equal to new Date(pauseUntil).getTime() after a rejected rate_limit event", async () => {
+    const taskId = randomUUID();
+    const sessionId = randomUUID();
+    await seedActiveSession(sessions, sessionId, taskId);
+
+    const agentClient = makeAgentClient(null);
+    const rateLimitSink = { onRateLimited: vi.fn(), onRateLimitResumed: vi.fn() };
+    const resetAt = new Date(Date.now() + 60_000).toISOString();
+    const expectedResumeAfter = new Date(resetAt).getTime();
+    const events: AgentEvent[] = [makeRateLimitRejectedEvent(resetAt)];
+
+    await spawnAndWait(apiClient, { events, taskId, sessionId, agentClient, rateLimitSink });
+
+    const session = sessions.read(sessionId);
+    expect(session).not.toBeNull();
+    expect(session!.resumeAfter).toBe(expectedResumeAfter);
+  });
+
+  it("persists resumeAfter equal to the later overage resetAt when overage resets later than main", async () => {
+    const taskId = randomUUID();
+    const sessionId = randomUUID();
+    await seedActiveSession(sessions, sessionId, taskId);
+
+    const agentClient = makeAgentClient(null);
+    const rateLimitSink = { onRateLimited: vi.fn(), onRateLimitResumed: vi.fn() };
+
+    // Overage resets later — pauseUntil should be the later time
+    const earlier = new Date(Date.now() + 30_000).toISOString();
+    const later = new Date(Date.now() + 120_000).toISOString();
+    const expectedResumeAfter = new Date(later).getTime();
+    const events: AgentEvent[] = [makeRateLimitRejectedWithOverageEvent(earlier, later)];
+
+    await spawnAndWait(apiClient, { events, taskId, sessionId, agentClient, rateLimitSink });
+
+    const session = sessions.read(sessionId);
+    expect(session).not.toBeNull();
+    expect(session!.resumeAfter).toBe(expectedResumeAfter);
+  });
+
+  it("persists resumeAfter approximately 1h from now when resetAt is missing and no overage", async () => {
+    const taskId = randomUUID();
+    const sessionId = randomUUID();
+    await seedActiveSession(sessions, sessionId, taskId);
+
+    const agentClient = makeAgentClient(null);
+    const rateLimitSink = { onRateLimited: vi.fn(), onRateLimitResumed: vi.fn() };
+
+    // No resetAt and no overage → resetTimestamps.length === 0 → fallback: Date.now() + 3_600_000
+    const before = Date.now();
+    const event: AgentEvent = {
+      type: "turn.rate_limit",
+      status: "rejected",
+      resetAt: undefined as any,
+      isUsingOverage: false,
+    } as AgentEvent;
+
+    await spawnAndWait(apiClient, { events: [event], taskId, sessionId, agentClient, rateLimitSink });
+
+    const session = sessions.read(sessionId);
+    expect(session).not.toBeNull();
+    const delta = session!.resumeAfter! - before;
+    // Should be approximately 1 hour (3_600_000ms), within ±5s tolerance
+    expect(delta).toBeGreaterThanOrEqual(3_600_000 - 5_000);
+    expect(delta).toBeLessThanOrEqual(3_600_000 + 5_000);
+  });
+});
+
 describe("routeRateLimit — allowed status", () => {
   it("clears rateLimited flag and calls onRateLimitResumed when allowed without overage", async () => {
     const taskId = randomUUID();
