pi-ext

import { readdirSync, readFileSync } from "node:fs";
import { join } from "node:path";
import { fileURLToPath } from "node:url";
import { fauxAssistantMessage, fauxToolCall } from "@earendil-works/pi-ai";
import {
	defineTool,
	type ExtensionAPI,
	type ExtensionFactory,
} from "@earendil-works/pi-coding-agent";
import type { Component, TUI } from "@earendil-works/pi-tui";
import { Type } from "typebox";
import { afterEach, describe, expect, it, vi } from "vitest";
import {
	createPiSession,
	type PiSession,
} from "../../../test/fixture/create-pi-session.ts";

let session: PiSession | undefined;

const finalReport = "Completed work with evidence and no unresolved blockers.";

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

function command(name: "afk" | "back") {
	const result = session?.extensions
		.flatMap((extension) => [...extension.commands.entries()])
		.find(([commandName]) => commandName === name)?.[1];
	if (!result) throw new Error(`missing /${name}`);
	return result;
}

async function runAfk(): Promise<void> {
	await command("afk").handler(
		"",
		session?.session.extensionRunner.createContext() as never,
	);
	await session?.session.waitForIdle();
}

const humanTool: ExtensionFactory = (pi: ExtensionAPI) => {
	pi.registerTool(
		defineTool({
			name: "ask_user_question",
			label: "ask user",
			description: "Ask a human a question.",
			parameters: Type.Object({}),
			async execute() {
				return {
					content: [{ type: "text", text: "human answered" }],
					details: undefined,
				};
			},
		}),
	);
};

function queueFollowUp(pi: ExtensionAPI): void {
	let queued = false;
	pi.on("agent_start", () => {
		if (queued) return;
		queued = true;
		pi.sendUserMessage("Finish queued follow-up work.", {
			deliverAs: "followUp",
		});
	});
}

function stopAfkWithoutAborting(pi: ExtensionAPI): void {
	let stopped = false;
	pi.on("agent_start", async (_event, ctx) => {
		if (stopped) return;
		stopped = true;
		await command("back").handler("", ctx as never);
	});
}

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

describe("afk contract", () => {
	it("loads through real Pi and leaves completion inactive", async () => {
		session = await createPiSession({ extensionPaths: [extensionPath()] });
		const extension = session.extensions[0]!;
		expect(session.errors).toEqual([]);
		expect([...extension.commands.keys()]).toEqual(["afk", "back"]);
		expect(extension.tools.get("afk_complete")?.definition.defaultActive).toBe(
			false,
		);
		expect(session.session.getActiveToolNames()).not.toContain("afk_complete");
	});

	it("uses real tool follow-up before napping after completion", async () => {
		session = await createPiSession({
			extensionPaths: [extensionPath()],
			responses: [
				fauxAssistantMessage("Investigated first approach."),
				fauxAssistantMessage(fauxToolCall("afk_complete", {}), {
					stopReason: "toolUse",
				}),
				fauxAssistantMessage(finalReport),
			],
		});
		await runAfk();

		expect(session.providerState.callCount).toBe(3);
		expect(session.session.getActiveToolNames()).toContain("afk_complete");
		expect(session.session.sessionManager.getBranch()).toContainEqual(
			expect.objectContaining({
				message: expect.objectContaining({
					role: "assistant",
					content: expect.any(Array),
				}),
			}),
		);
	});

	it("requires a fresh completion signal when reactivated after napping", async () => {
		session = await createPiSession({
			extensionPaths: [extensionPath()],
			responses: [
				fauxAssistantMessage(fauxToolCall("afk_complete", {}), {
					stopReason: "toolUse",
				}),
				fauxAssistantMessage(finalReport),
			],
		});
		await runAfk();
		session.setResponses([
			fauxAssistantMessage("More work remains."),
			fauxAssistantMessage(fauxToolCall("afk_complete", {}), {
				stopReason: "toolUse",
			}),
			fauxAssistantMessage("Reactivated work complete."),
		]);
		await runAfk();

		const branch = session.session.sessionManager.getBranch();
		expect(
			branch.filter(
				(entry) =>
					entry.type === "custom_message" && entry.customType === "afk-start",
			),
		).toHaveLength(2);
		expect(
			branch.filter(
				(entry) =>
					entry.type === "custom_message" &&
					entry.customType === "afk-continue",
			),
		).toHaveLength(1);
		expect(session.errors).toEqual([]);
	});

	it("ignores a completion call selected before /back and resets on reactivation", async () => {
		let returned = false;
		const returnBeforeCompletion: ExtensionFactory = (pi) => {
			pi.on("tool_call", async (event, ctx) => {
				if (event.toolName !== "afk_complete" || returned) return;
				returned = true;
				await command("back").handler("", ctx as never);
			});
		};
		session = await createPiSession({
			extensionPaths: [extensionPath()],
			extensionFactories: [returnBeforeCompletion],
			responses: [
				fauxAssistantMessage(fauxToolCall("afk_complete", {}), {
					stopReason: "toolUse",
				}),
				fauxAssistantMessage("Current work finished normally."),
			],
		});
		await runAfk();
		expect(returned).toBe(true);
		expect(session.session.getActiveToolNames()).not.toContain("afk_complete");
		expect(session.session.sessionManager.getBranch()).toContainEqual(
			expect.objectContaining({
				message: expect.objectContaining({
					role: "toolResult",
					toolName: "afk_complete",
					isError: false,
				}),
			}),
		);

		session.setResponses([
			fauxAssistantMessage("New work remains."),
			fauxAssistantMessage(fauxToolCall("afk_complete", {}), {
				stopReason: "toolUse",
			}),
			fauxAssistantMessage(finalReport),
		]);
		await runAfk();
		expect(
			session.session.sessionManager
				.getBranch()
				.filter(
					(entry) =>
						entry.type === "custom_message" &&
						entry.customType === "afk-continue",
				),
		).toHaveLength(1);
		expect(session.errors).toEqual([]);
	});

	it("returns blocked human tool result to model and receives a useful follow-up", async () => {
		session = await createPiSession({
			extensionPaths: [extensionPath()],
			extensionFactories: [humanTool],
			responses: [
				fauxAssistantMessage(fauxToolCall("ask_user_question", {}), {
					stopReason: "toolUse",
				}),
				fauxAssistantMessage(
					"Question unavailable, so I used the documented default.",
				),
				fauxAssistantMessage(fauxToolCall("afk_complete", {}), {
					stopReason: "toolUse",
				}),
				fauxAssistantMessage(finalReport),
			],
		});
		await runAfk();

		expect(session.providerState.callCount).toBe(4);
		expect(session.session.sessionManager.getBranch()).toContainEqual(
			expect.objectContaining({
				message: expect.objectContaining({
					role: "toolResult",
					isError: true,
					content: expect.arrayContaining([
						expect.objectContaining({
							text: expect.stringContaining(
								"unavailable while the user is away",
							),
						}),
					]),
				}),
			}),
		);
	});

	it("drains queued follow-up work before accepting completion", async () => {
		session = await createPiSession({
			extensionPaths: [extensionPath()],
			extensionFactories: [queueFollowUp],
			responses: [
				fauxAssistantMessage(fauxToolCall("afk_complete", {}), {
					stopReason: "toolUse",
				}),
				fauxAssistantMessage("Initial task report."),
				fauxAssistantMessage("Queued follow-up is still unfinished."),
				fauxAssistantMessage(fauxToolCall("afk_complete", {}), {
					stopReason: "toolUse",
				}),
				fauxAssistantMessage("Queued follow-up complete."),
			],
		});
		await runAfk();

		expect(session.providerState.callCount).toBe(5);
		expect(session.session.getActiveToolNames()).toContain("afk_complete");
	});

	it("activates during busy work without reordering already queued follow-ups", async () => {
		session = await createPiSession({
			extensionPaths: [extensionPath()],
			responses: [
				async () => {
					await session!.session.sendUserMessage("Follow-up A", {
						deliverAs: "followUp",
					});
					await session!.session.sendUserMessage("Follow-up B", {
						deliverAs: "followUp",
					});
					await command("afk").handler(
						"",
						session!.session.extensionRunner.createContext() as never,
					);
					return fauxAssistantMessage("Current task finished.");
				},
				fauxAssistantMessage("Follow-up A finished."),
				fauxAssistantMessage("Follow-up B finished."),
				fauxAssistantMessage(fauxToolCall("afk_complete", {}), {
					stopReason: "toolUse",
				}),
				fauxAssistantMessage(finalReport),
			],
		});
		await session.session.prompt("Start current task.");
		const branch = session.session.sessionManager.getBranch();
		const users = branch.flatMap((entry) =>
			entry.type === "message" && entry.message.role === "user"
				? [entry.message.content]
				: [],
		);
		expect(
			users.map((content) =>
				typeof content === "string"
					? content
					: content
							.filter((block) => block.type === "text")
							.map((block) => block.text)
							.join(""),
			),
		).toEqual(["Start current task.", "Follow-up A", "Follow-up B"]);
		expect(session.providerState.callCount).toBe(5);
		expect(session.errors).toEqual([]);
	});

	it("does not abort active work when /back disables continuation", async () => {
		session = await createPiSession({
			extensionPaths: [extensionPath()],
			extensionFactories: [stopAfkWithoutAborting],
			responses: [fauxAssistantMessage("Current work completed normally.")],
		});
		await runAfk();

		expect(session.providerState.callCount).toBe(1);
		expect(session.session.getActiveToolNames()).not.toContain("afk_complete");
	});

	it("stops after a bounded number of no-progress continuations", async () => {
		session = await createPiSession({
			extensionPaths: [extensionPath()],
			responses: [
				fauxAssistantMessage("Still considering it."),
				fauxAssistantMessage("Still considering it."),
				fauxAssistantMessage("Still considering it."),
			],
		});
		await runAfk();

		expect(session.providerState.callCount).toBe(3);
		expect(session.session.getActiveToolNames()).not.toContain("afk_complete");
		expect(session.session.sessionManager.getBranch()).toContainEqual(
			expect.objectContaining({ customType: "afk-incomplete" }),
		);
	});

	it.each(["aborted", "error"] as const)(
		"cleans up mode after %s without restarting",
		async (stopReason) => {
			session = await createPiSession({
				extensionPaths: [extensionPath()],
				responses: [fauxAssistantMessage("Stopped.", { stopReason })],
			});
			await runAfk();

			expect(session.providerState.callCount).toBe(1);
			expect(session.session.getActiveToolNames()).not.toContain(
				"afk_complete",
			);
		},
	);

	it("does not duplicate activation while AFK work is active", async () => {
		session = await createPiSession({ extensionPaths: [extensionPath()] });
		const ctx = session.session.extensionRunner.createContext();
		await command("afk").handler("", ctx as never);
		await command("afk").handler("", ctx as never);
		await session.session.waitForIdle();

		expect(session.providerState.callCount).toBe(1);
	});

	it("renders a napping creature through the real TUI binding and disposes it on /back", async () => {
		session = await createPiSession({
			extensionPaths: [extensionPath()],
			responses: [
				fauxAssistantMessage(fauxToolCall("afk_complete", {}), {
					stopReason: "toolUse",
				}),
				fauxAssistantMessage(finalReport),
			],
		});
		let creature: (Component & { dispose?(): void }) | undefined;
		const requestRender = vi.fn();
		await session.session.bindExtensions({
			mode: "tui",
			uiContext: {
				...session.session.extensionRunner.getUIContext(),
				setWidget: (_key, content) => {
					if (typeof content === "function") {
						creature = content(
							{ requestRender } as unknown as TUI,
							{ fg: (_color: string, text: string) => text } as never,
						);
					} else creature = undefined;
				},
			},
		});
		await runAfk();
		expect(creature?.render(80).join("\n")).toContain("zzz");
		expect(requestRender).toHaveBeenCalled();
		const dispose = vi.spyOn(creature!, "dispose");
		await command("back").handler(
			"",
			session.session.extensionRunner.createContext() as never,
		);
		expect(dispose).toHaveBeenCalledOnce();
		expect(creature).toBeUndefined();
		expect(session.errors).toEqual([]);
	});

	it.each(["alone", "first", "last"])(
		"recovers idle compaction with AFK loaded %s",
		async (order) => {
			const paths = [extensionPath()];
			if (order !== "alone")
				paths.push(
					fileURLToPath(new URL("../../continuity/index.ts", import.meta.url)),
				);
			if (order === "last") paths.reverse();
			session = await createPiSession({
				extensionPaths: paths,
				responses: [
					fauxAssistantMessage(fauxToolCall("afk_complete", {}), {
						stopReason: "toolUse",
					}),
					fauxAssistantMessage(finalReport),
					fauxAssistantMessage("Ordinary compaction recovery complete."),
				],
			});
			// Leave AFK enabled but idle to exercise recovery rather than the normal run loop.
			const start = vi
				.spyOn(session.session, "sendCustomMessage")
				.mockImplementationOnce(async () => {});
			await command("afk").handler(
				"",
				session.session.extensionRunner.createContext() as never,
			);
			start.mockRestore();
			expect(session.providerState.callCount).toBe(0);
			const compact = () =>
				session!.session.extensionRunner.emit({
					type: "session_compact",
					compactionEntry: { id: "afk-compact" },
					fromExtension: false,
					reason: "manual",
					willRetry: false,
				} as never);
			await compact();
			await vi.waitFor(() => expect(session!.providerState.callCount).toBe(2));
			await session.session.waitForIdle();
			const continuations = (type: string) =>
				session!.session.sessionManager
					.getBranch()
					.filter(
						(entry) =>
							entry.type === "custom_message" && entry.customType === type,
					);
			expect(continuations("afk-continue")).toHaveLength(1);
			expect(continuations("continuity-resume")).toHaveLength(0);
			await command("back").handler(
				"",
				session.session.extensionRunner.createContext() as never,
			);
			await compact();
			if (order !== "alone") {
				await vi.waitFor(() =>
					expect(session!.providerState.callCount).toBe(3),
				);
				await session.session.waitForIdle();
				expect(continuations("continuity-resume")).toHaveLength(1);
			} else {
				await new Promise((resolve) => setTimeout(resolve, 10));
				expect(session.providerState.callCount).toBe(2);
			}
			expect(continuations("afk-continue")).toHaveLength(1);
			expect(session.errors).toEqual([]);
		},
	);

	it.each(["first", "last"])(
		"does not restart napping AFK loaded %s, even without its completion tool active",
		async (order) => {
			const paths = [
				extensionPath(),
				fileURLToPath(new URL("../../continuity/index.ts", import.meta.url)),
			];
			if (order === "last") paths.reverse();
			session = await createPiSession({
				extensionPaths: paths,
				responses: [
					fauxAssistantMessage(fauxToolCall("afk_complete", {}), {
						stopReason: "toolUse",
					}),
					fauxAssistantMessage(finalReport),
				],
			});
			await runAfk();
			session.session.setActiveToolsByName(
				session.session
					.getActiveToolNames()
					.filter((name) => name !== "afk_complete"),
			);
			const send = vi.spyOn(session.session, "sendCustomMessage");
			await session.session.extensionRunner.emit({
				type: "session_compact",
				compactionEntry: { id: "afk-compact" },
				fromExtension: false,
				reason: "manual",
				willRetry: false,
			} as never);
			await new Promise((resolve) => setTimeout(resolve, 10));
			expect(send).not.toHaveBeenCalled();
			expect(session.providerState.callCount).toBe(2);
		},
	);

	it("writes only schema-bound lifecycle diagnostics when explicitly enabled", async () => {
		session = await createPiSession({
			extensionPaths: [extensionPath()],
			env: { PI_AFK_DEBUG: "1" },
		});
		await session.session.extensionRunner.emit({
			type: "session_start",
			reason: "startup",
		});
		await session.session.extensionRunner.emit({
			type: "session_shutdown",
			reason: "quit",
		});
		const directory = join(
			session.env["XDG_STATE_HOME"]!,
			"pi-ext",
			"debug",
			"afk",
		);
		const events = await vi.waitFor(() => {
			const records = readdirSync(directory).flatMap((file) =>
				readFileSync(join(directory, file), "utf8")
					.trim()
					.split("\n")
					.filter(Boolean)
					.map((line) => JSON.parse(line)),
			);
			expect(records).toContainEqual(
				expect.objectContaining({
					extension: "afk",
					event: "session.shutdown",
				}),
			);
			return records.map((record) => record.event);
		});
		expect(events).toContain("session.start");
	});

	it("cleans AFK mode on session-tree navigation and shutdown", async () => {
		session = await createPiSession({ extensionPaths: [extensionPath()] });
		await command("afk").handler(
			"",
			session.session.extensionRunner.createContext() as never,
		);
		await session.session.extensionRunner.emit({
			type: "session_tree",
			newLeafId: null,
			oldLeafId: null,
		});
		expect(session.session.getActiveToolNames()).not.toContain("afk_complete");
		await session.session.extensionRunner.emit({
			type: "session_shutdown",
			reason: "quit",
		});
		expect(session.errors).toEqual([]);
	});
});