feat: add Agent.keepAlive() to prevent idle eviction

.changeset/keep-alive-agent.md

@@ -0,0 +1,8 @@
+---
+"agents": minor
+"@cloudflare/ai-chat": patch
+---
+
+Add experimental `keepAlive()` method to the Agent class. Keeps the Durable Object alive via alarm heartbeats (every 30 seconds), preventing idle eviction during long-running work. Returns a disposer function to stop the heartbeat.
+
+`AIChatAgent` now automatically calls `keepAlive()` during `_reply()` streaming, preventing idle eviction during long LLM generations.

packages/agents/src/experimental/forever.ts

@@ -96,7 +96,6 @@ type RawFiberRow = {
 
 // ── Constants ─────────────────────────────────────────────────────────
 
-const KEEP_ALIVE_INTERVAL_MS = 10_000;
 const FIBER_CLEANUP_INTERVAL_MS = 10 * 60 * 1000;
 const FIBER_CLEANUP_COMPLETED_MS = 24 * 60 * 60 * 1000;
 const FIBER_CLEANUP_FAILED_MS = 7 * 24 * 60 * 60 * 1000;
@@ -111,7 +110,7 @@ type Constructor<T = object> = new (...args: any[]) => T;
 type AgentLike = Constructor<
   Pick<
     Agent<Cloudflare.Env>,
-    "sql" | "scheduleEvery" | "cancelSchedule" | "alarm"
+    "sql" | "scheduleEvery" | "cancelSchedule" | "alarm" | "keepAlive"
   >
 >;
 
@@ -161,37 +160,17 @@ export function withFibers<TBase extends AgentLike>(
       }
     }
 
-    // ── Heartbeat callback ────────────────────────────────────────
+    // ── Heartbeat callback override ───────────────────────────────
 
-    // Note: TypeScript `private` is compile-time only. The scheduler
-    // dispatches callbacks by string name (`this[row.callback]`),
-    // which works at runtime. The name is stable (stored in SQLite).
-    /** @internal */ async _cf_fiberHeartbeat() {
+    // Override the base Agent's no-op heartbeat to add fiber recovery.
+    // The scheduler dispatches by string name, so this override runs
+    // when the keepAlive schedule fires.
+    /** @internal */ async _cf_keepAliveHeartbeat() {
       await this._checkInterruptedFibers();
     }
 
     // ── Public API ────────────────────────────────────────────────
 
-    async keepAlive(): Promise<() => void> {
-      const heartbeatSeconds = Math.ceil(KEEP_ALIVE_INTERVAL_MS / 1000);
-      const schedule = await (
-        this as unknown as Agent<Cloudflare.Env>
-      ).scheduleEvery(
-        heartbeatSeconds,
-        "_cf_fiberHeartbeat" as keyof Agent<Cloudflare.Env>
-      );
-
-      this._fiberDebug("keepAlive started, schedule=%s", schedule.id);
-
-      let disposed = false;
-      return () => {
-        if (disposed) return;
-        disposed = true;
-        this._fiberDebug("keepAlive disposed, schedule=%s", schedule.id);
-        void this.cancelSchedule(schedule.id);
-      };
-    }
-
     spawnFiber(
       methodName: keyof this,
       payload?: unknown,
@@ -568,7 +547,7 @@ export function withFibers<TBase extends AgentLike>(
     /** @internal */ _cleanupOrphanedHeartbeats() {
       (this as unknown as Agent<Cloudflare.Env>).sql`
         DELETE FROM cf_agents_schedules
-        WHERE callback = '_cf_fiberHeartbeat'
+        WHERE callback = '_cf_keepAliveHeartbeat'
       `;
       this._fiberDebug("cleaned up orphaned heartbeat schedules");
     }
@@ -598,34 +577,3 @@ export function withFibers<TBase extends AgentLike>(
 
   return FiberAgent;
 }
-
-// ── Standalone keepAlive ──────────────────────────────────────────────
-
-/**
- * Keep a Durable Object alive via a scheduled heartbeat.
- * Returns a disposer function that cancels the heartbeat.
- *
- * Standalone version usable by any Agent subclass without requiring
- * the full fiber mixin. The agent must have a no-op method with the
- * given callbackName for the scheduler to invoke.
- *
- * @param agent - The agent instance (must have scheduleEvery and cancelSchedule)
- * @param callbackName - Name of a no-op method on the agent class (must exist)
- */
-export async function keepAlive(
-  agent: Pick<Agent<Cloudflare.Env>, "scheduleEvery" | "cancelSchedule">,
-  callbackName: string
-): Promise<() => void> {
-  const heartbeatSeconds = Math.ceil(KEEP_ALIVE_INTERVAL_MS / 1000);
-  const schedule = await agent.scheduleEvery(
-    heartbeatSeconds,
-    callbackName as keyof Agent<Cloudflare.Env>
-  );
-
-  let disposed = false;
-  return () => {
-    if (disposed) return;
-    disposed = true;
-    void agent.cancelSchedule(schedule.id);
-  };
-}

packages/agents/src/index.ts

@@ -333,6 +333,8 @@ export type AddRpcMcpServerOptions = {
   props?: Record<string, unknown>;
 };
 
+const KEEP_ALIVE_INTERVAL_MS = 30_000;
+
 const STATE_ROW_ID = "cf_state_row_id";
 const STATE_WAS_CHANGED = "cf_state_was_changed";
 
@@ -2281,6 +2283,51 @@ export class Agent<
     return true;
   }
 
+  /**
+   * Keep the Durable Object alive via alarm heartbeats.
+   * Returns a disposer function that stops the heartbeat when called.
+   *
+   * Use this when you have long-running work and need to prevent the
+   * DO from going idle (eviction after ~70-140s of inactivity).
+   * The heartbeat fires every 30 seconds via the scheduling system.
+   *
+   * @experimental This API may change between releases.
+   *
+   * @example
+   * ```ts
+   * const dispose = await this.keepAlive();
+   * try {
+   *   // ... long-running work ...
+   * } finally {
+   *   dispose();
+   * }
+   * ```
+   */
+  async keepAlive(): Promise<() => void> {
+    const heartbeatSeconds = Math.ceil(KEEP_ALIVE_INTERVAL_MS / 1000);
+    const schedule = await this.scheduleEvery(
+      heartbeatSeconds,
+      "_cf_keepAliveHeartbeat" as keyof this
+    );
+
+    let disposed = false;
+    return () => {
+      if (disposed) return;
+      disposed = true;
+      void this.cancelSchedule(schedule.id);
+    };
+  }
+
+  /**
+   * Internal no-op callback invoked by the keepAlive heartbeat schedule.
+   * Its only purpose is to keep the DO alive — the alarm machinery
+   * handles the rest.
+   * @internal
+   */
+  async _cf_keepAliveHeartbeat(): Promise<void> {
+    // intentionally empty — the alarm firing is what keeps the DO alive
+  }
+
   private async _scheduleNextAlarm() {
     // Find the next schedule that needs to be executed
     const result = this.sql`

packages/agents/src/tests/agents/fiber.ts

@@ -167,7 +167,7 @@ export class TestFiberAgent extends FiberAgent<Record<string, unknown>> {
   async getHeartbeatScheduleCount(): Promise<number> {
     const result = this.sql<{ count: number }>`
       SELECT COUNT(*) as count FROM cf_agents_schedules
-      WHERE callback = '_cf_fiberHeartbeat'
+      WHERE callback = '_cf_keepAliveHeartbeat'
     `;
     return result[0].count;
   }

packages/agents/src/tests/agents/index.ts

@@ -28,6 +28,7 @@ export { TestQueueAgent } from "./queue";
 export { TestRaceAgent } from "./race";
 export { TestRetryAgent, TestRetryDefaultsAgent } from "./retry";
 export { TestFiberAgent } from "./fiber";
+export { TestKeepAliveAgent } from "./keep-alive";
 export {
   TestSessionAgent,
   TestSessionAgentNoMicroCompaction,

packages/agents/src/tests/agents/keep-alive.ts

@@ -0,0 +1,58 @@
+import { Agent, callable } from "../../index.ts";
+
+export class TestKeepAliveAgent extends Agent<Record<string, unknown>> {
+  private _keepAliveDisposer: (() => void) | null = null;
+  keepAliveCallCount = 0;
+
+  @callable()
+  async startKeepAlive(): Promise<string> {
+    const dispose = await this.keepAlive();
+    this._keepAliveDisposer = dispose;
+    this.keepAliveCallCount++;
+    return "started";
+  }
+
+  @callable()
+  async stopKeepAlive(): Promise<string> {
+    if (this._keepAliveDisposer) {
+      this._keepAliveDisposer();
+      this._keepAliveDisposer = null;
+      this.keepAliveCallCount--;
+    }
+    return "stopped";
+  }
+
+  @callable()
+  async getKeepAliveCallCount(): Promise<number> {
+    return this.keepAliveCallCount;
+  }
+
+  @callable()
+  async getHeartbeatScheduleCount(): Promise<number> {
+    const result = this.sql<{ count: number }>`
+      SELECT COUNT(*) as count FROM cf_agents_schedules
+      WHERE callback = '_cf_keepAliveHeartbeat'
+    `;
+    return result[0].count;
+  }
+
+  @callable()
+  async getHeartbeatSchedule(): Promise<{
+    id: string;
+    callback: string;
+    type: string;
+    intervalSeconds: number;
+  } | null> {
+    const result = this.sql<{
+      id: string;
+      callback: string;
+      type: string;
+      intervalSeconds: number;
+    }>`
+      SELECT id, callback, type, intervalSeconds FROM cf_agents_schedules
+      WHERE callback = '_cf_keepAliveHeartbeat'
+      LIMIT 1
+    `;
+    return result[0] ?? null;
+  }
+}

packages/agents/src/tests/keep-alive.test.ts

@@ -0,0 +1,82 @@
+import { env } from "cloudflare:test";
+import { describe, expect, it } from "vitest";
+import type { Env } from "./worker";
+import { getAgentByName } from "..";
+
+declare module "cloudflare:test" {
+  interface ProvidedEnv extends Env {}
+}
+
+describe("keepAlive", () => {
+  it("should create a heartbeat schedule when started", async () => {
+    const agent = await getAgentByName(
+      env.TestKeepAliveAgent,
+      "create-heartbeat"
+    );
+
+    // No heartbeat schedules initially
+    expect(await agent.getHeartbeatScheduleCount()).toBe(0);
+
+    await agent.startKeepAlive();
+
+    // Should have created exactly one heartbeat schedule
+    expect(await agent.getHeartbeatScheduleCount()).toBe(1);
+
+    // Verify the schedule properties
+    const schedule = await agent.getHeartbeatSchedule();
+    expect(schedule).toBeDefined();
+    expect(schedule?.callback).toBe("_cf_keepAliveHeartbeat");
+    expect(schedule?.type).toBe("interval");
+    expect(schedule?.intervalSeconds).toBe(30);
+  });
+
+  it("should remove the heartbeat schedule when disposed", async () => {
+    const agent = await getAgentByName(
+      env.TestKeepAliveAgent,
+      "dispose-heartbeat"
+    );
+
+    await agent.startKeepAlive();
+    expect(await agent.getHeartbeatScheduleCount()).toBe(1);
+
+    await agent.stopKeepAlive();
+    expect(await agent.getHeartbeatScheduleCount()).toBe(0);
+  });
+
+  it("should be idempotent when disposed multiple times", async () => {
+    const agent = await getAgentByName(
+      env.TestKeepAliveAgent,
+      "double-dispose"
+    );
+
+    await agent.startKeepAlive();
+    expect(await agent.getHeartbeatScheduleCount()).toBe(1);
+
+    // First dispose removes the schedule
+    await agent.stopKeepAlive();
+    expect(await agent.getHeartbeatScheduleCount()).toBe(0);
+
+    // Second dispose is a no-op (doesn't throw)
+    await agent.stopKeepAlive();
+    expect(await agent.getHeartbeatScheduleCount()).toBe(0);
+  });
+
+  it("should support multiple concurrent keepAlive calls", async () => {
+    const agent = await getAgentByName(
+      env.TestKeepAliveAgent,
+      "multiple-keepalive"
+    );
+
+    await agent.startKeepAlive();
+    await agent.startKeepAlive();
+
+    // Each call creates its own schedule
+    expect(await agent.getHeartbeatScheduleCount()).toBe(2);
+    expect(await agent.getKeepAliveCallCount()).toBe(2);
+
+    // Stopping only cancels the latest disposer
+    await agent.stopKeepAlive();
+    expect(await agent.getHeartbeatScheduleCount()).toBe(1);
+    expect(await agent.getKeepAliveCallCount()).toBe(1);
+  });
+});

packages/agents/src/tests/worker.ts

@@ -32,6 +32,7 @@ export {
   TestRetryAgent,
   TestRetryDefaultsAgent,
   TestFiberAgent,
+  TestKeepAliveAgent,
   TestSessionAgent,
   TestSessionAgentNoMicroCompaction,
   TestSessionAgentCustomRules,
@@ -81,6 +82,7 @@ import type {
   TestRetryAgent,
   TestRetryDefaultsAgent,
   TestFiberAgent,
+  TestKeepAliveAgent,
   TestSessionAgent,
   TestSessionAgentNoMicroCompaction,
   TestSessionAgentCustomRules,
@@ -114,6 +116,7 @@ export type Env = {
   TestRetryAgent: DurableObjectNamespace<TestRetryAgent>;
   TestRetryDefaultsAgent: DurableObjectNamespace<TestRetryDefaultsAgent>;
   TestFiberAgent: DurableObjectNamespace<TestFiberAgent>;
+  TestKeepAliveAgent: DurableObjectNamespace<TestKeepAliveAgent>;
   TestSessionAgent: DurableObjectNamespace<TestSessionAgent>;
   TestSessionAgentNoMicroCompaction: DurableObjectNamespace<TestSessionAgentNoMicroCompaction>;
   TestSessionAgentCustomRules: DurableObjectNamespace<TestSessionAgentCustomRules>;

packages/agents/src/tests/wrangler.jsonc

@@ -115,6 +115,10 @@
         "class_name": "TestFiberAgent",
         "name": "TestFiberAgent"
       },
+      {
+        "class_name": "TestKeepAliveAgent",
+        "name": "TestKeepAliveAgent"
+      },
       {
         "class_name": "TestSessionAgent",
         "name": "TestSessionAgent"
@@ -196,6 +200,7 @@
         "TestRetryAgent",
         "TestRetryDefaultsAgent",
         "TestFiberAgent",
+        "TestKeepAliveAgent",
         "TestSessionAgent",
         "TestSessionAgentNoMicroCompaction",
         "TestSessionAgentCustomRules",