pi-ext

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

type Handler = (event: any, ctx: any) => unknown;
type ContinuationQuery = { sessionId: string; respond(state: unknown): void };

function compactionEvent(
	overrides: Record<string, unknown> = {},
): Record<string, unknown> {
	return {
		type: "session_compact",
		compactionEntry: { id: "compact-1" },
		fromExtension: false,
		reason: "manual",
		willRetry: false,
		...overrides,
	};
}

function createHarness() {
	const handlers = new Map<string, Handler[]>();
	const sent: Array<{ message: any; options: any }> = [];
	const notices: Array<[string, string]> = [];
	let activeTools: string[] = [];
	let idle = true;
	let pendingMessages = false;
	let sessionId = "session-1";
	let sessionFile: string | undefined = "/sessions/work.jsonl";
	let branch: any[] = [];
	let sendError: Error | undefined;

	const pi = {
		events: createEventBus(),
		on(event: string, handler: Handler) {
			handlers.set(event, [...(handlers.get(event) ?? []), handler]);
		},
		getActiveTools: () => activeTools,
		sendMessage(message: any, options: any) {
			if (sendError) throw sendError;
			sent.push({ message, options });
		},
	};
	const ctx = {
		isIdle: () => idle,
		hasPendingMessages: () => pendingMessages,
		sessionManager: {
			getSessionId: () => sessionId,
			getSessionFile: () => sessionFile,
			getBranch: () => branch,
		},
		ui: {
			notify(message: string, severity: string) {
				notices.push([message, severity]);
			},
		},
	};
	continuityExtension(pi as never);

	return {
		events: pi.events,
		sent,
		notices,
		ctx,
		emit(event: string, payload: any = {}) {
			for (const handler of handlers.get(event) ?? []) handler(payload, ctx);
		},
		setActiveTools(names: string[]) {
			activeTools = names;
		},
		setIdle(value: boolean) {
			idle = value;
		},
		setPendingMessages(value: boolean) {
			pendingMessages = value;
		},
		setSessionId(value: string) {
			sessionId = value;
		},
		setSessionFile(value: string | undefined) {
			sessionFile = value;
		},
		setBranch(value: any[]) {
			branch = value;
		},
		setSendError(error: Error | undefined) {
			sendError = error;
		},
	};
}

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

describe("continuation prompt", () => {
	it("uses the summary first and persisted active history only as fallback", () => {
		const prompt = buildContinuationPrompt(
			'/sessions/a "quoted" session.jsonl',
			"compact-7",
		);

		expect(prompt).toContain(
			"Recover the user's current intent from the compaction summary and current worktree.",
		);
		expect(prompt).toContain(
			'Persisted session JSONL: "/sessions/a \\"quoted\\" session.jsonl".',
		);
		expect(prompt).toContain("following parentId links");
		expect(prompt).toContain("Do not treat append order as branch order");
		expect(prompt).toContain('Compaction entry ID: "compact-7".');
		expect(prompt).toContain("If actionable work remains, continue it.");
		expect(prompt).toContain("all remaining work is blocked");
		expect(prompt).toContain("Do not rerun completed checks");
		expect(prompt).not.toContain("immediately perform");
	});

	it("handles ephemeral sessions without claiming history exists", () => {
		const prompt = buildContinuationPrompt(undefined, "compact-1");

		expect(prompt).toContain("This session is ephemeral");
		expect(prompt).not.toContain("Persisted session JSONL");
		expect(prompt).toContain('Compaction entry ID: "compact-1".');
	});
});

describe("continuity lifecycle", () => {
	it.each(["manual", "threshold"])(
		"continues once after successful %s compaction",
		async (reason) => {
			const h = createHarness();

			h.emit("session_compact", compactionEvent({ reason }));
			await vi.runAllTimersAsync();

			expect(h.sent).toHaveLength(1);
			expect(h.sent[0]?.message).toMatchObject({
				customType: __test.CONTINUATION_TYPE,
				display: false,
				details: { compactionEntryId: "compact-1", reason },
			});
			expect(h.sent[0]?.options).toEqual({
				triggerTurn: true,
				deliverAs: "followUp",
			});
		},
	);

	it("does not duplicate Pi's overflow retry", async () => {
		const h = createHarness();

		h.emit(
			"session_compact",
			compactionEvent({ reason: "overflow", willRetry: true }),
		);
		await vi.runAllTimersAsync();

		expect(h.sent).toEqual([]);
	});

	it("continues after non-retrying overflow compaction", async () => {
		const h = createHarness();

		h.emit(
			"session_compact",
			compactionEvent({ reason: "overflow", willRetry: false }),
		);
		await vi.runAllTimersAsync();

		expect(h.sent).toHaveLength(1);
	});

	it("leaves an already queued continuation or user message alone", async () => {
		const h = createHarness();
		h.setPendingMessages(true);

		h.emit("session_compact", compactionEvent());
		await vi.runAllTimersAsync();

		expect(h.sent).toEqual([]);
	});

	it("yields when a message becomes pending during deferral", async () => {
		const h = createHarness();

		h.emit("session_compact", compactionEvent());
		h.setPendingMessages(true);
		await vi.runAllTimersAsync();

		expect(h.sent).toEqual([]);
	});

	it("yields when another run starts during deferral", async () => {
		const h = createHarness();

		h.emit("session_compact", compactionEvent());
		h.emit("agent_start");
		await vi.runAllTimersAsync();

		expect(h.sent).toEqual([]);
	});

	it("does not dispatch while Pi remains busy", async () => {
		const h = createHarness();
		h.setIdle(false);

		h.emit("session_compact", compactionEvent());
		await vi.runAllTimersAsync();

		expect(h.sent).toEqual([]);
	});

	it("does not deliver into a replacement session", async () => {
		const h = createHarness();

		h.emit("session_compact", compactionEvent());
		h.setSessionId("session-2");
		await vi.runAllTimersAsync();

		expect(h.sent).toEqual([]);
	});

	it("replaces an obsolete scheduled continuation with the latest compaction", async () => {
		const h = createHarness();

		h.emit("session_compact", compactionEvent());
		h.emit(
			"session_compact",
			compactionEvent({ compactionEntry: { id: "compact-2" } }),
		);
		await vi.runAllTimersAsync();

		expect(h.sent).toHaveLength(1);
		expect(h.sent[0]?.message.details.compactionEntryId).toBe("compact-2");
	});

	it("clears scheduled work on session shutdown", async () => {
		const h = createHarness();

		h.emit("session_compact", compactionEvent());
		h.emit("session_shutdown");
		await vi.runAllTimersAsync();

		expect(h.sent).toEqual([]);
	});

	it.each(["update_goal"])(
		"yields to specialized continuation capability %s",
		async (tool) => {
			const h = createHarness();
			h.setActiveTools([tool]);

			h.emit("session_compact", compactionEvent());
			await vi.runAllTimersAsync();

			expect(h.sent).toEqual([]);
		},
	);

	it("yields to an active first-turn goal before update_goal becomes active", async () => {
		const h = createHarness();
		h.setBranch([
			{
				type: "custom",
				customType: __test.GOAL_STATE_TYPE,
				data: { status: "active", sessionId: "session-1", turnsUsed: 0 },
			},
		]);

		h.emit("session_compact", compactionEvent());
		await vi.runAllTimersAsync();

		expect(h.sent).toEqual([]);
	});

	it.each([
		["paused goal", { status: "paused", sessionId: "session-1" }],
		["foreign goal", { status: "active", sessionId: "session-2" }],
		["cleared goal", { cleared: true }],
	])("does not yield to %s state", async (_name, data) => {
		const h = createHarness();
		h.setBranch([{ type: "custom", customType: __test.GOAL_STATE_TYPE, data }]);

		h.emit("session_compact", compactionEvent());
		await vi.runAllTimersAsync();

		expect(h.sent).toHaveLength(1);
	});

	it("uses the latest goal state event", async () => {
		const h = createHarness();
		h.setBranch([
			{
				type: "custom",
				customType: __test.GOAL_STATE_TYPE,
				data: { status: "active", sessionId: "session-1" },
			},
			{
				type: "custom",
				customType: __test.GOAL_STATE_TYPE,
				data: { status: "paused", sessionId: "session-1" },
			},
		]);

		h.emit("session_compact", compactionEvent());
		await vi.runAllTimersAsync();

		expect(h.sent).toHaveLength(1);
	});

	it.each(["update_goal"])(
		"rechecks specialized continuation capability %s during deferral",
		async (tool) => {
			const h = createHarness();

			h.emit("session_compact", compactionEvent());
			h.setActiveTools([tool]);
			await vi.runAllTimersAsync();

			expect(h.sent).toEqual([]);
		},
	);

	it.each(["working", "done"])(
		"yields to synchronous ownership reply %s",
		async (state) => {
			const h = createHarness();
			h.events.on("continuation:query", (data) => {
				const query = data as ContinuationQuery;
				expect(query.sessionId).toBe("session-1");
				query.respond(state);
			});
			h.emit("session_compact", compactionEvent());
			await vi.runAllTimersAsync();
			expect(h.sent).toEqual([]);
		},
	);

	it.each(["working", "done"])(
		"rechecks ownership reply %s during deferral",
		async (state) => {
			const h = createHarness();
			h.emit("session_compact", compactionEvent());
			h.events.on("continuation:query", (data) =>
				(data as ContinuationQuery).respond(state),
			);
			await vi.runAllTimersAsync();
			expect(h.sent).toEqual([]);
		},
	);

	it("ignores late asynchronous ownership replies", async () => {
		const h = createHarness();
		h.events.on("continuation:query", async (data) => {
			await Promise.resolve();
			(data as ContinuationQuery).respond("done");
		});
		h.emit("session_compact", compactionEvent());
		await vi.runAllTimersAsync();
		expect(h.sent).toHaveLength(1);
	});

	it("ignores invalid ownership replies and completion-tool names", async () => {
		const h = createHarness();
		h.setActiveTools(["afk_complete"]);
		h.events.on("continuation:query", (data) =>
			(data as ContinuationQuery).respond("idle"),
		);
		h.emit("session_compact", compactionEvent());
		await vi.runAllTimersAsync();
		expect(h.sent).toHaveLength(1);
	});

	it("does not yield merely because inactive tools exist", async () => {
		const h = createHarness();
		h.setActiveTools(["read", "bash"]);

		h.emit("session_compact", compactionEvent());
		await vi.runAllTimersAsync();

		expect(h.sent).toHaveLength(1);
	});

	it("reports dispatch failures without throwing from the timer", async () => {
		const h = createHarness();
		h.setSendError(new Error("stale runtime"));

		h.emit("session_compact", compactionEvent());
		await vi.runAllTimersAsync();

		expect(h.sent).toEqual([]);
		expect(h.notices).toEqual([
			["Continuity could not resume: stale runtime", "error"],
		]);
	});

	it("uses ephemeral context from the live session", async () => {
		const h = createHarness();
		h.setSessionFile(undefined);

		h.emit("session_compact", compactionEvent());
		await vi.runAllTimersAsync();

		expect(h.sent[0]?.message.content).toContain("This session is ephemeral");
	});
});