repositories / pi-ext
pi-ext
bugabingas pi extensions
owned by admin
extensions/masks/__tests__/harness.test.ts
Rawimport { 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",
]);
});
});