pi-ext

import { mkdtempSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import path from "node:path";
import {
	type TestSession as BaseTestSession,
	type TestSessionOptions as BaseTestSessionOptions,
	createTestSession as createBaseTestSession,
	type SandboxOptions,
	type SandboxResult,
	verifySandboxInstall as verifyBaseSandboxInstall,
} from "@marcfargas/pi-test-harness";

export type {
	MockPi,
	MockPiCall,
	MockToolHandler,
	MockUIConfig,
	PlaybookAction,
	SandboxOptions,
	SandboxResult,
	TestEvents,
	ToolCallRecord,
	ToolResultRecord,
	Turn,
	UICallRecord,
} from "@marcfargas/pi-test-harness";
export {
	calls,
	createMockPi,
	safeRmSync,
	says,
	ToolBlockedError,
	when,
} from "@marcfargas/pi-test-harness";

export interface HarnessEnvOptions {
	/** Explicit env additions. Undefined deletes the key. */
	env?: Record<string, string | undefined>;
	/** Ambient env keys to copy intentionally. */
	inheritEnv?: string[];
}

export type HarnessTestSessionOptions = BaseTestSessionOptions &
	HarnessEnvOptions;
export type HarnessTestSession = BaseTestSession & { env: NodeJS.ProcessEnv };
export type TestSession = HarnessTestSession;
export type TestSessionOptions = HarnessTestSessionOptions;

const BASE_ENV_KEYS = [
	"PATH",
	"PATHEXT",
	"SystemRoot",
	"WINDIR",
	"COMSPEC",
	"TMPDIR",
	"TMP",
	"TEMP",
	"SHELL",
	"TERM",
	"COLORTERM",
	"NO_COLOR",
	"FORCE_COLOR",
	"CI",
	"VITEST",
	"VITEST_POOL_ID",
	"VITEST_WORKER_ID",
	"NODE_OPTIONS",
];

export function sandboxEnv(
	root: string,
	options: HarnessEnvOptions = {},
): NodeJS.ProcessEnv {
	const env: NodeJS.ProcessEnv = {};
	for (const key of BASE_ENV_KEYS) copyIfSet(env, process.env, key);
	for (const key of options.inheritEnv ?? []) copyIfSet(env, process.env, key);

	env.HOME = path.join(root, ".test-home");
	env.USERPROFILE = env.HOME;
	env.XDG_CONFIG_HOME = path.join(root, ".test-config");
	env.XDG_CACHE_HOME = path.join(root, ".test-cache");
	env.XDG_STATE_HOME = path.join(root, ".test-state");
	env.XDG_DATA_HOME = path.join(root, ".test-data");

	for (const [key, value] of Object.entries(options.env ?? {})) {
		setEnvValue(env, key, value);
	}
	return env;
}

export async function createTestSession(
	options: HarnessTestSessionOptions = {},
): Promise<HarnessTestSession> {
	const { env: envOverrides, inheritEnv, ...baseOptions } = options;
	const ownsCwd = !baseOptions.cwd;
	const cwd =
		baseOptions.cwd ?? mkdtempSync(path.join(tmpdir(), "pi-test-harness-"));
	const originalEnv = { ...process.env };
	const env = sandboxEnv(cwd, { env: envOverrides, inheritEnv });
	replaceProcessEnv(env);
	let session: BaseTestSession;
	try {
		session = await createBaseTestSession({ ...baseOptions, cwd });
	} catch (error) {
		replaceProcessEnv(originalEnv);
		if (ownsCwd) rmSync(cwd, { recursive: true, force: true });
		throw error;
	}
	const wrapped = Object.assign(session, {
		env: process.env,
	}) as HarnessTestSession;
	wrapSessionEnv(wrapped);
	const dispose = session.dispose.bind(session);
	wrapped.dispose = () => {
		dispose();
		if (ownsCwd) rmSync(cwd, { recursive: true, force: true });
		replaceProcessEnv(originalEnv);
	};
	return wrapped;
}

export async function verifySandboxInstall(
	options: SandboxOptions & HarnessEnvOptions,
): Promise<SandboxResult> {
	const { env, inheritEnv, ...baseOptions } = options;
	return await withProcessEnv(
		sandboxEnv(baseOptions.packageDir, { env, inheritEnv }),
		() => verifyBaseSandboxInstall(baseOptions),
	);
}

function wrapSessionEnv(session: HarnessTestSession): void {
	const run = session.run.bind(session);
	session.run = (...turns) => withProcessEnv(session.env, () => run(...turns));

	const prompt = session.session.prompt?.bind(session.session);
	if (prompt) {
		session.session.prompt = (...args: unknown[]) =>
			withProcessEnv(session.env, () => prompt(...args));
	}

	const extensionRunner = session.session.extensionRunner;
	const emit = extensionRunner?.emit?.bind(extensionRunner);
	if (emit) {
		extensionRunner.emit = (...args: unknown[]) =>
			withProcessEnv(session.env, () => emit(...args));
	}
	const emitBoundary = extensionRunner?.emitBoundary?.bind(extensionRunner);
	if (emitBoundary) {
		extensionRunner.emitBoundary = (...args: unknown[]) =>
			withProcessEnv(session.env, () => emitBoundary(...args));
	}

	for (const extension of extensionRunner?.extensions ?? []) {
		for (const command of extension.commands?.values?.() ?? []) {
			if (typeof command.handler !== "function") continue;
			const handler = command.handler.bind(command);
			command.handler = (...args: unknown[]) =>
				withProcessEnv(session.env, () => handler(...args));
		}
		for (const tool of extension.tools?.values?.() ?? []) {
			const definition = tool.definition;
			if (typeof definition?.execute !== "function") continue;
			const execute = definition.execute.bind(definition);
			definition.execute = (...args: unknown[]) =>
				withProcessEnv(session.env, () => execute(...args));
		}
	}
}

export async function withProcessEnv<T>(
	env: NodeJS.ProcessEnv,
	fn: () => T | Promise<T>,
): Promise<T> {
	const original = { ...process.env };
	const next = { ...env };
	replaceProcessEnv(next);
	try {
		return await fn();
	} finally {
		replaceProcessEnv(original);
	}
}

export function withProcessEnvSync<T>(env: NodeJS.ProcessEnv, fn: () => T): T {
	const original = { ...process.env };
	const next = { ...env };
	replaceProcessEnv(next);
	try {
		return fn();
	} finally {
		replaceProcessEnv(original);
	}
}

function replaceProcessEnv(env: NodeJS.ProcessEnv): void {
	for (const key of Object.keys(process.env)) delete process.env[key];
	Object.assign(process.env, env);
}

function copyIfSet(
	to: NodeJS.ProcessEnv,
	from: NodeJS.ProcessEnv,
	key: string,
): void {
	const sourceKey =
		process.platform === "win32"
			? Object.keys(from).find(
					(candidate) => candidate.toLowerCase() === key.toLowerCase(),
				)
			: key;
	if (!sourceKey) return;
	setEnvValue(to, sourceKey, from[sourceKey]);
}

function setEnvValue(
	env: NodeJS.ProcessEnv,
	key: string,
	value: string | undefined,
): void {
	if (process.platform === "win32") {
		for (const candidate of Object.keys(env)) {
			if (candidate.toLowerCase() === key.toLowerCase()) delete env[candidate];
		}
	} else {
		delete env[key];
	}
	if (value !== undefined) env[key] = value;
}