diff --git a/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.context-engine.test.ts b/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.context-engine.test.ts
index 734af10cec0..e3b9c4d43e8 100644
--- a/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.context-engine.test.ts
+++ b/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.context-engine.test.ts
@@ -257,7 +257,13 @@ describe("runEmbeddedAttempt context engine sessionKey forwarding", () => {
     hoisted.acquireSessionWriteLockMock.mockResolvedValue({
       release: releaseMock,
     });
-    hoisted.flushPendingToolResultsAfterIdleMock.mockRejectedValueOnce(new Error("flush failed"));
+    let flushCallCount = 0;
+    hoisted.flushPendingToolResultsAfterIdleMock.mockImplementation(async () => {
+      flushCallCount += 1;
+      if (flushCallCount >= 2) {
+        throw new Error("flush failed");
+      }
+    });
 
     const result = await createContextEngineAttemptRunner({
       contextEngine: {
diff --git a/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.test-support.ts b/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.test-support.ts
index 253d1c5662d..e5034c4efe0 100644
--- a/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.test-support.ts
+++ b/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.test-support.ts
@@ -157,6 +157,7 @@ vi.mock("../google.js", () => ({
   logToolSchemasForGoogle: () => {},
   sanitizeSessionHistory: async ({ messages }: { messages: unknown[] }) => messages,
   sanitizeToolsForGoogle: ({ tools }: { tools: unknown[] }) => tools,
+  validateReplayTurns: async ({ messages }: { messages: unknown[] }) => messages,
 }));
 
 vi.mock("../../session-file-repair.js", () => ({