pi-ext

import { createEventBus } from "@earendil-works/pi-coding-agent";
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
import afkExtension from "../index.ts";

type Handler = (event: never, ctx: never) => unknown;

function createHarness() {
	const handlers = new Map<string, Handler>();
	const commands = new Map<
		string,
		{ handler(args: string, ctx: never): Promise<void> }
	>();
	const sent: Array<{ customType: string; display: boolean }> = [];
	const events = createEventBus();
	let tool: { execute(): Promise<unknown> };
	let tools: string[] = [];
	let sessionId = "session-1";
	let idle = true;
	let pending = false;
	let sendError = false;
	const notify = vi.fn();
	const ctx = {
		mode: "print",
		sessionManager: { getSessionId: () => sessionId },
		isIdle: () => idle,
		hasPendingMessages: () => pending,
		ui: { notify },
	};
	afkExtension({
		events,
		on(name: string, handler: Handler) {
			handlers.set(name, handler);
		},
		registerCommand(
			name: string,
			command: { handler(args: string, ctx: never): Promise<void> },
		) {
			commands.set(name, command);
		},
		registerTool(definition: typeof tool) {
			tool = definition;
		},
		getActiveTools: () => tools,
		setActiveTools(names: string[]) {
			tools = names;
		},
		sendMessage(message: (typeof sent)[number]) {
			if (sendError) throw new Error("dispatch failed");
			sent.push(message);
		},
	} as never);
	const emit = (name: string, event: unknown = {}) =>
		handlers.get(name)?.(event as never, ctx as never);
	return {
		events,
		notify,
		emit,
		async command(name: string) {
			await commands.get(name)?.handler("", ctx as never);
		},
		async complete() {
			await tool.execute();
			emit("agent_before_settle", { outcome: "completed" });
		},
		compact(overrides: object = {}) {
			emit("session_compact", { willRetry: false, ...overrides });
		},
		continuations: () =>
			sent.filter((message) => message.customType === "afk-continue"),
		ownership(id = sessionId) {
			const states: unknown[] = [];
			events.emit("continuation:query", {
				sessionId: id,
				respond: (state: unknown) => states.push(state),
			});
			return states;
		},
		setIdle(value: boolean) {
			idle = value;
		},
		setPending(value: boolean) {
			pending = value;
		},
		setSessionId(value: string) {
			sessionId = value;
		},
		failDispatch() {
			sendError = true;
		},
	};
}

beforeEach(() => vi.useFakeTimers());
afterEach(() => vi.useRealTimers());

describe("AFK continuation ownership", () => {
	it("answers synchronously only for the active session, including napping state", async () => {
		const h = createHarness();
		expect(h.ownership()).toEqual([]);
		await h.command("afk");
		expect(h.ownership()).toEqual(["working"]);
		expect(h.ownership("foreign-session")).toEqual([]);
		for (const data of [
			null,
			{},
			{ sessionId: "session-1", respond: "invalid" },
		])
			h.events.emit("continuation:query", data);
		await h.complete();
		expect(h.ownership()).toEqual(["done"]);
		await h.command("back");
		expect(h.ownership()).toEqual([]);
	});

	it("owns new user work after napping and releases ownership on shutdown", async () => {
		const h = createHarness();
		await h.command("afk");
		await h.complete();
		h.emit("message_start", { message: { role: "user" } });
		expect(h.ownership()).toEqual(["working"]);
		h.emit("session_shutdown");
		expect(h.ownership()).toEqual([]);
		h.emit("session_start");
		await h.command("afk");
		expect(h.ownership()).toEqual(["working"]);
	});

	it("resumes idle compaction without Continuity, replacing duplicate schedules", async () => {
		const h = createHarness();
		await h.command("afk");
		h.compact();
		h.compact();
		await vi.runAllTimersAsync();
		expect(h.continuations()).toHaveLength(1);
		expect(h.continuations()[0]?.display).toBe(false);
	});

	it.each(["inactive", "done", "retry", "pending", "busy"])(
		"does not resume %s compaction",
		async (state) => {
			const h = createHarness();
			if (state !== "inactive") await h.command("afk");
			if (state === "done") await h.complete();
			if (state === "pending") h.setPending(true);
			if (state === "busy") h.setIdle(false);
			h.compact({ willRetry: state === "retry" });
			await vi.runAllTimersAsync();
			expect(h.continuations()).toEqual([]);
		},
	);

	it.each([
		"back",
		"start",
		"tree",
		"shutdown",
		"session",
		"pending",
		"completed",
		"abort",
	])("cancels stale compaction after %s", async (change) => {
		const h = createHarness();
		await h.command("afk");
		h.compact();
		if (change === "back") await h.command("back");
		if (change === "start") h.emit("agent_start");
		if (change === "tree") h.emit("session_tree");
		if (change === "shutdown") h.emit("session_shutdown");
		if (change === "session") h.setSessionId("replacement");
		if (change === "pending") h.setPending(true);
		if (change === "completed") await h.complete();
		if (change === "abort")
			h.emit("agent_before_settle", { outcome: "aborted" });
		await vi.runAllTimersAsync();
		expect(h.continuations()).toEqual([]);
	});

	it("disables AFK and reports a failed compaction dispatch", async () => {
		const h = createHarness();
		await h.command("afk");
		h.failDispatch();
		h.compact();
		await vi.runAllTimersAsync();
		expect(h.ownership()).toEqual([]);
		expect(h.notify).toHaveBeenLastCalledWith(
			"AFK could not resume: dispatch failed",
			"error",
		);
	});
});