1  /**
  2   * RP Native Tools Lock
  3   *
  4   * Disables Pi's native repo-file tools (read/write/edit/ls/find/grep) when RepoPrompt tools are available.
  5   *
  6   * Why:
  7   * - Some models will "reach" for the native tools because they appear first / are more familiar
  8   * - If RepoPrompt is available, we want to force repo-scoped work through rp (MCP) or rp_exec (CLI)
  9   *
 10   * Modes (user-facing):
 11   * - off     : no enforcement
 12   * - auto    : enforce via rp if available; else rp_exec if available; else off
 13   *            (auto mode only kicks in if the user has enabled `rp`/`rp_exec` in their active tools)
 14   *            Tip: use `/tools` to enable `rp` (and then toggle this lock with Alt+L or `/rp-tools-lock`)
 15   *
 16   * Advanced modes (set via config file):
 17   * - rp-mcp  : enforce when the `rp` tool exists
 18   * - rp-cli  : enforce when the `rp_exec` tool exists
 19   *
 20   * Hotkeys:
 21   * - Alt+L: toggle lock mode (off ↔ auto)
 22   *
 23   * Configuration precedence:
 24   * 1) Session branch override (via /rp-tools-lock)
 25   * 2) Global config file: ~/.pi/agent/extensions/rp-native-tools-lock/rp-native-tools-lock.json
 26   * 3) Default: auto
 27   */
 28  
 29  import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
 30  import { homedir } from "node:os";
 31  import { dirname, join } from "node:path";
 32  
 33  import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
 34  import { Key, type KeyId } from "@mariozechner/pi-tui";
 35  
 36  type Mode = "off" | "auto" | "rp-mcp" | "rp-cli";
 37  
 38  interface LockState {
 39  	mode: Mode;
 40  }
 41  
 42  const CUSTOM_TYPE = "rp-native-tools-lock";
 43  const CONFIG_PATH = join(homedir(), ".pi", "agent", "extensions", "rp-native-tools-lock", "rp-native-tools-lock.json");
 44  
 45  const REQUIRED_TOOL_BY_MODE: Record<Exclude<Mode, "off" | "auto">, string> = {
 46  	"rp-mcp": "rp",
 47  	"rp-cli": "rp_exec",
 48  };
 49  
 50  const NATIVE_FILE_TOOLS = ["read", "write", "edit", "ls", "find", "grep"];
 51  
 52  const TOGGLE_MODE_HOTKEY: KeyId = Key.alt("l");
 53  
 54  // Keep the interactive UX simple: users only toggle off/auto.
 55  // Advanced modes remain supported via the config file.
 56  const MODE_CYCLE_ORDER: Mode[] = ["off", "auto"];
 57  
 58  function normalizeMode(raw: string | undefined): Mode | undefined {
 59  	const value = (raw ?? "").trim().toLowerCase();
 60  	if (!value) return undefined;
 61  
 62  	if (value === "off" || value === "disabled" || value === "none") return "off";
 63  	if (value === "auto" || value === "aut" || value === "automatic") return "auto";
 64  	if (value === "rp-mcp" || value === "mcp" || value === "rp") return "rp-mcp";
 65  	if (value === "rp-cli" || value === "cli" || value === "rp_exec" || value === "rp-exec") return "rp-cli";
 66  
 67  	return undefined;
 68  }
 69  
 70  function loadGlobalConfig(): LockState | undefined {
 71  	if (!existsSync(CONFIG_PATH)) return undefined;
 72  	try {
 73  		const content = readFileSync(CONFIG_PATH, "utf-8");
 74  		const parsed = JSON.parse(content) as Partial<LockState> | undefined;
 75  		const mode = normalizeMode(parsed?.mode);
 76  		return mode ? { mode } : undefined;
 77  	} catch {
 78  		return undefined;
 79  	}
 80  }
 81  
 82  function saveGlobalConfig(state: LockState): void {
 83  	try {
 84  		const configDir = dirname(CONFIG_PATH);
 85  		if (!existsSync(configDir)) {
 86  			mkdirSync(configDir, { recursive: true });
 87  		}
 88  		writeFileSync(CONFIG_PATH, JSON.stringify(state, null, 2));
 89  	} catch (err) {
 90  		console.error(`Failed to save ${CONFIG_PATH}: ${err}`);
 91  	}
 92  }
 93  
 94  function restoreFromBranch(ctx: ExtensionContext, fallback: LockState): LockState {
 95  	const branchEntries = ctx.sessionManager.getBranch();
 96  	let restored: LockState | undefined;
 97  
 98  	for (const entry of branchEntries) {
 99  		if (entry.type === "custom" && entry.customType === CUSTOM_TYPE) {
100  			const data = entry.data as Partial<LockState> | undefined;
101  			const mode = normalizeMode(data?.mode);
102  			if (mode) restored = { mode };
103  		}
104  	}
105  
106  	return restored ?? fallback;
107  }
108  
109  function setStatus(ctx: ExtensionContext, text: string | undefined): void {
110  	if (!ctx.hasUI) return;
111  	ctx.ui.setStatus("rp-tools-lock", text);
112  }
113  
114  type EffectiveMode = Exclude<Mode, "auto">;
115  
116  function computeEffectiveMode(
117  	allToolNames: Set<string>,
118  	activeToolNames: Set<string>,
119  	requestedMode: Mode,
120  ): { effectiveMode: EffectiveMode; requiredTool: string | undefined } {
121  	if (requestedMode === "off") return { effectiveMode: "off", requiredTool: undefined };
122  
123  	if (requestedMode === "auto") {
124  		// In auto mode, respect the user's tool configuration.
125  		// We only prefer a RepoPrompt entrypoint if the user has enabled it.
126  		if (activeToolNames.has("rp")) return { effectiveMode: "rp-mcp", requiredTool: "rp" };
127  		if (activeToolNames.has("rp_exec")) return { effectiveMode: "rp-cli", requiredTool: "rp_exec" };
128  		return { effectiveMode: "off", requiredTool: undefined };
129  	}
130  
131  	// Advanced/maintainer modes: enforce based on tool availability (even if currently disabled)
132  	return {
133  		effectiveMode: requestedMode,
134  		requiredTool: REQUIRED_TOOL_BY_MODE[requestedMode],
135  	};
136  }
137  
138  function buildStatusText(effectiveMode: EffectiveMode): string | undefined {
139  	if (effectiveMode === "rp-mcp" || effectiveMode === "rp-cli") {
140  		return "RP 🔒";
141  	}
142  	return undefined;
143  }
144  
145  function enforceMode(
146  	pi: ExtensionAPI,
147  	ctx: ExtensionContext,
148  	requestedMode: Mode,
149  ): { enforced: boolean; reason?: string; effectiveMode: EffectiveMode; requiredTool?: string } {
150  	const allToolNames = new Set(pi.getAllTools().map((t) => t.name));
151  	const activeToolNames = new Set(pi.getActiveTools());
152  	const { effectiveMode, requiredTool } = computeEffectiveMode(allToolNames, activeToolNames, requestedMode);
153  
154  	if (effectiveMode === "off") {
155  		setStatus(ctx, undefined);
156  		return {
157  			enforced: false,
158  			reason: requestedMode === "auto" ? "auto:no-rp-tools" : "mode=off",
159  			effectiveMode,
160  		};
161  	}
162  
163  	if (!requiredTool || !allToolNames.has(requiredTool)) {
164  		setStatus(ctx, undefined);
165  		return {
166  			enforced: false,
167  			reason: `missing:${requiredTool ?? "unknown"}`,
168  			effectiveMode,
169  		};
170  	}
171  
172  	const active = pi.getActiveTools();
173  	const blocked = new Set(NATIVE_FILE_TOOLS.filter((t) => allToolNames.has(t)));
174  	const next = active.filter((t) => !blocked.has(t));
175  
176  	// Ensure the required RepoPrompt tool stays available
177  	if (!next.includes(requiredTool)) next.push(requiredTool);
178  
179  
180  	// Only apply when changed
181  	const activeSet = new Set(active);
182  	const nextSet = new Set(next);
183  	const changed =
184  		active.length !== next.length ||
185  		active.some((t) => !nextSet.has(t)) ||
186  		next.some((t) => !activeSet.has(t));
187  
188  	if (changed) {
189  		pi.setActiveTools(next);
190  	}
191  
192  	setStatus(ctx, buildStatusText(effectiveMode));
193  	return { enforced: true, effectiveMode, requiredTool };
194  }
195  
196  export default function rpNativeToolsLock(pi: ExtensionAPI): void {
197  	let state: LockState = { mode: "auto" };
198  
199  	function resolveState(ctx: ExtensionContext): LockState {
200  		const globalConfig = loadGlobalConfig();
201  		const fallback = globalConfig ?? { mode: "auto" };
202  		return restoreFromBranch(ctx, fallback);
203  	}
204  
205  	function apply(ctx: ExtensionContext): void {
206  		state = resolveState(ctx);
207  		enforceMode(pi, ctx, state.mode);
208  	}
209  
210  	function persistState(nextState: LockState): void {
211  		// Persist globally + in-session branch
212  		saveGlobalConfig(nextState);
213  		pi.appendEntry<LockState>(CUSTOM_TYPE, nextState);
214  	}
215  
216  	function setMode(
217  		ctx: ExtensionContext,
218  		mode: Mode,
219  	): { enforced: boolean; reason?: string; effectiveMode: EffectiveMode; requiredTool?: string } {
220  		state = { mode };
221  		persistState(state);
222  		return enforceMode(pi, ctx, state.mode);
223  	}
224  
225  	function notifyEnforcement(
226  		ctx: ExtensionContext,
227  		requestedMode: Mode,
228  		enforced: { enforced: boolean; reason?: string; effectiveMode: EffectiveMode },
229  	): void {
230  		if (!ctx.hasUI) return;
231  
232  		if (requestedMode === "off") {
233  			ctx.ui.notify("rp-tools-lock: off", "info");
234  			return;
235  		}
236  
237  		if (enforced.enforced) {
238  			// Keep user-facing messaging simple. Advanced detail in config.
239  			if (requestedMode === "auto") {
240  				ctx.ui.notify("rp-tools-lock: auto (native file tools disabled)", "info");
241  				return;
242  			}
243  
244  			ctx.ui.notify(`rp-tools-lock: ${requestedMode} (native file tools disabled)`, "info");
245  			return;
246  		}
247  
248  		if (requestedMode === "auto" && enforced.effectiveMode === "off") {
249  			ctx.ui.notify("rp-tools-lock: auto (no rp/rp_exec tools available)", "info");
250  			return;
251  		}
252  
253  		ctx.ui.notify(`rp-tools-lock: ${requestedMode} (not enforced: ${enforced.reason ?? "unknown"})`, "warning");
254  	}
255  
256  	function getNextMode(currentMode: Mode): Mode {
257  		const index = MODE_CYCLE_ORDER.indexOf(currentMode);
258  		const safeIndex = index >= 0 ? index : 0;
259  		return MODE_CYCLE_ORDER[(safeIndex + 1) % MODE_CYCLE_ORDER.length];
260  	}
261  
262  	pi.registerCommand("rp-tools-lock", {
263  		description:
264  			"RepoPrompt-first tooling: off | auto (disables read/write/edit/ls/find/grep). Advanced modes are available via config file.",
265  		handler: async (args, ctx) => {
266  			const raw = args?.trim();
267  
268  			const ALLOWED_MODES: Mode[] = ["off", "auto"];
269  
270  			// No args → interactive selector (if UI available)
271  			if (!raw) {
272  				if (!ctx.hasUI) {
273  					console.error("Usage: /rp-tools-lock <off|auto>");
274  					return;
275  				}
276  
277  				const choice = await ctx.ui.select("RepoPrompt tool policy", ALLOWED_MODES);
278  				if (!choice) return;
279  				state = { mode: choice as Mode };
280  			} else {
281  				const mode = normalizeMode(raw);
282  				if (!mode || !ALLOWED_MODES.includes(mode)) {
283  					const message =
284  						`Usage: /rp-tools-lock <off|auto> (got: ${raw})\n` +
285  						"Advanced modes (rp-mcp/rp-cli) can be set via: " +
286  						"~/.pi/agent/extensions/rp-native-tools-lock/rp-native-tools-lock.json";
287  
288  					if (ctx.hasUI) {
289  						ctx.ui.notify(message, "error");
290  					} else {
291  						console.error(message);
292  					}
293  					return;
294  				}
295  
296  				state = { mode };
297  			}
298  
299  			persistState(state);
300  
301  			const enforced = enforceMode(pi, ctx, state.mode);
302  			notifyEnforcement(ctx, state.mode, enforced);
303  		},
304  	});
305  
306  	pi.registerShortcut(TOGGLE_MODE_HOTKEY, {
307  		description: "Toggle rp-tools-lock mode (off ↔ auto)",
308  		handler: async (ctx) => {
309  			const current = resolveState(ctx).mode;
310  			const next = getNextMode(current);
311  			const enforced = setMode(ctx, next);
312  			notifyEnforcement(ctx, next, enforced);
313  		},
314  	});
315  
316  	// Apply early and often (covers /tools toggles, session navigation, etc.)
317  	pi.on("session_start", async (_event, ctx) => apply(ctx));
318  	pi.on("session_tree", async (_event, ctx) => apply(ctx));
319  
320  	// Enforce right when the user submits a prompt (before the agent starts)
321  	pi.on("input", async (_event, ctx) => apply(ctx));
322  
323  	// Safety backstop: even if a tool somehow remains active, block the call with a clear reason
324  	pi.on("tool_call", async (event, ctx) => {
325  		state = resolveState(ctx);
326  
327  		const allToolNames = new Set(pi.getAllTools().map((t) => t.name));
328  		const activeToolNames = new Set(pi.getActiveTools());
329  		const { effectiveMode, requiredTool } = computeEffectiveMode(allToolNames, activeToolNames, state.mode);
330  		if (effectiveMode === "off" || !requiredTool) return;
331  		if (!allToolNames.has(requiredTool)) return;
332  
333  		if (!NATIVE_FILE_TOOLS.includes(event.toolName)) return;
334  
335  		const suffix = state.mode === "auto" ? ` → ${effectiveMode}` : "";
336  		return {
337  			block: true,
338  			reason:
339  				`rp-tools-lock (${state.mode}${suffix}): native tool "${event.toolName}" is disabled. ` +
340  				`Use RepoPrompt instead (tool: ${requiredTool}). ` +
341  				`You can disable this lock with /rp-tools-lock off.`,
342  		};
343  	});
344  }