diff --git a/src/agents/pi-embedded-runner.run-embedded-pi-agent.auth-profile-rotation.test.ts b/src/agents/pi-embedded-runner.run-embedded-pi-agent.auth-profile-rotation.test.ts index de67dd10039..6c9038164af 100644 --- a/src/agents/pi-embedded-runner.run-embedded-pi-agent.auth-profile-rotation.test.ts +++ b/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(); - await withAgentWorkspace(async ({ agentDir, workspaceDir }) => { - await writeAuthStore(agentDir); - mockFailedThenSuccessfulAttempt(testCase.errorMessage); - await runAutoPinnedOpenAiTurn({ - agentDir, - workspaceDir, - sessionKey: testCase.sessionKey, - runId: testCase.runId, - }); - - expect(runEmbeddedAttemptMock).toHaveBeenCalledTimes(2); - await expectProfileP2UsageUpdated(agentDir); - }); + const { usageStats } = await runAutoPinnedRotationCase(testCase); + expect(typeof usageStats["openai:p2"]?.lastUsed).toBe("number"); } }); it("rotates on timeout without cooling down the timed-out profile", async () => { - await withAgentWorkspace(async ({ agentDir, workspaceDir }) => { - await writeAuthStore(agentDir); - mockFailedThenSuccessfulAttempt("request ended without sending any chunks"); - - 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(); + const { usageStats } = await runAutoPinnedRotationCase({ + errorMessage: "request ended without sending any chunks", + sessionKey: "agent:test:timeout-no-cooldown", + runId: "run:timeout-no-cooldown", }); + expect(typeof usageStats["openai:p2"]?.lastUsed).toBe("number"); + expect(usageStats["openai:p1"]?.cooldownUntil).toBeUndefined(); }); it("does not rotate for compaction timeouts", async () => { diff --git a/src/agents/pi-embedded-runner/run.ts b/src/agents/pi-embedded-runner/run.ts index 02e31b8febe..c60aaed0071 100644 --- a/src/agents/pi-embedded-runner/run.ts +++ b/src/agents/pi-embedded-runner/run.ts @@ -500,6 +500,22 @@ export async function runEmbeddedPiAgent( let lastRunPromptUsage: ReturnType | undefined; let autoCompactionCount = 0; let runLoopIterations = 0; + const maybeMarkAuthProfileFailure = async (params: { + profileId?: string; + reason?: Parameters[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 markAuthProfileFailure({ - store: authStore, - profileId: lastProfileId, - reason: promptFailoverReason, - cfg: params.config, - agentDir: params.agentDir, - }); - } + await maybeMarkAuthProfileFailure({ + profileId: lastProfileId, + reason: promptFailoverReason, + }); 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 markAuthProfileFailure({ - store: authStore, - profileId: lastProfileId, - reason, - cfg: params.config, - agentDir: params.agentDir, - }); - } + await maybeMarkAuthProfileFailure({ + profileId: lastProfileId, + reason, + }); if (timedOut && !isProbeSession) { log.warn(`Profile ${lastProfileId} timed out. Trying next account...`); }