1  /**
  2   * Inline Shell Extension - expands inline shell commands in user prompts.
  3   *
  4   * Start pi with this extension:
  5   *   pi -e ./extensions/inline-shell.ts
  6   *
  7   * Then type prompts with inline shell:
  8   *   What's in !{pwd}?
  9   *   The current branch is !{git branch --show-current} and status: !{git status --short}
 10   *   My node version is !{node --version}
 11   *
 12   * The !{command} patterns are executed and replaced with their output before
 13   * the prompt is sent to the agent.
 14   *
 15   * Shell selection:
 16   * - If the current shell is zsh and $PI_CODING_AGENT_DIR/shell/pi-inline.zsh exists,
 17   *   source that file in a fresh zsh before running the command
 18   * - Otherwise, if the current shell is zsh, run a fresh interactive zsh so aliases
 19   *   and functions from .zshrc are available
 20   * - If the current shell is not zsh, fall back to bash
 21   *
 22   * The spawned shell always gets PI_INLINE_SHELL=1 so your shell config can skip
 23   * noisy prompt/plugin setup while still loading aliases/functions.
 24   *
 25   * Note: Regular !command syntax (whole-line bash) is preserved and works as before.
 26   */
 27  import { spawn } from "node:child_process";
 28  import { randomUUID } from "node:crypto";
 29  import * as fs from "node:fs";
 30  import * as os from "node:os";
 31  import * as path from "node:path";
 32  
 33  import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
 34  
 35  type ShellType = "zsh" | "bash";
 36  type ShellMode = "shared-zsh-config" | "interactive-zsh" | "interactive-bash" | "bash-fallback";
 37  
 38  type ResolvedShell = {
 39  	path: string;
 40  	type: ShellType;
 41  	mode: ShellMode;
 42  	sharedConfigPath?: string;
 43  };
 44  
 45  const PATTERN = /!\{([^}]+)\}/g;
 46  const TIMEOUT_MS = 30_000;
 47  const COMMON_BASH_PATHS = [
 48  	"/bin/bash",
 49  	"/usr/bin/bash",
 50  	"/usr/local/bin/bash",
 51  	"/opt/homebrew/bin/bash",
 52  ];
 53  
 54  function getAgentDir(): string {
 55  	const envDir = process.env.PI_CODING_AGENT_DIR;
 56  	if (envDir) {
 57  		if (envDir === "~") {
 58  			return os.homedir();
 59  		}
 60  		if (envDir.startsWith("~/")) {
 61  			return path.join(os.homedir(), envDir.slice(2));
 62  		}
 63  		return envDir;
 64  	}
 65  	return path.join(os.homedir(), ".pi", "agent");
 66  }
 67  
 68  function detectShellType(shellPath: string | undefined): ShellType | null {
 69  	if (!shellPath) {
 70  		return null;
 71  	}
 72  
 73  	const baseName = path.basename(shellPath).toLowerCase();
 74  	if (baseName === "zsh" || baseName.startsWith("zsh")) {
 75  		return "zsh";
 76  	}
 77  	if (baseName === "bash" || baseName.startsWith("bash")) {
 78  		return "bash";
 79  	}
 80  	return null;
 81  }
 82  
 83  function findFirstExistingPath(paths: string[]): string | null {
 84  	for (const candidate of paths) {
 85  		if (fs.existsSync(candidate)) {
 86  			return candidate;
 87  		}
 88  	}
 89  	return null;
 90  }
 91  
 92  function getSharedInlineZshConfigPath(): string | null {
 93  	const candidate = path.join(getAgentDir(), "shell", "pi-inline.zsh");
 94  	return fs.existsSync(candidate) ? candidate : null;
 95  }
 96  
 97  function resolveExecutionShell(): ResolvedShell {
 98  	const userShellPath = process.env.SHELL;
 99  	const userShellType = detectShellType(userShellPath);
100  	const hasUserShell = Boolean(userShellPath && fs.existsSync(userShellPath));
101  
102  	if (hasUserShell && userShellType === "zsh") {
103  		const sharedConfigPath = getSharedInlineZshConfigPath();
104  		if (sharedConfigPath) {
105  			return {
106  				path: userShellPath as string,
107  				type: "zsh",
108  				mode: "shared-zsh-config",
109  				sharedConfigPath,
110  			};
111  		}
112  
113  		return {
114  			path: userShellPath as string,
115  			type: "zsh",
116  			mode: "interactive-zsh",
117  		};
118  	}
119  
120  	if (hasUserShell && userShellType === "bash") {
121  		return {
122  			path: userShellPath as string,
123  			type: "bash",
124  			mode: "interactive-bash",
125  		};
126  	}
127  
128  	return {
129  		path: findFirstExistingPath(COMMON_BASH_PATHS) ?? "/bin/bash",
130  		type: "bash",
131  		mode: "bash-fallback",
132  	};
133  }
134  
135  function shellQuote(value: string): string {
136  	return `'${value.replace(/'/g, `'"'"'`)}'`;
137  }
138  
139  function unwrapBraceWrappedCommand(command: string): string | null {
140  	const trimmed = command.trim();
141  	if (!trimmed.startsWith("{") || !trimmed.endsWith("}")) {
142  		return null;
143  	}
144  
145  	const inner = trimmed.slice(1, -1).trim();
146  	return inner.length > 0 ? inner : null;
147  }
148  
149  function buildShellScript(
150  	command: string,
151  	shell: ResolvedShell,
152  	options?: { startMarker?: string; endMarker?: string },
153  ): string {
154  	const scriptLines: string[] = ["export PI_INLINE_SHELL=1"];
155  
156  	if (options?.startMarker) {
157  		scriptLines.push(`printf '%s\\n' ${shellQuote(options.startMarker)}`);
158  		scriptLines.push(`printf '%s\\n' ${shellQuote(options.startMarker)} >&2`);
159  	}
160  
161  	if (shell.mode === "shared-zsh-config" && shell.sharedConfigPath) {
162  		scriptLines.push(`source ${shellQuote(shell.sharedConfigPath)}`);
163  	}
164  
165  	scriptLines.push("set +e");
166  	scriptLines.push(`eval -- ${shellQuote(command)}`);
167  	scriptLines.push("__pi_inline_status=$?");
168  
169  	if (options?.endMarker) {
170  		scriptLines.push(`printf '\\n%s\\n' ${shellQuote(options.endMarker)}`);
171  		scriptLines.push(`printf '\\n%s\\n' ${shellQuote(options.endMarker)} >&2`);
172  	}
173  
174  	scriptLines.push("exit $__pi_inline_status");
175  	return scriptLines.join("\n");
176  }
177  
178  function buildShellArgs(script: string, shell: ResolvedShell): string[] {
179  	return shell.mode === "shared-zsh-config"
180  		? ["-c", script]
181  		: ["-i", "-c", script];
182  }
183  
184  function createShellBackedBashOperations(shell: ResolvedShell, originalCommand?: string) {
185  	return {
186  		exec: async (
187  			command: string,
188  			cwd: string,
189  			options: {
190  				onData: (data: Buffer) => void;
191  				signal?: AbortSignal;
192  				timeout?: number;
193  				env?: NodeJS.ProcessEnv;
194  			},
195  		): Promise<{ exitCode: number | null }> => {
196  			const commandToRun = originalCommand ?? command;
197  			const effectiveCommand = unwrapBraceWrappedCommand(commandToRun) ?? commandToRun;
198  			const script = buildShellScript(effectiveCommand, shell);
199  			const child = spawn(shell.path, buildShellArgs(script, shell), {
200  				cwd,
201  				env: {
202  					...process.env,
203  					...options.env,
204  					PI_INLINE_SHELL: "1",
205  				},
206  				stdio: ["ignore", "pipe", "pipe"],
207  			});
208  
209  			child.stdout?.on("data", options.onData);
210  			child.stderr?.on("data", options.onData);
211  
212  			let timeoutHandle: NodeJS.Timeout | undefined;
213  			let settled = false;
214  
215  			const terminate = () => {
216  				if (child.killed) {
217  					return;
218  				}
219  				child.kill("SIGTERM");
220  				setTimeout(() => {
221  					if (!child.killed) {
222  						child.kill("SIGKILL");
223  					}
224  				}, 5000);
225  			};
226  
227  			const abortHandler = () => terminate();
228  			if (options.signal) {
229  				if (options.signal.aborted) {
230  					terminate();
231  				} else {
232  					options.signal.addEventListener("abort", abortHandler, { once: true });
233  				}
234  			}
235  
236  			if (options.timeout && options.timeout > 0) {
237  				timeoutHandle = setTimeout(() => terminate(), options.timeout * 1000);
238  			}
239  
240  			return await new Promise<{ exitCode: number | null }>((resolve, reject) => {
241  				const cleanup = () => {
242  					if (settled) {
243  						return;
244  					}
245  					settled = true;
246  					if (timeoutHandle) {
247  						clearTimeout(timeoutHandle);
248  					}
249  					options.signal?.removeEventListener("abort", abortHandler);
250  				};
251  
252  				child.on("error", (error) => {
253  					cleanup();
254  					reject(error);
255  				});
256  				child.on("close", (code) => {
257  					cleanup();
258  					resolve({ exitCode: code });
259  				});
260  			});
261  		},
262  	};
263  }
264  
265  function buildShellInvocation(command: string, shell: ResolvedShell): {
266  	args: string[];
267  	startMarker: string;
268  	endMarker: string;
269  } {
270  	const startMarker = `__PI_INLINE_START_${randomUUID()}__`;
271  	const endMarker = `__PI_INLINE_END_${randomUUID()}__`;
272  	const scriptLines = [
273  		"export PI_INLINE_SHELL=1",
274  		`printf '%s\\n' ${shellQuote(startMarker)}`,
275  		`printf '%s\\n' ${shellQuote(startMarker)} >&2`,
276  	];
277  
278  	if (shell.mode === "shared-zsh-config" && shell.sharedConfigPath) {
279  		scriptLines.push(`source ${shellQuote(shell.sharedConfigPath)}`);
280  	}
281  
282  	scriptLines.push("set +e");
283  	scriptLines.push(`eval -- ${shellQuote(command)}`);
284  	scriptLines.push("__pi_inline_status=$?");
285  	scriptLines.push(`printf '\\n%s\\n' ${shellQuote(endMarker)}`);
286  	scriptLines.push(`printf '\\n%s\\n' ${shellQuote(endMarker)} >&2`);
287  	scriptLines.push("exit $__pi_inline_status");
288  
289  	const args = shell.mode === "shared-zsh-config"
290  		? ["-c", scriptLines.join("\n")]
291  		: ["-i", "-c", scriptLines.join("\n")];
292  
293  	return { args, startMarker, endMarker };
294  }
295  
296  function extractMarkedOutput(text: string, startMarker: string, endMarker: string): string {
297  	const normalized = text.replace(/\r\n/g, "\n");
298  	const startIndex = normalized.indexOf(startMarker);
299  	if (startIndex === -1) {
300  		return normalized.trim();
301  	}
302  
303  	let contentStart = startIndex + startMarker.length;
304  	if (normalized[contentStart] === "\n") {
305  		contentStart += 1;
306  	}
307  
308  	const endIndex = normalized.lastIndexOf(`\n${endMarker}`);
309  	if (endIndex === -1 || endIndex < contentStart) {
310  		return normalized.slice(contentStart).trim();
311  	}
312  
313  	return normalized.slice(contentStart, endIndex).trim();
314  }
315  
316  function getReplacementText(stdout: string, stderr: string, exitCode: number): string {
317  	const trimmedStdout = stdout.trim();
318  	const trimmedStderr = stderr.trim();
319  	if (trimmedStdout.length > 0) {
320  		return trimmedStdout;
321  	}
322  	if (trimmedStderr.length > 0) {
323  		return trimmedStderr;
324  	}
325  	if (exitCode !== 0) {
326  		return `[error: exit code ${exitCode}]`;
327  	}
328  	return "";
329  }
330  
331  function describeShell(shell: ResolvedShell): string {
332  	switch (shell.mode) {
333  		case "shared-zsh-config":
334  			return "zsh (shell/pi-inline.zsh)";
335  		case "interactive-zsh":
336  			return "zsh (.zshrc)";
337  		case "interactive-bash":
338  			return "bash (.bashrc)";
339  		case "bash-fallback":
340  			return "bash fallback";
341  		default:
342  			return shell.type;
343  	}
344  }
345  
346  export default function (pi: ExtensionAPI) {
347  	pi.on("user_bash", async (event) => {
348  		const shell = resolveExecutionShell();
349  		return {
350  			operations: createShellBackedBashOperations(shell, event.command),
351  		};
352  	});
353  
354  	pi.on("input", async (event, ctx) => {
355  		const text = event.text;
356  
357  		if (text.trimStart().startsWith("!") && !text.trimStart().startsWith("!{")) {
358  			return { action: "continue" };
359  		}
360  
361  		if (!PATTERN.test(text)) {
362  			return { action: "continue" };
363  		}
364  
365  		PATTERN.lastIndex = 0;
366  
367  		let result = text;
368  		const shell = resolveExecutionShell();
369  		const expansions: Array<{ command: string; output: string; error?: string }> = [];
370  		const matches: Array<{ full: string; command: string }> = [];
371  
372  		let match = PATTERN.exec(text);
373  		while (match) {
374  			matches.push({ full: match[0], command: match[1] });
375  			match = PATTERN.exec(text);
376  		}
377  		PATTERN.lastIndex = 0;
378  
379  		for (const { full, command } of matches) {
380  			try {
381  				const invocation = buildShellInvocation(command, shell);
382  				const shellResult = await pi.exec(shell.path, invocation.args, {
383  					cwd: ctx.cwd,
384  					timeout: TIMEOUT_MS,
385  				});
386  
387  				const filteredStdout = extractMarkedOutput(shellResult.stdout, invocation.startMarker, invocation.endMarker);
388  				const filteredStderr = extractMarkedOutput(shellResult.stderr, invocation.startMarker, invocation.endMarker);
389  				const replacementText = getReplacementText(filteredStdout, filteredStderr, shellResult.code);
390  				const error = shellResult.code === 0 ? undefined : `exit code ${shellResult.code}`;
391  
392  				expansions.push({ command, output: replacementText, error });
393  				result = result.replace(full, () => replacementText);
394  			} catch (err) {
395  				const errorMsg = err instanceof Error ? err.message : String(err);
396  				expansions.push({ command, output: "", error: errorMsg });
397  				result = result.replace(full, () => `[error: ${errorMsg}]`);
398  			}
399  		}
400  
401  		if (ctx.hasUI && expansions.length > 0) {
402  			const summary = expansions
403  				.map((expansion) => {
404  					const status = expansion.error ? ` (${expansion.error})` : "";
405  					const preview = expansion.output.length > 50
406  						? `${expansion.output.slice(0, 50)}...`
407  						: expansion.output;
408  					return `!{${expansion.command}}${status} -> "${preview}"`;
409  				})
410  				.join("\n");
411  
412  			ctx.ui.notify(`Expanded ${expansions.length} inline command(s) via ${describeShell(shell)}:\n${summary}`, "info");
413  		}
414  
415  		return { action: "transform", text: result, images: event.images };
416  	});
417  }