pi-ext

import { mkdirSync, readdirSync, readFileSync, writeFileSync } from "node:fs";
import path from "node:path";
import {
	fauxAssistantMessage,
	fauxProvider,
	fauxToolCall,
	getCurrentTools,
} from "@earendil-works/pi-ai";
import { afterEach, describe, expect, it, vi } from "vitest";
import {
	createTestSession,
	type TestSession,
	withProcessEnv,
} from "../../../test/harness";
import {
	counts,
	goodJobDataDirectory,
	layout,
	loadRecords,
} from "../storage.js";

function writeUserSettings(t: TestSession, settings: unknown): void {
	const dir = path.join(t.cwd, ".test-home/.pi/agent");
	mkdirSync(dir, { recursive: true });
	writeFileSync(path.join(dir, "settings.json"), JSON.stringify(settings));
}

function extensionPath(): string {
	return path.resolve(import.meta.dirname, "../index.ts");
}

async function eventually<T>(
	read: () => Promise<T>,
	accept: (value: T) => boolean,
): Promise<T> {
	const deadline = Date.now() + 3_000;
	let value = await read();
	while (!accept(value) && Date.now() < deadline) {
		await new Promise((resolve) => setTimeout(resolve, 20));
		value = await read();
	}
	return value;
}

function learning(summary: string, slug = "shared-route-first") {
	return fauxAssistantMessage(
		fauxToolCall("submit_learning", {
			slug,
			summary,
			evidence: ["Focused tests passed"],
			behaviors: ["Inspected callers before editing"],
			candidate_rule: "Fix shared routes before leaf callers.",
		}),
		{ stopReason: "toolUse" },
	);
}

describe("good-job pi-test-harness", () => {
	let t: TestSession | undefined;

	afterEach(async () => {
		try {
			await t?.session.extensionRunner.emit({
				type: "session_shutdown",
				reason: "quit",
			});
		} finally {
			t?.dispose();
			t = undefined;
			vi.restoreAllMocks();
		}
	});

	it("registers only /gj with list-driven commands and no prompt or tools", async () => {
		t = await createTestSession({ extensions: [extensionPath()] });
		const extension = t.session.extensionRunner.extensions[0];
		expect(extension.path).toBe(extensionPath());
		expect([...extension.commands.keys()]).toEqual(["gj"]);
		expect([...extension.tools.keys()]).toEqual([]);
		expect(extension.handlers.get("resources_discover")).toBeUndefined();
		expect(
			t.session.resourceLoader
				.getPrompts()
				.prompts.find(({ name }) => name === "wtf"),
		).toBeUndefined();
		expect(extension.commands.get("gj").getArgumentCompletions("")).toEqual([
			{ value: "list", label: "list" },
			{ value: "status", label: "status" },
			{ value: "open", label: "open" },
			{ value: "learn", label: "learn" },
		]);
	});

	it("records praise while passing the original message to the agent", async () => {
		const originalDebug = process.env.PI_GOOD_JOB_DEBUG;
		let extensionModule: typeof import("../index.ts") | undefined;
		try {
			process.env.PI_GOOD_JOB_DEBUG = "1";
			vi.resetModules();
			extensionModule = await import("../index.ts");
		} finally {
			if (originalDebug === undefined) delete process.env.PI_GOOD_JOB_DEBUG;
			else process.env.PI_GOOD_JOB_DEBUG = originalDebug;
		}
		if (!extensionModule) throw new Error("could not load good-job extension");
		t = await createTestSession({
			extensionFactories: [extensionModule.default],
			env: { PI_GOOD_JOB_DEBUG: "1" },
		});
		const runner = t.session.extensionRunner;
		const registered = runner.extensions[0];
		const input = registered.handlers.get("input")[0];
		const base = runner.createContext();
		const faux = fauxProvider({
			provider: "good-job-test",
			models: [{ id: "worker", maxTokens: 4_096 }],
		});
		faux.setResponses([learning("Fixed the shared route")]);
		const stream = vi.spyOn(faux.provider, "streamSimple");
		base.modelRegistry.registerProvider(faux.provider);
		const model = base.modelRegistry.find("good-job-test", "worker");
		if (!model) throw new Error("missing faux model");
		const ctx = Object.create(base) as typeof base;
		Object.defineProperty(ctx, "model", { value: model });
		const manager = ctx.sessionManager as any;
		manager.appendMessage({
			role: "user",
			content: "Fix the shared route",
			timestamp: Date.now(),
		});
		manager.appendMessage(
			fauxAssistantMessage("Implemented and tested the shared fix."),
		);

		const result = await withProcessEnv(t.env, () =>
			input(
				{
					type: "input",
					text: "good job! excellent scope control",
					source: "interactive",
				},
				ctx,
			),
		);
		expect(result).toEqual({ action: "continue" });

		const root = goodJobDataDirectory(t.env, process.platform, t.env.HOME);
		const loaded = await eventually(
			() => loadRecords(layout(root)),
			(value) => value.records.length === 1,
		);
		expect(loaded.errors).toEqual([]);
		expect(loaded.records[0]).toMatchObject({
			kind: "praise",
			feedback: "excellent scope control",
			learning: {
				slug: "shared-route-first",
				summary: "Fixed the shared route",
			},
		});
		expect(loaded.records[0]?.id).toMatch(/^gj-[0-9a-hjkmnp-tv-z]{10}$/);
		expect(loaded.records[0]?.source).toMatchObject({
			type: "session",
			cwd: ctx.cwd,
		});
		const debugRoot = path.join(
			t.cwd,
			".test-state",
			"pi-ext",
			"debug",
			"good-job",
		);
		const events = readdirSync(debugRoot).flatMap((file) =>
			readFileSync(path.join(debugRoot, file), "utf8")
				.trim()
				.split("\n")
				.map((line) => JSON.parse(line)),
		);
		expect(events).toEqual(
			expect.arrayContaining([
				expect.objectContaining({
					event: "feedback.capture",
					kind: "praise",
				}),
				expect.objectContaining({
					event: "learning.outcome",
					status: "committed",
				}),
			]),
		);
		const sentMessages = stream.mock.calls[0]?.[1].messages ?? [];
		expect(JSON.stringify(sentMessages)).toContain(
			"Implemented and tested the shared fix.",
		);
		expect(JSON.stringify(sentMessages.at(-1))).toContain("kebab-case");
		expect(getCurrentTools(sentMessages).map(({ name }) => name)).toEqual([
			"submit_learning",
		]);

		await withProcessEnv(t.env, () =>
			registered.commands.get("gj").handler("status", ctx as never),
		);
		expect(t.events.ui.at(-1)?.args[0]).toContain("1 learnings");
	});

	it("rejects analysis that does not submit structured learning", async () => {
		t = await createTestSession({ extensions: [extensionPath()] });
		const runner = t.session.extensionRunner;
		const extension = runner.extensions[0];
		const input = extension.handlers.get("input")[0];
		const base = runner.createContext();
		const faux = fauxProvider({
			provider: "good-job-unstructured",
			models: [{ id: "worker", maxTokens: 4_096 }],
		});
		faux.setResponses([fauxAssistantMessage('{"summary":"not submitted"}')]);
		base.modelRegistry.registerProvider(faux.provider);
		const model = base.modelRegistry.find("good-job-unstructured", "worker");
		if (!model) throw new Error("missing unstructured faux model");
		const ctx = Object.create(base) as typeof base;
		Object.defineProperty(ctx, "model", { value: model });

		await withProcessEnv(t.env, () =>
			input({ type: "input", text: "gj", source: "interactive" }, ctx),
		);

		const paths = layout(
			goodJobDataDirectory(t.env, process.platform, t.env.HOME),
		);
		await eventually(
			() => counts(paths),
			(value) => value.failed === 1,
		);
		expect(await counts(paths)).toEqual({
			pending: 0,
			processing: 0,
			failed: 1,
			records: 0,
		});
		expect(
			t.events.ui
				.filter((call) => call.method === "notify")
				.map((call) => String(call.args[0]))
				.find((message) => message.includes("good-job failed")),
		).toContain("without calling submit_learning");
	});

	it("records first-word wtf feedback while passing it unchanged", async () => {
		t = await createTestSession({ extensions: [extensionPath()] });
		const runner = t.session.extensionRunner;
		const extension = runner.extensions[0];
		const input = extension.handlers.get("input")[0];
		const base = runner.createContext();
		const faux = fauxProvider({
			provider: "wtf-test",
			models: [{ id: "worker", maxTokens: 4_096 }],
		});
		faux.setResponses([
			learning("Expanded scope without approval", "scope-creep"),
		]);
		const stream = vi.spyOn(faux.provider, "streamSimple");
		base.modelRegistry.registerProvider(faux.provider);
		const model = base.modelRegistry.find("wtf-test", "worker");
		if (!model) throw new Error("missing wtf faux model");
		const ctx = Object.create(base) as typeof base;
		Object.defineProperty(ctx, "model", { value: model });
		const manager = ctx.sessionManager as any;
		manager.appendMessage({
			role: "user",
			content: "Implement only A",
			timestamp: Date.now(),
		});
		manager.appendMessage(
			fauxAssistantMessage("Implemented A and unrelated B."),
		);

		const result = await withProcessEnv(t.env, () =>
			input(
				{
					type: "input",
					text: "wtf you added unrelated B",
					source: "interactive",
				},
				ctx,
			),
		);
		expect(result).toEqual({ action: "continue" });

		const root = goodJobDataDirectory(t.env, process.platform, t.env.HOME);
		const loaded = await eventually(
			() => loadRecords(layout(root)),
			(value) => value.records.length === 1,
		);
		expect(loaded.records[0]).toMatchObject({
			kind: "problem",
			feedback: "you added unrelated B",
			learning: { slug: "scope-creep" },
		});
		expect(loaded.records[0]?.id).toMatch(/^wtf-[0-9a-hjkmnp-tv-z]{10}$/);
		expect(JSON.stringify(stream.mock.calls[0]?.[1].messages)).toContain(
			"Implemented A and unrelated B.",
		);
	});

	it("uses optional directory, model, and thinking settings", async () => {
		t = await createTestSession({ extensions: [extensionPath()] });
		const customRoot = path.join(t.cwd, "shared-gj");
		writeUserSettings(t, {
			"good-job": {
				directory: customRoot,
				model: "good-job-override/background",
				thinkingLevel: "xhigh",
			},
		});
		const runner = t.session.extensionRunner;
		const extension = runner.extensions[0];
		const input = extension.handlers.get("input")[0];
		const base = runner.createContext();
		const faux = fauxProvider({
			provider: "good-job-override",
			models: [
				{
					id: "active",
					maxTokens: 4_096,
					reasoning: true,
					thinkingLevelMap: { xhigh: "xhigh" },
				},
				{
					id: "background",
					maxTokens: 4_096,
					reasoning: true,
					thinkingLevelMap: { xhigh: "xhigh" },
				},
			],
		});
		faux.setResponses([learning("Used configured model", "configured-model")]);
		Object.assign(faux.getModel("active"), {
			thinkingLevelMap: { xhigh: "xhigh" },
		});
		Object.assign(faux.getModel("background"), {
			thinkingLevelMap: { xhigh: "xhigh" },
		});
		const stream = vi.spyOn(faux.provider, "streamSimple");
		base.modelRegistry.registerProvider(faux.provider);
		const active = base.modelRegistry.find("good-job-override", "active");
		if (!active) throw new Error("missing active faux model");
		const ctx = Object.create(base) as typeof base;
		Object.defineProperty(ctx, "model", { value: active });

		const result = await withProcessEnv(t.env, () =>
			input({ type: "input", text: "gj", source: "interactive" }, ctx),
		);
		expect(result).toEqual({ action: "continue" });

		const loaded = await eventually(
			() => loadRecords(layout(customRoot)),
			(value) => value.records.length === 1,
		);
		expect(loaded.records[0]?.analysisModel).toEqual({
			provider: "good-job-override",
			id: "background",
		});
		expect(stream.mock.calls[0]?.[0].id).toBe("background");
		expect(stream.mock.calls[0]?.[2]?.reasoning).toBe("xhigh");
	});

	it("sends only the records path and synthesis instructions to /gj learn", async () => {
		t = await createTestSession({ extensions: [extensionPath()] });
		const runner = t.session.extensionRunner;
		const extension = runner.extensions[0];
		const input = extension.handlers.get("input")[0];
		const base = runner.createContext();
		const faux = fauxProvider({
			provider: "learn-test",
			models: [{ id: "worker", maxTokens: 4_096 }],
		});
		faux.setResponses([
			learning("Private summary must not be embedded", "private-learning"),
			fauxAssistantMessage("Proposed changes."),
		]);
		base.modelRegistry.registerProvider(faux.provider);
		const model = base.modelRegistry.find("learn-test", "worker");
		if (!model) throw new Error("missing learn faux model");
		const ctx = Object.create(base) as typeof base;
		Object.defineProperty(ctx, "model", { value: model });
		await withProcessEnv(t.env, () =>
			input({ type: "input", text: "gj", source: "interactive" }, ctx),
		);
		const root = goodJobDataDirectory(t.env, process.platform, t.env.HOME);
		await eventually(
			() => loadRecords(layout(root)),
			(value) => value.records.length === 1,
		);

		const promptSpy = vi
			.spyOn(t.session, "prompt")
			.mockResolvedValue(undefined);
		await withProcessEnv(t.env, () =>
			extension.commands.get("gj").handler("learn", ctx as never),
		);
		await vi.waitFor(() => expect(promptSpy).toHaveBeenCalled());
		const prompt = String(promptSpy.mock.calls[0]?.[0]);
		expect(prompt).toContain(layout(root).records);
		expect(prompt).toContain("slug (id)");
		expect(prompt).toContain("praise and problem records together");
		expect(prompt).not.toContain("Private summary must not be embedded");
	});

	it("reports invalid settings during startup", async () => {
		t = await createTestSession({ extensions: [extensionPath()] });
		writeUserSettings(t, { "good-job": { thinkingLevel: "invalid" } });
		const runner = t.session.extensionRunner;
		const start = runner.extensions[0].handlers.get("session_start")[0];
		start({ type: "session_start", reason: "startup" }, runner.createContext());

		const notification = await eventually(
			async () =>
				t?.events.ui
					.filter((call) => call.method === "notify")
					.map((call) => String(call.args[0]))
					.find((message) => message.includes("startup failed")),
			(value) => value !== undefined,
		);
		expect(notification).toContain(
			"Invalid good-job.thinkingLevel from user settings",
		);
	});

	it("recovers in-flight work when a new session starts", async () => {
		t = await createTestSession({ extensions: [extensionPath()] });
		const runner = t.session.extensionRunner;
		const extension = runner.extensions[0];
		const input = extension.handlers.get("input")[0];
		const start = extension.handlers.get("session_start")[0];
		const base = runner.createContext();
		const faux = fauxProvider({
			provider: "good-job-restart",
			models: [{ id: "worker", maxTokens: 4_096 }],
		});
		let releaseFirst = () => {};
		let markStarted = () => {};
		const firstStarted = new Promise<void>((resolve) => {
			markStarted = resolve;
		});
		let firstAborted = false;
		faux.setResponses([
			async (_context, options) => {
				markStarted();
				await new Promise<void>((resolve) => {
					releaseFirst = resolve;
					options?.signal?.addEventListener(
						"abort",
						() => {
							firstAborted = true;
							resolve();
						},
						{ once: true },
					);
				});
				return learning("Stale response", "stale-response");
			},
			learning("Recovered after session change", "recovered-work"),
		]);
		base.modelRegistry.registerProvider(faux.provider);
		const model = base.modelRegistry.find("good-job-restart", "worker");
		if (!model) throw new Error("missing restart faux model");
		const ctx = Object.create(base) as typeof base;
		Object.defineProperty(ctx, "model", { value: model });

		await withProcessEnv(t.env, () =>
			input({ type: "input", text: "gj", source: "interactive" }, ctx),
		);
		await firstStarted;
		start({ type: "session_start", reason: "switch" }, ctx);
		await new Promise((resolve) => setTimeout(resolve, 50));
		releaseFirst();

		const root = goodJobDataDirectory(t.env, process.platform, t.env.HOME);
		const paths = layout(root);
		const loaded = await eventually(
			() => loadRecords(paths),
			(value) => value.records.length === 1,
		);
		expect(firstAborted).toBe(true);
		expect(loaded.records[0]?.learning.slug).toBe("recovered-work");
		expect(await counts(paths)).toEqual({
			pending: 0,
			processing: 0,
			failed: 0,
			records: 1,
		});
	});

	it("ignores unrelated and extension-injected text", async () => {
		t = await createTestSession({ extensions: [extensionPath()] });
		const input =
			t.session.extensionRunner.extensions[0].handlers.get("input")[0];
		for (const event of [
			{ type: "input", text: "great job", source: "interactive" },
			{ type: "input", text: "wtf happened", source: "extension" },
		])
			expect(
				await input(event, t.session.extensionRunner.createContext()),
			).toEqual({ action: "continue" });
	});
});