refactor(agents): dedupe embedded fallback e2e helpers

This commit is contained in:
Peter Steinberger
2026-03-07 17:22:17 +00:00
parent 2d4a0c79a3
commit 02b3e85eac

View File

@@ -228,6 +228,10 @@ async function runEmbeddedFallback(params: {
}
function mockPrimaryOverloadedThenFallbackSuccess() {
mockPrimaryErrorThenFallbackSuccess(OVERLOADED_ERROR_PAYLOAD);
}
function mockPrimaryErrorThenFallbackSuccess(errorMessage: string) {
runEmbeddedAttemptMock.mockImplementation(async (params: unknown) => {
const attemptParams = params as { provider: string; modelId: string; authProfileId?: string };
if (attemptParams.provider === "openai") {
@@ -237,7 +241,7 @@ function mockPrimaryOverloadedThenFallbackSuccess() {
provider: "openai",
model: "mock-1",
stopReason: "error",
errorMessage: OVERLOADED_ERROR_PAYLOAD,
errorMessage,
}),
});
}
@@ -256,6 +260,21 @@ function mockPrimaryOverloadedThenFallbackSuccess() {
});
}
function expectOpenAiThenGroqAttemptOrder(params?: { expectOpenAiAuthProfileId?: string }) {
expect(runEmbeddedAttemptMock).toHaveBeenCalledTimes(2);
const firstCall = runEmbeddedAttemptMock.mock.calls[0]?.[0] as
| { provider?: string; authProfileId?: string }
| undefined;
const secondCall = runEmbeddedAttemptMock.mock.calls[1]?.[0] as { provider?: string } | undefined;
expect(firstCall).toBeDefined();
expect(secondCall).toBeDefined();
expect(firstCall?.provider).toBe("openai");
if (params?.expectOpenAiAuthProfileId) {
expect(firstCall?.authProfileId).toBe(params.expectOpenAiAuthProfileId);
}
expect(secondCall?.provider).toBe("groq");
}
function mockAllProvidersOverloaded() {
runEmbeddedAttemptMock.mockImplementation(async (params: unknown) => {
const attemptParams = params as { provider: string; modelId: string; authProfileId?: string };
@@ -297,17 +316,7 @@ describe("runWithModelFallback + runEmbeddedPiAgent overload policy", () => {
expect(usageStats["openai:p1"]?.failureCounts).toMatchObject({ overloaded: 1 });
expect(typeof usageStats["groq:p1"]?.lastUsed).toBe("number");
expect(runEmbeddedAttemptMock).toHaveBeenCalledTimes(2);
const firstCall = runEmbeddedAttemptMock.mock.calls[0]?.[0] as
| { provider?: string }
| undefined;
const secondCall = runEmbeddedAttemptMock.mock.calls[1]?.[0] as
| { provider?: string }
| undefined;
expect(firstCall).toBeDefined();
expect(secondCall).toBeDefined();
expect(firstCall?.provider).toBe("openai");
expect(secondCall?.provider).toBe("groq");
expectOpenAiThenGroqAttemptOrder();
expect(computeBackoffMock).toHaveBeenCalledTimes(1);
expect(sleepWithAbortMock).toHaveBeenCalledTimes(1);
});
@@ -371,18 +380,7 @@ describe("runWithModelFallback + runEmbeddedPiAgent overload policy", () => {
});
expect(result.provider).toBe("groq");
expect(runEmbeddedAttemptMock).toHaveBeenCalledTimes(2);
const firstCall = runEmbeddedAttemptMock.mock.calls[0]?.[0] as
| { provider?: string; authProfileId?: string }
| undefined;
const secondCall = runEmbeddedAttemptMock.mock.calls[1]?.[0] as
| { provider?: string }
| undefined;
expect(firstCall).toBeDefined();
expect(secondCall).toBeDefined();
expect(firstCall?.provider).toBe("openai");
expect(firstCall?.authProfileId).toBe("openai:p1");
expect(secondCall?.provider).toBe("groq");
expectOpenAiThenGroqAttemptOrder({ expectOpenAiAuthProfileId: "openai:p1" });
});
});
@@ -414,19 +412,7 @@ describe("runWithModelFallback + runEmbeddedPiAgent overload policy", () => {
});
expect(secondResult.provider).toBe("groq");
expect(runEmbeddedAttemptMock).toHaveBeenCalledTimes(2);
const firstCall = runEmbeddedAttemptMock.mock.calls[0]?.[0] as
| { provider?: string; authProfileId?: string }
| undefined;
const secondCall = runEmbeddedAttemptMock.mock.calls[1]?.[0] as
| { provider?: string }
| undefined;
expect(firstCall).toBeDefined();
expect(secondCall).toBeDefined();
expect(firstCall?.provider).toBe("openai");
expect(firstCall?.authProfileId).toBe("openai:p1");
expect(secondCall?.provider).toBe("groq");
expectOpenAiThenGroqAttemptOrder({ expectOpenAiAuthProfileId: "openai:p1" });
const usageStats = await readUsageStats(agentDir);
expect(typeof usageStats["openai:p1"]?.cooldownUntil).toBe("number");
@@ -439,32 +425,7 @@ describe("runWithModelFallback + runEmbeddedPiAgent overload policy", () => {
it("keeps bare service-unavailable failures in the timeout lane without persisting cooldown", async () => {
await withAgentWorkspace(async ({ agentDir, workspaceDir }) => {
await writeAuthStore(agentDir);
runEmbeddedAttemptMock.mockImplementation(async (params: unknown) => {
const attemptParams = params as { provider: string };
if (attemptParams.provider === "openai") {
return makeAttempt({
assistantTexts: [],
lastAssistant: buildAssistant({
provider: "openai",
model: "mock-1",
stopReason: "error",
errorMessage: "LLM error: service unavailable",
}),
});
}
if (attemptParams.provider === "groq") {
return makeAttempt({
assistantTexts: ["fallback ok"],
lastAssistant: buildAssistant({
provider: "groq",
model: "mock-2",
stopReason: "stop",
content: [{ type: "text", text: "fallback ok" }],
}),
});
}
throw new Error(`Unexpected provider ${attemptParams.provider}`);
});
mockPrimaryErrorThenFallbackSuccess("LLM error: service unavailable");
const result = await runEmbeddedFallback({
agentDir,