import { mkdirSync, mkdtempSync, readdirSync, readFileSync, rmSync, writeFileSync, } from "node:fs"; import { tmpdir } from "node:os"; import { dirname, join, resolve } from "node:path"; import { fileURLToPath } from "node:url"; import { type Context, fauxAssistantMessage, fauxToolCall, type JsonObject, } from "@earendil-works/pi-ai"; import { DefaultPackageManager, type ExtensionCommandContext, type ExtensionContext, type ExtensionToolContext, initTheme, SessionManager, type Theme, } 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 { resolveSubagentExtensions } from "../implementation.ts"; import ultraExtension from "../index.ts"; import { DEFAULT_MODEL_TIERS, modelTierPrompt, parseUltraSettings, } from "../settings.ts"; import { WorkflowSpecSchema } from "../spec.ts"; let session: PiSession | undefined; let cwd: string | undefined; let notifySpy: ReturnType | undefined; afterEach(async () => { await session?.dispose(); if (cwd) rmSync(cwd, { recursive: true, force: true }); session = undefined; cwd = undefined; notifySpy = undefined; vi.unstubAllEnvs(); vi.restoreAllMocks(); }); function ultraPath(): string { return fileURLToPath(new URL("../index.ts", import.meta.url)); } function createToolContext( overrides: Partial = {}, toolCallId = "call", ): ExtensionToolContext { if (!session) throw new Error("missing session"); const context = Object.create( session.session.extensionRunner.createToolContext(toolCallId, undefined), ) as ExtensionToolContext; if (overrides.ui) Object.defineProperty(context, "ui", { value: overrides.ui }); if (overrides.sessionManager) Object.defineProperty(context, "sessionManager", { value: overrides.sessionManager, }); return context; } // The owned fixture binds in print mode with a no-op UI context. Replace it with // a recording UI (hasUI becomes true, so overlay-driven surfaces run) and spy on // notify to capture the UI's notify calls. function installUi(): void { if (!session) throw new Error("missing session"); const runner = session.session.extensionRunner; const ui = { ...runner.getUIContext() }; notifySpy = vi.spyOn(ui, "notify"); runner.setUIContext(ui, "print"); } function uiCallsFor(method: "notify"): Array<{ args: unknown[] }> { if (method !== "notify") return []; return (notifySpy?.mock.calls ?? []).map((args: unknown[]) => ({ args })); } async function start( options: { settings?: Record | null; env?: Record; factory?: boolean; extra?: string[]; } = {}, ) { cwd = mkdtempSync(join(tmpdir(), "ultra-contract-")); if (options.settings !== null) { mkdirSync(join(cwd, ".pi"), { recursive: true }); writeFileSync( join(cwd, ".pi", "settings.json"), JSON.stringify(options.settings ?? { ultra: { autoEnable: false } }), ); } session = await createPiSession({ cwd, ...(options.env ? { env: options.env } : {}), ...(options.factory ? { extensionFactories: [ultraExtension] } : { extensionPaths: [ultraPath(), ...(options.extra ?? [])] }), }); installUi(); return session.extensions[0]!; } function saveWorkflow(name: string, dynamicExtension?: string): void { if (!session) throw new Error("missing session"); const directory = join(session.cwd, ".pi", "workflows"); mkdirSync(directory, { recursive: true }); writeFileSync( join(directory, `${name}.json`), `${JSON.stringify({ name, phases: [ { id: "work", kind: "single", step: { summary: "Work", prompt: "Work", ...(dynamicExtension ? { dynamicExtensions: [dynamicExtension] } : {}), }, }, ], })}\n`, ); } function driveOverlay( ctx: ExtensionCommandContext, interact?: (overlay: { handleInput(data: string): void }) => void, ): void { Object.assign(ctx.ui, { custom: vi.fn( async ( factory: ( host: unknown, theme: unknown, keybindings: unknown, done: (value: unknown) => void, ) => { handleInput(data: string): void }, ) => new Promise((resolve) => { const overlay = factory( { requestRender: vi.fn() }, {}, { matches: () => false }, resolve, ); interact?.(overlay); }), ), }); } describe("ultra extension", () => { it("resolves explicit Pi sources through the SDK without installing them", async () => { const resolveSources = vi .spyOn(DefaultPackageManager.prototype, "resolveExtensionSources") .mockResolvedValue({ extensions: [ { path: "/enabled/index.ts", enabled: true }, { path: "/disabled/index.ts", enabled: false }, ], } as never); const selection = await resolveSubagentExtensions( parseUltraSettings({ subagentExtensions: [ "npm:@scope/pi-extension", "https://github.com/example/pi-extension.git", ], }).subagentExtensions, { cwd: process.cwd(), agentDir: join(process.cwd(), ".agent"), projectTrusted: false, }, ); expect(resolveSources).toHaveBeenNthCalledWith( 1, ["npm:@scope/pi-extension"], { temporary: true }, ); expect(resolveSources).toHaveBeenNthCalledWith( 2, ["https://github.com/example/pi-extension.git"], { temporary: true }, ); expect(selection.paths).toEqual(["/enabled/index.ts"]); const path = resolve("/enabled/index.ts").replaceAll("\\", "/"); expect(selection.pathKeys).toEqual( new Set([process.platform === "win32" ? path.toLowerCase() : path]), ); }); it("resolves relative sources from cwd regardless of project trust", async () => { cwd = mkdtempSync(join(tmpdir(), "ultra-relative-source-")); mkdirSync(join(cwd, "research")); const entry = join(cwd, "research", "evidence.ts"); writeFileSync(entry, "export default function () {}"); writeFileSync( join(cwd, "research", "package.json"), JSON.stringify({ pi: { extensions: ["evidence.ts"] } }), ); for (const projectTrusted of [false, true]) { const options = { cwd, agentDir: join(cwd, "unrelated-agent"), projectTrusted, }; expect( (await resolveSubagentExtensions(["./research"], options)).paths, ).toEqual([entry]); expect( (await resolveSubagentExtensions(["research"], options)).paths, ).toEqual([entry]); await expect( resolveSubagentExtensions(["./research", "./missing.ts"], options), ).rejects.toThrow("no enabled extensions resolved"); } }); it("rejects legacy bare extension nicknames with migration guidance", async () => { await expect( resolveSubagentExtensions(["nushell"], { cwd: process.cwd(), agentDir: join(process.cwd(), ".agent"), projectTrusted: true, }), ).rejects.toThrow(/no longer accepts bare extension nickname/u); }); it("loads only the command and renderer in the real Pi runtime", async () => { const indexPath = ultraPath(); const extension = await start(); expect(extension.path).toBe(indexPath); expect(extension.tools.has("run_workflow")).toBe(false); expect(readFileSync(indexPath, "utf8")).toContain( 'import("./implementation.ts")', ); expect(extension.commands.has("ultra")).toBe(true); expect(extension.commands.has("ultra-runs")).toBe(false); }); it("writes disabled command outcome in its debug sandbox", async () => { const extension = await start({ settings: null, env: { PI_ULTRA_DEBUG: "1" }, }); const runner = session?.session.extensionRunner; if (!runner || !session) throw new Error("missing session"); const context = runner.createCommandContext(); mkdirSync(join(session.cwd, ".pi"), { recursive: true }); writeFileSync( join(session.cwd, ".pi", "settings.json"), JSON.stringify({ ultra: { enabled: false } }), ); await extension.commands.get("ultra")!.handler("", context); await runner.emit({ type: "session_shutdown", reason: "quit" }); const directory = join( session.cwd, ".test-state", "pi-ext", "debug", "ultra", ); const records = readFileSync( join(directory, readdirSync(directory)[0]!), "utf8", ) .trim() .split("\n") .map((line) => JSON.parse(line)); expect( records.filter((record) => record.event === "command.handle.start"), ).toHaveLength(1); expect( records.filter((record) => record.event === "command.handle.finish"), ).toEqual([expect.objectContaining({ outcome: "disabled" })]); expect(records.map((record) => record.event)).toEqual( expect.arrayContaining(["session.start", "session.shutdown"]), ); }); it("auto-enables the workflow tool when configured", async () => { const extension = await start({ settings: { ultra: { autoEnable: true } }, }); if (!session) throw new Error("missing session"); expect(extension.tools.has("run_workflow")).toBe(true); expect(session.session.getActiveToolNames()).toContain("run_workflow"); expect(session.session.systemPrompt).toContain("- run_workflow:"); }); it.each([false, true])( "registers the authoring skill independently of tool activation (auto-enable: %s)", async (autoEnable) => { const extension = await start({ settings: { ultra: { autoEnable } } }); if (!session) throw new Error("missing session"); const skillPath = fileURLToPath( new URL("../prompts/ultra-authoring/SKILL.md", import.meta.url), ); const skill = session.session.resourceLoader .getSkills() .skills.find((candidate) => candidate.name === "ultra-authoring"); expect(skill?.filePath).toBe(skillPath); expect(session.session.systemPrompt).toContain(skillPath); expect(session.session.systemPrompt).not.toContain( readFileSync(skillPath, "utf8").trim(), ); expect(session.session.systemPrompt).not.toContain("## Worked example"); if (!autoEnable) { await extension.commands .get("ultra")! .handler("", session.session.extensionRunner.createCommandContext()); } expect(extension.tools.get("run_workflow")!.definition.exposure).toBe( "model-only", ); expect( extension.tools.get("run_workflow")!.definition.annotations, ).toEqual({ readOnlyHint: false, }); const rules: string[] = extension.tools.get("run_workflow")!.definition.promptGuidelines!; expect(rules).toEqual([ "Use run_workflow when decomposition into independent or sequential child-agent steps reduces context or verification risk.", "Use dynamic extensions when a workflow needs a reusable, task-specific capability; read the ultra-authoring skill for how to build and load them.", ]); for (const rule of rules) expect(session.session.systemPrompt.split(rule)).toHaveLength(2); session.session.setActiveToolsByName( session.session .getActiveToolNames() .filter((name) => name !== "run_workflow"), ); expect(session.session.systemPrompt).toContain(skillPath); for (const rule of rules) expect(session.session.systemPrompt).not.toContain(rule); await extension.commands .get("ultra")! .handler("", session.session.extensionRunner.createCommandContext()); for (const rule of rules) expect(session.session.systemPrompt.split(rule)).toHaveLength(2); expect(session.session.messages).toHaveLength(0); }, ); it("does not advertise the skill when Ultra is disabled, even after completion or activation attempts", async () => { const extension = await start({ settings: { ultra: { enabled: false, autoEnable: true } }, }); if (!session) throw new Error("missing session"); const command = extension.commands.get("ultra")!; await command.getArgumentCompletions!(""); await command.handler( "", session.session.extensionRunner.createCommandContext(), ); expect(session.session.getActiveToolNames()).not.toContain("run_workflow"); expect(session.session.systemPrompt).not.toContain( "ultra-authoring/SKILL.md", ); }); it("injects the effective tier policy only while active and refreshes settings each turn", async () => { const extension = await start(); if (!session) throw new Error("missing session"); const runner = session.session.extensionRunner; const emit = () => runner.emitBeforeAgentStart( { messageType: "user", prompt: "prompt" }, { cwd: session?.cwd ?? process.cwd(), }, ); const policy = async () => (await emit()).systemPromptOptions.sections; expect((await policy())["ultra_model_tier_policy"]).toBeUndefined(); expect((await policy())["ultra_fanout_tool_policy"]).toBeUndefined(); await extension.commands .get("ultra")! .handler("", runner.createCommandContext()); expect((await policy())["ultra_model_tier_policy"]).toBe( modelTierPrompt(DEFAULT_MODEL_TIERS), ); expect((await policy())["ultra_fanout_tool_policy"]).toContain( "read, grep, find, ls, bash", ); writeFileSync( join(session.cwd, ".pi", "settings.json"), JSON.stringify({ ultra: { fanoutToolAllowlist: ["web"], modelTiers: { large: { instructions: "Architecture only." }, fast_extract: { model: "custom/extractor", thinkingLevel: "low", instructions: "Select for literal extraction.", }, }, }, }), ); const updated = await policy(); expect(updated["ultra_model_tier_policy"]).toContain("Architecture only."); expect(updated["ultra_model_tier_policy"]).not.toContain( DEFAULT_MODEL_TIERS["large"]!.instructions, ); expect(updated["ultra_model_tier_policy"]).toContain( "- fast_extract: custom/extractor; low; Select for literal extraction.", ); expect( updated["ultra_model_tier_policy"]?.match(/# ultra model tier policy/g), ).toHaveLength(1); expect(updated["ultra_fanout_tool_policy"]).toContain("bash, web"); writeFileSync( join(session.cwd, ".pi", "settings.json"), JSON.stringify({ ultra: { enabled: false } }), ); expect((await policy())["ultra_model_tier_policy"]).toBeUndefined(); expect((await policy())["ultra_fanout_tool_policy"]).toBeUndefined(); }); it("advertises which step.flags enable sub-agent extension tools while active", async () => { const extension = await start({ settings: { ultra: { subagentExtensions: [ fileURLToPath(new URL("../../nushell/index.ts", import.meta.url)), fileURLToPath( new URL("../../chrome-cdp/index.ts", import.meta.url), ), fileURLToPath(new URL("../../ask/ask-tool.ts", import.meta.url)), ], }, }, }); if (!session) throw new Error("missing session"); const runner = session.session.extensionRunner; const sections = async () => ( await runner.emitBeforeAgentStart( { messageType: "user", prompt: "prompt" }, { cwd: session?.cwd ?? process.cwd(), }, ) ).systemPromptOptions.sections; expect((await sections())["ultra_subagent_flags"]).toBeUndefined(); await extension.commands .get("ultra")! .handler("", runner.createCommandContext()); const prompt = (await sections())["ultra_subagent_flags"] ?? ""; expect(prompt).toContain("# ultra sub-agent flags"); expect(prompt).toContain( "- nushell: nu; --nushell (boolean): Enable Nushell tools at startup instead of bash and powershell", ); expect(prompt).toContain("- chrome-cdp: chrome_cdp; --chrome (boolean)"); // ask registers tools but no flags, so it is not listed. expect(prompt).not.toContain("- ask:"); }); it("rejects unknown tiers before either surface starts work", async () => { const extension = await start(); if (!session) throw new Error("missing session"); const ctx = session.session.extensionRunner.createCommandContext(); await extension.commands.get("ultra")!.handler("", ctx); const spec = { name: "unknown-tier", phases: [ { id: "first", kind: "single", step: { summary: "First", prompt: "First", model: "medium" }, }, { id: "second", kind: "single", step: { summary: "Second", prompt: "Second", model: "missing_tier" }, }, ], }; const tool = extension.tools.get("run_workflow")!.definition; await expect( tool.execute( "call", { spec }, undefined, undefined, createToolContext({ ui: ctx.ui, sessionManager: ctx.sessionManager }), ), ).rejects.toThrow('unknown model tier "missing_tier"'); const directory = join(session.cwd, ".pi", "workflows"); mkdirSync(directory, { recursive: true }); writeFileSync(join(directory, "unknown-tier.json"), JSON.stringify(spec)); const custom = vi.fn(); Object.assign(ctx.ui, { custom }); await extension.commands.get("ultra")!.handler("run unknown-tier", ctx); expect(custom).not.toHaveBeenCalled(); expect(uiCallsFor("notify").at(-1)?.args).toEqual([ 'ultra: unknown model tier "missing_tier". Configure ultra.modelTiers or use provider/model.', "error", ]); }); it.each(["command", "tool"])( "inherits the parent's session directory through the %s surface", async (surface) => { const extension = await start({ factory: true }); if (!session) throw new Error("missing session"); saveWorkflow("session-directory"); const sessionDir = join(session.cwd, "custom-sessions"); const parent = SessionManager.create(session.cwd, sessionDir); const ctx = { ...session.session.extensionRunner.createCommandContext(), sessionManager: parent, }; const create = vi .spyOn(SessionManager, "create") .mockImplementation(() => { throw new Error("Stop before provider access"); }); if (surface === "command") { driveOverlay(ctx as unknown as ExtensionCommandContext); await extension.commands .get("ultra")! .handler("run session-directory", ctx); } else { await extension.commands.get("ultra")!.handler("", ctx); await extension.tools.get("run_workflow")!.definition.execute( "call", { name: "session-directory" }, undefined, undefined, createToolContext({ ui: ctx.ui, sessionManager: parent, }), ); } expect(create).toHaveBeenCalledWith(session.cwd, parent.getSessionDir(), { parentSession: parent.getSessionFile(), }); }, ); it.each([false, true])( "bounds orchestrator content while retaining UI details (oversized: %s)", async (oversized) => { const extension = await start(); const ctx = session?.session.extensionRunner.createCommandContext(); if (!ctx) throw new Error("missing session"); await extension.commands.get("ultra")!.handler("", ctx); const answer = oversized ? "界".repeat(30_000) : "Final answer"; const response = await extension.tools .get("run_workflow")! .definition.execute( "call", { spec: { name: "context-budget", phases: [ { id: "work", kind: "single", when: "{args.run}", step: { summary: "Unused", prompt: "Unused" }, }, ], return: "{args.answer}", }, args: { run: false, answer }, }, undefined, undefined, createToolContext(), ); const details = response.details as { workflowResult: { result: unknown; phaseResults: Record; }; fullOutputPath?: string; }; const body = JSON.parse((response.content[0] as { text: string }).text); try { expect(details.workflowResult.result).toBe(answer); expect(details.workflowResult.phaseResults).toEqual({ work: [], }); expect(body.phaseResults).toBeUndefined(); if (oversized) { expect(body.truncated).toBe(true); expect(body.fullOutputPath).toBe(details.fullOutputPath); expect({ ...JSON.parse(readFileSync(body.fullOutputPath, "utf8")), fullOutputPath: body.fullOutputPath, }).toEqual(details.workflowResult); } else { expect(body.fullOutputPath).toBe(details.fullOutputPath); expect( JSON.parse(readFileSync(body.fullOutputPath, "utf8")).result, ).toBe(answer); expect(body.result).toBe(answer); expect(body.phases).toEqual([ { id: "work", ran: 0, ok: 0, dropped: 0 }, ]); expect(body.truncated).toBeUndefined(); } } finally { if (details.fullOutputPath) rmSync(dirname(details.fullOutputPath), { recursive: true, force: true, }); } }, ); it.each(["command-output", "review", "research"])( "keeps %s command context bounded and its evidence retrievable", async (name) => { const extension = await start(); if (!session) throw new Error("missing session"); saveWorkflow(name); const ctx = session.session.extensionRunner.createCommandContext(); const result = { workflow: name, phases: [{ id: "work", ran: 1, ok: 1, dropped: 0 }], phaseResults: { work: [{ ok: true, value: "INTERMEDIATE".repeat(10000) }], }, phaseFailures: {}, result: [{ status: "blocked", blockers: ["Missing evidence"] }], report: { summary: "Human report" }, steered: false, dynamicExtensions: [], tokenUsage: { input: 1, output: 2, total: 3, cacheRead: 0, cacheWrite: 0, cost: 0, }, }; Object.assign(ctx.ui, { custom: vi.fn().mockResolvedValue({ status: "completed", result }), }); await extension.commands.get("ultra")!.handler(`run ${name}`, ctx); expect( uiCallsFor("notify").some( (call) => call.args[0] === `ultra: "${name}" complete · 1 phases · $0.0000 · 2 output.`, ), ).toBe(true); const message = session.session.messages.find( (message): message is Extract => "customType" in message && message.customType === "ultra-result", ); if (!message || typeof message.content !== "string") throw new Error("Missing command result"); const body = JSON.parse(message.content); try { expect(body.result).toEqual(result.result); expect(body.report).toEqual(result.report); expect(body.executionStatus).toBe("finished"); expect(body.phaseResults).toBeUndefined(); expect(message.content).not.toContain("INTERMEDIATE"); expect(message.details).toEqual({ ...result, fullOutputPath: body.fullOutputPath, }); expect(JSON.parse(readFileSync(body.fullOutputPath, "utf8"))).toEqual( result, ); } finally { rmSync(dirname(body.fullOutputPath), { recursive: true, force: true }); } }, ); it("completes subcommands first and saved workflows below run", async () => { const extension = await start(); if (!session) throw new Error("missing session"); const complete = extension.commands.get("ultra")!.getArgumentCompletions!; expect(extension.tools.has("run_workflow")).toBe(false); const subcommands = (await complete(""))!; const workflows = (await complete("run "))!; expect(extension.tools.has("run_workflow")).toBe(true); expect(session.session.getActiveToolNames()).not.toContain("run_workflow"); expect(subcommands.map((item: { value: string }) => item.value)).toEqual([ "exec", "run", ]); expect(workflows.map((item: { value: string }) => item.value)).toEqual( expect.arrayContaining(["run research", "run review"]), ); }); it("activates the workflow tool and notifies without starting a turn", async () => { const extension = await start(); if (!session) throw new Error("missing session"); const beforeMessages = session.session.messages.length; const beforeNotifications = uiCallsFor("notify").length; expect(session.session.getActiveToolNames()).not.toContain("run_workflow"); expect(session.session.systemPrompt).not.toContain("- run_workflow:"); await extension.commands .get("ultra")! .handler("", session.session.extensionRunner.createCommandContext()); expect(session.session.getActiveToolNames()).toContain("run_workflow"); expect(session.session.systemPrompt).toContain( "- run_workflow: run saved or inline multi-agent workflows.", ); expect(session.session.systemPrompt).toContain( "Use run_workflow when decomposition into independent or sequential child-agent steps reduces context or verification risk.", ); expect(session.session.messages).toHaveLength(beforeMessages); const notifications = uiCallsFor("notify"); expect(notifications).toHaveLength(beforeNotifications + 1); expect(notifications.at(-1)?.args).toEqual([ "ultra: run_workflow tool enabled.", "info", ]); }); it("sends exec instructions as the visible user message that starts the turn", async () => { const extension = await start(); if (!session) throw new Error("missing session"); const sendUserMessage = vi .spyOn(session.session, "sendUserMessage") .mockResolvedValue(); const beforeNotifications = uiCallsFor("notify").length; await extension.commands .get("ultra")! .handler( "exec inspect this failure once", session.session.extensionRunner.createCommandContext(), ); expect(session.session.getActiveToolNames()).toContain("run_workflow"); expect(sendUserMessage).toHaveBeenCalledOnce(); const prompt = sendUserMessage.mock.calls[0]?.[0]; expect(prompt).toBe( "# ultra exec\n\n# user request\n\ninspect this failure once", ); const guidance = extension.tools.get("run_workflow")!.definition.promptGuidelines!; for (const rule of guidance) { expect(session.session.systemPrompt.split(rule)).toHaveLength(2); expect(prompt).not.toContain(rule); } expect(session.session.systemPrompt).not.toContain( "# Authoring ultra workflows", ); expect(session.session.systemPrompt).not.toContain("## Worked example"); expect(prompt).not.toContain("The workflow tool is now available"); const injected = await session.session.extensionRunner.emitBeforeAgentStart( { messageType: "user", prompt: prompt as string }, { cwd: session.cwd }, ); expect( injected.systemPromptOptions.sections["ultra_model_tier_policy"], ).toBe(modelTierPrompt(DEFAULT_MODEL_TIERS)); expect(injected.systemPromptOptions.forceSystemPrompt).toBeUndefined(); expect(prompt).not.toContain(DEFAULT_MODEL_TIERS["large"]!.instructions); expect(uiCallsFor("notify")).toHaveLength(beforeNotifications); }); it("renders workflow progress as an overlay without replacing session content", async () => { const extension = await start(); if (!session) throw new Error("missing session"); const ctx = session.session.extensionRunner.createCommandContext(); const custom = vi.fn().mockResolvedValue(null); Object.assign(ctx.ui, { custom }); await extension.commands .get("ultra")! .handler("run review inspect scrolling", ctx); expect(custom).toHaveBeenCalledOnce(); expect(custom.mock.calls[0]![1]).toEqual({ overlay: true, overlayOptions: { anchor: "bottom-left", width: "100%", maxHeight: "100%", }, }); }); it("surfaces workflow errors instead of misreporting them as user aborts", async () => { const extension = await start(); if (!session) throw new Error("missing session"); saveWorkflow("runtime-error", "missing-live-capability"); const ctx = session.session.extensionRunner.createCommandContext(); driveOverlay(ctx); await extension.commands.get("ultra")!.handler("run runtime-error", ctx); expect(uiCallsFor("notify").at(-1)?.args).toEqual([ 'ultra: dynamic extension "missing-live-capability" has no canonical revision. Run a builder with all four dynamic-extension tools.', "error", ]); }); it("prepares dynamic extension runs on both command and tool surfaces", async () => { const message = 'ultra: dynamic extension "missing-live-capability" has no canonical revision. Run a builder with all four dynamic-extension tools.'; const extension = await start(); if (!session) throw new Error("missing session"); saveWorkflow("runtime-error-command", "missing-live-capability"); saveWorkflow("runtime-error-tool", "missing-live-capability"); const ctx = session.session.extensionRunner.createCommandContext(); driveOverlay(ctx); await extension.commands .get("ultra")! .handler("run runtime-error-command", ctx); expect(uiCallsFor("notify").at(-1)?.args).toEqual([message, "error"]); await extension.commands.get("ultra")!.handler("", ctx); await expect( extension.tools .get("run_workflow")! .definition.execute( "call", { name: "runtime-error-tool" }, undefined, undefined, createToolContext({ ui: ctx.ui }), ), ).rejects.toThrow(message); }); it("refuses project workflows on both surfaces when the project is untrusted", async () => { const extension = await start(); if (!session) throw new Error("missing session"); saveWorkflow("repo-owned"); const path = join(session.cwd, ".pi", "workflows", "repo-owned.json"); const message = `ultra: refused workflow "repo-owned" from untrusted project file ${path}. Project workflows are repo-controlled; trust this project to run it.`; session.session.settingsManager.setProjectTrusted(false); vi.spyOn(process, "cwd").mockReturnValue(session.cwd); const ctx = session.session.extensionRunner.createCommandContext(); expect(ctx.isProjectTrusted()).toBe(false); const command = extension.commands.get("ultra")!; const completions = await command.getArgumentCompletions!("run "); expect(completions?.map((item) => item.value)).not.toContain( "run repo-owned", ); expect(completions?.map((item) => item.value)).toContain("run review"); const custom = vi.fn(); Object.assign(ctx.ui, { custom }); await command.handler("run repo-owned", ctx); expect(custom).not.toHaveBeenCalled(); expect(uiCallsFor("notify").at(-1)?.args).toEqual([message, "error"]); await command.handler("", ctx); await expect( extension.tools .get("run_workflow")! .definition.execute( "call", { name: "repo-owned" }, undefined, undefined, createToolContext({ ui: ctx.ui }), ), ).rejects.toThrow(message); }); it("offers project workflows with their source when the project is trusted", async () => { const extension = await start(); if (!session) throw new Error("missing session"); saveWorkflow("repo-owned"); vi.spyOn(process, "cwd").mockReturnValue(session.cwd); const ctx = session.session.extensionRunner.createCommandContext(); expect(ctx.isProjectTrusted()).toBe(true); const completions = await extension.commands.get("ultra")!.getArgumentCompletions!("run re"); expect(completions).toEqual( expect.arrayContaining([ { value: "run repo-owned", label: "repo-owned", description: "project", }, { value: "run review", label: "review", description: "bundled" }, ]), ); }); it("reports confirmed whole-run cancellation as aborted while retaining its aggregate", async () => { const extension = await start(); if (!session) throw new Error("missing session"); saveWorkflow("cancelled-run"); const ctx = session.session.extensionRunner.createCommandContext(); driveOverlay(ctx, (overlay) => { overlay.handleInput("\u001b"); overlay.handleInput("\u001b"); }); await extension.commands.get("ultra")!.handler("run cancelled-run", ctx); expect(uiCallsFor("notify").at(-1)?.args).toEqual([ 'ultra: "cancelled-run" aborted · 1 phases · $0.0000 · 0 output.', "warning", ]); expect( uiCallsFor("notify").some((call) => String(call.args[0]).includes("complete"), ), ).toBe(false); const message = session.session.messages.find( (message): message is Extract => "customType" in message && message.customType === "ultra-result", ); const body = JSON.parse(String(message?.content)); try { expect(body.executionStatus).toBe("aborted"); expect( JSON.parse(readFileSync(body.fullOutputPath, "utf8")).aborted, ).toBe(true); } finally { rmSync(dirname(body.fullOutputPath), { recursive: true, force: true }); } }); it("refuses a legacy bare sub-agent extension nickname", async () => { const extension = await start({ settings: { ultra: { subagentExtensions: ["missing-extension"] } }, }); if (!session) throw new Error("missing session"); await extension.commands .get("ultra")! .handler( "run review", session.session.extensionRunner.createCommandContext(), ); expect(uiCallsFor("notify").at(-1)?.args).toEqual([ 'ultra: subagentExtensions no longer accepts bare extension nickname "missing-extension"; replace it with a Pi extension source path or package source.', "error", ]); }); it.each(["command", "tool"])( "names excluded tool owners and startup flags through the %s surface", async (surface) => { const extension = await start({ settings: { ultra: { subagentExtensions: [ fileURLToPath(new URL("../../ask/index.ts", import.meta.url)), ], }, }, env: { PI_NUSHELL_ENABLE: "false" }, extra: [ fileURLToPath(new URL("../../nushell/index.ts", import.meta.url)), ], }); if (!session) throw new Error("missing session"); const spec = { name: "excluded-nushell", phases: [ { id: "work", kind: "single", step: { summary: "Work", prompt: "Work", tools: ["nu"], flags: { nushell: true }, }, }, ], }; const ctx = session.session.extensionRunner.createCommandContext(); await extension.commands.get("ultra")!.handler("", ctx); let message: string; if (surface === "command") { const directory = join(session.cwd, ".pi", "workflows"); mkdirSync(directory, { recursive: true }); writeFileSync( join(directory, `${spec.name}.json`), JSON.stringify(spec), ); await extension.commands.get("ultra")!.handler(`run ${spec.name}`, ctx); message = String(uiCallsFor("notify").at(-1)?.args[0]); } else { message = await extension.tools .get("run_workflow")! .definition.execute( "call", { spec }, undefined, undefined, createToolContext(), ) .then( () => "unexpected success", (error: Error) => error.message, ); } expect(message).toContain("requires unavailable tools: nu"); expect(message).toContain( 'owning extension "nushell" is excluded by ultra.subagentExtensions', ); expect(message).toContain( `include ${JSON.stringify(resolve("extensions/nushell/index.ts"))} in that setting before using its step.flags`, ); expect(message).toContain("- nushell: nu; --nushell (boolean)"); expect(message).not.toContain("unexpected success"); }, ); it("diagnoses a loaded JS owner after package selection changes to TS without widening authorization", async () => { await start(); if (!session) throw new Error("missing session"); const extensionDir = join(session.cwd, "owner-fixture"); mkdirSync(extensionDir); const priorPath = join(extensionDir, "prior.js"); const replacementPath = join(extensionDir, "replacement.ts"); const otherPath = join(extensionDir, "other.ts"); const source = (flag: string) => [ 'import { Type } from "typebox";', "export default (pi) => {", ` pi.registerFlag(${JSON.stringify(flag)}, { type: "boolean" });`, ' pi.registerTool({ name: "fixture_tool", label: "fixture_tool", description: "fixture", parameters: Type.Object({}), execute: async () => ({ content: [{ type: "text", text: "ok" }] }) });', "};", "", ].join("\n"); writeFileSync(priorPath, source("prior-owner")); writeFileSync(replacementPath, source("replacement-owner")); writeFileSync(otherPath, source("other-owner")); writeFileSync( join(extensionDir, "package.json"), JSON.stringify({ pi: { extensions: ["prior.js"] } }), ); await session.dispose(); writeFileSync( join(session.cwd, ".pi", "settings.json"), JSON.stringify({ ultra: { subagentExtensions: [extensionDir] } }), ); session = await createPiSession({ cwd: cwd!, extensionPaths: [ultraPath(), priorPath], }); installUi(); writeFileSync( join(extensionDir, "package.json"), JSON.stringify({ pi: { extensions: ["replacement.ts"] } }), ); const spec = { name: "changed-entry", phases: [ { id: "work", kind: "single", step: { summary: "Work", prompt: "Work", tools: ["fixture_tool"], }, }, ], }; const createChild = vi.spyOn(SessionManager, "inMemory"); const staleExtension = session.extensions.find((item) => item.commands.has("ultra"), )!; await staleExtension.commands .get("ultra")! .handler("", session.session.extensionRunner.createCommandContext()); const mismatch = await staleExtension.tools .get("run_workflow")! .definition.execute( "call", { spec }, undefined, undefined, createToolContext(), ) .then( () => "unexpected success", (error: Error) => error.message, ); expect(mismatch).toContain("loaded owning extension"); expect(mismatch).toContain("does not match selected extension entries"); expect(mismatch).toContain( "runtime may be stale or explicit selection differs", ); expect(mismatch).toContain("Reload/restart Pi"); expect(mismatch).not.toContain("include"); expect(mismatch).toContain("replacement-owner"); expect(mismatch).not.toContain("prior-owner"); expect(createChild).not.toHaveBeenCalled(); createChild.mockRestore(); await session.dispose(); session = await createPiSession({ cwd: cwd!, extensionPaths: [ultraPath(), replacementPath], }); installUi(); const freshExtension = session.extensions.find((item) => item.commands.has("ultra"), )!; await freshExtension.commands .get("ultra")! .handler("", session.session.extensionRunner.createCommandContext()); const reachedChildCreation = vi .spyOn(SessionManager, "inMemory") .mockImplementation(() => { throw new Error("Stop before provider access"); }); const freshResponse = await freshExtension.tools .get("run_workflow")! .definition.execute( "call", { spec }, undefined, undefined, createToolContext(), ); expect(reachedChildCreation).toHaveBeenCalled(); expect(freshResponse.details).toBeDefined(); reachedChildCreation.mockRestore(); await session.dispose(); writeFileSync( join(session.cwd, ".pi", "settings.json"), JSON.stringify({ ultra: { subagentExtensions: [otherPath] } }), ); session = await createPiSession({ cwd: cwd!, extensionPaths: [ultraPath(), priorPath], }); installUi(); const excludedExtension = session.extensions.find((item) => item.commands.has("ultra"), )!; await excludedExtension.commands .get("ultra")! .handler("", session.session.extensionRunner.createCommandContext()); const noChild = vi.spyOn(SessionManager, "inMemory"); const excluded = await excludedExtension.tools .get("run_workflow")! .definition.execute( "call", { spec }, undefined, undefined, createToolContext(), ) .then( () => "unexpected success", (error: Error) => error.message, ); expect(excluded).toContain("does not match selected extension entries"); expect(excluded).toContain(JSON.stringify(otherPath)); expect(excluded).not.toContain("include"); expect(noChild).not.toHaveBeenCalled(); }); it("reports unknown tools without inventing an owning extension or flag", async () => { const extension = await start(); if (!session) throw new Error("missing session"); const ctx = session.session.extensionRunner.createCommandContext(); await extension.commands.get("ultra")!.handler("", ctx); await expect( extension.tools.get("run_workflow")!.definition.execute( "call", { spec: { name: "unknown-tool", phases: [ { id: "work", kind: "single", step: { summary: "Work", prompt: "Work", tools: ["not_installed"], }, }, ], }, }, undefined, undefined, createToolContext(), ), ).rejects.toThrow( "not_installed: no registered owning extension; check step.tools and installed extensions.", ); }); describe("inactive-tool preflight", () => { const nushellPath = () => fileURLToPath(new URL("../../nushell/index.ts", import.meta.url)); function specWith(step: Record) { return { name: "inactive-preflight", phases: [ { id: "work", kind: "fanout", over: ["a", "b", "c"], step: { summary: "Work {item}", prompt: "Work {item}", ...step }, }, ], }; } function deactivate(...names: string[]): void { if (!session) throw new Error("missing session"); session.session.setActiveToolsByName( session.session .getActiveToolNames() .filter((name) => !names.includes(name)), ); } it.each(["command", "tool"])( "refuses a step listing a registered-but-inactive tool before creating children through the %s surface", async (surface) => { const extension = await start({ settings: { ultra: { subagentExtensions: [nushellPath()] } }, }); if (!session) throw new Error("missing session"); // Like `pi --nushell`: bash is registered but switched off in this Pi. deactivate("bash"); expect(session.session.getActiveToolNames()).not.toContain("bash"); const spec = specWith({ tools: ["read", "bash"] }); const ctx = session.session.extensionRunner.createCommandContext(); const create = vi.spyOn(SessionManager, "create"); const inMemory = vi.spyOn(SessionManager, "inMemory"); let message: string; if (surface === "command") { const directory = join(session.cwd, ".pi", "workflows"); mkdirSync(directory, { recursive: true }); writeFileSync( join(directory, `${spec.name}.json`), JSON.stringify(spec), ); const custom = vi.fn(); Object.assign(ctx.ui, { custom }); await extension.commands .get("ultra")! .handler(`run ${spec.name}`, ctx); expect(custom).not.toHaveBeenCalled(); expect(uiCallsFor("notify").at(-1)?.args[1]).toBe("error"); message = String(uiCallsFor("notify").at(-1)?.args[0]); } else { await extension.commands.get("ultra")!.handler("", ctx); message = await extension.tools .get("run_workflow")! .definition.execute( "call", { spec }, undefined, undefined, createToolContext(), ) .then( () => "unexpected success", (error: Error) => error.message, ); } expect(message).toBe( "ultra: step tools are inactive in this Pi and children inherit its flags: bash. Possibly deactivated by an extension with startup flags: nushell (--nushell: Enable Nushell tools at startup instead of bash and powershell). Use an active alternative or set step.flags.", ); expect(create).not.toHaveBeenCalled(); expect(inMemory).not.toHaveBeenCalled(); }, ); it("recommends the active nu tool when bash is inactive", async () => { const extension = await start({ settings: { ultra: { subagentExtensions: [nushellPath()] } }, env: { PI_NUSHELL_ENABLE: "false" }, extra: [nushellPath()], }); if (!session) throw new Error("missing session"); deactivate("bash"); session.session.setActiveToolsByName([ ...session.session.getActiveToolNames(), "nu", ]); const ctx = session.session.extensionRunner.createCommandContext(); await extension.commands.get("ultra")!.handler("", ctx); await expect( extension.tools .get("run_workflow")! .definition.execute( "call", { spec: specWith({ tools: ["bash"] }) }, undefined, undefined, createToolContext(), ), ).rejects.toThrow( "Use an active alternative or set step.flags. The nu tool is active; prefer it.", ); }); // The parent runs in memory, so children are minted through `inMemory`; // stopping there proves the preflight let child creation begin. async function runPastPreflight(step: Record) { const extension = await start({ settings: { ultra: { subagentExtensions: [nushellPath()] } }, }); if (!session) throw new Error("missing session"); deactivate("bash"); const ctx = session.session.extensionRunner.createCommandContext(); await extension.commands.get("ultra")!.handler("", ctx); const inMemory = vi .spyOn(SessionManager, "inMemory") .mockImplementation(() => { throw new Error("Stop before provider access"); }); const response = await extension.tools .get("run_workflow")! .definition.execute( "call", { spec: specWith(step) }, undefined, undefined, createToolContext(), ); const details = response.details as { workflowResult: { phaseFailures: Record>; }; fullOutputPath?: string; }; if (details.fullOutputPath) rmSync(dirname(details.fullOutputPath), { recursive: true, force: true, }); expect(inMemory).toHaveBeenCalled(); const messages = (details.workflowResult.phaseFailures["work"] ?? []).map( (failure) => failure.message, ); expect(messages).toHaveLength(3); for (const message of messages) { expect(message).toContain("Stop before provider access"); expect(message).not.toContain("step tools are inactive"); } } it("skips the preflight for a step that declares its own flags", async () => { await runPastPreflight({ tools: ["bash"], flags: { nushell: false } }); }); // grep/find/ls are registered but inactive in a stock Pi; children activate // them explicitly, so they never trigger the preflight. it("passes a step listing active tools and Pi's optional builtins", async () => { await runPastPreflight({ tools: ["read", "grep", "find", "ls"] }); }); }); it("starts research without knowing which extension supplies research tools", async () => { const extension = await start(); if (!session) throw new Error("missing session"); const ctx = session.session.extensionRunner.createCommandContext(); const custom = vi.fn().mockResolvedValue(null); Object.assign(ctx.ui, { custom }); await extension.commands .get("ultra")! .handler("run research question", ctx); expect(custom).toHaveBeenCalledOnce(); expect( uiCallsFor("notify").some((call) => String(call.args[0]).includes("unavailable tools"), ), ).toBe(false); }); it("discovers a workflow saved after the extension loaded without /reload", async () => { const extension = await start(); if (!session) throw new Error("missing session"); const complete = () => extension.commands.get("ultra")!.getArgumentCompletions!("run "); await complete(); const previousCwd = process.cwd(); try { process.chdir(session.cwd); const directory = join(session.cwd, ".pi", "workflows"); mkdirSync(directory, { recursive: true }); writeFileSync( join(directory, "fresh.json"), `${JSON.stringify({ name: "fresh", phases: [ { id: "p", kind: "single", step: { summary: "hi", prompt: "hi" } }, ], })}\n`, ); const items = (await complete())!; expect(items.map((item: { value: string }) => item.value)).toContain( "run fresh", ); } finally { process.chdir(previousCwd); } }); it.each([ { params: {}, expected: ['provide either a workflow "name" or an inline "spec"'], }, { params: { name: "does-not-exist" }, expected: ['unknown workflow "does-not-exist"', "Available:", "review"], }, { params: { args: "not-an-object" }, expected: ["args", "object"] }, { params: { spec: { name: "invalid", phases: [{ id: "work", kind: "single", step: { prompt: "Work" } }], }, }, expected: ["summary"], }, ])( "delivers actionable argument errors to the agent and the tool UI: $params", async ({ params, expected }) => { const extension = await start(); if (!session) throw new Error("missing session"); const ctx = session.session.extensionRunner.createCommandContext(); await extension.commands.get("ultra")!.handler("", ctx); // Use the real tool boundary through the owned fixture's scripted // provider: the agent turns a thrown execute() into an isError // toolResult surfaced back to the model. const followUp = vi.fn((context: Context) => { const result = context.messages.findLast( (message) => message.role === "toolResult", ); expect(result?.isError).toBe(true); const error = result?.content .filter((part) => part.type === "text") .map((part) => part.text) .join("\n"); for (const text of expected) expect(error).toContain(text); return fauxAssistantMessage("Error received."); }); session.setResponses([ fauxAssistantMessage( fauxToolCall("run_workflow", params as JsonObject), { stopReason: "toolUse", }, ), followUp, ]); await session.session.prompt("Exercise invalid workflow arguments"); expect(followUp).toHaveBeenCalledOnce(); const toolResults = session.session.messages.filter( (message) => message.role === "toolResult", ); expect(toolResults).toHaveLength(1); const result = toolResults[0]!; expect(result.isError).toBe(true); const resultText = result.content .filter((part) => part.type === "text") .map((part) => part.text) .join("\n"); for (const text of expected) expect(resultText).toContain(text); const message = session.session.messages.find( (message): message is Extract => message.role === "toolResult", ); if (!message) throw new Error("missing tool result message"); expect(message).toMatchObject({ role: "toolResult", isError: true, content: result.content, }); expect(session.session.messages.at(-1)).toMatchObject({ role: "assistant", }); initTheme(); const theme = { fg: (_color: string, text: string) => text, bold: (text: string) => text, } as Theme; const tool = extension.tools.get("run_workflow")!.definition; const renderResult = tool.renderResult!; const toolResult = message as Parameters[0]; const renderContext = { isError: true, } as Parameters[3]; const collapsed = renderResult( toolResult, { expanded: false, isPartial: false }, theme, renderContext, ) .render(120) .join("\n"); expect(collapsed).toContain("ultra · failed:"); expect(collapsed).not.toContain("0/0 done"); const expanded = renderResult( toolResult, { expanded: true, isPartial: false }, theme, renderContext, ) .render(120) .join("\n"); for (const text of expected) expect(expanded).toContain(text); }, ); it("exposes the workflow schema while rejecting malformed specs before execution", async () => { const extension = await start(); if (!session) throw new Error("missing session"); await extension.commands.get("ultra")!.getArgumentCompletions!(""); const tool = extension.tools.get("run_workflow")!.definition; const parameters = tool.parameters as { properties: { spec: { type: string; properties: Record; additionalProperties?: unknown; }; background?: unknown; }; }; const specSchema = parameters.properties.spec; expect(tool.renderCall).toBeTypeOf("function"); expect(specSchema.type).toBe("object"); expect(specSchema.properties).toHaveProperty("name"); expect(specSchema.properties).toHaveProperty("phases"); expect(specSchema.properties).toHaveProperty("report"); expect(specSchema.additionalProperties).not.toBe(true); expect(tool.description).toContain( "run saved or inline multi-agent workflows.", ); expect(tool.promptSnippet).toBe( "run saved or inline multi-agent workflows.", ); expect(tool.promptGuidelines).toEqual([ "Use run_workflow when decomposition into independent or sequential child-agent steps reduces context or verification risk.", "Use dynamic extensions when a workflow needs a reusable, task-specific capability; read the ultra-authoring skill for how to build and load them.", ]); expect(parameters.properties.background).toBeUndefined(); expect(JSON.parse(JSON.stringify(specSchema))).toEqual( JSON.parse(JSON.stringify(WorkflowSpecSchema)), ); const step = { summary: "Work", prompt: "Work" }; for (const spec of [ {}, { name: "invalid", phases: [] }, { name: "invalid", phases: [{ id: "work", kind: "fanout", step }] }, { name: "invalid", phases: [{ id: "work", kind: "single", step: { prompt: "Work" } }], }, { name: "invalid", phases: [ { id: "work", kind: "single", step: { ...step, thinkingLevel: "invalid" }, }, ], }, { name: "invalid", phases: [ { id: "work", kind: "single", step: { ...step, schema: "missing" } }, ], }, { name: "invalid", phases: [{ id: "args", kind: "single", step }] }, ]) { await expect( tool.execute( "invalid", { spec }, undefined, undefined, createToolContext(), ), ).rejects.toThrow("Invalid workflow spec"); } }); });