pi-ext

// Repository-owned test fixture built only on public Pi APIs.
//
// It creates a real Pi session through `createAgentSession`, loads extension
// entrypoints through the real `DefaultResourceLoader`, and drives model turns
// with a deterministic scripted provider registered through the public
// `ModelRuntime.registerProvider`. It owns its temporary directories, a
// sandboxed environment (HOME, XDG_*, PI_CODING_AGENT_DIR), and shutdown, and
// collects extension load/bind errors so tests can fail on unexpected ones.
//
// Unlike the patched `@marcfargas/pi-test-harness`, this reaches no private
// agent internals (`streamFn`, `getApiKey`, `_modelRegistry`) and needs no
// install-time source patching.
//
// The sandboxed environment is applied to `process.env` for the session's
// lifetime and restored on `dispose`, so contract tests using this fixture must
// run in an isolated Vitest project (one file per worker process).
import { mkdirSync, mkdtempSync, writeFileSync } from "node:fs";
import { rm } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import {
	type CredentialStore,
	createFauxCore,
	type FauxProviderState,
	type FauxResponseStep,
	InMemoryModelsStore,
} from "@earendil-works/pi-ai";
import {
	createAgentSession,
	DefaultResourceLoader,
	type Extension,
	type ExtensionFactory,
	ModelRuntime,
	SessionManager,
} from "@earendil-works/pi-coding-agent";

const FAUX_PROVIDER = "faux";
const FAUX_MODEL = "faux-model";

export interface PiSessionOptions {
	/** Absolute paths to extension entrypoints to load (path-loaded contract). */
	extensionPaths?: string[];
	/**
	 * In-process extension factories loaded in the test's own module graph, so
	 * `vi.mock`/`vi.spyOn` of an extension's internal modules is observable.
	 * Prefer `extensionPaths`; use factories only for component tests that must
	 * intercept an extension's own lazy imports.
	 */
	extensionFactories?: ExtensionFactory[];
	/** Working directory; a temporary one is created and owned when omitted. */
	cwd?: string;
	/** User settings written to the isolated agent directory before load. */
	userSettings?: Record<string, unknown>;
	/** Credential storage override for tests that do not need file-backed auth. */
	credentials?: CredentialStore;
	/** Extra environment overrides applied on top of the sandbox. */
	env?: Record<string, string | undefined>;
	/** Initial scripted assistant responses for the deterministic provider. */
	responses?: FauxResponseStep[];
}

export interface PiSession {
	session: Awaited<ReturnType<typeof createAgentSession>>["session"];
	/** Loaded extensions via the public loader API (never runner internals). */
	extensions: Extension[];
	cwd: string;
	agentDir: string;
	/** The sandboxed environment applied to `process.env` for this session. */
	env: NodeJS.ProcessEnv;
	/** Extension load, bind, and shutdown errors. */
	errors: string[];
	/** Replace the scripted assistant responses before the next turn. */
	setResponses: (responses: FauxResponseStep[]) => void;
	/** Deterministic provider call state, e.g. `callCount`. */
	providerState: FauxProviderState;
	dispose: () => Promise<void>;
}

function replaceProcessEnv(env: NodeJS.ProcessEnv): void {
	for (const key of Object.keys(process.env)) delete process.env[key];
	for (const [key, value] of Object.entries(env))
		if (value !== undefined) process.env[key] = value;
}

export async function createPiSession(
	options: PiSessionOptions,
): Promise<PiSession> {
	const ownsCwd = !options.cwd;
	const cwd = options.cwd ?? mkdtempSync(join(tmpdir(), "pi-fixture-"));
	const agentDir = join(cwd, "agent");
	mkdirSync(agentDir, { recursive: true });
	if (options.userSettings)
		writeFileSync(
			join(agentDir, "settings.json"),
			`${JSON.stringify(options.userSettings, null, "\t")}\n`,
		);

	const originalEnv = { ...process.env };
	const env: NodeJS.ProcessEnv = {
		...process.env,
		HOME: join(cwd, ".home"),
		USERPROFILE: join(cwd, ".home"),
		APPDATA: options.env?.["XDG_CONFIG_HOME"] ?? join(cwd, ".config"),
		LOCALAPPDATA: options.env?.["XDG_STATE_HOME"] ?? join(cwd, ".test-state"),
		XDG_CONFIG_HOME: join(cwd, ".config"),
		XDG_CACHE_HOME: join(cwd, ".cache"),
		XDG_STATE_HOME: join(cwd, ".test-state"),
		XDG_DATA_HOME: join(cwd, ".data"),
		PI_CODING_AGENT_DIR: agentDir,
		...options.env,
	};
	for (const [key, value] of Object.entries(options.env ?? {}))
		if (value === undefined) delete env[key];
	replaceProcessEnv(env);

	const errors: string[] = [];
	let disposed = false;
	try {
		const loader = new DefaultResourceLoader({
			cwd,
			agentDir,
			noExtensions: true,
			noSkills: true,
			noPromptTemplates: true,
			noThemes: true,
			noContextFiles: true,
			additionalExtensionPaths: options.extensionPaths ?? [],
			...(options.extensionFactories
				? { extensionFactories: options.extensionFactories }
				: {}),
		});
		await loader.reload();
		const loaded = loader.getExtensions();
		errors.push(
			...loaded.errors.map((error) => `${error.path}: ${error.error}`),
		);

		const core = createFauxCore({
			api: "openai",
			provider: FAUX_PROVIDER,
			models: [{ id: FAUX_MODEL, name: "Faux Model" }],
		});
		if (options.responses) core.setResponses(options.responses);

		const modelRuntime = await ModelRuntime.create({
			...(options.credentials ? { credentials: options.credentials } : {}),
			modelsStore: new InMemoryModelsStore(),
			refreshOnCreate: false,
		});
		modelRuntime.registerProvider(FAUX_PROVIDER, {
			apiKey: "test-key",
			api: "openai",
			baseUrl: "https://faux.invalid",
			streamSimple: core.streamSimple,
			models: [
				{
					id: FAUX_MODEL,
					name: "Faux Model",
					reasoning: false,
					input: ["text"],
					cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
					contextWindow: 200_000,
					maxTokens: 8192,
				},
			],
		});
		const model = modelRuntime.getModel(FAUX_PROVIDER, FAUX_MODEL);
		if (!model)
			throw new Error("The owned fixture failed to register its model.");

		const { session } = await createAgentSession({
			cwd,
			agentDir,
			model,
			modelRuntime,
			resourceLoader: loader,
			sessionManager: SessionManager.inMemory(cwd),
		});
		await session.bindExtensions({
			mode: "print",
			onError: (error) =>
				errors.push(`${error.extensionPath}: ${error.event}: ${error.error}`),
		});

		return {
			session,
			extensions: loaded.extensions,
			cwd,
			agentDir,
			env,
			errors,
			setResponses: core.setResponses,
			providerState: core.state,
			dispose: async () => {
				if (disposed) return;
				disposed = true;
				const previousErrors = errors.length;
				try {
					await session.extensionRunner.emit({
						type: "session_shutdown",
						reason: "quit",
					});
					session.dispose();
					await session.settingsManager.flush();
					if (errors.length > previousErrors)
						throw new Error(
							`Fixture shutdown failed: ${errors.slice(previousErrors).join("; ")}`,
						);
				} finally {
					replaceProcessEnv(originalEnv);
					// Windows retains a child's cwd briefly after SDK shutdown returns.
					if (ownsCwd)
						await rm(cwd, {
							recursive: true,
							force: true,
							maxRetries: 10,
							retryDelay: 100,
						});
				}
			},
		};
	} catch (error) {
		replaceProcessEnv(originalEnv);
		if (ownsCwd)
			await rm(cwd, {
				recursive: true,
				force: true,
				maxRetries: 10,
				retryDelay: 100,
			});
		throw error;
	}
}