import fs from "node:fs/promises"; import os from "node:os"; import path from "node:path"; import type { AssistantMessage } from "@mariozechner/pi-ai"; import { beforeEach, describe, expect, it, vi } from "vitest"; import type { ClawdbotConfig } from "../config/config.js"; import type { EmbeddedRunAttemptResult } from "./pi-embedded-runner/run/types.js"; const runEmbeddedAttemptMock = vi.fn, [unknown]>(); vi.mock("./pi-embedded-runner/run/attempt.js", () => ({ runEmbeddedAttempt: (params: unknown) => runEmbeddedAttemptMock(params), })); let runEmbeddedPiAgent: typeof import("./pi-embedded-runner.js").runEmbeddedPiAgent; beforeEach(async () => { vi.useRealTimers(); vi.resetModules(); runEmbeddedAttemptMock.mockReset(); ({ runEmbeddedPiAgent } = await import("./pi-embedded-runner.js")); }); const baseUsage = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }; const buildAssistant = (overrides: Partial): AssistantMessage => ({ role: "assistant", content: [], api: "openai-responses", provider: "openai", model: "mock-1", usage: baseUsage, stopReason: "stop", timestamp: Date.now(), ...overrides, }); const makeAttempt = (overrides: Partial): EmbeddedRunAttemptResult => ({ aborted: false, timedOut: false, promptError: null, sessionIdUsed: "session:test", systemPromptReport: undefined, messagesSnapshot: [], assistantTexts: [], toolMetas: [], lastAssistant: undefined, didSendViaMessagingTool: false, messagingToolSentTexts: [], messagingToolSentTargets: [], cloudCodeAssistFormatError: false, ...overrides, }); const makeConfig = (): ClawdbotConfig => ({ agents: { defaults: { model: { fallbacks: [], }, }, }, models: { providers: { openai: { api: "openai-responses", apiKey: "sk-test", baseUrl: "https://example.com", models: [ { id: "mock-1", name: "Mock 1", reasoning: false, input: ["text"], cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, contextWindow: 16_000, maxTokens: 2048, }, ], }, }, }, }) satisfies ClawdbotConfig; const writeAuthStore = async (agentDir: string, opts?: { includeAnthropic?: boolean }) => { 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" }, ...(opts?.includeAnthropic ? { "anthropic:default": { type: "api_key", provider: "anthropic", key: "sk-anth" } } : {}), }, usageStats: { "openai:p1": { lastUsed: 1 }, "openai:p2": { lastUsed: 2 }, }, }; await fs.writeFile(authPath, JSON.stringify(payload)); }; describe("runEmbeddedPiAgent auth profile rotation", () => { it("rotates for auto-pinned profiles", async () => { const agentDir = await fs.mkdtemp(path.join(os.tmpdir(), "clawdbot-agent-")); const workspaceDir = await fs.mkdtemp(path.join(os.tmpdir(), "clawdbot-workspace-")); try { await writeAuthStore(agentDir); runEmbeddedAttemptMock .mockResolvedValueOnce( makeAttempt({ assistantTexts: [], lastAssistant: buildAssistant({ stopReason: "error", errorMessage: "rate limit", }), }), ) .mockResolvedValueOnce( makeAttempt({ assistantTexts: ["ok"], lastAssistant: buildAssistant({ stopReason: "stop", content: [{ type: "text", text: "ok" }], }), }), ); await runEmbeddedPiAgent({ sessionId: "session:test", sessionKey: "agent:test:auto", sessionFile: path.join(workspaceDir, "session.jsonl"), workspaceDir, agentDir, config: makeConfig(), prompt: "hello", provider: "openai", model: "mock-1", authProfileId: "openai:p1", authProfileIdSource: "auto", timeoutMs: 5_000, runId: "run:auto", }); expect(runEmbeddedAttemptMock).toHaveBeenCalledTimes(2); const stored = JSON.parse( await fs.readFile(path.join(agentDir, "auth-profiles.json"), "utf-8"), ) as { usageStats?: Record }; expect(typeof stored.usageStats?.["openai:p2"]?.lastUsed).toBe("number"); } 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(), "clawdbot-agent-")); const workspaceDir = await fs.mkdtemp(path.join(os.tmpdir(), "clawdbot-workspace-")); try { await writeAuthStore(agentDir); runEmbeddedAttemptMock.mockResolvedValueOnce( makeAttempt({ assistantTexts: [], lastAssistant: buildAssistant({ stopReason: "error", errorMessage: "rate limit", }), }), ); await runEmbeddedPiAgent({ sessionId: "session:test", sessionKey: "agent:test:user", sessionFile: path.join(workspaceDir, "session.jsonl"), workspaceDir, agentDir, config: makeConfig(), prompt: "hello", provider: "openai", model: "mock-1", authProfileId: "openai:p1", authProfileIdSource: "user", timeoutMs: 5_000, runId: "run:user", }); expect(runEmbeddedAttemptMock).toHaveBeenCalledTimes(1); const stored = JSON.parse( await fs.readFile(path.join(agentDir, "auth-profiles.json"), "utf-8"), ) as { usageStats?: Record }; expect(stored.usageStats?.["openai:p2"]?.lastUsed).toBe(2); } finally { await fs.rm(agentDir, { recursive: true, force: true }); await fs.rm(workspaceDir, { recursive: true, force: true }); } }); it("ignores user-locked profile when provider mismatches", async () => { const agentDir = await fs.mkdtemp(path.join(os.tmpdir(), "clawdbot-agent-")); const workspaceDir = await fs.mkdtemp(path.join(os.tmpdir(), "clawdbot-workspace-")); try { await writeAuthStore(agentDir, { includeAnthropic: true }); runEmbeddedAttemptMock.mockResolvedValueOnce( makeAttempt({ assistantTexts: ["ok"], lastAssistant: buildAssistant({ stopReason: "stop", content: [{ type: "text", text: "ok" }], }), }), ); await runEmbeddedPiAgent({ sessionId: "session:test", sessionKey: "agent:test:mismatch", sessionFile: path.join(workspaceDir, "session.jsonl"), workspaceDir, agentDir, config: makeConfig(), prompt: "hello", provider: "openai", model: "mock-1", authProfileId: "anthropic:default", authProfileIdSource: "user", timeoutMs: 5_000, runId: "run:mismatch", }); expect(runEmbeddedAttemptMock).toHaveBeenCalledTimes(1); } finally { await fs.rm(agentDir, { recursive: true, force: true }); await fs.rm(workspaceDir, { recursive: true, force: true }); } }); });