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 a2f311ca7..a054377d7 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 @@ -12,6 +12,14 @@ vi.mock("./pi-embedded-runner/run/attempt.js", () => ({ runEmbeddedAttempt: (params: unknown) => runEmbeddedAttemptMock(params), })); +vi.mock("./models-config.js", async (importOriginal) => { + const mod = await importOriginal(); + return { + ...mod, + ensureOpenClawModelsJson: vi.fn(async () => ({ wrote: false })), + }; +}); + let runEmbeddedPiAgent: typeof import("./pi-embedded-runner.js").runEmbeddedPiAgent; beforeAll(async () => { @@ -235,6 +243,19 @@ async function withTimedAgentWorkspace( } } +async function withAgentWorkspace( + run: (ctx: { agentDir: string; workspaceDir: string }) => Promise, +) { + const agentDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-agent-")); + const workspaceDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-workspace-")); + try { + return await run({ agentDir, workspaceDir }); + } finally { + await fs.rm(agentDir, { recursive: true, force: true }); + await fs.rm(workspaceDir, { recursive: true, force: true }); + } +} + async function runTurnWithCooldownSeed(params: { sessionKey: string; runId: string; @@ -288,9 +309,7 @@ describe("runEmbeddedPiAgent auth profile rotation", () => { for (const testCase of cases) { runEmbeddedAttemptMock.mockClear(); - const agentDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-agent-")); - const workspaceDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-workspace-")); - try { + await withAgentWorkspace(async ({ agentDir, workspaceDir }) => { await writeAuthStore(agentDir); mockFailedThenSuccessfulAttempt(testCase.errorMessage); await runAutoPinnedOpenAiTurn({ @@ -302,17 +321,12 @@ describe("runEmbeddedPiAgent auth profile rotation", () => { expect(runEmbeddedAttemptMock).toHaveBeenCalledTimes(2); await expectProfileP2UsageUpdated(agentDir); - } finally { - await fs.rm(agentDir, { recursive: true, force: true }); - await fs.rm(workspaceDir, { recursive: true, force: true }); - } + }); } }); it("does not rotate for compaction timeouts", async () => { - const agentDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-agent-")); - const workspaceDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-workspace-")); - try { + await withAgentWorkspace(async ({ agentDir, workspaceDir }) => { await writeAuthStore(agentDir); runEmbeddedAttemptMock.mockResolvedValueOnce( @@ -348,16 +362,11 @@ describe("runEmbeddedPiAgent auth profile rotation", () => { expect(result.meta.aborted).toBe(true); await expectProfileP2UsageUnchanged(agentDir); - } finally { - await fs.rm(agentDir, { recursive: true, force: true }); - await fs.rm(workspaceDir, { recursive: true, force: true }); - } + }); }); it("does not rotate for user-pinned profiles", async () => { - const agentDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-agent-")); - const workspaceDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-workspace-")); - try { + await withAgentWorkspace(async ({ agentDir, workspaceDir }) => { await writeAuthStore(agentDir); mockSingleErrorAttempt({ errorMessage: "rate limit" }); @@ -380,10 +389,7 @@ describe("runEmbeddedPiAgent auth profile rotation", () => { expect(runEmbeddedAttemptMock).toHaveBeenCalledTimes(1); await expectProfileP2UsageUnchanged(agentDir); - } finally { - await fs.rm(agentDir, { recursive: true, force: true }); - await fs.rm(workspaceDir, { recursive: true, force: true }); - } + }); }); it("honors user-pinned profiles even when in cooldown", async () => { @@ -400,9 +406,7 @@ describe("runEmbeddedPiAgent auth profile rotation", () => { }); it("ignores user-locked profile when provider mismatches", async () => { - const agentDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-agent-")); - const workspaceDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-workspace-")); - try { + await withAgentWorkspace(async ({ agentDir, workspaceDir }) => { await writeAuthStore(agentDir, { includeAnthropic: true }); runEmbeddedAttemptMock.mockResolvedValueOnce( @@ -432,10 +436,7 @@ describe("runEmbeddedPiAgent auth profile rotation", () => { }); expect(runEmbeddedAttemptMock).toHaveBeenCalledTimes(1); - } finally { - await fs.rm(agentDir, { recursive: true, force: true }); - await fs.rm(workspaceDir, { recursive: true, force: true }); - } + }); }); it("skips profiles in cooldown during initial selection", async () => { @@ -486,47 +487,43 @@ describe("runEmbeddedPiAgent auth profile rotation", () => { }); it("fails over when auth is unavailable and fallbacks are configured", async () => { - const agentDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-agent-")); - const workspaceDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-workspace-")); const previousOpenAiKey = process.env.OPENAI_API_KEY; delete process.env.OPENAI_API_KEY; try { - const authPath = path.join(agentDir, "auth-profiles.json"); - await fs.writeFile(authPath, JSON.stringify({ version: 1, profiles: {}, usageStats: {} })); + await withAgentWorkspace(async ({ agentDir, workspaceDir }) => { + const authPath = path.join(agentDir, "auth-profiles.json"); + await fs.writeFile(authPath, JSON.stringify({ version: 1, profiles: {}, usageStats: {} })); - await expect( - runEmbeddedPiAgent({ - sessionId: "session:test", - sessionKey: "agent:test:auth-unavailable", - sessionFile: path.join(workspaceDir, "session.jsonl"), - workspaceDir, - agentDir, - config: makeConfig({ fallbacks: ["openai/mock-2"], apiKey: "" }), - prompt: "hello", - provider: "openai", - model: "mock-1", - authProfileIdSource: "auto", - timeoutMs: 5_000, - runId: "run:auth-unavailable", - }), - ).rejects.toMatchObject({ name: "FailoverError", reason: "auth" }); + await expect( + runEmbeddedPiAgent({ + sessionId: "session:test", + sessionKey: "agent:test:auth-unavailable", + sessionFile: path.join(workspaceDir, "session.jsonl"), + workspaceDir, + agentDir, + config: makeConfig({ fallbacks: ["openai/mock-2"], apiKey: "" }), + prompt: "hello", + provider: "openai", + model: "mock-1", + authProfileIdSource: "auto", + timeoutMs: 5_000, + runId: "run:auth-unavailable", + }), + ).rejects.toMatchObject({ name: "FailoverError", reason: "auth" }); - expect(runEmbeddedAttemptMock).not.toHaveBeenCalled(); + expect(runEmbeddedAttemptMock).not.toHaveBeenCalled(); + }); } finally { if (previousOpenAiKey === undefined) { delete process.env.OPENAI_API_KEY; } else { process.env.OPENAI_API_KEY = previousOpenAiKey; } - await fs.rm(agentDir, { recursive: true, force: true }); - await fs.rm(workspaceDir, { recursive: true, force: true }); } }); it("uses the active erroring model in billing failover errors", async () => { - const agentDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-agent-")); - const workspaceDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-workspace-")); - try { + await withAgentWorkspace(async ({ agentDir, workspaceDir }) => { await writeAuthStore(agentDir); mockSingleErrorAttempt({ errorMessage: "insufficient credits", @@ -564,56 +561,40 @@ describe("runEmbeddedPiAgent auth profile rotation", () => { expect(thrown).toBeInstanceOf(Error); expect((thrown as Error).message).toContain("openai (mock-rotated) returned a billing error"); expect(runEmbeddedAttemptMock).toHaveBeenCalledTimes(1); - } finally { - await fs.rm(agentDir, { recursive: true, force: true }); - await fs.rm(workspaceDir, { recursive: true, force: true }); - } + }); }); it("skips profiles in cooldown when rotating after failure", async () => { - vi.useFakeTimers(); - try { - const agentDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-agent-")); - const workspaceDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-workspace-")); - const now = Date.now(); - vi.setSystemTime(now); + await withTimedAgentWorkspace(async ({ agentDir, workspaceDir, now }) => { + const authPath = path.join(agentDir, "auth-profiles.json"); + const payload = { + version: 1, + profiles: { + "openai:p1": { type: "api_key", provider: "openai", key: "sk-one" }, + "openai:p2": { type: "api_key", provider: "openai", key: "sk-two" }, + "openai:p3": { type: "api_key", provider: "openai", key: "sk-three" }, + }, + usageStats: { + "openai:p1": { lastUsed: 1 }, + "openai:p2": { cooldownUntil: now + 60 * 60 * 1000 }, // p2 in cooldown + "openai:p3": { lastUsed: 3 }, + }, + }; + await fs.writeFile(authPath, JSON.stringify(payload)); - try { - const authPath = path.join(agentDir, "auth-profiles.json"); - const payload = { - version: 1, - profiles: { - "openai:p1": { type: "api_key", provider: "openai", key: "sk-one" }, - "openai:p2": { type: "api_key", provider: "openai", key: "sk-two" }, - "openai:p3": { type: "api_key", provider: "openai", key: "sk-three" }, - }, - usageStats: { - "openai:p1": { lastUsed: 1 }, - "openai:p2": { cooldownUntil: now + 60 * 60 * 1000 }, // p2 in cooldown - "openai:p3": { lastUsed: 3 }, - }, - }; - await fs.writeFile(authPath, JSON.stringify(payload)); + mockFailedThenSuccessfulAttempt("rate limit"); + await runAutoPinnedOpenAiTurn({ + agentDir, + workspaceDir, + sessionKey: "agent:test:rotate-skip-cooldown", + runId: "run:rotate-skip-cooldown", + }); - mockFailedThenSuccessfulAttempt("rate limit"); - await runAutoPinnedOpenAiTurn({ - agentDir, - workspaceDir, - sessionKey: "agent:test:rotate-skip-cooldown", - runId: "run:rotate-skip-cooldown", - }); - - expect(runEmbeddedAttemptMock).toHaveBeenCalledTimes(2); - const usageStats = await readUsageStats(agentDir); - expect(typeof usageStats["openai:p1"]?.lastUsed).toBe("number"); - expect(typeof usageStats["openai:p3"]?.lastUsed).toBe("number"); - expect(usageStats["openai:p2"]?.cooldownUntil).toBe(now + 60 * 60 * 1000); - } finally { - await fs.rm(agentDir, { recursive: true, force: true }); - await fs.rm(workspaceDir, { recursive: true, force: true }); - } - } finally { - vi.useRealTimers(); - } + expect(runEmbeddedAttemptMock).toHaveBeenCalledTimes(2); + const usageStats = await readUsageStats(agentDir); + expect(typeof usageStats["openai:p1"]?.lastUsed).toBe("number"); + expect(typeof usageStats["openai:p3"]?.lastUsed).toBe("number"); + expect(usageStats["openai:p2"]?.cooldownUntil).toBe(now + 60 * 60 * 1000); + }); }); });