cloudflare · mattzcarey · Sep 1, 2026 · Sep 1, 2026
diff --git a/packages/think/src/tests/action-pause-recovery.test.ts b/packages/think/src/tests/action-pause-recovery.test.ts
@@ -76,10 +76,11 @@ describe("action-pause during an active recovery incident", () => {
       user: null
     });
 
-    await agent.triggerFiberRecovery();
-    expect(
-      await agent.getScheduledChatRecoveryCountForTest("_chatRecoveryRetry")
-    ).toBe(1);
+    // Assert on the counts returned from inside the trigger RPC: the
+    // zero-delay recovery alarm may fire (and consume the job row) as soon as
+    // the RPC releases the DO, so a follow-up count read can race it.
+    const scheduled = await agent.triggerFiberRecovery();
+    expect(scheduled.scheduledRetryCount).toBe(1);
     await agent.runScheduledRecoveryRetryForTest();
 
     // Recovery ran the interrupted turn to completion (no leaked fiber)...
@@ -155,10 +156,8 @@ describe("action-pause during an active recovery incident", () => {
       }
     );
 
-    await agent.triggerFiberRecovery();
-    expect(
-      await agent.getScheduledChatRecoveryCountForTest("_chatRecoveryContinue")
-    ).toBe(1);
+    const scheduled = await agent.triggerFiberRecovery();
+    expect(scheduled.scheduledContinueCount).toBe(1);
     await agent.runScheduledRecoveryContinueForTest();
 
     // Recovery settled the interrupted turn without disturbing the pause.

diff --git a/packages/think/src/tests/agent-tool-reattach-recovery.test.ts b/packages/think/src/tests/agent-tool-reattach-recovery.test.ts
@@ -77,10 +77,11 @@ describe("agent-tool child re-attach: request_id rebinding across recovery", ()
       }
     );
 
-    await agent.triggerFiberRecovery();
-    expect(
-      await agent.getScheduledChatRecoveryCountForTest("_chatRecoveryContinue")
-    ).toBe(1);
+    // Assert on the counts returned from inside the trigger RPC: the
+    // zero-delay recovery alarm may fire (and consume the job row) as soon as
+    // the RPC releases the DO, so a follow-up count read can race it.
+    const scheduled = await agent.triggerFiberRecovery();
+    expect(scheduled.scheduledContinueCount).toBe(1);
     await agent.runScheduledRecoveryContinueForTest();
 
     // The row's request_id moved off the pre-eviction turn to the recovery
@@ -124,10 +125,8 @@ describe("agent-tool child re-attach: request_id rebinding across recovery", ()
       }
     );
 
-    await agent.triggerFiberRecovery();
-    expect(
-      await agent.getScheduledChatRecoveryCountForTest("_chatRecoveryRetry")
-    ).toBe(1);
+    const scheduled = await agent.triggerFiberRecovery();
+    expect(scheduled.scheduledRetryCount).toBe(1);
     await agent.runScheduledRecoveryRetryForTest();
 
     const reboundReqId =

diff --git a/packages/think/src/tests/agent-tools.test.ts b/packages/think/src/tests/agent-tools.test.ts
@@ -4,7 +4,7 @@ import {
   AGENT_TOOL_PROGRESS_PART,
   getAgentByName
 } from "agents";
-import { describe, expect, it } from "vitest";
+import { describe, expect, it, vi } from "vitest";
 import type { ThinkAgentToolParent, ThinkTestAgent } from "./agents";
 import type {
   AgentToolEventMessage,
@@ -23,7 +23,9 @@ type ThinkAgentToolTestStub = {
   setAgentToolOutputForTest(runId: string, output: unknown): Promise<void>;
   clearAgentToolOutputForTest(runId: string): Promise<void>;
   setStripTextResponseForTest(strip: boolean): Promise<void>;
-  setBeforeStepAsyncDelay(ms: number): Promise<void>;
+  holdBeforeStepForTest(): Promise<void>;
+  hasEnteredBeforeStepForTest(): Promise<boolean>;
+  releaseBeforeStepForTest(): Promise<void>;
   resetTurnStateForTest(): Promise<void>;
   startAgentToolRun(
     input: unknown,
@@ -191,18 +193,18 @@ async function waitForAgentToolRun(
   agent: ThinkAgentToolTestStub,
   runId: string
 ): Promise<AgentToolInspection> {
-  for (let attempt = 0; attempt < 20; attempt++) {
-    const inspection = await agent.inspectAgentToolRun(runId);
-    if (
-      inspection?.status === "completed" ||
-      inspection?.status === "error" ||
-      inspection?.status === "aborted"
-    ) {
+  // The child turn runs detached (`startAgentToolRun` returns immediately),
+  // so terminal status is only observable by polling. Use a long deadline —
+  // it costs nothing when the run is fast, and fails with a clear timeout
+  // instead of handing callers a misleading non-terminal snapshot.
+  return vi.waitFor(
+    async () => {
+      const inspection = await agent.inspectAgentToolRun(runId);
+      expect(["completed", "error", "aborted"]).toContain(inspection?.status);
       return inspection;
-    }
-    await new Promise((resolve) => setTimeout(resolve, 10));
-  }
-  return agent.inspectAgentToolRun(runId);
+    },
+    { timeout: 8000, interval: 25 }
+  );
 }
 
 describe("Think agent tools", () => {
@@ -269,10 +271,21 @@ describe("Think agent tools", () => {
     const agent = await freshAgent();
     const runId = crypto.randomUUID();
 
-    await agent.setBeforeStepAsyncDelay(50);
+    // Park the child turn inside `beforeStep` on a promise gate so the reset
+    // deterministically lands while the turn is in flight (a wall-clock sleep
+    // here could lose the race and observe a completed turn instead).
+    await agent.holdBeforeStepForTest();
     await agent.startAgentToolRun("skipped probe", { runId });
-    await new Promise((resolve) => setTimeout(resolve, 5));
+    await vi.waitFor(
+      async () => {
+        expect(await agent.hasEnteredBeforeStepForTest()).toBe(true);
+      },
+      { timeout: 8000, interval: 25 }
+    );
     await agent.resetTurnStateForTest();
+    // Release AFTER the reset so the resumed turn observes the generation
+    // bump and seals the child run promptly.
+    await agent.releaseBeforeStepForTest();
 
     const inspection = await waitForAgentToolRun(agent, runId);
 
@@ -287,9 +300,17 @@ describe("Think agent tools", () => {
     const agent = await freshAgent();
     const runId = crypto.randomUUID();
 
-    await agent.setBeforeStepAsyncDelay(50);
+    // Park the child turn inside `beforeStep` so the cancel deterministically
+    // lands while the run is still cancellable (a wall-clock sleep here could
+    // lose the race and observe a completed turn instead).
+    await agent.holdBeforeStepForTest();
     await agent.startAgentToolRun("cancelled probe", { runId });
-    await new Promise((resolve) => setTimeout(resolve, 5));
+    await vi.waitFor(
+      async () => {
+        expect(await agent.hasEnteredBeforeStepForTest()).toBe(true);
+      },
+      { timeout: 8000, interval: 25 }
+    );
     await agent.cancelAgentToolRun(runId, "stop");
 
     const inspection = await waitForAgentToolRun(agent, runId);
@@ -299,7 +320,19 @@ describe("Think agent tools", () => {
       error: "stop"
     });
 
-    await new Promise((resolve) => setTimeout(resolve, 100));
+    // Let the parked turn resume and run its finish path to the end (the
+    // cleanup maps empty only when it has), then re-assert: the finalizer's
+    // guarded UPDATE must NOT clobber the aborted seal.
+    await agent.releaseBeforeStepForTest();
+    await vi.waitFor(
+      async () => {
+        expect(await agent.getAgentToolCleanupMapSizesForTest()).toEqual({
+          lastErrors: 0,
+          preTurnAssistantIds: 0
+        });
+      },
+      { timeout: 8000, interval: 25 }
+    );
     await expect(agent.inspectAgentToolRun(runId)).resolves.toMatchObject({
       runId,
       status: "aborted",

diff --git a/packages/think/src/tests/agents/think-session.ts b/packages/think/src/tests/agents/think-session.ts
@@ -743,6 +743,9 @@ export class ThinkTestAgent extends Think {
   private _turnConfigOverride: TurnConfig | null = null;
   private _stepConfigOverride: StepConfig | null = null;
   private _beforeStepAsyncDelayMs = 0;
+  private _beforeStepGate: Promise<void> | null = null;
+  private _releaseBeforeStepGate: (() => void) | null = null;
+  private _beforeStepGateEntered = false;
   private _lastModelCallSettings: CapturedModelCallSettings | null = null;
   private _reasoningResponse: { response: string; reasoning: string } | null =
     null;
@@ -902,6 +905,10 @@ export class ThinkTestAgent extends Think {
     if (this._beforeStepAsyncDelayMs > 0) {
       await new Promise((r) => setTimeout(r, this._beforeStepAsyncDelayMs));
     }
+    if (this._beforeStepGate) {
+      this._beforeStepGateEntered = true;
+      await this._beforeStepGate;
+    }
     if (this._stepConfigOverride) return this._stepConfigOverride;
   }
 
@@ -917,6 +924,31 @@ export class ThinkTestAgent extends Think {
     this._beforeStepAsyncDelayMs = ms;
   }
 
+  /**
+   * Arm a promise gate that parks the next `beforeStep` until released. Tests
+   * use it to hold a turn deterministically in flight — instead of racing a
+   * wall-clock delay — while they reset or cancel from outside the turn.
+   */
+  async holdBeforeStepForTest(): Promise<void> {
+    this._beforeStepGateEntered = false;
+    this._beforeStepGate = new Promise((resolve) => {
+      this._releaseBeforeStepGate = resolve;
+    });
+  }
+
+  /** Whether a turn is currently parked inside the armed `beforeStep` gate. */
+  async hasEnteredBeforeStepForTest(): Promise<boolean> {
+    return this._beforeStepGateEntered;
+  }
+
+  /** Release the parked turn (and disarm the gate for later steps). */
+  async releaseBeforeStepForTest(): Promise<void> {
+    const release = this._releaseBeforeStepGate;
+    this._beforeStepGate = null;
+    this._releaseBeforeStepGate = null;
+    release?.();
+  }
+
   async resetTurnStateForTest(): Promise<void> {
     this.resetTurnState();
   }
@@ -4803,21 +4835,39 @@ export class ThinkToolsTestAgent extends Think {
     `;
   }
 
-  async triggerFiberRecovery(): Promise<void> {
+  async triggerFiberRecovery(): Promise<{
+    scheduledContinueCount: number;
+    scheduledRetryCount: number;
+  }> {
     await (
       this as unknown as { _checkRunFibers(): Promise<void> }
     )._checkRunFibers();
+    // Recovery arms ZERO-delay schedules whose alarms the workers test pool
+    // can fire as soon as this RPC releases the DO, consuming the job rows.
+    // Read the counts synchronously inside the same invocation so callers can
+    // assert on them without racing that alarm.
+    return {
+      scheduledContinueCount: this._countScheduledChatRecovery(
+        "_chatRecoveryContinue"
+      ),
+      scheduledRetryCount:
+        this._countScheduledChatRecovery("_chatRecoveryRetry")
+    };
   }
 
-  async getScheduledChatRecoveryCountForTest(
-    callback = "_chatRecoveryContinue"
-  ): Promise<number> {
+  private _countScheduledChatRecovery(callback: string): number {
     const rows = this.sql<{ count: number }>`
       SELECT COUNT(*) as count FROM cf_agents_jobs WHERE capability = 'scheduler' AND fn = ${callback}
     `;
     return rows[0]?.count ?? 0;
   }
 
+  async getScheduledChatRecoveryCountForTest(
+    callback = "_chatRecoveryContinue"
+  ): Promise<number> {
+    return this._countScheduledChatRecovery(callback);
+  }
+
   async runScheduledRecoveryRetryForTest(): Promise<void> {
     const rows = this.sql<{ payload: string }>`
       SELECT json_extract(payload, '$.payload') AS payload FROM cf_agents_jobs
@@ -5540,8 +5590,11 @@ export class ThinkProgrammaticTestAgent extends Think {
       const submission = await this.testSubmitMessages("alarm owned", {
         submissionId
       });
-      for (let attempt = 0; attempt < 20 && alarmDrainCalls === 0; attempt++) {
-        await new Promise((resolve) => setTimeout(resolve, 10));
+      // The drain is delivered by a platform-scheduled alarm whose firing
+      // latency is not bounded by our timer ticks — give it a generous (~5s)
+      // deadline; the happy path still exits on the first tick after it fires.
+      for (let attempt = 0; attempt < 200 && alarmDrainCalls === 0; attempt++) {
+        await new Promise((resolve) => setTimeout(resolve, 25));
       }
       return { alarmDrainCalls, inlineDrainCalls, submission };
     } finally {
@@ -7710,12 +7763,7 @@ export class ThinkRecoveryTestAgent extends Think {
   async getScheduledChatRecoveryCountForTest(
     callback = "_chatRecoveryContinue"
   ): Promise<number> {
-    const rows = this.sql<{ count: number }>`
-      SELECT COUNT(*) as count
-      FROM cf_agents_jobs
-      WHERE capability = 'scheduler' AND fn = ${callback}
-    `;
-    return rows[0]?.count ?? 0;
+    return this._countScheduledChatRecovery(callback);
   }
 
   /** Insert a stream-metadata row aged `ageMs` in the past (for cleanup tests). */
@@ -7826,10 +7874,33 @@ export class ThinkRecoveryTestAgent extends Think {
     `;
   }
 
-  async triggerFiberRecovery(): Promise<void> {
+  async triggerFiberRecovery(): Promise<{
+    scheduledContinueCount: number;
+    scheduledRetryCount: number;
+  }> {
     await (
       this as unknown as { _checkRunFibers(): Promise<void> }
     )._checkRunFibers();
+    // Recovery arms ZERO-delay schedules whose alarms the workers test pool
+    // can fire as soon as this RPC releases the DO, consuming the job rows.
+    // Read the counts synchronously inside the same invocation so callers can
+    // assert on them without racing that alarm.
+    return {
+      scheduledContinueCount: this._countScheduledChatRecovery(
+        "_chatRecoveryContinue"
+      ),
+      scheduledRetryCount:
+        this._countScheduledChatRecovery("_chatRecoveryRetry")
+    };
+  }
+
+  private _countScheduledChatRecovery(callback: string): number {
+    const rows = this.sql<{ count: number }>`
+      SELECT COUNT(*) as count
+      FROM cf_agents_jobs
+      WHERE capability = 'scheduler' AND fn = ${callback}
+    `;
+    return rows[0]?.count ?? 0;
   }
 
   async persistTestMessage(msg: UIMessage): Promise<void> {

diff --git a/packages/think/src/tests/assistant-agent.test.ts b/packages/think/src/tests/assistant-agent.test.ts
@@ -1,5 +1,5 @@
 import { env, exports } from "cloudflare:workers";
-import { describe, expect, it } from "vitest";
+import { describe, expect, it, vi } from "vitest";
 import { getAgentByName } from "agents";
 import type { UIMessage } from "ai";
 
@@ -235,14 +235,21 @@ describe("Think — clear", () => {
 
     await waitForMessageOfType(ws, MSG_CHAT_MESSAGES);
 
-    let messages = (await agent.getMessages()) as unknown as UIMessage[];
+    const messages = (await agent.getMessages()) as unknown as UIMessage[];
     expect(messages.length).toBe(2);
 
     ws.send(JSON.stringify({ type: MSG_CHAT_CLEAR }));
-    await new Promise((r) => setTimeout(r, 200));
-
-    messages = (await agent.getMessages()) as unknown as UIMessage[];
-    expect(messages.length).toBe(0);
+    // The clear frame (WS) and the read (RPC) travel independent transports
+    // with no ordering guarantee, so poll for the effect instead of sleeping.
+    // Messages only reach 0 via the clear, so this cannot pass spuriously.
+    await vi.waitFor(
+      async () => {
+        expect(
+          ((await agent.getMessages()) as unknown as UIMessage[]).length
+        ).toBe(0);
+      },
+      { timeout: 5000, interval: 50 }
+    );
 
     await closeWS(ws);
   });

diff --git a/packages/think/src/tests/channel-recovery.test.ts b/packages/think/src/tests/channel-recovery.test.ts
@@ -74,10 +74,11 @@ describe("recovery × channels", () => {
       }
     );
 
-    await agent.triggerFiberRecovery();
-    expect(
-      await agent.getScheduledChatRecoveryCountForTest("_chatRecoveryContinue")
-    ).toBe(1);
+    // Assert on the counts returned from inside the trigger RPC: the
+    // zero-delay recovery alarm may fire (and consume the job row) as soon as
+    // the RPC releases the DO, so a follow-up count read can race it.
+    const scheduled = await agent.triggerFiberRecovery();
+    expect(scheduled.scheduledContinueCount).toBe(1);
     await agent.runScheduledRecoveryContinueForTest();
 
     // (a) The channel stamp survives recovery.
@@ -117,10 +118,8 @@ describe("recovery × channels", () => {
       user: null
     });
 
-    await agent.triggerFiberRecovery();
-    expect(
-      await agent.getScheduledChatRecoveryCountForTest("_chatRecoveryRetry")
-    ).toBe(1);
+    const scheduled = await agent.triggerFiberRecovery();
+    expect(scheduled.scheduledRetryCount).toBe(1);
     await agent.runScheduledRecoveryRetryForTest();
 
     // (a) The channel stamp survives recovery.