repositories / pi-ext
pi-ext
bugabingas pi extensions
owned by admin
extensions/masks/__tests__/contract.test.ts
Rawimport {
mkdirSync,
mkdtempSync,
readdirSync,
readFileSync,
writeFileSync,
} from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { fileURLToPath } from "node:url";
import {
createFauxCore,
type FauxResponseStep,
fauxAssistantMessage,
getCurrentSystemPrompt,
getCurrentTools,
} from "@earendil-works/pi-ai";
import {
type ExtensionRuntime,
initTheme,
parseFrontmatter,
} 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 { __test } from "../index.ts";
let session: PiSession | undefined;
let cwd: string | undefined;
let selectIndex: number | undefined;
interface RecordedUI {
notify: Array<{ args: unknown[] }>;
select: Array<{ args: unknown[] }>;
setStatus: Array<{ args: unknown[] }>;
[kind: string]: Array<{ args: unknown[] }>;
}
let recorded: RecordedUI = { notify: [], select: [], setStatus: [] };
function extensionPath(): string {
return fileURLToPath(new URL("../index.ts", import.meta.url));
}
// The owned fixture disables prompt-template discovery; re-expose the project's
// prompts directory as live prompt templates so masks resolves /review, /plan.
function scanPrompts() {
const dir = join(cwd ?? "", "prompts");
const prompts: unknown[] = [];
let entries: string[];
try {
entries = readdirSync(dir);
} catch {
return prompts;
}
for (const entry of entries) {
if (!entry.endsWith(".md")) continue;
const filePath = join(dir, entry);
const { frontmatter, body } = parseFrontmatter(
readFileSync(filePath, "utf8"),
);
prompts.push({
name: entry.replace(/\.md$/, ""),
description:
typeof frontmatter["description"] === "string"
? frontmatter["description"]
: "",
content: body,
sourceInfo: {
path: filePath,
source: "project",
scope: "project",
origin: "top-level",
},
filePath,
});
}
return prompts;
}
function fauxModelDefs() {
return [
{
id: "gpt-4o",
name: "gpt-4o",
reasoning: true,
input: ["text", "image"] as ("text" | "image")[],
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
contextWindow: 200_000,
maxTokens: 8192,
},
{
id: "gpt-4o-mini",
name: "gpt-4o-mini",
reasoning: true,
input: ["text", "image"] as ("text" | "image")[],
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
contextWindow: 200_000,
maxTokens: 8192,
},
];
}
function installUI(): void {
recorded = { notify: [], select: [], setStatus: [] };
if (!session) throw new Error("missing test session");
const ui = session.session.extensionRunner.createCommandContext().ui;
vi.spyOn(ui, "notify").mockImplementation((...args: unknown[]) => {
recorded.notify.push({ args });
});
vi.spyOn(ui, "select").mockImplementation(async (...args: unknown[]) => {
recorded.select.push({ args });
const entries = args[1] as string[];
return selectIndex === undefined ? undefined : entries[selectIndex];
});
vi.spyOn(ui, "setStatus").mockImplementation((...args: unknown[]) => {
recorded.setStatus.push({ args });
});
}
function uiCallsFor(kind: string): Array<{ args: unknown[] }> {
return recorded[kind] ?? [];
}
async function startSession(
settings: Record<string, unknown>,
options: { selectIndex?: number } = {},
): Promise<ReturnType<typeof createFauxCore>> {
cwd = mkdtempSync(join(tmpdir(), "masks-test-"));
selectIndex = options.selectIndex;
// Print-mode fixture leaves the theme singleton uninitialized.
initTheme(undefined, false);
session = await createPiSession({
cwd,
extensionPaths: [extensionPath()],
userSettings: settings,
});
(
session.session.resourceLoader as unknown as { getPrompts: () => unknown }
).getPrompts = () => ({ prompts: scanPrompts(), diagnostics: [] });
const core = createFauxCore({
api: "openai",
provider: "openai",
models: [
{ id: "gpt-4o", reasoning: true, input: ["text", "image"] },
{ id: "gpt-4o-mini", reasoning: true, input: ["text", "image"] },
],
});
session.session.modelRuntime.registerProvider("openai", {
apiKey: "test-key",
api: "openai",
baseUrl: "https://faux.invalid",
streamSimple: core.streamSimple,
models: fauxModelDefs(),
});
const model = session.session.modelRuntime.getModel("openai", "gpt-4o");
if (!model) throw new Error("missing openai model");
await session.session.setModel(model);
installUI();
return core;
}
function writePrompt(name: string, content: string): void {
if (!cwd) throw new Error("missing test root");
const promptsDir = join(cwd, "prompts");
mkdirSync(promptsDir, { recursive: true });
writeFileSync(join(promptsDir, `${name}.md`), content);
}
function extension() {
if (!session) throw new Error("missing test session");
const first = session.extensions[0];
if (!first) throw new Error("missing test extension");
return first;
}
function getRuntime(): ExtensionRuntime {
if (!session) throw new Error("missing test session");
return (
session.session.extensionRunner as unknown as {
runtime: ExtensionRuntime;
}
).runtime;
}
function handler(type: string) {
const handlers = extension().handlers.get(type);
if (!handlers?.[0]) throw new Error(`missing handler ${type}`);
return handlers[0];
}
type ShortcutKey = Parameters<
ReturnType<typeof extension>["shortcuts"]["get"]
>[0];
function shortcut(key: ShortcutKey) {
const found = extension().shortcuts.get(key);
if (!found) throw new Error(`missing shortcut ${key}`);
return found;
}
function command(name: string) {
const found = extension().commands.get(name);
if (!found) throw new Error(`missing command ${name}`);
return found;
}
function context() {
if (!session) throw new Error("missing test session");
return session.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 core = await startSession({
masks: {
items: [
{ name: "focused", model: "openai/gpt-4o-mini", thinkingLevel: "high" },
{ name: "quick", model: "openai/gpt-4o", thinkingLevel: "low" },
{
name: "auditor",
model: "openai/gpt-4o-mini",
thinkingLevel: "low",
tools: ["read", "grep"],
instructions: "Audit only.",
},
],
},
});
writePrompt("review", "---\nmask: focused\n---\nReview $1.");
writePrompt("plan", "---\nmask: quick\n---\nPlan $1.");
writePrompt("audit", "---\nmask: auditor\n---\nAudit $1.");
writeFileSync(join(cwd ?? "", "fixture.txt"), "fixture");
for (const name of ["gpt-4o", "gpt-4o-mini"]) {
const model = context().modelRegistry.find("openai", name);
if (model) model.reasoning = true;
}
session?.session.settingsManager.setRetryEnabled(false);
// Fixture images are not decodable; keep Pi from resizing or omitting them.
session?.session.settingsManager.setImageAutoResize(false);
const requests: Request[] = [];
const recorder = (async (ctx: any, options: any, _state: any, model: any) => {
if (options?.signal?.aborted)
return fauxAssistantMessage([{ type: "text", text: "Done." }], {
stopReason: "aborted",
});
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,
};
requests.push(request);
const reply = await onRequest(
request,
requests.length - 1,
getCurrentTools(ctx.messages).map((tool) => tool.name),
getCurrentSystemPrompt(ctx.messages),
);
return fauxAssistantMessage(
(reply?.["content"] as never) ?? [{ type: "text", text: "Done." }],
{
stopReason: (reply?.["stopReason"] as never) ?? "stop",
...(reply?.["errorMessage"] !== undefined
? { errorMessage: reply["errorMessage"] as string }
: {}),
},
);
}) as FauxResponseStep;
core.setResponses(Array(100).fill(recorder));
return requests;
}
async function queue(text: string): Promise<void> {
await session?.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(() => session?.session.isIdle).toBe(true);
}
afterEach(async () => {
vi.restoreAllMocks();
await session?.dispose();
session = undefined;
cwd = undefined;
selectIndex = undefined;
});
describe("masks contract", () => {
it("switches complete masks from direct shortcut, picker, and command", async () => {
await startSession(
{
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",
},
],
},
},
{ selectIndex: 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 shortcut("alt+1").handler(context());
expect(session?.session.model?.provider).toBe("openai");
expect(session?.session.model?.id).toBe("gpt-4o");
expect(session?.session.thinkingLevel).toBe("low");
await shortcut("alt+m").handler(context());
expect(session?.session.thinkingLevel).toBe("high");
expect(uiCallsFor("select")[0]?.args).toEqual([
"Masks",
[
"low openai/gpt-4o:low \u00b7 alt+1",
"high openai/gpt-4o:high \u00b7 alt+2",
"missing absent/nope:max \u00b7 alt+3",
],
]);
await command("masks").handler("low", context());
expect(session?.session.thinkingLevel).toBe("low");
await shortcut("alt+3").handler(context());
expect(session?.session.thinkingLevel).toBe("low");
expect(uiCallsFor("notify").at(-1)?.args).toEqual([
'Mask "missing": model absent/nope not found',
"error",
]);
});
it("applies manual mask tools, warns on unknown names, and injects instructions", async () => {
await startSession({
masks: {
items: [
{
name: "reader",
model: "openai/gpt-4o",
thinkingLevel: "low",
tools: ["read", "absent", "ls"],
instructions: "Read only.\nNever edit.",
},
{ name: "plain", model: "openai/gpt-4o", thinkingLevel: "off" },
],
},
});
const runtime = getRuntime();
const before = handler("before_agent_start");
const sectionsFor = async () => {
const event = {
type: "before_agent_start",
prompt: "hi",
systemPrompt: "",
systemPromptOptions: { sections: {} as Record<string, string> },
};
await before(event, context());
return event.systemPromptOptions.sections;
};
expect(await sectionsFor()).toEqual({});
await command("masks").handler("reader", context());
expect(runtime.getActiveTools()).toEqual(["read", "ls"]);
expect(uiCallsFor("notify").at(-1)?.args).toEqual([
'Mask "reader": unknown tools: absent',
"warning",
]);
expect(await sectionsFor()).toEqual({ masks: "Read only.\nNever edit." });
await command("masks").handler("plain", context());
expect(runtime.getActiveTools()).toEqual(["read", "ls"]);
expect(await sectionsFor()).toEqual({});
});
it("leaves tools unchanged when no mask tool is known", async () => {
await startSession({
masks: {
items: [
{
name: "ghost",
model: "openai/gpt-4o",
thinkingLevel: "low",
tools: ["absent"],
},
],
},
});
const runtime = getRuntime();
const originalTools = runtime.getActiveTools();
await command("masks").handler("ghost", context());
expect(runtime.getActiveTools()).toEqual(originalTools);
expect(session?.session.thinkingLevel).toBe("low");
expect(uiCallsFor("notify").at(-1)?.args).toEqual([
'Mask "ghost": no known tools; tool state unchanged',
"error",
]);
});
it("describes tools and truncated instructions in selector entries", async () => {
await startSession({
masks: {
items: [
{
name: "reader",
model: "openai/gpt-4o",
thinkingLevel: "low",
shortcut: "alt+1",
tools: ["read", "ls"],
instructions:
"You are in review mode.\n\nNever edit files, only report.",
},
],
},
});
await shortcut("alt+m").handler(context());
expect(uiCallsFor("select")[0]?.args).toEqual([
"Masks",
[
'reader openai/gpt-4o:low \u00b7 alt+1 \u00b7 tools:read,ls \u00b7 "You are in review mode. Never\u2026"',
],
]);
});
it("applies mask tools and instructions for a masked prompt, then restores", async () => {
const toolSets: string[][] = [];
const systemPrompts: string[] = [];
const requests = await queuedSession(
async (_request, _index, tools, systemPrompt) => {
toolSets.push(tools);
systemPrompts.push(systemPrompt);
return undefined;
},
);
await session?.session.prompt("/audit scoped");
await settled(1, requests);
await session?.session.prompt("ordinary");
await settled(2, requests);
expect(toolSets[0]).toEqual(["read", "grep"]);
expect(systemPrompts[0]).toContain("<masks>\nAudit only.\n</masks>");
expect(systemPrompts[1]).not.toContain("Audit only.");
expect(getRuntime().getActiveTools()).toEqual([
"read",
"bash",
"edit",
"write",
]);
});
it("lets prompt frontmatter tools override mask tools", async () => {
const toolSets: string[][] = [];
const requests = await queuedSession(async (_request, _index, tools) => {
toolSets.push(tools);
return undefined;
});
writePrompt("audit", "---\nmask: auditor\ntools: [ls]\n---\nAudit $1.");
await session?.session.prompt("/audit scoped");
await settled(1, requests);
expect(toolSets).toEqual([["ls"]]);
});
it("applies prompt frontmatter for one agent run, then restores", async () => {
await startSession({
masks: {
items: [
{ name: "focused", model: "openai/gpt-4o", thinkingLevel: "high" },
],
},
});
writePrompt(
"review",
"---\ndescription: Review changes\nmask: focused\n---\nReview it.",
);
const model = context().modelRegistry.find("openai", "gpt-4o");
if (!model) throw new Error("missing harness model");
model.reasoning = true;
const inputHandler = handler("input");
const endHandler = handler("agent_settled");
expect(
await inputHandler(
{ type: "input", text: "/review", source: "interactive" },
context(),
),
).toEqual({ action: "continue" });
expect(session?.session.thinkingLevel).toBe("high");
await endHandler({ type: "agent_settled", aborted: false }, context());
expect(session?.session.model?.provider).toBe("openai");
expect(session?.session.model?.id).toBe("gpt-4o");
expect(session?.session.thinkingLevel).toBe("off");
});
it("applies a masked prompt tool profile and restores the exact snapshot", async () => {
await startSession({
masks: {
items: [
{ name: "focused", model: "openai/gpt-4o", thinkingLevel: "high" },
],
},
});
writePrompt(
"review",
"---\nmask: focused\ntools: [read, ls]\n---\nReview it.",
);
const model = context().modelRegistry.find("openai", "gpt-4o");
if (!model) throw new Error("missing harness model");
model.reasoning = true;
const runtime = getRuntime();
const originalTools = runtime.getActiveTools();
expect(
await handler("input")(
{ type: "input", text: "/review", source: "interactive" },
context(),
),
).toEqual({ action: "continue" });
expect(runtime.getActiveTools()).toEqual(["read", "ls"]);
await handler("agent_start")({ type: "agent_start" }, context());
runtime.setActiveTools(["read", "ls", "bash"]);
await handler("agent_settled")(
{ type: "agent_settled", aborted: false },
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 session?.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(getRuntime().getActiveTools()).toEqual([
"read",
"bash",
"edit",
"write",
]);
});
it("leaves tools alone when masked prompt frontmatter omits them", async () => {
await startSession({
masks: {
items: [
{ name: "focused", model: "openai/gpt-4o", thinkingLevel: "high" },
],
},
});
writePrompt("review", "---\nmask: focused\n---\nReview it.");
const runtime = getRuntime();
await handler("input")(
{ type: "input", text: "/review", source: "interactive" },
context(),
);
await handler("agent_start")({ type: "agent_start" }, context());
runtime.setActiveTools(["read", "ls"]);
await handler("agent_settled")(
{ type: "agent_settled", aborted: false },
context(),
);
expect(runtime.getActiveTools()).toEqual(["read", "ls"]);
});
it("rejects invalid prompt tool profiles before changing the mask", async () => {
await startSession({
masks: {
items: [
{
name: "focused",
model: "openai/gpt-4o-mini",
thinkingLevel: "high",
},
],
},
});
writePrompt(
"review",
"---\nmask: focused\ntools: [read, absent]\n---\nReview it.",
);
const runtime = getRuntime();
const originalTools = runtime.getActiveTools();
const originalModel = session?.session.model?.id;
expect(
await handler("input")(
{ type: "input", text: "/review", source: "interactive" },
context(),
),
).toEqual({ action: "handled" });
expect(runtime.getActiveTools()).toEqual(originalTools);
expect(session?.session.model?.id).toBe(originalModel);
expect(uiCallsFor("notify").at(-1)?.args).toEqual([
'Prompt /review: unknown tool "absent"',
"error",
]);
});
it("rolls back model, thinking, and tools when tool application fails", async () => {
await startSession({
masks: {
items: [
{
name: "focused",
model: "openai/gpt-4o-mini",
thinkingLevel: "high",
},
],
},
});
writePrompt(
"review",
"---\nmask: focused\ntools: [read, ls]\n---\nReview it.",
);
const runtime = getRuntime();
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 handler("input")(
{ type: "input", text: "/review", source: "interactive" },
context(),
),
).toEqual({ action: "handled" });
expect(runtime.getActiveTools()).toEqual(originalTools);
expect(session?.session.model?.id).toBe("gpt-4o");
expect(session?.session.thinkingLevel).toBe("off");
});
it("restores tools even when model restoration fails", async () => {
await startSession({
masks: {
items: [
{
name: "focused",
model: "openai/gpt-4o-mini",
thinkingLevel: "high",
},
],
},
});
writePrompt(
"review",
"---\nmask: focused\ntools: [read, ls]\n---\nReview it.",
);
const runtime = getRuntime();
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 handler("input")(
{ type: "input", text: "/review", source: "interactive" },
context(),
);
await handler("agent_start")({ type: "agent_start" }, context());
failRestore = true;
await handler("agent_settled")(
{ type: "agent_settled", aborted: false },
context(),
);
expect(runtime.getActiveTools()).toEqual(originalTools);
expect(uiCallsFor("notify").at(-1)?.args[0]).toContain(
"Failed to restore prompt mask",
);
});
it("restores after tree navigation races temporary profile application", async () => {
await startSession({
masks: {
items: [
{
name: "focused",
model: "openai/gpt-4o-mini",
thinkingLevel: "high",
},
],
},
});
writePrompt(
"review",
"---\nmask: focused\ntools: [read, ls]\n---\nReview it.",
);
const runtime = getRuntime();
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 = handler("input")(
{ type: "input", text: "/review", source: "interactive" },
context(),
);
await entered.promise;
const resetting = handler("session_tree")(
{ type: "session_tree" },
context(),
);
release.resolve();
await Promise.all([applying, resetting]);
expect(session?.session.model?.id).toBe("gpt-4o");
expect(session?.session.thinkingLevel).toBe("off");
expect(runtime.getActiveTools()).toEqual(originalTools);
});
it("keeps reset joined to profile application after a manual switch is rejected", async () => {
await startSession({
masks: {
items: [
{
name: "focused",
model: "openai/gpt-4o-mini",
thinkingLevel: "high",
},
],
},
});
writePrompt(
"review",
"---\nmask: focused\ntools: [read, ls]\n---\nReview it.",
);
const runtime = getRuntime();
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 = handler("input")(
{ type: "input", text: "/review", source: "interactive" },
context(),
);
await entered.promise;
await command("masks").handler("focused", context());
const resetting = handler("session_tree")(
{ type: "session_tree" },
context(),
);
release.resolve();
await Promise.all([applying, resetting]);
expect(session?.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 session?.session.prompt("start");
await settled(1, requests);
expect(
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 = getRuntime();
const setModel = runtime.setModel.bind(runtime);
vi.spyOn(runtime, "setModel").mockImplementation(async (model) =>
model.id === "gpt-4o-mini" ? false : setModel(model),
);
await session?.session.prompt("start");
await settled(1, requests);
expect(
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 () => {
await startSession({
masks: {
items: [
{ name: "focused", model: "openai/gpt-4o", thinkingLevel: "high" },
],
},
});
writePrompt("review", "---\nmask: focused\n---\nReview it.");
const busyContext = { ...context(), isIdle: () => false };
const result = await handler("input")(
{
type: "input",
text: "/review later",
source: "interactive",
streamingBehavior: "followUp",
},
busyContext,
);
expect(result).toEqual({ action: "handled" });
expect(session?.session.thinkingLevel).toBe("off");
expect(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(session?.session.model?.id).toBe("gpt-4o");
}
return undefined;
});
await session?.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(session?.session.model?.id).toBe("gpt-4o");
expect(session?.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 session?.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 session?.session.prompt("/review now");
await settled(1, requests);
expect(session?.session.model?.id).toBe("gpt-4o");
expect(session?.session.thinkingLevel).toBe("off");
expect(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 session?.session.prompt("/masks-cancel");
expect(context().signal?.aborted).toBe(false);
return undefined;
});
await session?.session.prompt("/review now");
await settled(1, requests);
expect(session?.session.thinkingLevel).toBe("off");
});
it("rejects steering without changing the active mask", async () => {
const requests = await queuedSession(async () => {
await session?.session.sendUserMessage("/plan later", {
expandPromptTemplates: true,
deliverAs: "steer",
});
expect(session?.session.thinkingLevel).toBe("high");
return undefined;
});
await session?.session.prompt("/review now");
await settled(1, requests);
expect(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 handler("input")(
{
type: "input",
text: "/plan later",
source: "extension",
streamingBehavior: "followUp",
},
{ ...context(), isIdle: () => false },
);
await handler(type)({ type, reason: "reload" }, context());
await handler("agent_settled")(
{ type: "agent_settled", aborted: false },
context(),
);
expect(requests).toHaveLength(0);
expect(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(getRuntime(), "setModel").mockResolvedValue(false);
await session?.session.prompt("/review now");
expect(requests).toHaveLength(0);
expect(
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(getRuntime(), "setModel").mockRejectedValue(
new Error("offline"),
);
return undefined;
});
await session?.session.prompt("/review now");
await settled(1, requests);
expect(
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,
);
if (!session) throw new Error("missing test session");
vi.spyOn(
session.session.settingsManager,
"getRetrySettings",
).mockReturnValue({
enabled: true,
maxRetries: 1,
baseDelayMs: 1,
maxAgentDelayMs: 0,
});
await session?.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(session?.session.model?.id).toBe("gpt-4o");
expect(session?.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(session?.session as any, "_runAutoCompaction")
.mockImplementation(async () => {
expect(session?.session.model?.id).toBe("gpt-4o-mini");
return true;
});
await session?.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(session?.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 = getRuntime();
const setModel = runtime.setModel.bind(runtime);
vi.spyOn(runtime, "setModel").mockImplementation(async (model) => {
entered.resolve();
await release.promise;
return setModel(model);
});
const first = session?.session.prompt("/review first");
await entered.promise;
try {
await session?.session.prompt("ordinary concurrent prompt");
await session?.session.prompt("/plan second");
} finally {
release.resolve();
}
await first;
expect(requests).toEqual([
{ text: "Review first.", model: "gpt-4o-mini", thinking: "high" },
]);
expect(session?.session.model?.id).toBe("gpt-4o");
expect(session?.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 = getRuntime();
const setModel = runtime.setModel.bind(runtime);
vi.spyOn(runtime, "setModel").mockImplementation(async (model) => {
entered.resolve();
await release.promise;
return setModel(model);
});
const manual = session?.session.prompt("/masks focused");
await entered.promise;
try {
await session?.session.prompt("/plan blocked");
} finally {
release.resolve();
}
await manual;
expect(requests).toHaveLength(0);
expect(session?.session.model?.id).toBe("gpt-4o-mini");
expect(session?.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 = getRuntime();
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 session?.session.prompt("/review first");
await session?.session.prompt("ordinary blocked prompt");
await session?.session.prompt("/plan blocked");
expect(requests).toHaveLength(1);
failRestore = false;
await session?.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 = getRuntime();
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 session?.session.prompt("/review now");
await settled(1, requests);
expect(uiCallsFor("notify").at(-1)?.args).toEqual([
"Cancelled 1 deferred masked prompts",
"warning",
]);
expect(session?.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 session?.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" as const,
mimeType: "image/png",
data: "Zml4dHVyZS1pbWFnZQ==",
};
const requests = await queuedSession(async (_request, index) => {
if (index === 0) {
await session?.session.sendUserMessage(
[{ type: "text", text: "/review same" }, image],
{
expandPromptTemplates: true,
deliverAs: "followUp",
},
);
await queue("/review same");
}
return undefined;
});
await session?.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(session?.session.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 handler("input")(
{ type: "input", text: "/plan later", source: "extension" },
{ ...context(), isIdle: () => false },
);
await handler("agent_settled")(
{ type: "agent_settled", aborted: false },
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" },
{
name: "c",
model: "provider/c",
thinkingLevel: "low",
tools: [],
instructions: " ",
},
{
name: "d",
model: "provider/d",
thinkingLevel: "low",
tools: ["read", "read"],
},
{
name: "e",
model: "provider/e",
thinkingLevel: "low",
tools: ["read", "ls"],
instructions: " Be brief. ",
},
],
});
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,
},
{
name: "e",
model: "provider/e",
thinkingLevel: "low",
shortcut: undefined,
tools: ["read", "ls"],
instructions: "Be brief.",
},
],
});
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",
"masks.items[4].tools must be a non-empty array",
"masks.items[4].instructions must be a non-empty string",
'masks.items[5].tools duplicates "read"',
]);
});
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("truncates long instructions in mask descriptions", () => {
expect(
__test.describeMask({
name: "x",
model: "p/m",
thinkingLevel: "low",
instructions: "short",
}),
).toBe('p/m:low \u00b7 "short"');
expect(
__test.describeMask({
name: "x",
model: "p/m",
thinkingLevel: "low",
instructions: "a".repeat(40),
}),
).toBe(`p/m:low \u00b7 "${"a".repeat(29)}\u2026"`);
});
it("keeps a model id containing slashes", () => {
expect(__test.splitModelRef("provider/org/model")).toEqual([
"provider",
"org/model",
]);
});
});