pi-ext

import {
	chmodSync,
	mkdirSync,
	mkdtempSync,
	readdirSync,
	readFileSync,
	rmSync,
	writeFileSync,
} from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { pathToFileURL } from "node:url";
import { Agent } from "@earendil-works/pi-agent-core";
import { createAssistantMessageEventStream } from "@earendil-works/pi-ai";
import {
	DefaultResourceLoader,
	estimateTokens,
	formatSkillsForPrompt,
	loadSkills,
	SettingsManager,
} from "@earendil-works/pi-coding-agent";
import { afterEach, expect, it, vi } from "vitest";
import {
	CASES,
	comparable,
	formatOverview,
	measure,
	PROMPT,
	snapshot,
	summarize,
} from "./context-startup.mjs";
import {
	estimateSourceTokens,
	installedPiDist,
	observeRequests,
	publish,
	resourceOptions,
} from "./context-startup-session.mjs";

const message = {
	role: "assistant",
	provider: "test",
	model: "test-model",
	stopReason: "stop",
	usage: { input: 100, cacheRead: 200, cacheWrite: 30, output: 2 },
};
const event = (value = message) =>
	JSON.stringify({ type: "message_end", message: value });
const metadata = {
	cwd: "/project",
	version: "test",
	provider: "test",
	model: "test-model",
	thinking: "high",
	contextWindow: 100000,
};
let directory: string | undefined;
afterEach(() => {
	vi.unstubAllEnvs();
	if (directory) rmSync(directory, { recursive: true, force: true });
	directory = undefined;
});

it("locates installed Pi from PATH without executing it", async () => {
	directory = mkdtempSync(join(tmpdir(), "pi-locator-test-"));
	const bin = join(directory, "bin");
	mkdirSync(bin);
	writeFileSync(
		join(directory, "package.json"),
		JSON.stringify({ name: "@earendil-works/pi-coding-agent" }),
	);
	const executable = join(bin, "pi");
	writeFileSync(executable, "This file must never execute.");
	chmodSync(executable, 0o755);
	vi.stubEnv("PATH", bin);
	expect(await installedPiDist()).toBe(join(directory, "dist"));
});
it("reports missing installed Pi without a repository SDK fallback", async () => {
	directory = mkdtempSync(join(tmpdir(), "pi-locator-test-"));
	vi.stubEnv("PATH", directory);
	await expect(installedPiDist()).rejects.toThrow(
		"Could not locate installed Pi",
	);
});

it("estimates global appendix and visible skill catalog, excluding bodies and rejecting incomplete catalogs", async () => {
	directory = mkdtempSync(join(tmpdir(), "pi-overview-test-"));
	const sdk = { estimateTokens, formatSkillsForPrompt };
	const measureSources = () =>
		estimateSourceTokens(
			directory,
			sdk,
			loadSkills({
				cwd: directory,
				agentDir: directory,
				skillPaths: [join(directory, "skills")],
				includeDefaults: false,
			}),
		);
	writeFileSync(join(directory, "APPEND_SYSTEM.md"), "漢漢漢漢");
	mkdirSync(join(directory, "skills"));
	expect(await measureSources()).toEqual({
		appendix: 1,
		descriptions: [],
		piAgent: { tokens: 0, count: 0 },
		extensions: { tokens: 0, count: 0 },
		shared: 0,
		ignoredProjectSkills: 0,
		catalog: 0,
		total: 1,
	});
	const visible = join(directory, "skills", "visible.md");
	const frontmatter = "---\nname: visible\ndescription: Read & write\n---\n";
	writeFileSync(visible, frontmatter + "Body");
	const result = await measureSources();
	const expectedCatalog = formatSkillsForPrompt([
		{
			name: "visible",
			description: "Read & write",
			filePath: visible,
			baseDir: join(directory, "skills"),
			disableModelInvocation: false,
			sourceInfo: {
				path: visible,
				source: "user",
				scope: "user",
				origin: "top-level",
			},
		},
	]);
	const expectedEntry = expectedCatalog.match(
		/ {2}<skill>[\s\S]*? {2}<\/skill>\n/,
	)?.[0];
	if (!expectedEntry) throw new Error("Fixture catalog missing skill entry");
	expect(result).toEqual({
		appendix: 1,
		descriptions: [{ name: "visible", description: "Read & write", tokens: 3 }],
		piAgent: { tokens: Math.ceil(expectedEntry.length / 4), count: 1 },
		extensions: { tokens: 0, count: 0 },
		shared: Math.ceil((expectedCatalog.length - expectedEntry.length) / 4),
		ignoredProjectSkills: 0,
		catalog: Math.ceil(expectedCatalog.length / 4),
		total: 1 + Math.ceil(expectedCatalog.length / 4),
	});
	writeFileSync(visible, frontmatter + "Long skill body ".repeat(1000));
	writeFileSync(
		join(directory, "skills", "hidden.md"),
		"---\nname: hidden\ndescription: Not injected\ndisable-model-invocation: true\n---\nHidden body",
	);
	expect(await measureSources()).toEqual(result);
	writeFileSync(
		join(directory, "skills", "invalid.md"),
		"---\nname: INVALID\ndescription: Invalid uppercase name\n---\nBody",
	);
	await expect(measureSources()).rejects.toThrow(
		"Skill catalog has diagnostics",
	);
});

it("ranks only visible agent-dir descriptions by length, then name", async () => {
	directory = mkdtempSync(join(tmpdir(), "pi-overview-test-"));
	writeFileSync(join(directory, "APPEND_SYSTEM.md"), "");
	const skill = (
		name: string,
		description: string,
		sourceInfo = { source: "user", scope: "user", origin: "top-level" },
	) => ({
		name,
		description,
		filePath: join(directory, "skills", `${name}.md`),
		disableModelInvocation: false,
		sourceInfo,
	});
	const result = await estimateSourceTokens(
		directory,
		{ estimateTokens, formatSkillsForPrompt },
		{
			diagnostics: [],
			skills: [
				skill("z", "漢漢漢漢"),
				skill("short", "abc"),
				skill("long", "abcde"),
				skill("a", "1234"),
				{ ...skill("hidden", "x".repeat(100)), disableModelInvocation: true },
				skill("package", "x".repeat(100), {
					source: "fixture",
					scope: "user",
					origin: "package",
				}),
				{
					...skill("project", "x".repeat(100), {
						source: "project",
						scope: "project",
						origin: "top-level",
					}),
					filePath: join(directory, "project", "SKILL.md"),
				},
			],
		},
	);
	expect(
		result.descriptions.map(({ name, tokens }) => ({ name, tokens })),
	).toEqual([
		{ name: "long", tokens: 2 },
		{ name: "a", tokens: 1 },
		{ name: "z", tokens: 1 },
		{ name: "short", tokens: 1 },
	]);
});

it("counts cached input exactly once and ignores notices, output, and other events", () => {
	expect(summarize(`notice\nnull\n{}\n${event()}\n`)).toEqual({
		provider: "test",
		model: "test-model",
		inputTokens: 330,
		uncachedInputTokens: 100,
		cacheReadTokens: 200,
		cacheWriteTokens: 30,
		outputTokens: 2,
	});
	expect(
		summarize(
			event({
				...message,
				usage: { input: 100, cacheRead: 0, cacheWrite: 0, output: 2 },
			}),
		).inputTokens,
	).toBe(100);
});
it.each(["", `${event()}\n${event()}`])(
	"rejects missing or multiple responses",
	(stdout) => {
		expect(() => summarize(stdout)).toThrow("Expected one assistant response");
	},
);
it.each(["error", "aborted", "toolUse", "length"])(
	"rejects %s responses",
	(stopReason) => {
		expect(() => summarize(event({ ...message, stopReason }))).toThrow(
			"measurement rejected",
		);
	},
);
it.each([
	undefined,
	{ input: 0, cacheRead: 0, cacheWrite: 0, output: 1 },
	{ ...message.usage, input: -1 },
	{ ...message.usage, cacheRead: "200" },
])("rejects invalid usage", (usage) => {
	expect(() =>
		summarize(event({ ...message, usage } as typeof message)),
	).toThrow("provider token usage");
});
it("uses Pi's character heuristic for system and tool definitions, excluding implementation and display fields", () => {
	const parameters = { type: "object" };
	expect(
		snapshot(
			{
				systemPrompt: "漢漢漢漢",
				tools: [
					{
						name: "read",
						description: "Read",
						parameters,
						label: "SECRET",
						execute() {},
					},
				],
				messages: [],
			},
			estimateTokens,
		),
	).toEqual({
		system: 1,
		tools: Math.ceil(
			JSON.stringify([{ name: "read", description: "Read", parameters }])
				.length / 4,
		),
		messages: 0,
		toolCount: 1,
	});
});
it("uses Pi's per-message rounding and image allowance, excluding message metadata", () => {
	expect(snapshot({}, estimateTokens)).toEqual({
		system: 0,
		tools: 0,
		messages: 0,
		toolCount: 0,
	});
	const messages = [
		{ role: "user", content: "a", timestamp: 0 },
		{
			role: "user",
			content: [
				{ type: "text", text: "abcde" },
				{ type: "image", data: "ignored", mimeType: "image/png" },
			],
			timestamp: 0,
		},
	];
	expect(snapshot({ messages }, estimateTokens).messages).toBe(1203);
	expect(snapshot({ systemPrompt: "abcde" }, estimateTokens).system).toBe(2);
});
it("baseline prevents every discovery route including SYSTEM and APPEND, configured preserves defaults", () => {
	const extensions = [];
	expect(resourceOptions("baseline+1msg", extensions)).toEqual({
		noExtensions: true,
		noSkills: true,
		noPromptTemplates: true,
		noThemes: true,
		noContextFiles: true,
		systemPrompt: "",
		appendSystemPrompt: [],
		extensionFactories: extensions,
	});
	expect(resourceOptions("headless+1msg", extensions)).toEqual({
		extensionFactories: extensions,
	});
});
it.each([
	{ projectTrusted: true, projectPrompts: false },
	{ projectTrusted: true, projectPrompts: true },
	{ projectTrusted: false, projectPrompts: true },
])(
	"removes only agent-dir text resources through Pi's real loader: %j",
	async ({ projectTrusted, projectPrompts }) => {
		directory = mkdtempSync(join(tmpdir(), "pi-overview-test-"));
		vi.stubEnv("HOME", directory);
		vi.stubEnv("USERPROFILE", directory);
		const cwd = join(directory, "project");
		const agentDir = join(directory, "agent");
		const packageDir = join(agentDir, "packages", "fixture");
		const put = (path: string, content: string) => {
			mkdirSync(join(path, ".."), { recursive: true });
			writeFileSync(path, content);
		};
		const skill = (name: string) =>
			`---\nname: ${name}\ndescription: Fixture skill\n---\nFixture body\n`;
		put(join(agentDir, "AGENTS.md"), "Global instructions");
		put(join(agentDir, "SYSTEM.md"), "Global system");
		put(join(agentDir, "APPEND_SYSTEM.md"), "Global append");
		put(join(agentDir, "skills", "global-fixture.md"), skill("global-fixture"));
		put(join(agentDir, "prompts", "global-fixture.md"), "Global template");
		put(join(cwd, "AGENTS.md"), "Project instructions");
		mkdirSync(join(cwd, ".git"));
		put(
			join(cwd, ".pi", "skills", "project-fixture.md"),
			skill("project-fixture"),
		);
		put(join(cwd, ".pi", "prompts", "project-fixture.md"), "Project template");
		if (projectPrompts) {
			put(join(cwd, ".pi", "SYSTEM.md"), "Project system");
			put(join(cwd, ".pi", "APPEND_SYSTEM.md"), "Project append");
		}
		const settingsPath = join(agentDir, "settings.json");
		const settingsText = JSON.stringify({
			packages: [packageDir],
			contextFixture: "unchanged",
		});
		put(settingsPath, settingsText);
		put(
			join(packageDir, "package.json"),
			JSON.stringify({
				name: "context-fixture",
				pi: {
					extensions: ["index.mjs"],
					skills: ["skills"],
					prompts: ["prompts"],
				},
			}),
		);
		put(
			join(packageDir, "index.mjs"),
			`
		import { readFileSync } from 'node:fs';
		export default pi => pi.registerFlag('fixture-setting', {
			type: 'boolean', description: JSON.parse(readFileSync(${JSON.stringify(settingsPath)}, 'utf8')).contextFixture
		});
	`,
		);
		put(
			join(packageDir, "skills", "extension-fixture.md"),
			skill("extension-fixture"),
		);
		put(
			join(packageDir, "prompts", "extension-fixture.md"),
			"Extension template",
		);
		const settingsManager = SettingsManager.create(cwd, agentDir, {
			projectTrusted,
		});
		const load = async (kind: string) => {
			const loader = new DefaultResourceLoader({
				cwd,
				agentDir,
				settingsManager,
				...resourceOptions(kind, [], { cwd, agentDir, projectTrusted }),
			});
			await loader.reload();
			expect(loader.getExtensions().errors).toEqual([]);
			return loader;
		};
		const normal = await load("headless+1msg");
		const stripped = await load("headless+1msg+no_pi_config");
		const dynamicSkill = join(directory, "dynamic-fixture.md");
		put(dynamicSkill, skill("dynamic-fixture"));
		normal.extendResources({
			skillPaths: [
				{
					path: dynamicSkill,
					metadata: {
						source: "extension:fixture",
						scope: "temporary",
						origin: "top-level",
					},
				},
			],
		});
		const sources = await estimateSourceTokens(
			agentDir,
			{ estimateTokens, formatSkillsForPrompt },
			normal.getSkills(),
		);
		expect(sources.piAgent.count).toBe(1);
		expect(sources.extensions.count).toBe(2);
		expect(sources.ignoredProjectSkills).toBe(projectTrusted ? 1 : 0);
		const included = normal
			.getSkills()
			.skills.filter((skill) => skill.sourceInfo.scope !== "project");
		expect(sources.catalog).toBe(
			Math.ceil(formatSkillsForPrompt(included).length / 4),
		);
		expect(sources.shared).toBeGreaterThan(0);
		expect(sources.extensions.tokens).toBeGreaterThan(0);
		const extensionOnly = await estimateSourceTokens(
			agentDir,
			{ estimateTokens, formatSkillsForPrompt },
			{
				skills: included.filter(
					(skill) =>
						skill.sourceInfo.origin === "package" ||
						skill.sourceInfo.source.startsWith("extension:"),
				),
				diagnostics: [],
			},
		);
		expect(extensionOnly.shared).toBe(sources.shared);
		expect(extensionOnly.piAgent).toEqual({ tokens: 0, count: 0 });
		expect(extensionOnly.extensions).toEqual(sources.extensions);
		await expect(
			estimateSourceTokens(
				agentDir,
				{ estimateTokens, formatSkillsForPrompt: () => "changed" },
				normal.getSkills(),
			),
		).rejects.toThrow("Pi skill catalog format changed");
		const unknown = {
			...included[0],
			filePath: join(directory, "unknown.md"),
			sourceInfo: undefined,
		};
		await expect(
			estimateSourceTokens(
				agentDir,
				{ estimateTokens, formatSkillsForPrompt },
				{ skills: [unknown], diagnostics: [] },
			),
		).rejects.toThrow("Unattributed skill source");
		for (const loader of [normal, stripped]) {
			expect(
				loader
					.getExtensions()
					.extensions.find((ext) => ext.flags.has("fixture-setting"))
					?.flags.get("fixture-setting")?.description,
			).toBe("unchanged");
			expect(
				loader
					.getSkills()
					.skills.some((skill) => skill.name === "extension-fixture"),
			).toBe(true);
			expect(
				loader
					.getPrompts()
					.prompts.some((prompt) => prompt.name === "extension-fixture"),
			).toBe(true);
			expect(
				loader
					.getAgentsFiles()
					.agentsFiles.some((file) => file.path === join(cwd, "AGENTS.md")),
			).toBe(true);
			expect(
				loader
					.getSkills()
					.skills.some((skill) => skill.name === "project-fixture"),
			).toBe(projectTrusted);
			expect(
				loader
					.getPrompts()
					.prompts.some((prompt) => prompt.name === "project-fixture"),
			).toBe(projectTrusted);
		}
		expect(stripped.getExtensions().extensions.map((ext) => ext.path)).toEqual(
			normal.getExtensions().extensions.map((ext) => ext.path),
		);
		expect(
			normal
				.getSkills()
				.skills.some((skill) => skill.name === "global-fixture"),
		).toBe(true);
		expect(
			normal
				.getPrompts()
				.prompts.some((prompt) => prompt.name === "global-fixture"),
		).toBe(true);
		expect(
			normal
				.getAgentsFiles()
				.agentsFiles.some((file) => file.path === join(agentDir, "AGENTS.md")),
		).toBe(true);
		for (let i = 0; i < 2; i++) {
			expect(
				stripped
					.getSkills()
					.skills.some((skill) => skill.name === "global-fixture"),
			).toBe(false);
			expect(
				stripped
					.getPrompts()
					.prompts.some((prompt) => prompt.name === "global-fixture"),
			).toBe(false);
			expect(
				stripped
					.getAgentsFiles()
					.agentsFiles.some(
						(file) => file.path === join(agentDir, "AGENTS.md"),
					),
			).toBe(false);
			expect(stripped.getSystemPrompt()).toBe(
				projectTrusted && projectPrompts ? "Project system" : undefined,
			);
			expect(stripped.getAppendSystemPrompt()).toEqual(
				projectTrusted && projectPrompts ? ["Project append"] : [],
			);
			if (i === 0) await stripped.reload();
		}
		expect(readFileSync(settingsPath, "utf8")).toBe(settingsText);
	},
);
it("shows only tokens with explicit heuristic labels and a first-turn estimate diff", () => {
	const rows = CASES.map((kind, i) => ({
		...metadata,
		kind,
		local: {
			system: 1024 * (i + 1),
			tools: 2048,
			messages: i === 1 ? 2 : 1024,
		},
		...(i === 1 ? {} : { inputTokens: [100, 0, 400, 250][i] }),
	}));
	const output = formatOverview(rows);
	expect(output).toContain("Config: +300 tok (+300%)");
	expect(output).toContain("Pi text resources: +150 tok (+60%)");
	expect(output).toContain("headless+1msg+no_pi_config");
	expect(output).not.toMatch(/TUI|tui|WezTerm|KiB|bytes/);
	expect(output).toContain("~sys tok ~tools tok");
	expect(output).toContain("~ = Pi heuristic (chars/4)");
	expect(output).toContain(
		"First turn (~tok): system +1,024 · tools +0 · messages +1,022",
	);
	expect(output).toContain("n/a");
	expect(
		new Set(
			output
				.split("\n")
				.slice(3, 9)
				.map((line) => line.length),
		).size,
	).toBe(1);
	rows[2].model = "different";
	expect(formatOverview(rows)).toContain(
		"Config: n/a (failed or mismatched runs)",
	);
	expect(comparable(rows[0], { ...rows[0], error: "failure" })).toBe(false);
});
it("publishes complete JSON atomically", async () => {
	directory = mkdtempSync(join(tmpdir(), "pi-overview-test-"));
	const result = join(directory, "result.json");
	await publish(result, { error: "startup failed" });
	expect(JSON.parse(readFileSync(result, "utf8"))).toEqual({
		error: "startup failed",
	});
	expect(readdirSync(directory)).toEqual(["result.json"]);
});

it("observes the real Agent stream boundary after context hooks without changing its response", async () => {
	const response = {
		...message,
		role: "assistant" as const,
		api: "openai-responses" as const,
		stopReason: "stop" as const,
		content: [{ type: "text" as const, text: "OK" }],
		timestamp: 0,
		usage: {
			...message.usage,
			totalTokens: 332,
			cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
		},
	};
	const agent = new Agent({
		initialState: {
			model: {
				id: "test-model",
				name: "Test",
				provider: "test",
				api: "openai-responses",
				baseUrl: "https://example.invalid",
				reasoning: false,
				input: ["text"],
				contextWindow: 100000,
				maxTokens: 1000,
				cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
			},
			systemPrompt: "stock",
		},
		transformContext: async (messages) => [
			...messages,
			{ role: "user", content: "injected", timestamp: 0 },
		],
		streamFn: () => {
			const stream = createAssistantMessageEventStream();
			stream.push({ type: "done", reason: "stop", message: response });
			stream.end();
			return stream;
		},
	});
	const observe = vi.fn();
	observeRequests(agent, observe);
	await agent.prompt(PROMPT);
	expect(observe).toHaveBeenCalledTimes(1);
	expect(observe.mock.calls[0][0].messages).toHaveLength(3);
	expect(agent.state.messages.at(-1)?.role).toBe("assistant");
	expect(() => observeRequests({}, observe)).toThrow(
		"Pi stream API unavailable",
	);
});

function fixtures() {
	directory = mkdtempSync(join(tmpdir(), "pi-overview-test-"));
	const cwd = join(directory, "project with spaces");
	mkdirSync(cwd);
	const shared = join(directory, "shared.cjs");
	writeFileSync(
		shared,
		`
		const fs = require('node:fs');
		const path = require('node:path');
		const root = ${JSON.stringify(directory)};
		exports.log = value => fs.appendFileSync(path.join(root, 'calls.jsonl'), JSON.stringify(value) + '\\n');
		exports.result = options => ({ ...${JSON.stringify(metadata)}, cwd: options.cwd, kind: options.kind,
			mode: 'json', hasUI: false,
			local: { system: 1024, tools: 2048, messages: 2, toolCount: 4 },
			...(options.kind === 'headless+0msg' ? {} : { inputTokens: options.kind === 'baseline+1msg' ? 100 : 400 }),
		});
	`,
	);
	const worker = join(directory, "worker.mjs");
	writeFileSync(
		worker,
		`
		import fs from 'node:fs';
		import assert from 'node:assert/strict';
		import shared from ${JSON.stringify(pathToFileURL(shared).href)};
		const options = JSON.parse(fs.readFileSync(process.argv[2], 'utf8'));
		shared.log({ options });
		assert.equal(process.cwd(), options.cwd);
		if (options.kind === 'baseline+1msg') {
			assert.notEqual(options.agentDir, options.sourceAgentDir);
			assert.deepEqual(fs.readdirSync(options.agentDir), []);
		}
		const fail = process.env.CONTEXT_TEST_FAIL === options.kind;
		const result = fail ? { kind: options.kind, error: 'controlled failure' } : shared.result(options);
		if (process.env.CONTEXT_TEST_FAIL === 'mismatch' && options.kind === 'baseline+1msg') result.model = 'other';
		fs.writeFileSync(options.resultPath, JSON.stringify(result));
	`,
	);
	// No terminal executable or control context is available to these subprocesses.
	vi.stubEnv("PATH", directory);
	vi.stubEnv("WEZTERM_PANE", "");
	const callsPath = join(directory, "calls.jsonl");
	return {
		cwd,
		worker,
		calls: () =>
			readFileSync(callsPath, "utf8")
				.trim()
				.split("\n")
				.map((line) => JSON.parse(line)),
	};
}

it("runs four isolated headless cases without WezTerm, pins selection, and cleans measurement files", async () => {
	const { cwd, worker, calls } = fixtures();
	const result = await measure(cwd, worker);
	expect(result.map((row) => row.kind)).toEqual([
		CASES[1],
		CASES[0],
		CASES[2],
		CASES[3],
	]);
	expect(result.some((row) => row.error)).toBe(false);
	expect(result.every((row) => row.mode === "json" && !row.hasUI)).toBe(true);
	const jobs = calls()
		.filter((call) => call.options)
		.map((call) => call.options);
	expect(
		jobs.slice(1).every((job) => job.selection.model === "test-model"),
	).toBe(true);
	for (const job of jobs) expect(() => readFileSync(job.resultPath)).toThrow();
	expect(jobs).toHaveLength(4);
	expect(jobs[3].agentDir).toBe(jobs[3].sourceAgentDir);
});
it.each([
	"baseline+1msg",
	"headless+1msg",
	"headless+1msg+no_pi_config",
	"mismatch",
])("reports %s failure without fabricated diffs", async (failure) => {
	const { cwd, worker } = fixtures();
	vi.stubEnv("CONTEXT_TEST_FAIL", failure);
	const result = await measure(cwd, worker);
	expect(result.some((row) => row.error)).toBe(true);
	expect(formatOverview(result)).toContain("n/a (failed or mismatched runs)");
});
it("blocks all paid runs when configured startup fails", async () => {
	const { cwd, worker, calls } = fixtures();
	vi.stubEnv("CONTEXT_TEST_FAIL", "headless+0msg");
	const result = await measure(cwd, worker);
	expect(result.every((row) => row.error)).toBe(true);
	expect(calls().filter((call) => call.options)).toHaveLength(1);
});