pi-ext

import { readdirSync, readFileSync, writeFileSync } from "node:fs";
import path from "node:path";
import { fileURLToPath } from "node:url";
import {
	type FauxModelDefinition,
	fauxAssistantMessage,
	fauxProvider,
	fauxToolCall,
	getCurrentTools,
} from "@earendil-works/pi-ai";
import { initTheme } from "@earendil-works/pi-coding-agent";
import { afterEach, describe, expect, it, vi } from "vitest";
import {
	createPiSession,
	type PiSession,
} from "../../../test/fixture/create-pi-session.ts";
import {
	counts,
	goodJobDataDirectory,
	layout,
	loadRecords,
} from "../storage.js";

function extensionPath(): string {
	return fileURLToPath(new URL("../index.ts", import.meta.url));
}

function writeUserSettings(session: PiSession, settings: unknown): void {
	writeFileSync(
		path.join(session.agentDir, "settings.json"),
		JSON.stringify(settings),
	);
}

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 contract", () => {
	let session: PiSession | undefined;

	afterEach(async () => {
		await session?.dispose();
		session = undefined;
		vi.restoreAllMocks();
	});

	it("registers only /gj with list-driven commands and no prompt or tools", async () => {
		session = await createPiSession({ extensionPaths: [extensionPath()] });
		const extension = session.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(
			session.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 () => {
		session = await createPiSession({
			extensionPaths: [extensionPath()],
			env: { PI_GOOD_JOB_DEBUG: "1" },
		});
		const runner = session.session.extensionRunner;
		const registered = session.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.defineProperties(ctx, {
			model: { value: model },
			hasUI: { value: true },
		});
		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."),
		);
		initTheme(undefined, false);
		const notify = vi.spyOn(base.ui, "notify").mockImplementation(() => {});

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

		const root = goodJobDataDirectory(
			session.env,
			process.platform,
			session.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(
			session.cwd,
			".test-state",
			"pi-ext",
			"debug",
			"good-job",
		);
		const readDebugEvents = async () =>
			readdirSync(debugRoot).flatMap((file) =>
				readFileSync(path.join(debugRoot, file), "utf8")
					.trim()
					.split("\n")
					.map((line) => JSON.parse(line)),
			);
		// Debug records flush asynchronously; poll until the praise events land.
		const events = await eventually(
			readDebugEvents,
			(value) =>
				value.some(
					(record) =>
						record.event === "feedback.capture" && record.kind === "praise",
				) &&
				value.some(
					(record) =>
						record.event === "learning.outcome" &&
						record.status === "committed",
				),
		);
		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 registered.commands.get("gj")!.handler("status", ctx as never);
		expect(notify.mock.calls.at(-1)?.[0]).toContain("1 learnings");
	});

	it("rejects analysis that does not submit structured learning", async () => {
		session = await createPiSession({ extensionPaths: [extensionPath()] });
		const runner = session.session.extensionRunner;
		const extension = session.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.defineProperties(ctx, {
			model: { value: model },
			hasUI: { value: true },
		});
		initTheme(undefined, false);
		const notify = vi.spyOn(base.ui, "notify").mockImplementation(() => {});

		await input({ type: "input", text: "gj", source: "interactive" }, ctx);

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

	it("records first-word wtf feedback while passing it unchanged", async () => {
		session = await createPiSession({ extensionPaths: [extensionPath()] });
		const runner = session.session.extensionRunner;
		const extension = session.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 input(
			{
				type: "input",
				text: "wtf you added unrelated B",
				source: "interactive",
			},
			ctx,
		);
		expect(result).toEqual({ action: "continue" });

		const root = goodJobDataDirectory(
			session.env,
			process.platform,
			session.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 () => {
		session = await createPiSession({ extensionPaths: [extensionPath()] });
		const customRoot = path.join(session.cwd, "shared-gj");
		writeUserSettings(session, {
			"good-job": {
				directory: customRoot,
				model: "good-job-override/background",
				thinkingLevel: "xhigh",
			},
		});
		const runner = session.session.extensionRunner;
		const extension = session.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" },
				},
			] as (FauxModelDefinition & {
				thinkingLevelMap: Record<string, string>;
			})[],
		});
		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 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 () => {
		session = await createPiSession({ extensionPaths: [extensionPath()] });
		const runner = session.session.extensionRunner;
		const extension = session.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 input({ type: "input", text: "gj", source: "interactive" }, ctx);
		const root = goodJobDataDirectory(
			session.env,
			process.platform,
			session.env["HOME"],
		);
		await eventually(
			() => loadRecords(layout(root)),
			(value) => value.records.length === 1,
		);

		const promptSpy = vi
			.spyOn(session.session, "prompt")
			.mockResolvedValue(undefined);
		await 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(JSON.stringify(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 () => {
		session = await createPiSession({ extensionPaths: [extensionPath()] });
		writeUserSettings(session, { "good-job": { thinkingLevel: "invalid" } });
		const runner = session.session.extensionRunner;
		const start = session.extensions[0]!.handlers.get("session_start")![0]!;
		const base = runner.createContext();
		const ctx = Object.create(base) as typeof base;
		Object.defineProperty(ctx, "hasUI", { value: true });
		initTheme(undefined, false);
		const notify = vi.spyOn(base.ui, "notify").mockImplementation(() => {});
		start({ type: "session_start", reason: "startup" }, ctx);

		const notification = await eventually(
			async () =>
				notify.mock.calls
					.map((call) => String(call[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 () => {
		session = await createPiSession({ extensionPaths: [extensionPath()] });
		const runner = session.session.extensionRunner;
		const extension = session.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 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(
			session.env,
			process.platform,
			session.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 eventually(
				() => counts(paths),
				(value) => value.processing === 0,
			),
		).toEqual({
			pending: 0,
			processing: 0,
			failed: 0,
			records: 1,
		});
	});

	it("continues steering input before the capture is queued", async () => {
		session = await createPiSession({ extensionPaths: [extensionPath()] });
		const runner = session.session.extensionRunner;
		const input = session.extensions[0]!.handlers.get("input")![0]!;
		const base = runner.createContext();
		const faux = fauxProvider({
			provider: "steer-test",
			models: [{ id: "worker", maxTokens: 4_096 }],
		});
		faux.setResponses([learning("Steered mid-turn", "steered-mid-turn")]);
		base.modelRegistry.registerProvider(faux.provider);
		const model = base.modelRegistry.find("steer-test", "worker");
		if (!model) throw new Error("missing steer faux model");
		const ctx = Object.create(base) as typeof base;
		Object.defineProperty(ctx, "model", { value: model });
		const paths = layout(
			goodJobDataDirectory(session.env, process.platform, session.env["HOME"]),
		);

		const result = await input(
			{
				type: "input",
				text: "wtf stop touching B",
				source: "interactive",
				streamingBehavior: "steer",
			},
			ctx,
		);
		expect(result).toEqual({ action: "continue" });
		// The handler resolved before any queue write could complete.
		expect(await counts(paths)).toEqual({
			pending: 0,
			processing: 0,
			failed: 0,
			records: 0,
		});

		const loaded = await eventually(
			() => loadRecords(paths),
			(value) => value.records.length === 1,
		);
		expect(loaded.records[0]).toMatchObject({
			kind: "problem",
			feedback: "stop touching B",
			learning: { slug: "steered-mid-turn" },
		});
	});

	it("still reports steering capture failures", async () => {
		session = await createPiSession({ extensionPaths: [extensionPath()] });
		writeUserSettings(session, { "good-job": { thinkingLevel: "invalid" } });
		const runner = session.session.extensionRunner;
		const input = session.extensions[0]!.handlers.get("input")![0]!;
		const base = runner.createContext();
		const ctx = Object.create(base) as typeof base;
		Object.defineProperty(ctx, "hasUI", { value: true });
		initTheme(undefined, false);
		const notify = vi.spyOn(base.ui, "notify").mockImplementation(() => {});

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

		const notification = await eventually(
			async () =>
				notify.mock.calls
					.map((call) => String(call[0]))
					.find((message) => message.includes("gj could not queue")),
			(value) => value !== undefined,
		);
		expect(notification).toContain(
			"Invalid good-job.thinkingLevel from user settings",
		);
	});

	it("ignores unrelated and extension-injected text", async () => {
		session = await createPiSession({ extensionPaths: [extensionPath()] });
		const input = session.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, session.session.extensionRunner.createContext()),
			).toEqual({ action: "continue" });
	});
});