/ extensions / plan-mode.ts
plan-mode.ts
  1  /**
  2   * Plan Mode Extension
  3   *
  4   * Provides a Claude Code-style "plan mode" read-only sandbox for safe code exploration.
  5   * When enabled, Pi-native write tools are removed from the active Pi tool list and
  6   * write-capable bash/RepoPrompt operations are blocked.
  7   *
  8   * Features:
  9   * - /plan command (and Ctrl+Alt/Option+P shortcut) to toggle plan mode
 10   * - --plan flag to start in plan mode
 11   * - Removes Pi-native write tools (`edit`, `write`) from the active Pi tool list while enabled
 12   * - Blocks destructive bash commands while plan mode is enabled (including redirects)
 13   * - Blocks RepoPrompt write commands (edit/file/file_actions/apply-edits), even via bash rp-cli -e, rp_exec, or rp (repoprompt-mcp)
 14   * - Blocks rp-cli interactive REPL (-i/--interactive) to prevent bypassing the sandbox
 15   * - Adds plan-mode instructions via the system prompt only while plan mode is enabled
 16   * - Shows a "plan" indicator in the status line when active
 17   * - Persists plan mode state only when toggled (and once at startup if --plan is used)
 18   *
 19   * Usage:
 20   * 1. Copy this file to ~/.pi/agent/extensions/ or your project's .pi/extensions/
 21   * 2. Use /plan (or the ctrl+opt+P / ctrl+alt+P hotkey) to toggle plan mode on/off
 22   * 3. Or start in plan mode with --plan
 23   */
 24  
 25  import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
 26  import { Key } from "@mariozechner/pi-tui";
 27  
 28  let parseBash: ((input: string) => any) | null = null;
 29  let justBashLoadPromise: Promise<void> | null = null;
 30  let justBashLoadDone = false;
 31  
 32  async function ensureJustBashLoaded(): Promise<void> {
 33  	if (justBashLoadDone) return;
 34  
 35  	if (!justBashLoadPromise) {
 36  		justBashLoadPromise = import("just-bash")
 37  			.then((mod: any) => {
 38  				parseBash = typeof mod?.parse === "function" ? mod.parse : null;
 39  			})
 40  			.catch(() => {
 41  				parseBash = null;
 42  			})
 43  			.finally(() => {
 44  				justBashLoadDone = true;
 45  			});
 46  	}
 47  
 48  	await justBashLoadPromise;
 49  }
 50  
 51  let warnedAstUnavailable = false;
 52  function maybeWarnAstUnavailable(ctx: ExtensionContext): void {
 53  	if (warnedAstUnavailable) return;
 54  	if (parseBash) return;
 55  	if (!ctx.hasUI) return;
 56  
 57  	warnedAstUnavailable = true;
 58  	ctx.ui.notify(
 59  		"plan-mode: bash AST parser unavailable; falling back to best-effort regex command checks",
 60  		"warning",
 61  	);
 62  }
 63  
 64  type BashInvocation = {
 65  	commandNameRaw: string;
 66  	commandName: string;
 67  	effectiveCommandName: string;
 68  	effectiveArgs: string[];
 69  	hasWriteRedirection: boolean;
 70  };
 71  
 72  const WRAPPER_COMMANDS = new Set(["command", "builtin", "exec", "nohup"]);
 73  const WRITE_REDIRECTION_OPERATORS = new Set([">", ">>", ">|", "<>", "&>", "&>>", ">&"]);
 74  
 75  function commandBaseName(value: string): string {
 76  	const normalized = value.replace(/\\+/g, "/");
 77  	const idx = normalized.lastIndexOf("/");
 78  	const base = idx >= 0 ? normalized.slice(idx + 1) : normalized;
 79  	return base.toLowerCase();
 80  }
 81  
 82  function partToText(part: any): string {
 83  	if (!part || typeof part !== "object") return "";
 84  
 85  	switch (part.type) {
 86  		case "Literal":
 87  		case "SingleQuoted":
 88  		case "Escaped":
 89  			return typeof part.value === "string" ? part.value : "";
 90  		case "DoubleQuoted":
 91  			return Array.isArray(part.parts) ? part.parts.map(partToText).join("") : "";
 92  		case "Glob":
 93  			return typeof part.pattern === "string" ? part.pattern : "";
 94  		case "TildeExpansion":
 95  			return typeof part.user === "string" && part.user.length > 0 ? `~${part.user}` : "~";
 96  		case "ParameterExpansion":
 97  			return typeof part.parameter === "string" && part.parameter.length > 0
 98  				? "${" + part.parameter + "}"
 99  				: "${}";
100  		case "CommandSubstitution":
101  			return "$(...)";
102  		case "ProcessSubstitution":
103  			return part.direction === "output" ? ">(...)" : "<(...)";
104  		case "ArithmeticExpansion":
105  			return "$((...))";
106  		default:
107  			return "";
108  	}
109  }
110  
111  function wordToText(word: any): string {
112  	if (!word || typeof word !== "object" || !Array.isArray(word.parts)) return "";
113  	return word.parts.map(partToText).join("");
114  }
115  
116  function resolveEffectiveCommand(commandNameRaw: string, args: string[]): {
117  	effectiveCommandName: string;
118  	effectiveArgs: string[];
119  } {
120  	const primary = commandNameRaw.trim();
121  	const primaryBase = commandBaseName(primary);
122  
123  	if (WRAPPER_COMMANDS.has(primaryBase)) {
124  		const next = args[0] ?? "";
125  		return {
126  			effectiveCommandName: commandBaseName(next),
127  			effectiveArgs: args.slice(1),
128  		};
129  	}
130  
131  	if (primaryBase === "env") {
132  		let idx = 0;
133  		while (idx < args.length) {
134  			const token = args[idx] ?? "";
135  			if (token === "--") {
136  				idx += 1;
137  				break;
138  			}
139  			if (token.startsWith("-") || /^[A-Za-z_][A-Za-z0-9_]*=.*/.test(token)) {
140  				idx += 1;
141  				continue;
142  			}
143  			break;
144  		}
145  
146  		const next = args[idx] ?? "";
147  		return {
148  			effectiveCommandName: commandBaseName(next),
149  			effectiveArgs: args.slice(idx + 1),
150  		};
151  	}
152  
153  	if (primaryBase === "sudo") {
154  		let idx = 0;
155  		while (idx < args.length) {
156  			const token = args[idx] ?? "";
157  			if (token === "--") {
158  				idx += 1;
159  				break;
160  			}
161  			if (token.startsWith("-")) {
162  				idx += 1;
163  				continue;
164  			}
165  			break;
166  		}
167  
168  		const next = args[idx] ?? "";
169  		return {
170  			effectiveCommandName: commandBaseName(next),
171  			effectiveArgs: args.slice(idx + 1),
172  		};
173  	}
174  
175  	return {
176  		effectiveCommandName: primaryBase,
177  		effectiveArgs: args,
178  	};
179  }
180  
181  function collectNestedScriptsFromWord(word: any, collect: (script: any) => void): void {
182  	if (!word || typeof word !== "object" || !Array.isArray(word.parts)) return;
183  
184  	for (const part of word.parts) {
185  		if (!part || typeof part !== "object") continue;
186  
187  		if (part.type === "DoubleQuoted") {
188  			collectNestedScriptsFromWord(part, collect);
189  			continue;
190  		}
191  
192  		if ((part.type === "CommandSubstitution" || part.type === "ProcessSubstitution") && part.body) {
193  			collect(part.body);
194  		}
195  	}
196  }
197  
198  function analyzeBashScript(command: string): { parseError?: string; invocations: BashInvocation[] } {
199  	try {
200  		if (!parseBash) {
201  			return { parseError: "just-bash parse unavailable", invocations: [] };
202  		}
203  
204  		const ast: any = parseBash(command);
205  		const invocations: BashInvocation[] = [];
206  
207  		const visitScript = (script: any) => {
208  			if (!script || typeof script !== "object" || !Array.isArray(script.statements)) return;
209  
210  			for (const statement of script.statements) {
211  				if (!statement || typeof statement !== "object" || !Array.isArray(statement.pipelines)) continue;
212  
213  				for (const pipeline of statement.pipelines) {
214  					if (!pipeline || typeof pipeline !== "object" || !Array.isArray(pipeline.commands)) continue;
215  
216  					for (const commandNode of pipeline.commands) {
217  						if (!commandNode || typeof commandNode !== "object") continue;
218  
219  						if (commandNode.type === "SimpleCommand") {
220  							const commandNameRaw = wordToText(commandNode.name).trim();
221  							const commandName = commandBaseName(commandNameRaw);
222  							const args = Array.isArray(commandNode.args)
223  								? commandNode.args.map((arg: any) => wordToText(arg)).filter(Boolean)
224  								: [];
225  							const redirections = Array.isArray(commandNode.redirections)
226  								? commandNode.redirections.map((r: any) => typeof r?.operator === "string" ? r.operator : "")
227  								: [];
228  							const effective = resolveEffectiveCommand(commandNameRaw, args);
229  
230  							invocations.push({
231  								commandNameRaw,
232  								commandName,
233  								effectiveCommandName: effective.effectiveCommandName,
234  								effectiveArgs: effective.effectiveArgs,
235  								hasWriteRedirection: redirections.some((op) => WRITE_REDIRECTION_OPERATORS.has(op)),
236  							});
237  
238  							if (commandNode.name) collectNestedScriptsFromWord(commandNode.name, visitScript);
239  							if (Array.isArray(commandNode.args)) {
240  								for (const arg of commandNode.args) {
241  									collectNestedScriptsFromWord(arg, visitScript);
242  								}
243  							}
244  							continue;
245  						}
246  
247  						if (Array.isArray(commandNode.body)) visitScript({ statements: commandNode.body });
248  						if (Array.isArray(commandNode.condition)) visitScript({ statements: commandNode.condition });
249  						if (Array.isArray(commandNode.clauses)) {
250  							for (const clause of commandNode.clauses) {
251  								if (Array.isArray(clause?.condition)) visitScript({ statements: clause.condition });
252  								if (Array.isArray(clause?.body)) visitScript({ statements: clause.body });
253  							}
254  						}
255  						if (Array.isArray(commandNode.elseBody)) visitScript({ statements: commandNode.elseBody });
256  						if (Array.isArray(commandNode.items)) {
257  							for (const item of commandNode.items) {
258  								if (Array.isArray(item?.body)) visitScript({ statements: item.body });
259  							}
260  						}
261  						if (commandNode.word) collectNestedScriptsFromWord(commandNode.word, visitScript);
262  						if (Array.isArray(commandNode.words)) {
263  							for (const word of commandNode.words) {
264  								collectNestedScriptsFromWord(word, visitScript);
265  							}
266  						}
267  					}
268  				}
269  			}
270  		};
271  
272  		visitScript(ast);
273  		return { invocations };
274  	} catch (error: any) {
275  		return { parseError: error?.message ?? String(error), invocations: [] };
276  	}
277  }
278  
279  const PLAN_MODE_DISABLED_TOOLS = ["edit", "write"] as const;
280  const PLAN_MODE_DISABLED_TOOL_SET = new Set<string>(PLAN_MODE_DISABLED_TOOLS);
281  
282  function removePlanModeWriteTools(toolNames: string[]): string[] {
283  	return toolNames.filter((toolName) => !PLAN_MODE_DISABLED_TOOL_SET.has(toolName));
284  }
285  
286  function restorePlanModeWriteTools(toolNames: string[], toolsBeforePlanMode: string[]): string[] {
287  	const activeWithoutWrites = removePlanModeWriteTools(toolNames);
288  	const remainingActiveTools = new Set(activeWithoutWrites);
289  	const restored: string[] = [];
290  
291  	for (const toolName of toolsBeforePlanMode) {
292  		if (PLAN_MODE_DISABLED_TOOL_SET.has(toolName)) {
293  			restored.push(toolName);
294  			continue;
295  		}
296  
297  		if (remainingActiveTools.has(toolName)) {
298  			restored.push(toolName);
299  			remainingActiveTools.delete(toolName);
300  		}
301  	}
302  
303  	for (const toolName of activeWithoutWrites) {
304  		if (remainingActiveTools.has(toolName)) {
305  			restored.push(toolName);
306  		}
307  	}
308  
309  	return restored;
310  }
311  
312  function toolListsMatch(current: string[], next: string[]): boolean {
313  	return current.length === next.length && current.every((toolName, index) => toolName === next[index]);
314  }
315  
316  // Patterns for destructive bash commands that should be blocked in plan mode
317  const DESTRUCTIVE_PATTERNS = [
318  	/\brm\b/i,
319  	/\brmdir\b/i,
320  	/\bmv\b/i,
321  	/\bcp\b/i,
322  	/\bmkdir\b/i,
323  	/\btouch\b/i,
324  	/\bchmod\b/i,
325  	/\bchown\b/i,
326  	/\bchgrp\b/i,
327  	/\bln\b/i,
328  	/\btee\b/i,
329  	/\btruncate\b/i,
330  	/\bdd\b/i,
331  	/\bshred\b/i,
332  	/[^<]>(?![>&])/, // redirect stdout to a file
333  	/>>/, // append redirect
334  	/\bnpm\s+(install|uninstall|update|ci|link|publish)/i,
335  	/\byarn\s+(add|remove|install|publish)/i,
336  	/\bpnpm\s+(add|remove|install|publish)/i,
337  	/\bpip\s+(install|uninstall)/i,
338  	/\bapt(-get)?\s+(install|remove|purge|update|upgrade)/i,
339  	/\bbrew\s+(install|uninstall|upgrade)/i,
340  	/\bgit\s+(add|commit|push|pull|merge|rebase|reset|checkout\s+-b|branch\s+-[dD]|stash|cherry-pick|revert|tag|init|clone)/i,
341  	/\bsudo\b/i,
342  	/\bsu\b/i,
343  	/\bkill\b/i,
344  	/\bpkill\b/i,
345  	/\bkillall\b/i,
346  	/\breboot\b/i,
347  	/\bshutdown\b/i,
348  	/\bsystemctl\s+(start|stop|restart|enable|disable)/i,
349  	/\bservice\s+\S+\s+(start|stop|restart)/i,
350  	/\b(vim?|nano|emacs|code|subl)\b/i,
351  ];
352  
353  // Read-only commands that are always safe
354  const SAFE_COMMANDS = [
355  	/^\s*cat\b/,
356  	/^\s*head\b/,
357  	/^\s*tail\b/,
358  	/^\s*less\b/,
359  	/^\s*more\b/,
360  	/^\s*grep\b/,
361  	/^\s*find\b/,
362  	/^\s*ls\b/,
363  	/^\s*pwd\b/,
364  	/^\s*echo\b/,
365  	/^\s*printf\b/,
366  	/^\s*wc\b/,
367  	/^\s*sort\b/,
368  	/^\s*uniq\b/,
369  	/^\s*diff\b/,
370  	/^\s*file\b/,
371  	/^\s*stat\b/,
372  	/^\s*du\b/,
373  	/^\s*df\b/,
374  	/^\s*tree\b/,
375  	/^\s*which\b/,
376  	/^\s*whereis\b/,
377  	/^\s*type\b/,
378  	/^\s*env\b/,
379  	/^\s*printenv\b/,
380  	/^\s*uname\b/,
381  	/^\s*whoami\b/,
382  	/^\s*id\b/,
383  	/^\s*date\b/,
384  	/^\s*cal\b/,
385  	/^\s*uptime\b/,
386  	/^\s*ps\b/,
387  	/^\s*top\b/,
388  	/^\s*htop\b/,
389  	/^\s*free\b/,
390  	/^\s*git\s+(status|log|diff|show|branch|remote|config\s+--get)/i,
391  	/^\s*git\s+ls-/i,
392  	/^\s*npm\s+(list|ls|view|info|search|outdated|audit)/i,
393  	/^\s*yarn\s+(list|info|why|audit)/i,
394  	/^\s*node\s+--version/i,
395  	/^\s*python\s+--version/i,
396  	/^\s*curl\s/i,
397  	/^\s*wget\s+-O\s*-/i,
398  	/^\s*jq\b/,
399  	/^\s*sed\s+-n/i,
400  	/^\s*awk\b/,
401  	/^\s*rg\b/,
402  	/^\s*fd\b/,
403  	/^\s*bat\b/,
404  	/^\s*exa\b/,
405  	/^\s*rp-cli\b/,
406  	/^\s*rp_exec\b/,
407  	/^\s*rp_bind\b/,
408  ];
409  
410  const REPROMPT_WRITE_PATTERNS = [
411  	/(^|&&\s*)\s*edit\b/i,
412  	/(^|&&\s*)\s*file\s+(create|delete|move)\b/i,
413  	/(^|&&\s*)\s*file_actions\b/i,
414  	/(^|&&\s*)\s*call\s+(apply-edits|file_actions)\b/i,
415  ];
416  
417  const RP_CLI_INTERACTIVE_PATTERN =
418  	/(^|\s)rp-cli\b.*(?:\s)(?:-i|--interactive)(?:\s|$)/i;
419  
420  const RP_CLI_EXEC_WRITE_PATTERN =
421  	/(^|\s)rp-cli\b.*(?:\s)(?:-e|--exec)(?:\s*)[\s\S]*\b(edit|file_actions|file\s+(create|delete|move)|call\s+(apply-edits|file_actions))\b/i;
422  
423  function isRepoPromptWriteCommand(command: string): boolean {
424  	return REPROMPT_WRITE_PATTERNS.some((pattern) => pattern.test(command));
425  }
426  
427  function isRepoPromptMcpWriteRequest(input: unknown): boolean {
428  	if (input === null || typeof input !== "object") {
429  		return false;
430  	}
431  
432  	const request = input as { call?: unknown };
433  	const call = request.call;
434  	if (typeof call !== "string") {
435  		return false;
436  	}
437  
438  	// `rp` (repoprompt-mcp) proxies RepoPrompt MCP tools. Treat these as write-capable.
439  	// Be tolerant of tool name prefixing, e.g. RepoPrompt_apply_edits
440  	const normalizedCall = call.trim();
441  	return /(^|_)(apply[-_]edits)$/.test(normalizedCall) || /(^|_)(file_actions)$/.test(normalizedCall);
442  }
443  
444  const AST_READ_ONLY_COMMANDS = new Set([
445  	"cat", "head", "tail", "less", "more", "grep", "find", "ls", "pwd", "echo", "printf", "wc", "sort", "uniq",
446  	"diff", "file", "stat", "du", "df", "tree", "which", "whereis", "type", "env", "printenv", "uname", "whoami",
447  	"id", "date", "cal", "uptime", "ps", "top", "htop", "free", "jq", "awk", "rg", "fd", "bat", "exa", "rp-cli",
448  	"rp_exec", "rp_bind", "curl",
449  ]);
450  
451  const AST_BLOCKED_COMMANDS = new Set([
452  	"rm", "rmdir", "mv", "cp", "mkdir", "touch", "chmod", "chown", "chgrp", "ln", "tee", "truncate", "dd", "shred",
453  	"sudo", "su", "kill", "pkill", "killall", "reboot", "shutdown", "systemctl", "service", "vim", "vi", "nano", "emacs",
454  	"code", "subl", "apt", "apt-get", "brew", "pip",
455  ]);
456  
457  const ALLOWED_GIT_SUBCOMMANDS = new Set(["status", "log", "diff", "show", "branch", "remote", "config", "ls-files", "ls-tree", "ls-remote"]);
458  const ALLOWED_NPM_SUBCOMMANDS = new Set(["list", "ls", "view", "info", "search", "outdated", "audit"]);
459  const ALLOWED_YARN_SUBCOMMANDS = new Set(["list", "info", "why", "audit"]);
460  const ALLOWED_PNPM_SUBCOMMANDS = new Set(["list", "ls", "view", "info", "search", "outdated", "audit"]);
461  
462  function isInvocationReadOnly(invocation: { effectiveCommandName: string; effectiveArgs: string[]; commandName: string; commandNameRaw: string }): boolean {
463  	const commandName = invocation.effectiveCommandName || invocation.commandName;
464  	const args = invocation.effectiveArgs;
465  
466  	if (!commandName) {
467  		return true;
468  	}
469  
470  	if (AST_BLOCKED_COMMANDS.has(commandName)) {
471  		return false;
472  	}
473  
474  	if (commandName === "git") {
475  		const sub = (args[0] ?? "").toLowerCase();
476  		if (!sub) return true;
477  		if (sub === "config") {
478  			return args[1] === "--get";
479  		}
480  		return ALLOWED_GIT_SUBCOMMANDS.has(sub) || sub.startsWith("ls-");
481  	}
482  
483  	if (commandName === "npm") {
484  		const sub = (args[0] ?? "").toLowerCase();
485  		return !sub || ALLOWED_NPM_SUBCOMMANDS.has(sub);
486  	}
487  
488  	if (commandName === "yarn") {
489  		const sub = (args[0] ?? "").toLowerCase();
490  		return !sub || ALLOWED_YARN_SUBCOMMANDS.has(sub);
491  	}
492  
493  	if (commandName === "pnpm") {
494  		const sub = (args[0] ?? "").toLowerCase();
495  		return !sub || ALLOWED_PNPM_SUBCOMMANDS.has(sub);
496  	}
497  
498  	if (commandName === "node" || commandName === "python" || commandName === "python3") {
499  		return args.length > 0 && args.every((arg) => arg === "--version");
500  	}
501  
502  	if (commandName === "wget") {
503  		for (let i = 0; i < args.length; i += 1) {
504  			if (args[i] === "-O") {
505  				return args[i + 1] === "-";
506  			}
507  		}
508  		return false;
509  	}
510  
511  	if (commandName === "sed") {
512  		return args.includes("-n");
513  	}
514  
515  	return AST_READ_ONLY_COMMANDS.has(commandName);
516  }
517  
518  function isSafeCommand(command: string): boolean {
519  	// Prevent using rp-cli via bash to enter interactive REPL while in plan mode
520  	if (RP_CLI_INTERACTIVE_PATTERN.test(command)) {
521  		return false;
522  	}
523  
524  	// Prevent using rp-cli via bash to perform edits/file actions while in plan mode
525  	if (RP_CLI_EXEC_WRITE_PATTERN.test(command)) {
526  		return false;
527  	}
528  
529  	const analysis = analyzeBashScript(command);
530  	if (!analysis.parseError) {
531  		if (analysis.invocations.some((invocation) => invocation.hasWriteRedirection)) {
532  			return false;
533  		}
534  
535  		return analysis.invocations.every((invocation) => isInvocationReadOnly(invocation));
536  	}
537  
538  	// Fallback: original regex policy if parsing fails
539  	if (DESTRUCTIVE_PATTERNS.some((pattern) => pattern.test(command))) {
540  		return false;
541  	}
542  
543  	return SAFE_COMMANDS.some((pattern) => pattern.test(command));
544  }
545  
546  type PlanModeState = {
547  	enabled: boolean;
548  	activeToolsBeforePlan?: string[];
549  };
550  
551  export default function planModeExtension(pi: ExtensionAPI) {
552  	let planModeEnabled = false;
553  	let activeToolsBeforePlan: string[] | null = null;
554  
555  	// Register --plan CLI flag
556  	pi.registerFlag("plan", {
557  		description: "Start in plan mode (read-only exploration)",
558  		type: "boolean",
559  		default: false,
560  	});
561  
562  	function applyToolMode(): void {
563  		const currentTools = pi.getActiveTools();
564  		const nextTools = planModeEnabled
565  			? removePlanModeWriteTools(currentTools)
566  			: activeToolsBeforePlan
567  				? restorePlanModeWriteTools(currentTools, activeToolsBeforePlan)
568  				: currentTools;
569  
570  		if (!toolListsMatch(currentTools, nextTools)) {
571  			pi.setActiveTools(nextTools);
572  		}
573  	}
574  
575  	function updateStatus(ctx: ExtensionContext): void {
576  		if (!ctx.hasUI) return;
577  
578  		if (planModeEnabled) {
579  			ctx.ui.setStatus("plan-mode", ctx.ui.theme.fg("warning", "⏸ plan"));
580  		} else {
581  			ctx.ui.setStatus("plan-mode", undefined);
582  		}
583  
584  		// Backward-compat cleanup: older versions used a widget for todo display
585  		ctx.ui.setWidget("plan-todos", undefined);
586  	}
587  
588  	function togglePlanMode(ctx: ExtensionContext): void {
589  		if (!planModeEnabled) {
590  			activeToolsBeforePlan = pi.getActiveTools();
591  			planModeEnabled = true;
592  			applyToolMode();
593  			persistPlanModeState();
594  
595  			if (ctx.hasUI) {
596  				ctx.ui.notify("Plan mode enabled. Pi write tools disabled; bash and RepoPrompt writes are blocked.");
597  			}
598  
599  			updateStatus(ctx);
600  			return;
601  		}
602  
603  		planModeEnabled = false;
604  		applyToolMode();
605  		activeToolsBeforePlan = null;
606  		persistPlanModeState();
607  
608  		if (ctx.hasUI) {
609  			ctx.ui.notify("Plan mode disabled. Full access restored.");
610  		}
611  
612  		updateStatus(ctx);
613  	}
614  
615  	// Register /plan command
616  	pi.registerCommand("plan", {
617  		description: "Toggle plan mode (read-only exploration)",
618  		handler: async (_args, ctx) => {
619  			togglePlanMode(ctx);
620  		},
621  	});
622  
623  	// Register Ctrl+Option+P shortcut
624  	pi.registerShortcut(Key.ctrlAlt("p"), {
625  		description: "Toggle plan mode",
626  		handler: async (ctx) => {
627  			togglePlanMode(ctx);
628  		},
629  	});
630  
631  	// Block write operations in plan mode (bash + RepoPrompt + native file tools as a backstop)
632  	pi.on("tool_call", async (event, ctx) => {
633  		if (!planModeEnabled) return;
634  
635  		// Backstop: even if another extension (e.g. /tools) re-enables these, plan mode must remain read-only
636  		if (event.toolName === "edit" || event.toolName === "write") {
637  			return {
638  				block: true,
639  				reason: `Plan mode: native tool "${event.toolName}" is blocked. Use /plan to disable plan mode first.`,
640  			};
641  		}
642  
643  		if (event.toolName === "bash") {
644  			await ensureJustBashLoaded();
645  			maybeWarnAstUnavailable(ctx);
646  			const command = event.input.command as string;
647  			if (!isSafeCommand(command)) {
648  				return {
649  					block: true,
650  					reason: `Plan mode: command blocked. Use /plan to disable plan mode first.\nCommand: ${command}`,
651  				};
652  			}
653  			return;
654  		}
655  
656  		if (event.toolName === "rp_exec" || event.toolName === "rp-cli") {
657  			const input = event.input as { cmd?: unknown; command?: unknown };
658  			const command = (input.cmd ?? input.command) as string | undefined;
659  			if (typeof command !== "string") return;
660  
661  			if (isRepoPromptWriteCommand(command)) {
662  				return {
663  					block: true,
664  					reason: `Plan mode: RepoPrompt write command blocked. Use /plan to disable plan mode first.\nCommand: ${command}`,
665  				};
666  			}
667  		}
668  
669  		if (event.toolName === "rp") {
670  			if (isRepoPromptMcpWriteRequest(event.input)) {
671  				const call = (event.input as { call?: unknown } | undefined)?.call;
672  				return {
673  					block: true,
674  					reason: `Plan mode: RepoPrompt write tool blocked. Use /plan to disable plan mode first.\nTool: rp(call=${String(call)})`,
675  				};
676  			}
677  		}
678  	});
679  
680  	// Re-apply tool restrictions right before the agent starts, in case other extensions mutate tool state
681  	pi.on("input", async (_event, ctx) => {
682  		if (!planModeEnabled) return;
683  		applyToolMode();
684  		updateStatus(ctx);
685  	});
686  
687  	// Filter out legacy plan-mode custom messages from older sessions so only current mode applies
688  	pi.on("context", async (event) => {
689  		const filtered = event.messages.filter((message) => {
690  			const customMessage = message as { role?: string; customType?: string };
691  			return !(
692  				customMessage.role === "custom"
693  				&& (customMessage.customType === "plan-mode-context" || customMessage.customType === "plan-mode-exit")
694  			);
695  		});
696  
697  		return { messages: filtered };
698  	});
699  
700  	// Add plan-mode instructions through the system prompt only while plan mode is active
701  	pi.on("before_agent_start", async (event) => {
702  		if (!planModeEnabled) {
703  			return;
704  		}
705  
706  		return {
707  			systemPrompt: `${event.systemPrompt}
708  
709  You are in plan mode (read-only). Describe what you would change rather than making changes directly.`,
710  		};
711  	});
712  
713  	function persistPlanModeState(): void {
714  		const data: PlanModeState = {
715  			enabled: planModeEnabled,
716  		};
717  
718  		if (planModeEnabled && activeToolsBeforePlan) {
719  			data.activeToolsBeforePlan = [...activeToolsBeforePlan];
720  		}
721  
722  		pi.appendEntry("plan-mode", data);
723  	}
724  
725  	function restorePlanModeFromBranch(
726  		ctx: ExtensionContext,
727  		options?: { preferStartFlag?: boolean },
728  	): void {
729  		const previousPlanModeEnabled = planModeEnabled;
730  		const previousActiveToolsBeforePlan = activeToolsBeforePlan ? [...activeToolsBeforePlan] : null;
731  		activeToolsBeforePlan = null;
732  
733  		// Optionally force plan mode on at startup
734  		if (options?.preferStartFlag && pi.getFlag("plan") === true) {
735  			activeToolsBeforePlan = pi.getActiveTools();
736  			planModeEnabled = true;
737  			// Persist once so /tree navigation remains branch-consistent even before the first turn starts
738  			persistPlanModeState();
739  			return;
740  		}
741  
742  		planModeEnabled = false;
743  
744  		for (const entry of ctx.sessionManager.getBranch()) {
745  			if (entry.type !== "custom" || entry.customType !== "plan-mode") {
746  				continue;
747  			}
748  
749  			const data = entry.data as PlanModeState | undefined;
750  			if (typeof data?.enabled === "boolean") {
751  				planModeEnabled = data.enabled;
752  				activeToolsBeforePlan = Array.isArray(data.activeToolsBeforePlan)
753  					? data.activeToolsBeforePlan.filter((toolName): toolName is string => typeof toolName === "string")
754  					: null;
755  			}
756  		}
757  
758  		if (!planModeEnabled && previousPlanModeEnabled && previousActiveToolsBeforePlan) {
759  			activeToolsBeforePlan = previousActiveToolsBeforePlan;
760  		}
761  	}
762  
763  	function applyRestoredState(ctx: ExtensionContext): void {
764  		applyToolMode();
765  		if (!planModeEnabled) {
766  			activeToolsBeforePlan = null;
767  		}
768  		updateStatus(ctx);
769  	}
770  
771  	// Initialize state on session start and post-transition runtime recreation
772  	pi.on("session_start", async (event, ctx) => {
773  		restorePlanModeFromBranch(ctx, { preferStartFlag: event.reason === "startup" });
774  		applyRestoredState(ctx);
775  	});
776  
777  	pi.on("session_tree", async (_event, ctx) => {
778  		restorePlanModeFromBranch(ctx);
779  		applyRestoredState(ctx);
780  	});
781  
782  }