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

This commit is contained in:
Peter Steinberger
2026-02-22 15:43:57 +01:00
parent 44dfbd23df
commit d0b59270a7
2 changed files with 54 additions and 52 deletions

View File

@@ -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 () => {

View File

@@ -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 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...`);
}