From c55bee5ec72d55794f11f52820bd3eeb30bd147b Mon Sep 17 00:00:00 2001
From: Vincent Koc <vincentkoc@ieee.org>
Date: Mon, 25 May 2026 22:03:00 +0200
Subject: [PATCH] fix(test): model active assistant failover attempts

---
 ...pi-agent.auth-profile-rotation.e2e.test.ts | 34 ++++++++++++-------
 1 file changed, 22 insertions(+), 12 deletions(-)

diff --git a/src/agents/pi-embedded-runner.run-embedded-pi-agent.auth-profile-rotation.e2e.test.ts b/src/agents/pi-embedded-runner.run-embedded-pi-agent.auth-profile-rotation.e2e.test.ts
index f0766217094..8c47e335c89 100644
--- a/src/agents/pi-embedded-runner.run-embedded-pi-agent.auth-profile-rotation.e2e.test.ts
+++ b/src/agents/pi-embedded-runner.run-embedded-pi-agent.auth-profile-rotation.e2e.test.ts
@@ -385,15 +385,26 @@ const writeOpenAiCodexAuthStore = async (agentDir: string) => {
 const buildCopilotAssistant = (overrides: Partial<AssistantMessage> = {}) =>
   buildAssistant({ provider: "github-copilot", model: copilotModelId, ...overrides });
 
+const makeErrorAttempt = (
+  overrides: Partial<AssistantMessage> = {},
+  opts?: { currentAttempt?: boolean },
+) => {
+  const assistant = buildAssistant({
+    stopReason: "error",
+    ...overrides,
+  });
+  return makeAttempt({
+    assistantTexts: [],
+    lastAssistant: assistant,
+    ...(opts?.currentAttempt ? { currentAttemptAssistant: assistant } : {}),
+  });
+};
+
 const mockFailedThenSuccessfulAttempt = (errorMessage = "rate limit") => {
   runEmbeddedAttemptMock
     .mockResolvedValueOnce(
-      makeAttempt({
-        assistantTexts: [],
-        lastAssistant: buildAssistant({
-          stopReason: "error",
-          errorMessage,
-        }),
+      makeErrorAttempt({
+        errorMessage,
       }),
     )
     .mockResolvedValueOnce(
@@ -540,15 +551,14 @@ function mockSingleErrorAttempt(params: {
   model?: string;
 }) {
   runEmbeddedAttemptMock.mockResolvedValueOnce(
-    makeAttempt({
-      assistantTexts: [],
-      lastAssistant: buildAssistant({
-        stopReason: "error",
+    makeErrorAttempt(
+      {
         errorMessage: params.errorMessage,
         ...(params.provider ? { provider: params.provider } : {}),
         ...(params.model ? { model: params.model } : {}),
-      }),
-    }),
+      },
+      { currentAttempt: true },
+    ),
   );
 }