refactor: dedupe auth-profile failure marking and rotation test setup

src/agents/pi-embedded-runner.run-embedded-pi-agent.auth-profile-rotation.test.ts

@@ -188,16 +188,33 @@ async function readUsageStats(agentDir: string) {
   return stored.usageStats ?? {};
 }
 
-async function expectProfileP2UsageUpdated(agentDir: string) {
-  const usageStats = await readUsageStats(agentDir);
-  expect(typeof usageStats["openai:p2"]?.lastUsed).toBe("number");
-}
-
 async function expectProfileP2UsageUnchanged(agentDir: string) {
   const usageStats = await readUsageStats(agentDir);
   expect(usageStats["openai:p2"]?.lastUsed).toBe(2);
 }
 
+async function runAutoPinnedRotationCase(params: {
+  errorMessage: string;
+  sessionKey: string;
+  runId: string;
+}) {
+  runEmbeddedAttemptMock.mockClear();
+  return withAgentWorkspace(async ({ agentDir, workspaceDir }) => {
+    await writeAuthStore(agentDir);
+    mockFailedThenSuccessfulAttempt(params.errorMessage);
+    await runAutoPinnedOpenAiTurn({
+      agentDir,
+      workspaceDir,
+      sessionKey: params.sessionKey,
+      runId: params.runId,
+    });
+
+    expect(runEmbeddedAttemptMock).toHaveBeenCalledTimes(2);
+    const usageStats = await readUsageStats(agentDir);
+    return { usageStats };
+  });
+}
+
 function mockSingleSuccessfulAttempt() {
   runEmbeddedAttemptMock.mockResolvedValueOnce(
     makeAttempt({
@@ -314,40 +331,19 @@ describe("runEmbeddedPiAgent auth profile rotation", () => {
     ] as const;
 
     for (const testCase of cases) {
-      runEmbeddedAttemptMock.mockClear();
+      const { usageStats } = await runAutoPinnedRotationCase(testCase);
-      await withAgentWorkspace(async ({ agentDir, workspaceDir }) => {
+      expect(typeof usageStats["openai:p2"]?.lastUsed).toBe("number");
-        await writeAuthStore(agentDir);
-        mockFailedThenSuccessfulAttempt(testCase.errorMessage);
-        await runAutoPinnedOpenAiTurn({
-          agentDir,
-          workspaceDir,
-          sessionKey: testCase.sessionKey,
-          runId: testCase.runId,
-        });
-
-        expect(runEmbeddedAttemptMock).toHaveBeenCalledTimes(2);
-        await expectProfileP2UsageUpdated(agentDir);
-      });
     }
   });
 
   it("rotates on timeout without cooling down the timed-out profile", async () => {
-    await withAgentWorkspace(async ({ agentDir, workspaceDir }) => {
+    const { usageStats } = await runAutoPinnedRotationCase({
-      await writeAuthStore(agentDir);
+      errorMessage: "request ended without sending any chunks",
-      mockFailedThenSuccessfulAttempt("request ended without sending any chunks");
+      sessionKey: "agent:test:timeout-no-cooldown",
-
+      runId: "run:timeout-no-cooldown",
-      await runAutoPinnedOpenAiTurn({
-        agentDir,
-        workspaceDir,
-        sessionKey: "agent:test:timeout-no-cooldown",
-        runId: "run:timeout-no-cooldown",
-      });
-
-      expect(runEmbeddedAttemptMock).toHaveBeenCalledTimes(2);
-      const usageStats = await readUsageStats(agentDir);
-      expect(typeof usageStats["openai:p2"]?.lastUsed).toBe("number");
-      expect(usageStats["openai:p1"]?.cooldownUntil).toBeUndefined();
     });
+    expect(typeof usageStats["openai:p2"]?.lastUsed).toBe("number");
+    expect(usageStats["openai:p1"]?.cooldownUntil).toBeUndefined();
   });
 
   it("does not rotate for compaction timeouts", async () => {

src/agents/pi-embedded-runner/run.ts

@@ -500,6 +500,22 @@ export async function runEmbeddedPiAgent(
       let lastRunPromptUsage: ReturnType<typeof normalizeUsage> | undefined;
       let autoCompactionCount = 0;
       let runLoopIterations = 0;
+      const maybeMarkAuthProfileFailure = async (params: {
+        profileId?: string;
+        reason?: Parameters<typeof markAuthProfileFailure>[0]["reason"] | null;
+      }) => {
+        const { profileId, reason } = params;
+        if (!profileId || !reason || reason === "timeout") {
+          return;
+        }
+        await markAuthProfileFailure({
+          store: authStore,
+          profileId,
+          reason,
+          cfg: params.config,
+          agentDir: params.agentDir,
+        });
+      };
       try {
         while (true) {
           if (runLoopIterations >= MAX_RUN_LOOP_ITERATIONS) {
@@ -869,15 +885,10 @@ export async function runEmbeddedPiAgent(
               };
             }
             const promptFailoverReason = classifyFailoverReason(errorText);
-            if (promptFailoverReason && promptFailoverReason !== "timeout" && lastProfileId) {
+            await maybeMarkAuthProfileFailure({
-              await markAuthProfileFailure({
+              profileId: lastProfileId,
-                store: authStore,
+              reason: promptFailoverReason,
-                profileId: lastProfileId,
+            });
-                reason: promptFailoverReason,
-                cfg: params.config,
-                agentDir: params.agentDir,
-              });
-            }
             if (
               isFailoverErrorMessage(errorText) &&
               promptFailoverReason !== "timeout" &&
@@ -963,15 +974,10 @@ export async function runEmbeddedPiAgent(
               // Skip cooldown for timeouts: a timeout is model/network-specific,
               // not an auth issue. Marking the profile would poison fallback models
               // on the same provider (e.g. gpt-5.3 timeout blocks gpt-5.2).
-              if (reason !== "timeout") {
+              await maybeMarkAuthProfileFailure({
-                await markAuthProfileFailure({
+                profileId: lastProfileId,
-                  store: authStore,
+                reason,
-                  profileId: lastProfileId,
+              });
-                  reason,
-                  cfg: params.config,
-                  agentDir: params.agentDir,
-                });
-              }
               if (timedOut && !isProbeSession) {
                 log.warn(`Profile ${lastProfileId} timed out. Trying next account...`);
               }