Schedule fixes (#653)

* add destroyed flag and yield abort

* Implement destroyed flag and context abort

Add a destroyed flag and yield context abort functionality.

* add test and fix immediate schedules

* remove manual alarm call in constructor

* update test + better changeset

Author: deathbyknowledge

.changeset/neat-beds-cheat.md

@@ -0,0 +1,7 @@
+---
+"agents": patch
+---
+
+Allow `this.destroy` inside a schedule by including a `destroyed` flag and yielding `ctx.abort` instead of calling it directly
+Fix issue where schedules would not be able to run for more 30 seconds due to `blockConccurencyWhile`. `alarm()` isn't manually called anymore, getting rid of the bCW.
+Fix an issue where immediate schedules (e.g. `this.schedule(0, "foo"))`) would not get immediately scheduled.

packages/agents/src/index.ts

@@ -321,6 +321,7 @@ export class Agent<
   private _state = DEFAULT_STATE as State;
   private _disposables = new DisposableStore();
   private _mcpStateRestored = false;
+  private _destroyed = false;
 
   private _ParentClass: typeof Agent<Env, State> =
     Object.getPrototypeOf(this).constructor;
@@ -457,26 +458,18 @@ export class Agent<
       )
     `;
 
-    void this.ctx.blockConcurrencyWhile(async () => {
-      return this._tryCatch(async () => {
-        // Create alarms table if it doesn't exist
-        this.sql`
-        CREATE TABLE IF NOT EXISTS cf_agents_schedules (
-          id TEXT PRIMARY KEY NOT NULL DEFAULT (randomblob(9)),
-          callback TEXT,
-          payload TEXT,
-          type TEXT NOT NULL CHECK(type IN ('scheduled', 'delayed', 'cron')),
-          time INTEGER,
-          delayInSeconds INTEGER,
-          cron TEXT,
-          created_at INTEGER DEFAULT (unixepoch())
-        )
-      `;
-
-        // execute any pending alarms and schedule the next alarm
-        await this.alarm();
-      });
-    });
+    this.sql`
+      CREATE TABLE IF NOT EXISTS cf_agents_schedules (
+        id TEXT PRIMARY KEY NOT NULL DEFAULT (randomblob(9)),
+        callback TEXT,
+        payload TEXT,
+        type TEXT NOT NULL CHECK(type IN ('scheduled', 'delayed', 'cron')),
+        time INTEGER,
+        delayInSeconds INTEGER,
+        cron TEXT,
+        created_at INTEGER DEFAULT (unixepoch())
+      )
+    `;
 
     this.sql`
       CREATE TABLE IF NOT EXISTS cf_agents_mcp_servers (
@@ -1251,7 +1244,7 @@ export class Agent<
     // Find the next schedule that needs to be executed
     const result = this.sql`
       SELECT time FROM cf_agents_schedules
-      WHERE time > ${Math.floor(Date.now() / 1000)}
+      WHERE time >= ${Math.floor(Date.now() / 1000)}
       ORDER BY time ASC
       LIMIT 1
     `;
@@ -1321,6 +1314,7 @@ export class Agent<
           }
         );
         if (row.type === "cron") {
+          if (this._destroyed) return;
           // Update next execution time for cron schedules
           const nextExecutionTime = getNextCronTime(row.cron);
           const nextTimestamp = Math.floor(nextExecutionTime.getTime() / 1000);
@@ -1329,13 +1323,15 @@ export class Agent<
           UPDATE cf_agents_schedules SET time = ${nextTimestamp} WHERE id = ${row.id}
         `;
         } else {
+          if (this._destroyed) return;
           // Delete one-time schedules after execution
           this.sql`
           DELETE FROM cf_agents_schedules WHERE id = ${row.id}
         `;
         }
       }
     }
+    if (this._destroyed) return;
 
     // Schedule the next alarm
     await this._scheduleNextAlarm();
@@ -1356,7 +1352,13 @@ export class Agent<
     await this.ctx.storage.deleteAll();
     this._disposables.dispose();
     await this.mcp.dispose?.();
-    this.ctx.abort("destroyed"); // enforce that the agent is evicted
+    this._destroyed = true;
+
+    // `ctx.abort` throws an uncatchable error, so we yield to the event loop
+    // to avoid capturing it and let handlers finish cleaning up
+    setTimeout(() => {
+      this.ctx.abort("destroyed");
+    }, 0);
 
     this.observability?.emit(
       {

packages/agents/src/tests/alarms.test.ts

@@ -0,0 +1,31 @@
+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("scheduled destroys", () => {
+  it("should not throw when a scheduled callback nukes storage", async () => {
+    let agentStub = await getAgentByName(
+      env.TestDestroyScheduleAgent,
+      "alarm-destroy-repro"
+    );
+
+    // Alarm should fire immediately
+    await agentStub.scheduleSelfDestructingAlarm();
+    await expect(agentStub.getStatus()).resolves.toBe("scheduled");
+
+    // Let the alarm run
+    await new Promise((resolve) => setTimeout(resolve, 50));
+
+    agentStub = await getAgentByName(
+      env.TestDestroyScheduleAgent,
+      "alarm-destroy-repro"
+    );
+
+    await expect(agentStub.getStatus()).resolves.toBe("unscheduled");
+  });
+});

packages/agents/src/tests/worker.ts

@@ -30,6 +30,7 @@ export type Env = {
   TestChatAgent: DurableObjectNamespace<TestChatAgent>;
   TestOAuthAgent: DurableObjectNamespace<TestOAuthAgent>;
   TEST_MCP_JURISDICTION: DurableObjectNamespace<TestMcpJurisdiction>;
+  TestDestroyScheduleAgent: DurableObjectNamespace<TestDestroyScheduleAgent>;
 };
 
 type State = unknown;
@@ -39,6 +40,7 @@ type Props = {
 };
 
 export class TestMcpAgent extends McpAgent<Env, State, Props> {
+  observability = undefined;
   private tempToolHandle?: { remove: () => void };
 
   server = new McpServer(
@@ -130,6 +132,7 @@ export class TestMcpAgent extends McpAgent<Env, State, Props> {
 
 // Test email agents
 export class TestEmailAgent extends Agent<Env> {
+  observability = undefined;
   emailsReceived: AgentEmail[] = [];
 
   async onEmail(email: AgentEmail) {
@@ -144,6 +147,7 @@ export class TestEmailAgent extends Agent<Env> {
 }
 
 export class TestCaseSensitiveAgent extends Agent<Env> {
+  observability = undefined;
   emailsReceived: AgentEmail[] = [];
 
   async onEmail(email: AgentEmail) {
@@ -156,6 +160,7 @@ export class TestCaseSensitiveAgent extends Agent<Env> {
 }
 
 export class TestUserNotificationAgent extends Agent<Env> {
+  observability = undefined;
   emailsReceived: AgentEmail[] = [];
 
   async onEmail(email: AgentEmail) {
@@ -167,12 +172,30 @@ export class TestUserNotificationAgent extends Agent<Env> {
   }
 }
 
+export class TestDestroyScheduleAgent extends Agent<Env, { status: string }> {
+  observability = undefined;
+  initialState = {
+    status: "unscheduled"
+  };
+
+  async scheduleSelfDestructingAlarm() {
+    this.setState({ status: "scheduled" });
+    await this.schedule(0, "destroy");
+  }
+
+  getStatus() {
+    return this.state.status;
+  }
+}
+
 // An Agent that tags connections in onConnect,
 // then echoes whether the tag was observed in onMessage
 export class TestRaceAgent extends Agent<Env> {
   initialState = { hello: "world" };
   static options = { hibernate: true };
 
+  observability = undefined;
+
   async onConnect(conn: Connection<{ tagged: boolean }>) {
     // Simulate real async setup to widen the window a bit
     conn.setState({ tagged: true });
@@ -187,6 +210,8 @@ export class TestRaceAgent extends Agent<Env> {
 
 // Test Agent for OAuth client side flows
 export class TestOAuthAgent extends Agent<Env> {
+  observability = undefined;
+
   async onRequest(_request: Request): Promise<Response> {
     return new Response("Test OAuth Agent");
   }
@@ -311,6 +336,8 @@ export class TestOAuthAgent extends Agent<Env> {
 }
 
 export class TestChatAgent extends AIChatAgent<Env> {
+  observability = undefined;
+
   async onChatMessage() {
     // Simple echo response for testing
     return new Response("Hello from chat agent!", {
@@ -373,6 +400,8 @@ export class TestChatAgent extends AIChatAgent<Env> {
 
 // Test MCP Agent for jurisdiction feature
 export class TestMcpJurisdiction extends McpAgent<Env> {
+  observability = undefined;
+
   server = new McpServer(
     { name: "test-jurisdiction-server", version: "1.0.0" },
     { capabilities: { tools: {} } }

packages/agents/src/tests/wrangler.jsonc

@@ -38,6 +38,10 @@
       {
         "class_name": "TestMcpJurisdiction",
         "name": "TEST_MCP_JURISDICTION"
+      },
+      {
+        "class_name": "TestDestroyScheduleAgent",
+        "name": "TestDestroyScheduleAgent"
       }
     ]
   },
@@ -52,7 +56,8 @@
         "TestRaceAgent",
         "TestChatAgent",
         "TestOAuthAgent",
-        "TestMcpJurisdiction"
+        "TestMcpJurisdiction",
+        "TestDestroyScheduleAgent"
       ],
       "tag": "v1"
     }