fix(coding-agent): fix plan-mode example

closes #5940

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

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
 

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<string>(["edit", "write"]);
+const PLAN_MANAGED_TOOLS = new Set<string>([...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);
 	});

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> | void;
+type AgentEndHandler = (
+	event: { type: "agent_end"; messages: AgentMessage[] },
+	ctx: ExtensionContext,
+) => Promise<void> | 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<string, CommandHandler>();
+	let agentEndHandler: AgentEndHandler | undefined;
+
+	const sendMessage = vi.fn<ExtensionAPI["sendMessage"]>();
+	const sendUserMessage = vi.fn<ExtensionAPI["sendUserMessage"]>();
+	const setActiveTools = vi.fn<ExtensionAPI["setActiveTools"]>((toolNames) => {
+		activeTools = [...toolNames];
+	});
+	const appendEntry = vi.fn<ExtensionAPI["appendEntry"]>();
+
+	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<void> {
+		const command = commands.get(name);
+		if (!command) throw new Error(`Missing command: ${name}`);
+		await command("", ctx);
+	}
+
+	async function triggerAgentEnd(text: string): Promise<void> {
+		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",
+		});
+	});
+});