Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
11 kB · 357 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358/** * 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 "@earendil-works/pi-agent-core";import { getModel, type OAuthCredentials, type OAuthProvider } from "@earendil-works/pi-ai";import { getOAuthApiKey } from "@earendil-works/pi-ai/oauth";import { AgentSession } from "../src/core/agent-session.js";import { AuthStorage } from "../src/core/auth-storage.js";import { createEventBus } from "../src/core/event-bus.js";import type { Extension, ExtensionFactory, LoadExtensionsResult } from "../src/core/extensions/index.js";import { createExtensionRuntime, loadExtensionFromFactory } from "../src/core/extensions/loader.js";import { ModelRegistry } from "../src/core/model-registry.js";import type { ResourceLoader } from "../src/core/resource-loader.js";import { SessionManager } from "../src/core/session-manager.js";import { SettingsManager } from "../src/core/settings-manager.js";import type { Skill } from "../src/core/skills.js";import { createIpythonTool } from "../src/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;
/** * Overlay for fixture spawns that must not inherit this process's worker identity. A child that * keeps `PRIME_AGENT_INTERNAL_DAEMON_WORKER=1` boots as a daemon worker and speaks private * framing, so a suite that spawns a real supervisor can never complete its public handshake; * `RLM_*` bind the child to this session's artifacts instead of its fixture directory. */export function fixtureEnv(): NodeJS.ProcessEnv { return { PRIME_AGENT_INTERNAL_LEGACY_OWNED_WORKER_FRONTEND: "0", PRIME_AGENT_INTERNAL_DAEMON_WORKER: undefined, PRIME_AGENT_INTERNAL_DAEMON_WORKER_TOKEN: undefined, PRIME_AGENT_INTERNAL_DAEMON_WORKER_ACTIVE_SESSION_ID: undefined, PRIME_AGENT_INTERNAL_DAEMON_SUPERVISOR_SOCKET: undefined, PRIME_AGENT_INTERNAL_DAEMON_WORKER_RECOVERY_JOURNAL: undefined, PRIME_AGENT_INTERNAL_ORPHAN_PROCESS_JOURNAL: undefined, PRIME_AGENT_INTERNAL_SESSION_LEASES: undefined, PRIME_AGENT_INTERNAL_SESSION_LEASE_OWNER_ID: undefined, PRIME_AGENT_KERNEL_OWNER_PID: undefined, RLM_DEPTH: undefined, RLM_MAX_DEPTH: undefined, RLM_SESSION_DIR: undefined, RLM_HARNESS_STATE_DIR: undefined, RLM_GLOBAL_HARNESS_STATE_DIR: undefined, };}
// ============================================================================// 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<string, AuthCredential>;
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). * */export async function resolveApiKey(provider: string): Promise<string | undefined> { 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<string, OAuthCredentials> = {}; 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 AuthStorage.create(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<string, unknown>;}
/** * Resources returned by createTestSession that need cleanup. */export interface TestSessionContext { session: AgentSession; sessionManager: SessionManager; tempDir: string; cleanup: () => void;}
export interface CreateTestExtensionsResultInput { factory: ExtensionFactory; path?: string;}
export async function createTestExtensionsResult( inputs: Array<ExtensionFactory | CreateTestExtensionsResultInput>, cwd = process.cwd(),): Promise<LoadExtensionsResult> { const runtime = createExtensionRuntime(); const eventBus = createEventBus(); const extensions: Extension[] = [];
for (const [index, input] of inputs.entries()) { const factory = typeof input === "function" ? input : input.factory; const extensionPath = typeof input === "function" ? `<inline:${index + 1}>` : (input.path ?? `<inline:${index + 1}>`); extensions.push(await loadExtensionFromFactory(factory, cwd, eventBus, runtime, extensionPath)); }
return { extensions, errors: [], runtime, };}
export interface CreateTestResourceLoaderOptions { extensionsResult?: LoadExtensionsResult; skills?: Skill[];}
export function createTestResourceLoader(options: CreateTestResourceLoaderOptions = {}): ResourceLoader { const extensionsResult = options.extensionsResult ?? { extensions: [], errors: [], runtime: createExtensionRuntime(), };
return { getExtensions: () => extensionsResult, getSkills: () => ({ skills: options.skills ?? [], diagnostics: [] }), getPrompts: () => ({ prompts: [], diagnostics: [] }), getThemes: () => ({ themes: [], diagnostics: [] }), getAgentsFiles: () => ({ agentsFiles: [] }), getSystemPrompt: () => undefined, getAppendSystemPrompt: () => [], extendResources: () => {}, reload: async () => {}, };}
/** * 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: [createIpythonTool(process.cwd())], }, });
const sessionManager = options.inMemory ? SessionManager.inMemory() : SessionManager.create(tempDir); const settingsManager = SettingsManager.create(tempDir, tempDir);
if (options.settingsOverrides) { settingsManager.applyOverrides(options.settingsOverrides); }
const authStorage = AuthStorage.create(join(tempDir, "auth.json")); const modelRegistry = ModelRegistry.create(authStorage, tempDir);
const session = new AgentSession({ agent, sessionManager, settingsManager, cwd: tempDir, modelRegistry, resourceLoader: createTestResourceLoader(), });
// 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<string, string> { const ids = new Map<string, string>();
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;}
/** * A permissive role-routing policy for tests that exercise spawn lifecycle rather * than routing: every listed selector carries every tag any role asks for, and the * price and context floors are set so nothing is filtered by accident. */export function writeTestRlmModels(dir: string, selectors: readonly string[]): string { const toml = `scouts = ${JSON.stringify(selectors)}\ncoders = ${JSON.stringify(selectors)}\n`; const path = join(dir, "rlm-models.toml"); writeFileSync(path, toml, "utf8"); return path;}
export const writeTestRlmPolicy = writeTestRlmModels;