test: cover recovery after forced Durable Object eviction

packages/agents/src/tests/agents-core-eviction.test.ts

@@ -0,0 +1,159 @@
+/**
+ * Forced Durable Object eviction coverage for the base Agent.
+ *
+ * `evictDurableObject()` tears down a running test actor while preserving its
+ * durable storage. These tests prove that selected Agent state is reconstructed
+ * from storage on the next access. They do not assert natural idle hibernation,
+ * the absence of pending timers, or hibernation eligibility.
+ */
+import { env } from "cloudflare:workers";
+import {
+  evictDurableObject,
+  runDurableObjectAlarm,
+  runInDurableObject
+} from "cloudflare:test";
+import { describe, expect, it } from "vitest";
+import { getAgentByName } from "..";
+import type { FiberRecoveryContext } from "..";
+import type { TestRunFiberAgent } from "./agents/run-fiber";
+import type { TestScheduleAgent } from "./agents/schedule";
+
+function uniqueName(prefix: string): string {
+  return `${prefix}-${crypto.randomUUID()}`;
+}
+
+describe("Agent recovery after forced Durable Object eviction", () => {
+  it("drops instance state while restoring persisted Agent state", async () => {
+    const agent = await getAgentByName(
+      env.TestStateAgent,
+      uniqueName("evict-state")
+    );
+    const stored = {
+      count: 314,
+      items: ["before-eviction"],
+      lastUpdated: "pre-evict"
+    };
+
+    await agent.updateState(stored);
+
+    let calls = await agent.getStateUpdateCalls();
+    const deadline = Date.now() + 1000;
+    while (calls.length === 0 && Date.now() < deadline) {
+      await new Promise((resolve) => setTimeout(resolve, 10));
+      calls = await agent.getStateUpdateCalls();
+    }
+    expect(calls.length).toBeGreaterThanOrEqual(1);
+
+    await evictDurableObject(agent);
+
+    // The hook log is instance-only, while the Agent state is SQL-backed.
+    expect(await agent.getStateUpdateCalls()).toEqual([]);
+    expect(await agent.getState()).toEqual(stored);
+  });
+
+  it("runs a persisted interval schedule on the reconstructed instance", async () => {
+    const agent = await getAgentByName(
+      env.TestScheduleAgent,
+      uniqueName("evict-schedule")
+    );
+    const scheduleId = await agent.createIntervalSchedule(86_400);
+
+    try {
+      await agent.clearStoredAlarm();
+      await agent.backdateSchedule(
+        scheduleId,
+        Math.floor(Date.now() / 1000) - 1
+      );
+
+      await evictDurableObject(agent);
+
+      expect(
+        await runInDurableObject(
+          agent,
+          (instance: TestScheduleAgent) => instance.intervalCallbackCount
+        )
+      ).toBe(0);
+      expect((await agent.getStoredScheduleById(scheduleId))?.id).toBe(
+        scheduleId
+      );
+
+      await agent.setStoredAlarm(Date.now() + 1000);
+      await runDurableObjectAlarm(agent);
+
+      expect(
+        await runInDurableObject(
+          agent,
+          (instance: TestScheduleAgent) => instance.intervalCallbackCount
+        )
+      ).toBe(1);
+      expect(await agent.getScheduleCount()).toBe(1);
+    } finally {
+      await agent.cancelScheduleById(scheduleId);
+      await agent.clearStoredAlarm();
+    }
+  });
+
+  it("recovers an interrupted unmanaged fiber from SQL", async () => {
+    const stub = await getAgentByName(
+      env.TestRunFiberAgent,
+      uniqueName("evict-fiber")
+    );
+    const agent = stub as unknown as TestRunFiberAgent;
+
+    await agent.insertInterruptedFiber("evicted-fiber", "research", {
+      step: 7
+    });
+    await agent.runSimple("warm");
+    expect(await agent.getExecutionLog()).toContain("executed:warm");
+
+    await evictDurableObject(stub);
+
+    expect(await agent.getExecutionLog()).not.toContain("executed:warm");
+    expect(await agent.getRunningFiberCount()).toBe(1);
+
+    await agent.triggerRecoveryCheck();
+
+    const recovered =
+      (await agent.getRecoveredFibers()) as unknown as FiberRecoveryContext[];
+    expect(recovered).toEqual([
+      expect.objectContaining({
+        id: "evicted-fiber",
+        name: "research",
+        snapshot: { step: 7 },
+        recoveryReason: "interrupted"
+      })
+    ]);
+    expect(await agent.getRunningFiberCount()).toBe(0);
+  });
+
+  it("applies managed-fiber recovery on a fresh instance", async () => {
+    const stub = await getAgentByName(
+      env.TestRunFiberAgent,
+      uniqueName("evict-managed-fiber")
+    );
+    const agent = stub as unknown as TestRunFiberAgent;
+
+    await agent.insertInterruptedManagedFiber(
+      "evicted-managed",
+      "managed-recovery-complete",
+      { progress: 42 }
+    );
+
+    await evictDurableObject(stub);
+
+    expect(await agent.getRecoveredFibers()).toEqual([]);
+    await agent.simulateAlarmCycle();
+
+    const recovered =
+      (await agent.getRecoveredFibers()) as unknown as FiberRecoveryContext[];
+    expect(recovered).toEqual([
+      expect.objectContaining({
+        id: "evicted-managed",
+        snapshot: { progress: 42 }
+      })
+    ]);
+    expect((await agent.inspectManagedFiber("evicted-managed"))?.status).toBe(
+      "completed"
+    );
+  });
+});

packages/agents/src/tests/agents/wait-connections.ts

@@ -94,6 +94,26 @@ export class TestWaitConnectionsAgent extends Agent {
     };
   }
 
+  /** Wait for the lifecycle-started restores without triggering another one. */
+  async waitAndReport(timeout?: number): Promise<{
+    connectionIds: string[];
+    connectionStates: Record<string, string>;
+  }> {
+    await this.mcp.waitForConnections(
+      timeout != null ? { timeout } : undefined
+    );
+
+    return {
+      connectionIds: Object.keys(this.mcp.mcpConnections),
+      connectionStates: Object.fromEntries(
+        Object.entries(this.mcp.mcpConnections).map(([id, connection]) => [
+          id,
+          connection.connectionState
+        ])
+      )
+    };
+  }
+
   /**
    * Trigger restore WITHOUT waiting, then immediately check states.
    * This simulates the race condition (old behavior).

packages/agents/src/tests/experimental/memory/agents-memory-eviction.test.ts

@@ -0,0 +1,97 @@
+/**
+ * Forced Durable Object eviction coverage for Session storage.
+ *
+ * These tests tear down running test actors and verify that a new Session
+ * instance reconstructs its active-leaf cache from SQLite. They do not assert
+ * natural idle hibernation or hibernation eligibility.
+ */
+import type { UIMessage } from "ai";
+import { env } from "cloudflare:workers";
+import { evictAllDurableObjects, evictDurableObject } from "cloudflare:test";
+import { describe, expect, it } from "vitest";
+import { getAgentByName } from "../../..";
+
+interface SessionAgentStub {
+  appendMessage(message: UIMessage, parentId?: string | null): Promise<void>;
+  getHistory(): Promise<UIMessage[]>;
+  getLatestLeaf(): Promise<UIMessage | null>;
+  getBranches(messageId: string): Promise<UIMessage[]>;
+  appendLinearChainForTest(count: number, prefix?: string): Promise<void>;
+  getAntiJoinCountForTest(): Promise<number>;
+  resetAntiJoinCountForTest(): Promise<void>;
+}
+
+async function getAgent(name: string): Promise<SessionAgentStub> {
+  return getAgentByName(
+    env.TestSessionAgent,
+    name
+  ) as unknown as Promise<SessionAgentStub>;
+}
+
+async function evict(agent: SessionAgentStub): Promise<void> {
+  await evictDurableObject(agent as unknown as DurableObjectStub);
+}
+
+function uniqueName(prefix: string): string {
+  return `${prefix}-${crypto.randomUUID()}`;
+}
+
+describe("Session recovery after forced Durable Object eviction", () => {
+  it("rebuilds and rewarms the active-leaf cache from SQLite", async () => {
+    const agent = await getAgent(uniqueName("evict-session"));
+    await agent.appendLinearChainForTest(20);
+    expect((await agent.getLatestLeaf())?.id).toBe("m19");
+
+    await agent.resetAntiJoinCountForTest();
+    expect((await agent.getLatestLeaf())?.id).toBe("m19");
+    expect(await agent.getAntiJoinCountForTest()).toBe(0);
+
+    await evict(agent);
+
+    // A reconstructed provider has no activeLeafId cache. Its first lookup must
+    // discover the tip once; subsequent auto-parenting uses the warmed cache.
+    expect((await agent.getLatestLeaf())?.id).toBe("m19");
+    expect(await agent.getAntiJoinCountForTest()).toBe(1);
+
+    await agent.resetAntiJoinCountForTest();
+    await agent.appendMessage({
+      id: "after",
+      role: "user",
+      parts: [{ type: "text", text: "after eviction" }]
+    });
+
+    expect(
+      (await agent.getBranches("m19")).map((message) => message.id)
+    ).toEqual(["after"]);
+    expect((await agent.getHistory()).map((message) => message.id)).toEqual([
+      ...Array.from({ length: 20 }, (_, index) => `m${index}`),
+      "after"
+    ]);
+    expect(await agent.getAntiJoinCountForTest()).toBe(0);
+  });
+
+  it("keeps named Session histories isolated after a global forced eviction", async () => {
+    const agentA = await getAgent(uniqueName("evict-session-a"));
+    const agentB = await getAgent(uniqueName("evict-session-b"));
+
+    await agentA.appendMessage({
+      id: "a1",
+      role: "user",
+      parts: [{ type: "text", text: "A" }]
+    });
+    await agentB.appendMessage({
+      id: "b1",
+      role: "user",
+      parts: [{ type: "text", text: "B" }]
+    });
+
+    await evictAllDurableObjects();
+
+    expect((await agentA.getHistory()).map((message) => message.id)).toEqual([
+      "a1"
+    ]);
+    expect((await agentB.getHistory()).map((message) => message.id)).toEqual([
+      "b1"
+    ]);
+  });
+});

packages/agents/src/tests/mcp/eviction-mcp-rehydration.test.ts

@@ -0,0 +1,52 @@
+/**
+ * Forced eviction coverage for the stateful McpAgent transport.
+ *
+ * This verifies reconstruction after a test-requested actor teardown. It does
+ * not assert natural idle hibernation or hibernation eligibility.
+ */
+import { env } from "cloudflare:workers";
+import { createExecutionContext, evictDurableObject } from "cloudflare:test";
+import { describe, expect, it } from "vitest";
+import { getAgentByName } from "../..";
+import worker from "../worker";
+import { initializeStreamableHTTPServer } from "../shared/test-utils";
+
+describe("McpAgent recovery after forced Durable Object eviction", () => {
+  it("restores the initialize request and continues the HTTP session", async () => {
+    const ctx = createExecutionContext();
+    const baseUrl = "http://example.com/mcp";
+    const sessionId = await initializeStreamableHTTPServer(ctx, baseUrl);
+    const name = `streamable-http:${sessionId}`;
+    let stub = await getAgentByName(env.MCP_OBJECT, name);
+
+    const initializeRequest = await stub.getInitializeRequest();
+    expect(initializeRequest).toBeDefined();
+
+    await evictDurableObject(stub);
+
+    stub = await getAgentByName(env.MCP_OBJECT, name);
+    expect(await stub.getInitializeRequest()).toEqual(initializeRequest);
+
+    const response = await worker.fetch(
+      new Request(baseUrl, {
+        method: "POST",
+        headers: {
+          "Content-Type": "application/json",
+          Accept: "application/json, text/event-stream",
+          "mcp-session-id": sessionId
+        },
+        body: JSON.stringify({
+          jsonrpc: "2.0",
+          id: 1234,
+          method: "tools/call",
+          params: { name: "greet", arguments: { name: "Evicted" } }
+        })
+      }),
+      env,
+      ctx
+    );
+
+    expect(response.status).toBe(200);
+    expect(await response.text()).toContain("Hello, Evicted!");
+  });
+});

packages/agents/src/tests/mcp/wait-connections-e2e.test.ts

@@ -1,5 +1,7 @@
 import { describe, it, expect } from "vitest";
 import { env } from "cloudflare:workers";
+import { evictDurableObject } from "cloudflare:test";
+import { getAgentByName } from "../..";
 
 /**
  * E2E tests for waitForConnections() through the full Agent lifecycle.
@@ -235,6 +237,35 @@ describe("waitForConnections E2E", () => {
       // OAuth server should be in authenticating state
       expect(result.connectionStates["oauth-server"]).toBe("authenticating");
     });
+
+    // This forces a running actor out of memory. It verifies reconstruction,
+    // not natural idle-hibernation eligibility.
+    it("automatically restores MCP connections after forced eviction", async () => {
+      const name = `forced-evict-${crypto.randomUUID()}`;
+      let stub = await getAgentByName(env.TestWaitConnectionsAgent, name);
+
+      await stub.insertMcpServer(
+        "evict-roundtrip-server",
+        "Evict Round Trip Server",
+        "http://nonexistent-mcp.example.com",
+        "http://localhost:3000/callback",
+        null
+      );
+      await stub.resetRestoredFlag();
+      await stub.restoreAndWait(5000);
+      expect(await stub.hasMcpConnection("evict-roundtrip-server")).toBe(true);
+
+      await evictDurableObject(stub);
+
+      // Re-routing through getAgentByName runs the normal Agent startup wrapper.
+      // waitAndReport only waits for that lifecycle-started restore; it does not
+      // manually invoke onStart() or restoreConnectionsFromStorage().
+      stub = await getAgentByName(env.TestWaitConnectionsAgent, name);
+      const result = await stub.waitAndReport(5000);
+
+      expect(result.connectionIds).toContain("evict-roundtrip-server");
+      expect(result.connectionStates["evict-roundtrip-server"]).toBe("failed");
+    });
   });
 
   it("should handle timeout when connections are slow", async () => {