From 542683b29ab2865976dddb006b4d70cffe315e25 Mon Sep 17 00:00:00 2001 From: Armin Ronacher Date: Sun, 21 Jun 2026 22:10:17 +0200 Subject: [PATCH] fix(coding-agent): fix plan-mode example closes #5940 --- packages/coding-agent/CHANGELOG.md | 1 + .../examples/extensions/plan-mode/README.md | 5 +- .../examples/extensions/plan-mode/index.ts | 128 ++++++++++---- .../test/plan-mode-extension.test.ts | 167 ++++++++++++++++++ 4 files changed, 260 insertions(+), 41 deletions(-) create mode 100644 packages/coding-agent/test/plan-mode-extension.test.ts diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 92bdd6ba..20d13deb 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -4,6 +4,7 @@ ### Fixed +- Fixed the plan-mode example to preserve active custom tools, skip the action prompt when no plan is found, and queue refinement/execution follow-ups correctly from `agent_end` ([#5940](https://github.com/earendil-works/pi/issues/5940)). - Fixed `pi update` to install the exact version returned by the Pi update check, make `--force` reinstall that checked version, fail instead of falling back to an unversioned reinstall when no version is available, and report both the old and updated versions. ## [0.79.9] - 2026-06-20 diff --git a/packages/coding-agent/examples/extensions/plan-mode/README.md b/packages/coding-agent/examples/extensions/plan-mode/README.md index 549e3473..2568a684 100644 --- a/packages/coding-agent/examples/extensions/plan-mode/README.md +++ b/packages/coding-agent/examples/extensions/plan-mode/README.md @@ -4,7 +4,7 @@ Read-only exploration mode for safe code analysis. ## Features -- **Read-only tools**: Restricts available tools to read, bash, grep, find, ls, question +- **Built-in write tools disabled**: Disables edit/write while preserving other active tools - **Bash allowlist**: Only read-only bash commands are allowed - **Plan extraction**: Extracts numbered steps from `Plan:` sections - **Progress tracking**: Widget shows completion status during execution @@ -37,7 +37,8 @@ Plan: ## How It Works ### Plan Mode (Read-Only) -- Only read-only tools available +- Built-in edit/write tools disabled +- Other active tools remain available - Bash commands filtered through allowlist - Agent creates a plan without making changes diff --git a/packages/coding-agent/examples/extensions/plan-mode/index.ts b/packages/coding-agent/examples/extensions/plan-mode/index.ts index 40db408c..737ce56a 100644 --- a/packages/coding-agent/examples/extensions/plan-mode/index.ts +++ b/packages/coding-agent/examples/extensions/plan-mode/index.ts @@ -2,7 +2,7 @@ * Plan Mode Extension * * Read-only exploration mode for safe code analysis. - * When enabled, only read-only tools are available. + * When enabled, built-in write tools are disabled. * * Features: * - /plan command or Ctrl+Alt+P to toggle @@ -21,6 +21,15 @@ import { extractTodoItems, isSafeCommand, markCompletedSteps, type TodoItem } fr // Tools const PLAN_MODE_TOOLS = ["read", "bash", "grep", "find", "ls", "questionnaire"]; const NORMAL_MODE_TOOLS = ["read", "bash", "edit", "write"]; +const PLAN_MODE_DISABLED_TOOLS = new Set(["edit", "write"]); +const PLAN_MANAGED_TOOLS = new Set([...PLAN_MODE_TOOLS, ...NORMAL_MODE_TOOLS]); + +interface PlanModeState { + enabled: boolean; + todos?: TodoItem[]; + executing?: boolean; + toolsBeforePlanMode?: string[]; +} // Type guard for assistant messages function isAssistantMessage(m: AgentMessage): m is AssistantMessage { @@ -39,6 +48,7 @@ export default function planModeExtension(pi: ExtensionAPI): void { let planModeEnabled = false; let executionMode = false; let todoItems: TodoItem[] = []; + let toolsBeforePlanMode: string[] | undefined; pi.registerFlag("plan", { description: "Start in plan mode (read-only exploration)", @@ -73,19 +83,34 @@ export default function planModeExtension(pi: ExtensionAPI): void { } } - function togglePlanMode(ctx: ExtensionContext): void { - planModeEnabled = !planModeEnabled; - executionMode = false; - todoItems = []; + function uniqueToolNames(toolNames: string[]): string[] { + return [...new Set(toolNames)]; + } - if (planModeEnabled) { - pi.setActiveTools(PLAN_MODE_TOOLS); - ctx.ui.notify(`Plan mode enabled. Tools: ${PLAN_MODE_TOOLS.join(", ")}`); - } else { - pi.setActiveTools(NORMAL_MODE_TOOLS); - ctx.ui.notify("Plan mode disabled. Full access restored."); + function getPlanModeTools(activeToolNames: string[]): string[] { + return uniqueToolNames([ + ...activeToolNames.filter((name) => !PLAN_MODE_DISABLED_TOOLS.has(name)), + ...PLAN_MODE_TOOLS, + ]); + } + + function getNormalModeTools(activeToolNames: string[]): string[] { + return uniqueToolNames([ + ...NORMAL_MODE_TOOLS, + ...activeToolNames.filter((name) => !PLAN_MANAGED_TOOLS.has(name)), + ]); + } + + function enablePlanModeTools(): void { + if (toolsBeforePlanMode === undefined) { + toolsBeforePlanMode = pi.getActiveTools(); } - updateStatus(ctx); + pi.setActiveTools(getPlanModeTools(toolsBeforePlanMode)); + } + + function restoreNormalModeTools(): void { + pi.setActiveTools(toolsBeforePlanMode ?? getNormalModeTools(pi.getActiveTools())); + toolsBeforePlanMode = undefined; } function persistState(): void { @@ -93,9 +118,26 @@ export default function planModeExtension(pi: ExtensionAPI): void { enabled: planModeEnabled, todos: todoItems, executing: executionMode, + toolsBeforePlanMode, }); } + function togglePlanMode(ctx: ExtensionContext): void { + planModeEnabled = !planModeEnabled; + executionMode = false; + todoItems = []; + + if (planModeEnabled) { + enablePlanModeTools(); + ctx.ui.notify("Plan mode enabled. Built-in write tools disabled."); + } else { + restoreNormalModeTools(); + ctx.ui.notify("Plan mode disabled. Full access restored."); + } + updateStatus(ctx); + persistState(); + } + pi.registerCommand("plan", { description: "Toggle plan mode (read-only exploration)", handler: async (_args, ctx) => togglePlanMode(ctx), @@ -165,8 +207,8 @@ export default function planModeExtension(pi: ExtensionAPI): void { You are in plan mode - a read-only exploration mode for safe code analysis. Restrictions: -- You can only use: read, bash, grep, find, ls, questionnaire -- You CANNOT use: edit, write (file modifications are disabled) +- Built-in edit and write tools are disabled +- Other currently active tools remain available - Bash is restricted to an allowlist of read-only commands Ask clarifying questions using the questionnaire tool. @@ -228,7 +270,6 @@ After completing a step, include a [DONE:n] tag in your response.`, ); executionMode = false; todoItems = []; - pi.setActiveTools(NORMAL_MODE_TOOLS); updateStatus(ctx); persistState(); // Save cleared state so resume doesn't restore old execution mode } @@ -246,43 +287,51 @@ After completing a step, include a [DONE:n] tag in your response.`, } } + if (todoItems.length === 0) return; + persistState(); + // Show plan steps and prompt for next action - if (todoItems.length > 0) { - const todoListText = todoItems.map((t, i) => `${i + 1}. ☐ ${t.text}`).join("\n"); - pi.sendMessage( - { - customType: "plan-todo-list", - content: `**Plan Steps (${todoItems.length}):**\n\n${todoListText}`, - display: true, - }, - { triggerTurn: false }, - ); - } + const todoListText = todoItems.map((t, i) => `${i + 1}. ☐ ${t.text}`).join("\n"); + const planTodoListMessage = { + customType: "plan-todo-list", + content: `**Plan Steps (${todoItems.length}):**\n\n${todoListText}`, + display: true, + }; const choice = await ctx.ui.select("Plan mode - what next?", [ - todoItems.length > 0 ? "Execute the plan (track progress)" : "Execute the plan", + "Execute the plan (track progress)", "Stay in plan mode", "Refine the plan", ]); if (choice?.startsWith("Execute")) { - planModeEnabled = false; - executionMode = todoItems.length > 0; - pi.setActiveTools(NORMAL_MODE_TOOLS); - updateStatus(ctx); + const firstTodoItem = todoItems[0]; + if (!firstTodoItem) return; - const execMessage = - todoItems.length > 0 - ? `Execute the plan. Start with: ${todoItems[0].text}` - : "Execute the plan you just created."; + planModeEnabled = false; + executionMode = true; + restoreNormalModeTools(); + updateStatus(ctx); + persistState(); + + const remainingList = todoItems.map((t) => `${t.step}. ${t.text}`).join("\n"); + const execMessage = `Execute the plan. + +Remaining steps: +${remainingList} + +Start with: ${firstTodoItem.text} +After completing a step, include a [DONE:n] tag in your response.`; + pi.sendMessage(planTodoListMessage, { deliverAs: "followUp" }); pi.sendMessage( { customType: "plan-mode-execute", content: execMessage, display: true }, - { triggerTurn: true }, + { triggerTurn: true, deliverAs: "followUp" }, ); } else if (choice === "Refine the plan") { const refinement = await ctx.ui.editor("Refine the plan:", ""); if (refinement?.trim()) { - pi.sendUserMessage(refinement.trim()); + pi.sendMessage(planTodoListMessage, { deliverAs: "followUp" }); + pi.sendUserMessage(refinement.trim(), { deliverAs: "followUp" }); } } }); @@ -298,12 +347,13 @@ After completing a step, include a [DONE:n] tag in your response.`, // Restore persisted state const planModeEntry = entries .filter((e: { type: string; customType?: string }) => e.type === "custom" && e.customType === "plan-mode") - .pop() as { data?: { enabled: boolean; todos?: TodoItem[]; executing?: boolean } } | undefined; + .pop() as { data?: PlanModeState } | undefined; if (planModeEntry?.data) { planModeEnabled = planModeEntry.data.enabled ?? planModeEnabled; todoItems = planModeEntry.data.todos ?? todoItems; executionMode = planModeEntry.data.executing ?? executionMode; + toolsBeforePlanMode = planModeEntry.data.toolsBeforePlanMode ?? toolsBeforePlanMode; } // On resume: re-scan messages to rebuild completion state @@ -333,7 +383,7 @@ After completing a step, include a [DONE:n] tag in your response.`, } if (planModeEnabled) { - pi.setActiveTools(PLAN_MODE_TOOLS); + enablePlanModeTools(); } updateStatus(ctx); }); diff --git a/packages/coding-agent/test/plan-mode-extension.test.ts b/packages/coding-agent/test/plan-mode-extension.test.ts new file mode 100644 index 00000000..419a707e --- /dev/null +++ b/packages/coding-agent/test/plan-mode-extension.test.ts @@ -0,0 +1,167 @@ +import type { AgentMessage } from "@earendil-works/pi-agent-core"; +import type { AssistantMessage } from "@earendil-works/pi-ai"; +import { describe, expect, it, vi } from "vitest"; +import planModeExtension from "../examples/extensions/plan-mode/index.ts"; +import type { ExtensionAPI, ExtensionContext } from "../src/core/extensions/index.ts"; + +type CommandHandler = (args: string, ctx: ExtensionContext) => Promise | void; +type AgentEndHandler = ( + event: { type: "agent_end"; messages: AgentMessage[] }, + ctx: ExtensionContext, +) => Promise | void; + +function createAssistantMessage(text: string): AssistantMessage { + return { + role: "assistant", + content: [{ type: "text", text }], + api: "anthropic-messages", + provider: "anthropic", + model: "mock", + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop", + timestamp: Date.now(), + }; +} + +function setup(options: { activeTools?: string[]; selectChoice?: string; editorText?: string } = {}) { + let activeTools = options.activeTools ?? ["read", "bash", "edit", "write"]; + const commands = new Map(); + let agentEndHandler: AgentEndHandler | undefined; + + const sendMessage = vi.fn(); + const sendUserMessage = vi.fn(); + const setActiveTools = vi.fn((toolNames) => { + activeTools = [...toolNames]; + }); + const appendEntry = vi.fn(); + + const api = { + registerFlag: vi.fn(), + registerCommand(name: string, command: { handler: CommandHandler }) { + commands.set(name, command.handler); + }, + registerShortcut: vi.fn(), + on(event: string, handler: unknown) { + if (event === "agent_end") agentEndHandler = handler as AgentEndHandler; + }, + getFlag: vi.fn(() => false), + getActiveTools: vi.fn(() => [...activeTools]), + setActiveTools, + sendMessage, + sendUserMessage, + appendEntry, + } as unknown as ExtensionAPI; + + planModeExtension(api); + + const ctx = { + hasUI: true, + ui: { + notify: vi.fn(), + select: vi.fn(async () => options.selectChoice), + editor: vi.fn(async () => options.editorText), + setStatus: vi.fn(), + setWidget: vi.fn(), + theme: { + fg: (_name: string, text: string) => text, + strikethrough: (text: string) => text, + }, + }, + sessionManager: { getEntries: () => [] }, + isIdle: () => false, + hasPendingMessages: () => false, + } as unknown as ExtensionContext; + + async function runCommand(name: string): Promise { + const command = commands.get(name); + if (!command) throw new Error(`Missing command: ${name}`); + await command("", ctx); + } + + async function triggerAgentEnd(text: string): Promise { + if (!agentEndHandler) throw new Error("Missing agent_end handler"); + await agentEndHandler({ type: "agent_end", messages: [createAssistantMessage(text)] }, ctx); + } + + return { + activeTools: () => activeTools, + appendEntry, + ctx, + runCommand, + sendMessage, + sendUserMessage, + setActiveTools, + triggerAgentEnd, + }; +} + +describe("plan-mode example extension", () => { + it("preserves custom active tools while toggling plan mode", async () => { + const { activeTools, runCommand, setActiveTools } = setup({ + activeTools: ["read", "bash", "edit", "write", "echo_tool"], + }); + + await runCommand("plan"); + + expect(activeTools()).toEqual(["read", "bash", "echo_tool", "grep", "find", "ls", "questionnaire"]); + expect(setActiveTools).toHaveBeenLastCalledWith([ + "read", + "bash", + "echo_tool", + "grep", + "find", + "ls", + "questionnaire", + ]); + + await runCommand("plan"); + + expect(activeTools()).toEqual(["read", "bash", "edit", "write", "echo_tool"]); + expect(setActiveTools).toHaveBeenLastCalledWith(["read", "bash", "edit", "write", "echo_tool"]); + }); + + it("does not prompt when the assistant response contains no plan", async () => { + const { ctx, runCommand, sendMessage, triggerAgentEnd } = setup(); + + await runCommand("plan"); + await triggerAgentEnd("This file defines the command-line argument parser."); + + expect(ctx.ui.select).not.toHaveBeenCalled(); + expect(sendMessage).not.toHaveBeenCalled(); + }); + + it("queues plan refinement as a follow-up user message", async () => { + const { runCommand, sendUserMessage, triggerAgentEnd } = setup({ + selectChoice: "Refine the plan", + editorText: "Add a regression test.", + }); + + await runCommand("plan"); + await triggerAgentEnd("Plan:\n1. Inspect the current implementation\n2. Add a regression test"); + + expect(sendUserMessage).toHaveBeenCalledWith("Add a regression test.", { deliverAs: "followUp" }); + }); + + it("queues plan execution as a follow-up custom message", async () => { + const { activeTools, runCommand, sendMessage, triggerAgentEnd } = setup({ + activeTools: ["read", "bash", "edit", "write", "echo_tool"], + selectChoice: "Execute the plan (track progress)", + }); + + await runCommand("plan"); + await triggerAgentEnd("Plan:\n1. Inspect the current implementation\n2. Add a regression test"); + + expect(activeTools()).toEqual(["read", "bash", "edit", "write", "echo_tool"]); + expect(sendMessage).toHaveBeenCalledWith(expect.objectContaining({ customType: "plan-mode-execute" }), { + triggerTurn: true, + deliverAs: "followUp", + }); + }); +});