/** * Shared test utilities for coding-agent tests. */ import { chmodSync, existsSync, mkdirSync, readFileSync, rmSync, writeFileSync } from "node:fs"; import { homedir, tmpdir } from "node:os"; import { dirname, join } from "node:path"; import { Agent } from "@mariozechner/pi-agent-core"; import { getModel, getOAuthApiKey, type OAuthCredentials, type OAuthProvider } from "@mariozechner/pi-ai"; import { AgentSession } from "../src/core/agent-session.js"; import { AuthStorage } from "../src/core/auth-storage.js"; import { ModelRegistry } from "../src/core/model-registry.js"; import { SessionManager } from "../src/core/session-manager.js"; import { SettingsManager } from "../src/core/settings-manager.js"; import { codingTools } from "../src/core/tools/index.js"; /** * API key for authenticated tests. Tests using this should be wrapped in * describe.skipIf(!API_KEY) */ export const API_KEY = process.env.ANTHROPIC_OAUTH_TOKEN || process.env.ANTHROPIC_API_KEY; // ============================================================================ // OAuth API key resolution from ~/.pi/agent/auth.json // ============================================================================ const AUTH_PATH = join(homedir(), ".pi", "agent", "auth.json"); type ApiKeyCredential = { type: "api_key"; key: string; }; type OAuthCredentialEntry = { type: "oauth"; } & OAuthCredentials; type AuthCredential = ApiKeyCredential | OAuthCredentialEntry; type AuthStorageData = Record; function loadAuthStorage(): AuthStorageData { if (!existsSync(AUTH_PATH)) { return {}; } try { const content = readFileSync(AUTH_PATH, "utf-8"); return JSON.parse(content); } catch { return {}; } } function saveAuthStorage(storage: AuthStorageData): void { const configDir = dirname(AUTH_PATH); if (!existsSync(configDir)) { mkdirSync(configDir, { recursive: true, mode: 0o700 }); } writeFileSync(AUTH_PATH, JSON.stringify(storage, null, 2), "utf-8"); chmodSync(AUTH_PATH, 0o600); } /** * Resolve API key for a provider from ~/.pi/agent/auth.json * * For API key credentials, returns the key directly. * For OAuth credentials, returns the access token (refreshing if expired and saving back). * * For google-gemini-cli and google-antigravity, returns JSON-encoded { token, projectId } */ export async function resolveApiKey(provider: string): Promise { const storage = loadAuthStorage(); const entry = storage[provider]; if (!entry) return undefined; if (entry.type === "api_key") { return entry.key; } if (entry.type === "oauth") { // Build OAuthCredentials record for getOAuthApiKey const oauthCredentials: Record = {}; for (const [key, value] of Object.entries(storage)) { if (value.type === "oauth") { const { type: _, ...creds } = value; oauthCredentials[key] = creds; } } const result = await getOAuthApiKey(provider as OAuthProvider, oauthCredentials); if (!result) return undefined; // Save refreshed credentials back to auth.json storage[provider] = { type: "oauth", ...result.newCredentials }; saveAuthStorage(storage); return result.apiKey; } return undefined; } /** * Check if a provider has credentials in ~/.pi/agent/auth.json */ export function hasAuthForProvider(provider: string): boolean { const storage = loadAuthStorage(); return provider in storage; } /** Path to the real pi agent config directory */ export const PI_AGENT_DIR = join(homedir(), ".pi", "agent"); /** * Get an AuthStorage instance backed by ~/.pi/agent/auth.json * Use this for tests that need real OAuth credentials. */ export function getRealAuthStorage(): AuthStorage { return new AuthStorage(AUTH_PATH); } /** * Create a minimal user message for testing. */ export function userMsg(text: string) { return { role: "user" as const, content: text, timestamp: Date.now() }; } /** * Create a minimal assistant message for testing. */ export function assistantMsg(text: string) { return { role: "assistant" as const, content: [{ type: "text" as const, text }], api: "anthropic-messages" as const, provider: "anthropic", model: "test", usage: { input: 1, output: 1, cacheRead: 0, cacheWrite: 0, totalTokens: 2, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, stopReason: "stop" as const, timestamp: Date.now(), }; } /** * Options for creating a test session. */ export interface TestSessionOptions { /** Use in-memory session (no file persistence) */ inMemory?: boolean; /** Custom system prompt */ systemPrompt?: string; /** Custom settings overrides */ settingsOverrides?: Record; } /** * Resources returned by createTestSession that need cleanup. */ export interface TestSessionContext { session: AgentSession; sessionManager: SessionManager; tempDir: string; cleanup: () => void; } /** * Create an AgentSession for testing with proper setup and cleanup. * Use this for e2e tests that need real LLM calls. */ export function createTestSession(options: TestSessionOptions = {}): TestSessionContext { const tempDir = join(tmpdir(), `pi-test-${Date.now()}-${Math.random().toString(36).slice(2)}`); mkdirSync(tempDir, { recursive: true }); const model = getModel("anthropic", "claude-sonnet-4-5")!; const agent = new Agent({ getApiKey: () => API_KEY, initialState: { model, systemPrompt: options.systemPrompt ?? "You are a helpful assistant. Be extremely concise.", tools: codingTools, }, }); const sessionManager = options.inMemory ? SessionManager.inMemory() : SessionManager.create(tempDir); const settingsManager = SettingsManager.create(tempDir, tempDir); if (options.settingsOverrides) { settingsManager.applyOverrides(options.settingsOverrides); } const authStorage = new AuthStorage(join(tempDir, "auth.json")); const modelRegistry = new ModelRegistry(authStorage, tempDir); const session = new AgentSession({ agent, sessionManager, settingsManager, modelRegistry, }); // Must subscribe to enable session persistence session.subscribe(() => {}); const cleanup = () => { session.dispose(); if (tempDir && existsSync(tempDir)) { rmSync(tempDir, { recursive: true }); } }; return { session, sessionManager, tempDir, cleanup }; } /** * Build a session tree for testing using SessionManager. * Returns the IDs of all created entries. * * Example tree structure: * ``` * u1 -> a1 -> u2 -> a2 * -> u3 -> a3 (branch from a1) * u4 -> a4 (another root) * ``` */ export function buildTestTree( session: SessionManager, structure: { messages: Array<{ role: "user" | "assistant"; text: string; branchFrom?: string }>; }, ): Map { const ids = new Map(); for (const msg of structure.messages) { if (msg.branchFrom) { const branchFromId = ids.get(msg.branchFrom); if (!branchFromId) { throw new Error(`Cannot branch from unknown entry: ${msg.branchFrom}`); } session.branch(branchFromId); } const id = msg.role === "user" ? session.appendMessage(userMsg(msg.text)) : session.appendMessage(assistantMsg(msg.text)); ids.set(msg.text, id); } return ids; }