Luigit
repositories / pi-ext

pi-ext

bugabingas pi extensions

owned by admin

extensions/masks/__tests__/harness.test.ts

Raw
import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { fileURLToPath } from "node:url";
import { getCurrentSystemPrompt, getCurrentTools } from "@earendil-works/pi-ai";
import { afterEach, describe, expect, it, vi } from "vitest";
import { createTestSession, type TestSession } from "../../../test/harness";
import { __test } from "../index.ts";

let root: string | undefined;
let t: TestSession | undefined;

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

function writeSettings(settings: Record<string, unknown>): string {
	root = mkdtempSync(join(tmpdir(), "masks-test-"));
	const agentDir = join(root, ".test-home", ".pi", "agent");
	mkdirSync(agentDir, { recursive: true });
	writeFileSync(join(agentDir, "settings.json"), JSON.stringify(settings));
	return root;
}

function writePrompt(name: string, content: string): void {
	if (!root) throw new Error("missing test root");
	const promptsDir = join(root, "prompts");
	mkdirSync(promptsDir, { recursive: true });
	writeFileSync(join(promptsDir, `${name}.md`), content);
}

function extension() {
	if (!t) throw new Error("missing test session");
	return t.session.extensionRunner.extensions[0];
}

function context() {
	if (!t) throw new Error("missing test session");
	return t.session.extensionRunner.createCommandContext();
}

type Request = { text: string; model: string; thinking: string | undefined };

async function queuedSession(
	onRequest: (
		request: Request,
		index: number,
		tools: string[],
		systemPrompt: string,
	) => Promise<Record<string, unknown> | undefined>,
): Promise<Request[]> {
	const cwd = writeSettings({
		masks: {
			items: [
				{ name: "focused", model: "openai/gpt-4o-mini", thinkingLevel: "high" },
				{ name: "quick", model: "openai/gpt-4o", thinkingLevel: "low" },
			],
		},
	});
	writePrompt("review", "---\nmask: focused\n---\nReview $1.");
	writePrompt("plan", "---\nmask: quick\n---\nPlan $1.");
	writeFileSync(join(cwd, "fixture.txt"), "fixture");
	t = await createTestSession({ cwd, extensions: [extensionPath()] });
	for (const name of ["gpt-4o", "gpt-4o-mini"]) {
		context().modelRegistry.find("openai", name).reasoning = true;
	}
	t.session.settingsManager.setRetryEnabled(false);
	// Fixture images are not decodable; keep Pi from resizing or omitting them.
	t.session.settingsManager.setImageAutoResize(false);
	const requests: Request[] = [];
	t.session.agent.streamFn = async (model: any, ctx: any, options: any) => {
		const request = {
			text: ctx.messages
				.filter((m: any) => m.role === "user")
				.at(-1)
				?.content.filter((p: any) => p.type === "text")
				.map((p: any) => p.text)
				.join(""),
			model: model.id,
			thinking: options.reasoning,
		};
		let reply: Record<string, unknown> | undefined;
		if (!options.signal.aborted) {
			requests.push(request);
			reply = await onRequest(
				request,
				requests.length - 1,
				getCurrentTools(ctx.messages).map((tool) => tool.name),
				getCurrentSystemPrompt(ctx.messages),
			);
		}
		const message = {
			role: "assistant",
			content: [{ type: "text", text: "Done." }],
			api: model.api,
			provider: model.provider,
			model: model.id,
			usage: {
				input: 1,
				output: 1,
				cacheRead: 0,
				cacheWrite: 0,
				totalTokens: 2,
				cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
			},
			stopReason: options.signal.aborted ? "aborted" : "stop",
			timestamp: Date.now(),
			...reply,
		};
		return {
			async *[Symbol.asyncIterator]() {
				yield { type: "done", reason: message.stopReason, message };
			},
			result: async () => message,
		};
	};
	return requests;
}

async function queue(text: string): Promise<void> {
	await t?.session.sendUserMessage(text, {
		expandPromptTemplates: true,
		deliverAs: "followUp",
	});
}

async function settled(count: number, requests: Request[]): Promise<void> {
	await expect.poll(() => requests.length).toBe(count);
	await expect.poll(() => t?.session.isIdle).toBe(true);
}

afterEach(() => {
	vi.restoreAllMocks();
	t?.dispose();
	t = undefined;
	if (root) rmSync(root, { recursive: true, force: true });
	root = undefined;
});

describe("masks pi-test-harness", () => {
	it("switches complete masks from direct shortcut, picker, and command", async () => {
		const cwd = writeSettings({
			masks: {
				pickerShortcut: "alt+m",
				items: [
					{
						name: "low",
						model: "openai/gpt-4o",
						thinkingLevel: "low",
						shortcut: "alt+1",
					},
					{
						name: "high",
						model: "openai/gpt-4o",
						thinkingLevel: "high",
						shortcut: "alt+2",
					},
					{
						name: "missing",
						model: "absent/nope",
						thinkingLevel: "max",
						shortcut: "alt+3",
					},
				],
			},
		});
		t = await createTestSession({
			cwd,
			extensions: [extensionPath()],
			mockUI: { select: 1 },
		});

		const model = context().modelRegistry.find("openai", "gpt-4o");
		if (!model) throw new Error("missing harness model");
		model.reasoning = true;

		expect([...extension().commands.keys()]).toEqual(["masks-cancel", "masks"]);
		expect([...extension().shortcuts.keys()]).toEqual([
			"alt+m",
			"alt+1",
			"alt+2",
			"alt+3",
		]);

		await extension().shortcuts.get("alt+1").handler(context());
		expect(t.session.model.provider).toBe("openai");
		expect(t.session.model.id).toBe("gpt-4o");
		expect(t.session.thinkingLevel).toBe("low");

		await extension().shortcuts.get("alt+m").handler(context());
		expect(t.session.thinkingLevel).toBe("high");
		expect(t.events.uiCallsFor("select")[0].args).toEqual([
			"Masks",
			[
				"low  openai/gpt-4o:low",
				"high  openai/gpt-4o:high",
				"missing  absent/nope:max",
			],
		]);

		await extension().commands.get("masks").handler("low", context());
		expect(t.session.thinkingLevel).toBe("low");

		await extension().shortcuts.get("alt+3").handler(context());
		expect(t.session.thinkingLevel).toBe("low");
		expect(t.events.uiCallsFor("notify").at(-1)?.args).toEqual([
			'Mask "missing": model absent/nope not found',
			"error",
		]);
	});

	it("applies prompt frontmatter for one agent run, then restores", async () => {
		const cwd = writeSettings({
			masks: {
				items: [
					{
						name: "focused",
						model: "openai/gpt-4o",
						thinkingLevel: "high",
					},
				],
			},
		});
		writePrompt(
			"review",
			"---\ndescription: Review changes\nmask: focused\n---\nReview it.",
		);
		t = await createTestSession({ cwd, extensions: [extensionPath()] });
		const model = context().modelRegistry.find("openai", "gpt-4o");
		if (!model) throw new Error("missing harness model");
		model.reasoning = true;

		const inputHandler = extension().handlers.get("input")[0];
		const endHandler = extension().handlers.get("agent_settled")[0];
		expect(
			await inputHandler(
				{ type: "input", text: "/review", source: "interactive" },
				context(),
			),
		).toEqual({ action: "continue" });

		expect(t.session.thinkingLevel).toBe("high");

		await endHandler({ type: "agent_settled" }, context());
		expect(t.session.model.provider).toBe("openai");
		expect(t.session.model.id).toBe("gpt-4o");
		expect(t.session.thinkingLevel).toBe("off");
	});

	it("applies a masked prompt tool profile and restores the exact snapshot", async () => {
		const cwd = writeSettings({
			masks: {
				items: [
					{
						name: "focused",
						model: "openai/gpt-4o",
						thinkingLevel: "high",
					},
				],
			},
		});
		writePrompt(
			"review",
			"---\nmask: focused\ntools: [read, ls]\n---\nReview it.",
		);
		t = await createTestSession({ cwd, extensions: [extensionPath()] });
		const model = context().modelRegistry.find("openai", "gpt-4o");
		if (!model) throw new Error("missing harness model");
		model.reasoning = true;
		const runtime = t.session.extensionRunner.runtime;
		const originalTools = runtime.getActiveTools();

		expect(
			await extension().handlers.get("input")[0](
				{ type: "input", text: "/review", source: "interactive" },
				context(),
			),
		).toEqual({ action: "continue" });
		expect(runtime.getActiveTools()).toEqual(["read", "ls"]);
		await extension().handlers.get("agent_start")[0](
			{ type: "agent_start" },
			context(),
		);

		runtime.setActiveTools(["read", "ls", "bash"]);
		await extension().handlers.get("agent_settled")[0](
			{ type: "agent_settled" },
			context(),
		);
		expect(runtime.getActiveTools()).toEqual(originalTools);
	});

	it("sends a consistent tool profile and system prompt to the provider", async () => {
		const toolSets: string[][] = [];
		const systemPrompts: string[] = [];
		const requests = await queuedSession(
			async (_request, _index, tools, systemPrompt) => {
				toolSets.push(tools);
				systemPrompts.push(systemPrompt);
				return undefined;
			},
		);
		writePrompt(
			"review",
			"---\nmask: focused\ntools: [read, ls]\n---\nReview $1.",
		);

		await t?.session.prompt("/review profiled");
		await settled(1, requests);
		expect(toolSets).toEqual([["read", "ls"]]);
		expect(systemPrompts[0]).toContain("- read:");
		expect(systemPrompts[0]).toContain("- ls:");
		expect(systemPrompts[0]).not.toContain("- bash:");
		expect(t?.session.extensionRunner.runtime.getActiveTools()).toEqual([
			"read",
			"bash",
			"edit",
			"write",
		]);
	});

	it("leaves tools alone when masked prompt frontmatter omits them", async () => {
		const cwd = writeSettings({
			masks: {
				items: [
					{
						name: "focused",
						model: "openai/gpt-4o",
						thinkingLevel: "high",
					},
				],
			},
		});
		writePrompt("review", "---\nmask: focused\n---\nReview it.");
		t = await createTestSession({ cwd, extensions: [extensionPath()] });
		const runtime = t.session.extensionRunner.runtime;
		await extension().handlers.get("input")[0](
			{ type: "input", text: "/review", source: "interactive" },
			context(),
		);
		await extension().handlers.get("agent_start")[0](
			{ type: "agent_start" },
			context(),
		);
		runtime.setActiveTools(["read", "ls"]);
		await extension().handlers.get("agent_settled")[0](
			{ type: "agent_settled" },
			context(),
		);
		expect(runtime.getActiveTools()).toEqual(["read", "ls"]);
	});

	it("rejects invalid prompt tool profiles before changing the mask", async () => {
		const cwd = writeSettings({
			masks: {
				items: [
					{
						name: "focused",
						model: "openai/gpt-4o-mini",
						thinkingLevel: "high",
					},
				],
			},
		});
		writePrompt(
			"review",
			"---\nmask: focused\ntools: [read, absent]\n---\nReview it.",
		);
		t = await createTestSession({ cwd, extensions: [extensionPath()] });
		const runtime = t.session.extensionRunner.runtime;
		const originalTools = runtime.getActiveTools();
		const originalModel = t.session.model.id;

		expect(
			await extension().handlers.get("input")[0](
				{ type: "input", text: "/review", source: "interactive" },
				context(),
			),
		).toEqual({ action: "handled" });
		expect(runtime.getActiveTools()).toEqual(originalTools);
		expect(t.session.model.id).toBe(originalModel);
		expect(t.events.uiCallsFor("notify").at(-1)?.args).toEqual([
			'Prompt /review: unknown tool "absent"',
			"error",
		]);
	});

	it("rolls back model, thinking, and tools when tool application fails", async () => {
		const cwd = writeSettings({
			masks: {
				items: [
					{
						name: "focused",
						model: "openai/gpt-4o-mini",
						thinkingLevel: "high",
					},
				],
			},
		});
		writePrompt(
			"review",
			"---\nmask: focused\ntools: [read, ls]\n---\nReview it.",
		);
		t = await createTestSession({ cwd, extensions: [extensionPath()] });
		const runtime = t.session.extensionRunner.runtime;
		const originalTools = runtime.getActiveTools();
		const setActiveTools = runtime.setActiveTools.bind(runtime);
		vi.spyOn(runtime, "setActiveTools").mockImplementation((tools) => {
			setActiveTools(tools);
			if (tools.join(",") === "read,ls") throw new Error("profile failure");
		});

		expect(
			await extension().handlers.get("input")[0](
				{ type: "input", text: "/review", source: "interactive" },
				context(),
			),
		).toEqual({ action: "handled" });
		expect(runtime.getActiveTools()).toEqual(originalTools);
		expect(t.session.model.id).toBe("gpt-4o");
		expect(t.session.thinkingLevel).toBe("off");
	});

	it("restores tools even when model restoration fails", async () => {
		const cwd = writeSettings({
			masks: {
				items: [
					{
						name: "focused",
						model: "openai/gpt-4o-mini",
						thinkingLevel: "high",
					},
				],
			},
		});
		writePrompt(
			"review",
			"---\nmask: focused\ntools: [read, ls]\n---\nReview it.",
		);
		t = await createTestSession({ cwd, extensions: [extensionPath()] });
		const runtime = t.session.extensionRunner.runtime;
		const originalTools = runtime.getActiveTools();
		const setModel = runtime.setModel.bind(runtime);
		let failRestore = false;
		vi.spyOn(runtime, "setModel").mockImplementation(async (model) => {
			if (failRestore && model.id === "gpt-4o") return false;
			return setModel(model);
		});

		await extension().handlers.get("input")[0](
			{ type: "input", text: "/review", source: "interactive" },
			context(),
		);
		await extension().handlers.get("agent_start")[0](
			{ type: "agent_start" },
			context(),
		);
		failRestore = true;
		await extension().handlers.get("agent_settled")[0](
			{ type: "agent_settled" },
			context(),
		);
		expect(runtime.getActiveTools()).toEqual(originalTools);
		expect(t.events.uiCallsFor("notify").at(-1)?.args[0]).toContain(
			"Failed to restore prompt mask",
		);
	});

	it("restores after tree navigation races temporary profile application", async () => {
		const cwd = writeSettings({
			masks: {
				items: [
					{
						name: "focused",
						model: "openai/gpt-4o-mini",
						thinkingLevel: "high",
					},
				],
			},
		});
		writePrompt(
			"review",
			"---\nmask: focused\ntools: [read, ls]\n---\nReview it.",
		);
		t = await createTestSession({ cwd, extensions: [extensionPath()] });
		const runtime = t.session.extensionRunner.runtime;
		const originalTools = runtime.getActiveTools();
		const setModel = runtime.setModel.bind(runtime);
		const entered = Promise.withResolvers<void>();
		const release = Promise.withResolvers<void>();
		vi.spyOn(runtime, "setModel").mockImplementation(async (model) => {
			if (model.id === "gpt-4o-mini") {
				entered.resolve();
				await release.promise;
			}
			return setModel(model);
		});
		const applying = extension().handlers.get("input")[0](
			{ type: "input", text: "/review", source: "interactive" },
			context(),
		);
		await entered.promise;
		const resetting = extension().handlers.get("session_tree")[0](
			{ type: "session_tree" },
			context(),
		);
		release.resolve();
		await Promise.all([applying, resetting]);
		expect(t.session.model.id).toBe("gpt-4o");
		expect(t.session.thinkingLevel).toBe("off");
		expect(runtime.getActiveTools()).toEqual(originalTools);
	});

	it("keeps reset joined to profile application after a manual switch is rejected", async () => {
		const cwd = writeSettings({
			masks: {
				items: [
					{
						name: "focused",
						model: "openai/gpt-4o-mini",
						thinkingLevel: "high",
					},
				],
			},
		});
		writePrompt(
			"review",
			"---\nmask: focused\ntools: [read, ls]\n---\nReview it.",
		);
		t = await createTestSession({ cwd, extensions: [extensionPath()] });
		const runtime = t.session.extensionRunner.runtime;
		const originalTools = runtime.getActiveTools();
		const setModel = runtime.setModel.bind(runtime);
		const entered = Promise.withResolvers<void>();
		const release = Promise.withResolvers<void>();
		vi.spyOn(runtime, "setModel").mockImplementation(async (model) => {
			if (model.id === "gpt-4o-mini") {
				entered.resolve();
				await release.promise;
			}
			return setModel(model);
		});
		const applying = extension().handlers.get("input")[0](
			{ type: "input", text: "/review", source: "interactive" },
			context(),
		);
		await entered.promise;
		await extension().commands.get("masks").handler("focused", context());
		const resetting = extension().handlers.get("session_tree")[0](
			{ type: "session_tree" },
			context(),
		);
		release.resolve();
		await Promise.all([applying, resetting]);
		expect(t.session.model.id).toBe("gpt-4o");
		expect(runtime.getActiveTools()).toEqual(originalTools);
	});

	it("cancels later deferred masks when a queued tool profile is invalid", async () => {
		const requests = await queuedSession(async (_request, index) => {
			if (index === 0) {
				await queue("/review invalid");
				await queue("/plan later");
			}
			return undefined;
		});
		writePrompt(
			"review",
			"---\nmask: focused\ntools: [read, absent]\n---\nReview $1.",
		);
		await t?.session.prompt("start");
		await settled(1, requests);
		expect(
			t?.events
				.uiCallsFor("notify")
				.some((call) => String(call.args[0]).includes("Cancelled 1 deferred")),
		).toBe(true);
	});

	it("cancels later deferred masks when queued profile application fails", async () => {
		const requests = await queuedSession(async (_request, index) => {
			if (index === 0) {
				await queue("/review unavailable");
				await queue("/plan stale");
			}
			return undefined;
		});
		const runtime = t?.session.extensionRunner.runtime;
		const setModel = runtime.setModel.bind(runtime);
		vi.spyOn(runtime, "setModel").mockImplementation(async (model) =>
			model.id === "gpt-4o-mini" ? false : setModel(model),
		);
		await t?.session.prompt("start");
		await settled(1, requests);
		expect(
			t?.events
				.uiCallsFor("notify")
				.some((call) => String(call.args[0]).includes("Cancelled 1 deferred")),
		).toBe(true);
	});

	it("defers masked follow-ups without changing the running model", async () => {
		const cwd = writeSettings({
			masks: {
				items: [
					{
						name: "focused",
						model: "openai/gpt-4o",
						thinkingLevel: "high",
					},
				],
			},
		});
		writePrompt("review", "---\nmask: focused\n---\nReview it.");
		t = await createTestSession({ cwd, extensions: [extensionPath()] });
		const busyContext = { ...context(), isIdle: () => false };
		const result = await extension().handlers.get("input")[0](
			{
				type: "input",
				text: "/review later",
				source: "interactive",
				streamingBehavior: "followUp",
			},
			busyContext,
		);

		expect(result).toEqual({ action: "handled" });
		expect(t.session.thinkingLevel).toBe("off");
		expect(t.events.uiCallsFor("setStatus").at(-1)?.args).toEqual([
			"masks-queue",
			"masks: 1 deferred · /masks-cancel",
		]);
	});

	it("runs deferred masks FIFO after native follow-ups and restores between them", async () => {
		const requests = await queuedSession(async (_request, index) => {
			if (index === 0) {
				await queue("/review first");
				await queue("ordinary follow-up");
				await queue("/plan second");
				expect(t?.session.model.id).toBe("gpt-4o");
			}
			return undefined;
		});
		await t?.session.prompt("start");
		await settled(4, requests);
		expect(requests).toEqual([
			{ text: "start", model: "gpt-4o", thinking: undefined },
			{ text: "ordinary follow-up", model: "gpt-4o", thinking: undefined },
			{ text: "Review first.", model: "gpt-4o-mini", thinking: "high" },
			{ text: "Plan second.", model: "gpt-4o", thinking: "low" },
		]);
		expect(t?.session.model.id).toBe("gpt-4o");
		expect(t?.session.thinkingLevel).toBe("off");
	});

	it("keeps the mask through tools, then restores before an ordinary follow-up", async () => {
		const requests = await queuedSession(async (_request, index) => {
			if (index === 0) {
				await queue("ordinary follow-up");
				return {
					content: [
						{
							type: "toolCall",
							id: "read-fixture",
							name: "read",
							arguments: { path: "fixture.txt" },
						},
					],
					stopReason: "toolUse",
				};
			}
			return undefined;
		});
		await t?.session.prompt("/review first");
		await settled(3, requests);
		expect(
			requests.map(({ model, thinking }) => ({ model, thinking })),
		).toEqual([
			{ model: "gpt-4o-mini", thinking: "high" },
			{ model: "gpt-4o-mini", thinking: "high" },
			{ model: "gpt-4o", thinking: undefined },
		]);
	});

	it.each(["aborted", "error"])(
		"discards deferred masks on %s and restores",
		async (stopReason) => {
			const requests = await queuedSession(async () => {
				await queue("/plan later");
				return { stopReason, errorMessage: "test failure" };
			});
			await t?.session.prompt("/review now");
			await settled(1, requests);
			expect(t?.session.model.id).toBe("gpt-4o");
			expect(t?.session.thinkingLevel).toBe("off");
			expect(t?.events.uiCallsFor("notify").at(-1)?.args).toEqual([
				"Cancelled 1 deferred masked prompts",
				"warning",
			]);
		},
	);

	it("cancels deferred prompts explicitly without stopping the current run", async () => {
		const requests = await queuedSession(async () => {
			await queue("/plan later");
			await t?.session.prompt("/masks-cancel");
			expect(context().signal?.aborted).toBe(false);
			return undefined;
		});
		await t?.session.prompt("/review now");
		await settled(1, requests);
		expect(t?.session.thinkingLevel).toBe("off");
	});

	it("rejects steering without changing the active mask", async () => {
		const requests = await queuedSession(async () => {
			await t?.session.sendUserMessage("/plan later", {
				expandPromptTemplates: true,
				deliverAs: "steer",
			});
			expect(t?.session.thinkingLevel).toBe("high");
			return undefined;
		});
		await t?.session.prompt("/review now");
		await settled(1, requests);
		expect(t?.events.uiCallsFor("notify").at(-1)?.args).toEqual([
			"Prompt /plan: use follow-up delivery, not steering",
			"warning",
		]);
	});

	it.each(["session_tree", "session_shutdown"])(
		"clears deferred prompts on %s",
		async (type) => {
			const requests = await queuedSession(async () => undefined);
			await extension().handlers.get("input")[0](
				{
					type: "input",
					text: "/plan later",
					source: "extension",
					streamingBehavior: "followUp",
				},
				{ ...context(), isIdle: () => false },
			);
			await extension().handlers.get(type)[0](
				{ type, reason: "reload" },
				context(),
			);
			await extension().handlers.get("agent_settled")[0](
				{ type: "agent_settled" },
				context(),
			);
			expect(requests).toHaveLength(0);
			expect(t?.events.uiCallsFor("setStatus").at(-1)?.args).toEqual([
				"masks-queue",
				undefined,
			]);
		},
	);

	it("does not run the prompt under the wrong model when equipping fails", async () => {
		const requests = await queuedSession(async () => undefined);
		vi.spyOn(t?.session.extensionRunner.runtime, "setModel").mockResolvedValue(
			false,
		);
		await t?.session.prompt("/review now");
		expect(requests).toHaveLength(0);
		expect(
			t?.events
				.uiCallsFor("notify")
				.some((call) => String(call.args[0]).includes("no credentials")),
		).toBe(true);
	});

	it("cancels deferred prompts and aborts when restoration fails", async () => {
		const requests = await queuedSession(async () => {
			await queue("/plan later");
			vi.spyOn(
				t?.session.extensionRunner.runtime,
				"setModel",
			).mockRejectedValue(new Error("offline"));
			return undefined;
		});
		await t?.session.prompt("/review now");
		await settled(1, requests);
		expect(
			t?.events
				.uiCallsFor("notify")
				.some((call) => String(call.args[0]).includes("Failed to restore")),
		).toBe(true);
	});

	it("retains the mask through automatic retry", async () => {
		const requests = await queuedSession(async (_request, index) =>
			index === 0
				? { stopReason: "error", errorMessage: "503 service unavailable" }
				: undefined,
		);
		vi.spyOn(t?.session.settingsManager, "getRetrySettings").mockReturnValue({
			enabled: true,
			maxRetries: 1,
			baseDelayMs: 1,
		});
		await t?.session.prompt("/review retry");
		await settled(2, requests);
		expect(
			requests.map(({ model, thinking }) => ({ model, thinking })),
		).toEqual([
			{ model: "gpt-4o-mini", thinking: "high" },
			{ model: "gpt-4o-mini", thinking: "high" },
		]);
		expect(t?.session.model.id).toBe("gpt-4o");
		expect(t?.session.thinkingLevel).toBe("off");
	});

	it("keeps the mask for overflow recovery instead of skipping the same-model check", async () => {
		const requests = await queuedSession(async (_request, index) =>
			index === 0
				? {
						stopReason: "error",
						errorMessage: "maximum context length exceeded",
					}
				: undefined,
		);
		const compact = vi
			.spyOn(t?.session, "_runAutoCompaction")
			.mockImplementation(async () => {
				expect(t?.session.model.id).toBe("gpt-4o-mini");
				return true;
			});
		await t?.session.prompt("/review overflow");
		await settled(2, requests);
		expect(compact).toHaveBeenCalledWith("overflow", true);
		expect(requests[1]).toEqual({
			text: "Review overflow.",
			model: "gpt-4o-mini",
			thinking: "high",
		});
		expect(t?.session.thinkingLevel).toBe("off");
	});

	it("rejects competing prompts during profile application without losing the mask", async () => {
		const requests = await queuedSession(async () => undefined);
		const entered = Promise.withResolvers<void>();
		const release = Promise.withResolvers<void>();
		const runtime = t?.session.extensionRunner.runtime;
		const setModel = runtime.setModel.bind(runtime);
		vi.spyOn(runtime, "setModel").mockImplementation(async (model) => {
			entered.resolve();
			await release.promise;
			return setModel(model);
		});
		const first = t?.session.prompt("/review first");
		await entered.promise;
		try {
			await t?.session.prompt("ordinary concurrent prompt");
			await t?.session.prompt("/plan second");
		} finally {
			release.resolve();
		}
		await first;
		expect(requests).toEqual([
			{ text: "Review first.", model: "gpt-4o-mini", thinking: "high" },
		]);
		expect(t?.session.model.id).toBe("gpt-4o");
		expect(t?.session.thinkingLevel).toBe("off");
	});

	it("blocks prompt dispatch while a manual mask switch is applying", async () => {
		const requests = await queuedSession(async () => undefined);
		const entered = Promise.withResolvers<void>();
		const release = Promise.withResolvers<void>();
		const runtime = t?.session.extensionRunner.runtime;
		const setModel = runtime.setModel.bind(runtime);
		vi.spyOn(runtime, "setModel").mockImplementation(async (model) => {
			entered.resolve();
			await release.promise;
			return setModel(model);
		});
		const manual = t?.session.prompt("/masks focused");
		await entered.promise;
		try {
			await t?.session.prompt("/plan blocked");
		} finally {
			release.resolve();
		}
		await manual;
		expect(requests).toHaveLength(0);
		expect(t?.session.model.id).toBe("gpt-4o-mini");
		expect(t?.session.thinkingLevel).toBe("high");
	});

	it("blocks later prompts until a failed restoration succeeds", async () => {
		let failRestore = true;
		const requests = await queuedSession(async () => undefined);
		const runtime = t?.session.extensionRunner.runtime;
		const setModel = runtime.setModel.bind(runtime);
		vi.spyOn(runtime, "setModel").mockImplementation(async (model: any) => {
			if (model.id === "gpt-4o" && failRestore) return false;
			return setModel(model);
		});
		await t?.session.prompt("/review first");
		await t?.session.prompt("ordinary blocked prompt");
		await t?.session.prompt("/plan blocked");
		expect(requests).toHaveLength(1);
		failRestore = false;
		await t?.session.prompt("ordinary recovered prompt");
		expect(requests[1]).toEqual({
			text: "ordinary recovered prompt",
			model: "gpt-4o",
			thinking: undefined,
		});
	});

	it("cancels deferred prompts on a late abort after a successful response", async () => {
		const requests = await queuedSession(async () => {
			await queue("/plan later");
			return undefined;
		});
		const runtime = t?.session.extensionRunner.runtime;
		const setModel = runtime.setModel.bind(runtime);
		vi.spyOn(runtime, "setModel").mockImplementation(async (model: any) => {
			if (model.id === "gpt-4o") context().abort();
			return setModel(model);
		});
		await t?.session.prompt("/review now");
		await settled(1, requests);
		expect(t?.events.uiCallsFor("notify").at(-1)?.args).toEqual([
			"Cancelled 1 deferred masked prompts",
			"warning",
		]);
		expect(t?.session.thinkingLevel).toBe("off");
	});

	it("restores before a native follow-up after a terminating tool batch", async () => {
		const requests = await queuedSession(async (_request, index) => {
			if (index !== 0) return undefined;
			await queue("ordinary follow-up");
			return {
				content: [
					{
						type: "toolCall",
						id: "blocked-read",
						name: "read",
						arguments: { path: "fixture.txt" },
					},
				],
				stopReason: "toolUse",
			};
		});
		extension().handlers.set("tool_call", [
			async () => ({ block: true, terminate: true, reason: "Finished" }),
		]);
		await t?.session.prompt("/review now");
		await settled(2, requests);
		expect(requests[1]).toEqual({
			text: "ordinary follow-up",
			model: "gpt-4o",
			thinking: undefined,
		});
	});

	it("preserves images and repeated masked follow-ups", async () => {
		const image = {
			type: "image",
			mimeType: "image/png",
			data: "Zml4dHVyZS1pbWFnZQ==",
		};
		const requests = await queuedSession(async (_request, index) => {
			if (index === 0) {
				await t?.session.sendUserMessage(
					[{ type: "text", text: "/review same" }, image],
					{
						expandPromptTemplates: true,
						deliverAs: "followUp",
					},
				);
				await queue("/review same");
			}
			return undefined;
		});
		await t?.session.prompt("start");
		await settled(3, requests);
		expect(requests.slice(1)).toEqual([
			{ text: "Review same.", model: "gpt-4o-mini", thinking: "high" },
			{ text: "Review same.", model: "gpt-4o-mini", thinking: "high" },
		]);
		expect(t?.events.messages).toContainEqual(
			expect.objectContaining({
				role: "user",
				content: [{ type: "text", text: "Review same." }, image],
			}),
		);
	});

	it("defers busy input during recovery when Pi omits streamingBehavior", async () => {
		const requests = await queuedSession(async () => undefined);
		await extension().handlers.get("input")[0](
			{ type: "input", text: "/plan later", source: "extension" },
			{ ...context(), isIdle: () => false },
		);
		await extension().handlers.get("agent_settled")[0](
			{ type: "agent_settled" },
			context(),
		);
		await settled(1, requests);
		expect(requests[0]).toEqual({
			text: "Plan later.",
			model: "gpt-4o",
			thinking: "low",
		});
	});

	it("fuzzy-matches non-prefix mask queries", () => {
		const parsed = __test.parseMasksSettings({
			items: [
				{
					name: "terra",
					model: "openai-codex/gpt-5.6-terra",
					thinkingLevel: "medium",
				},
				{
					name: "sol-xhigh",
					model: "openai-codex/gpt-5.6-sol",
					thinkingLevel: "xhigh",
				},
			],
		});

		expect(
			__test.fuzzyMasks(parsed.config.items, "sxh").map((mask) => mask.name),
		).toEqual(["sol-xhigh"]);
	});

	it("equips the fuzzy picker's selected result", async () => {
		const parsed = __test.parseMasksSettings({
			items: [
				{
					name: "terra",
					model: "openai-codex/gpt-5.6-terra",
					thinkingLevel: "medium",
				},
				{
					name: "sol-xhigh",
					model: "openai-codex/gpt-5.6-sol",
					thinkingLevel: "xhigh",
				},
			],
		});
		const selected = await __test.chooseFuzzyMask(parsed.config.items, {
			ui: {
				custom: async (factory: any) => {
					let result: string | undefined;
					const component = await factory(
						{ requestRender() {} },
						{
							bold: (text: string) => text,
							fg: (_color: string, text: string) => text,
						},
						{
							matches: (data: string, action: string) =>
								action === "tui.select.confirm" && data === "\r",
						},
						(value: string | undefined) => {
							result = value;
						},
					);
					for (const key of ["s", "x", "h", "\r"]) component.handleInput(key);
					return result;
				},
			},
		} as any);

		expect(selected?.name).toBe("sol-xhigh");
	});

	it("validates schema while preserving usable picker entries", () => {
		const parsed = __test.parseMasksSettings({
			pickerShortcut: "alt+m",
			items: [
				{
					name: "a",
					model: "provider/a",
					thinkingLevel: "high",
					shortcut: "alt+1",
				},
				{
					name: "a",
					model: "provider/b",
					thinkingLevel: "low",
				},
				{
					name: "b",
					model: "provider/b",
					thinkingLevel: "low",
					shortcut: "alt+1",
				},
				{ name: "bad", model: "missing-provider", thinkingLevel: "huge" },
			],
		});

		expect(parsed.config).toEqual({
			pickerShortcut: "alt+m",
			items: [
				{
					name: "a",
					model: "provider/a",
					thinkingLevel: "high",
					shortcut: "alt+1",
				},
				{
					name: "b",
					model: "provider/b",
					thinkingLevel: "low",
					shortcut: undefined,
				},
			],
		});
		expect(parsed.errors).toEqual([
			'masks.items[1].name duplicates "a"',
			'masks.items[2].shortcut duplicates "alt+1"',
			"masks.items[3].model must be provider/model",
			"masks.items[3].thinkingLevel is invalid",
		]);
	});

	it("validates prompt tool profile shape and uniqueness", () => {
		const registered = new Set(["read", "ls"]);
		expect(__test.parsePromptTools(["read", "ls"], registered)).toEqual({
			tools: ["read", "ls"],
		});
		expect(__test.parsePromptTools([], registered).error).toMatch(
			/non-empty array/,
		);
		expect(__test.parsePromptTools([" "], registered).error).toMatch(
			/non-empty string/,
		);
		expect(__test.parsePromptTools(["read", "read"], registered).error).toMatch(
			/duplicates/,
		);
	});

	it("keeps a model id containing slashes", () => {
		expect(__test.splitModelRef("provider/org/model")).toEqual([
			"provider",
			"org/model",
		]);
	});
});