SSH Remote Extension for pi - routes file operations and bash to a remote server over SSH using SSHFS

ssh-remote.ts

/**
 * SSH Remote Extension for pi
 *
 * Routes file operations (read, write, edit) and bash to a remote server over SSH.
 * Uses SSHFS for file operations and direct SSH for bash execution.
 *
 * Usage:
 *   pi -e ./ssh-remote.ts --ssh user@host
 *   pi -e ./ssh-remote.ts --ssh user@host:/remote/path --ssh-key ~/.ssh/key
 *   pi -e ./ssh-remote.ts --ssh user@host --ssh-port 2222
 */

import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import {
	createReadTool,
	createWriteTool,
	createEditTool,
	createBashTool,
	type ReadOperations,
	type WriteOperations,
	type EditOperations,
	type BashOperations,
} from "@mariozechner/pi-coding-agent";
import { spawn } from "node:child_process";
import { mkdir, rm, access, writeFile as fsWriteFile, constants as fsConstants } from "node:fs/promises";
import { homedir } from "node:os";
import { join, isAbsolute } from "node:path";

// Configuration from CLI flags
interface SshConfig {
	user: string;
	host: string;
	key?: string;
	port: number;
	remoteCwd: string;
	mountPoint: string;
}

// Track active SSH config for system prompt and cleanup
let activeConfig: SshConfig | null = null;

/**
 * Parse SSH connection string: user@host or user@host:/path
 */
function parseSshArg(arg: string): { user: string; host: string; remoteCwd?: string } {
	const match = arg.match(/^([^@]+)@([^:]+)(?::(.+))?$/);
	if (!match) {
		throw new Error(
			`Invalid SSH format. Expected: user@host or user@host:/path, got: ${arg}`,
		);
	}
	return {
		user: match[1],
		host: match[2],
		remoteCwd: match[3],
	};
}

/**
 * Build SSH argument array with options
 */
function buildSshArgs(config: Pick<SshConfig, "port" | "key">): string[] {
	const args: string[] = [];
	if (config.port !== 22) {
		args.push("-p", config.port.toString());
	}
	if (config.key) {
		args.push("-i", config.key);
	}
	args.push("-o", "StrictHostKeyChecking=accept-new");
	return args;
}

/**
 * Execute a command via SSH
 */
async function sshExec(
	command: string,
	config: Pick<SshConfig, "user" | "host" | "port" | "key">,
): Promise<Buffer> {
	const args = buildSshArgs(config);
	args.push(`${config.user}@${config.host}`, command);

	return new Promise((resolve, reject) => {
		const child = spawn("ssh", args, { stdio: ["ignore", "pipe", "pipe"] });
		const chunks: Buffer[] = [];
		const errChunks: Buffer[] = [];

		child.stdout.on("data", (data) => chunks.push(data));
		child.stderr.on("data", (data) => errChunks.push(data));
		child.on("error", reject);
		child.on("close", (code) => {
			if (code !== 0) {
				reject(
					new Error(
						`SSH failed (${code}): ${Buffer.concat(errChunks).toString()}`,
					),
				);
			} else {
				resolve(Buffer.concat(chunks));
			}
		});
	});
}

/**
 * Detect remote home directory via SSH
 */
async function detectRemoteCwd(
	user: string,
	host: string,
	port: number,
	key?: string,
): Promise<string> {
	const result = await sshExec("pwd", { user, host, port, key });
	return result.toString().trim();
}

/**
 * Mount remote directory via SSHFS
 */
async function mountRemote(config: SshConfig): Promise<void> {
	const args: string[] = [];

	if (config.port !== 22) {
		args.push("-o", `port=${config.port}`);
	}
	if (config.key) {
		args.push("-o", `IdentityFile=${config.key}`);
	}

	args.push(
		"-o",
		"StrictHostKeyChecking=accept-new",
		"-o",
		"cache=yes",
		"-o",
		"kernel_cache",
		"-o",
		"compression=yes",
	);

	args.push(
		`${config.user}@${config.host}:${config.remoteCwd}`,
		config.mountPoint,
	);

	return new Promise((resolve, reject) => {
		const child = spawn("sshfs", args, { stdio: ["ignore", "pipe", "pipe"] });
		const errChunks: Buffer[] = [];

		child.stderr.on("data", (data) => errChunks.push(data));
		child.on("error", (err) => {
			if ((err as any).code === "ENOENT") {
				reject(new Error("sshfs not found. Please install SSHFS."));
			} else {
				reject(err);
			}
		});
		child.on("close", (code) => {
			if (code !== 0) {
				reject(
					new Error(
						`SSHFS mount failed (${code}): ${Buffer.concat(errChunks).toString()}`,
					),
				);
			} else {
				resolve();
			}
		});
	});
}

/**
 * Unmount SSHFS directory
 */
async function unmountRemote(mountPoint: string): Promise<void> {
	const platform = process.platform;
	const command = platform === "linux" ? "fusermount" : "umount";
	const args = platform === "linux" ? ["-u", mountPoint] : [mountPoint];

	return new Promise((resolve) => {
		const child = spawn(command, args, { stdio: "ignore" });
		child.on("error", () => resolve());
		child.on("close", () => resolve());
	});
}

/**
 * Convert a logical path to the mounted path
 */
function toMountPath(logicalPath: string, mountPoint: string, remoteCwd: string): string {
	// If absolute path, make it relative to remoteCwd first
	if (isAbsolute(logicalPath)) {
		// Check if the path starts with remoteCwd
		if (logicalPath.startsWith(remoteCwd)) {
			const relativePath = logicalPath.slice(remoteCwd.length);
			return join(mountPoint, relativePath.startsWith("/") ? relativePath.slice(1) : relativePath);
		}
		// Path is absolute but not under remoteCwd, just join it
		return join(mountPoint, logicalPath);
	}
	// Relative path - join directly with mountPoint
	return join(mountPoint, logicalPath);
}

/**
 * Create read operations for SSHFS-mounted remote
 */
function createSshfsReadOps(mountPoint: string, remoteCwd: string): ReadOperations {
	return {
		readFile: async (p: string) => {
			const { readFile } = await import("node:fs/promises");
			return readFile(toMountPath(p, mountPoint, remoteCwd));
		},
		access: async (p: string) => {
			const { access } = await import("node:fs/promises");
			await access(toMountPath(p, mountPoint, remoteCwd), fsConstants.R_OK);
		},
		detectImageMimeType: async (p: string) => {
			const { exec } = await import("node:child_process");
			try {
				const result = await new Promise<string>((resolve, reject) => {
					exec(
						`file --mime-type -b ${JSON.stringify(toMountPath(p, mountPoint, remoteCwd))}`,
						(err, stdout) => {
							if (err) reject(err);
							else resolve(stdout.trim());
						},
					);
				});
				return ["image/jpeg", "image/png", "image/gif", "image/webp"].includes(result)
					? result
					: null;
			} catch {
				return null;
			}
		},
	};
}

/**
 * Create write operations for SSHFS-mounted remote
 */
async function createSshfsWriteOps(mountPoint: string, remoteCwd: string): Promise<WriteOperations> {
	return {
		writeFile: async (p: string, content: string) => {
			await fsWriteFile(toMountPath(p, mountPoint, remoteCwd), content, "utf-8");
		},
		mkdir: async (dir: string) => {
			await mkdir(toMountPath(dir, mountPoint, remoteCwd), { recursive: true });
		},
	};
}

/**
 * Create edit operations for SSHFS-mounted remote
 */
function createSshfsEditOps(mountPoint: string, remoteCwd: string): EditOperations {
	return {
		readFile: async (p: string) => {
			const { readFile } = await import("node:fs/promises");
			return readFile(toMountPath(p, mountPoint, remoteCwd));
		},
		access: async (p: string) => {
			const { access } = await import("node:fs/promises");
			await access(toMountPath(p, mountPoint, remoteCwd), fsConstants.R_OK);
		},
		writeFile: async (p: string, content: string) => {
			await fsWriteFile(toMountPath(p, mountPoint, remoteCwd), content, "utf-8");
		},
	};
}

/**
 * Create bash operations for remote SSH execution
 */
function createSshBashOps(config: SshConfig): BashOperations {
	return {
		exec: (command, cwd, { onData, signal, timeout }) =>
			new Promise((resolve, reject) => {
				const args = buildSshArgs(config);
				const remoteCommand = `cd ${JSON.stringify(config.remoteCwd)} && ${command}`;
				args.push(`${config.user}@${config.host}`, remoteCommand);

				const child = spawn("ssh", args, { stdio: ["ignore", "pipe", "pipe"] });
				let timedOut = false;

				const timer = timeout
					? setTimeout(() => {
							timedOut = true;
							child.kill();
						}, timeout * 1000)
					: undefined;

				child.stdout.on("data", onData);
				child.stderr.on("data", onData);

				child.on("error", (e) => {
					if (timer) clearTimeout(timer);
					reject(e);
				});

				const onAbort = () => child.kill();
				signal?.addEventListener("abort", onAbort, { once: true });

				child.on("close", (code) => {
					if (timer) clearTimeout(timer);
					signal?.removeEventListener("abort", onAbort);

					if (signal?.aborted) {
						reject(new Error("aborted"));
					} else if (timedOut) {
						reject(new Error(`timeout:${timeout}`));
					} else {
						resolve({ exitCode: code });
					}
				});
			}),
	};
}

export default function (pi: ExtensionAPI) {
	// 1. Register CLI flags
	pi.registerFlag("ssh", {
		type: "string",
		description: "SSH connection: user@host or user@host:/path",
	});
	pi.registerFlag("ssh-key", {
		type: "string",
		description: "Path to SSH private key",
	});
	pi.registerFlag("ssh-port", {
		type: "number",
		default: 22,
	});

	// 2. Session lifecycle - conditionally register tools if --ssh provided
	pi.on("session_start", async (_event, ctx) => {
		const sshArg = pi.getFlag("ssh") as string | undefined;
		if (!sshArg) {
			return;
		}

		try {
			// Parse connection string
			const { user, host, remoteCwd: explicitCwd } = parseSshArg(sshArg);
			const port = (pi.getFlag("ssh-port") as number) ?? 22;
			const key = pi.getFlag("ssh-key") as string | undefined;

			// Verify SSHFS is available
			const sshfsPaths = [
				"/usr/bin/sshfs",
				"/opt/homebrew/bin/sshfs",
				"/usr/local/bin/sshfs",
			];
			let sshfsFound = false;
			for (const path of sshfsPaths) {
				try {
					await access(path);
					sshfsFound = true;
					break;
				} catch {
					// Try next path
				}
			}
			if (!sshfsFound) {
				throw new Error(
					"SSHFS not found. Install with: brew install macfuse sshfs (macOS) or apt install sshfs (Linux)",
				);
			}

			// If no remote path specified, detect home directory via SSH
			const remoteCwd =
				explicitCwd ?? (await detectRemoteCwd(user, host, port, key));

			// Setup mount point
			const mountPoint = join(
				homedir(),
				".pi",
				"ssh-mounts",
				`${host}_${Date.now()}`,
			);
			await mkdir(mountPoint, { recursive: true });

			const config: SshConfig = {
				user,
				host,
				port,
				key,
				remoteCwd,
				mountPoint,
			};

			// Mount SSHFS
			ctx.ui.notify(`Mounting ${user}@${host}:${remoteCwd}...`, "info");
			await mountRemote(config);
			activeConfig = config;

			// Create tools with SSHFS operations
			const localCwd = process.cwd();

			const readTool = createReadTool(localCwd, {
				operations: createSshfsReadOps(mountPoint, remoteCwd),
			});
			const writeTool = createWriteTool(localCwd, {
				operations: createSshfsWriteOps(mountPoint, remoteCwd),
			});
			const editTool = createEditTool(localCwd, {
				operations: createSshfsEditOps(mountPoint, remoteCwd),
			});
			const bashTool = createBashTool(localCwd, {
				operations: createSshBashOps(config),
			});

			// Register tool overrides
			pi.registerTool(readTool);
			pi.registerTool(writeTool);
			pi.registerTool(editTool);
			pi.registerTool(bashTool);

			ctx.ui.notify(`SSH mode: ${user}@${host}:${remoteCwd}`, "success");
			ctx.ui.setStatus("ssh", `SSH: ${user}@${host}:${remoteCwd}`);
		} catch (error: any) {
			ctx.ui.notify(`SSH setup failed: ${error.message}`, "error");
			throw error;
		}
	});

	// 3. Modify system prompt when SSH is active
	pi.on("before_agent_start", async (event) => {
		if (activeConfig) {
			const { user, host, remoteCwd } = activeConfig;
			return {
				systemPrompt: event.systemPrompt.replace(
					`Current working directory: ${process.cwd()}`,
					`Current working directory: ${remoteCwd} (SSH: ${user}@${host})`,
				),
			};
		}
	});

	// 4. Cleanup on shutdown
	pi.on("session_shutdown", async () => {
		if (activeConfig) {
			try {
				await unmountRemote(activeConfig.mountPoint);
				await rm(activeConfig.mountPoint, { recursive: true, force: true });
			} catch {
				// Ignore cleanup errors
			}
			activeConfig = null;
		}
	});
}