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 { 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 | undefined>, ): Promise { 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 | 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 { await t?.session.sendUserMessage(text, { expandPromptTemplates: true, deliverAs: "followUp", }); } async function settled(count: number, requests: Request[]): Promise { 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(); const release = Promise.withResolvers(); 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(); const release = Promise.withResolvers(); 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(); const release = Promise.withResolvers(); 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(); const release = Promise.withResolvers(); 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", ]); }); });