@oh-my-pi/pi-coding-agent 3.33.0 → 3.34.0

package/src/core/ssh/ssh-executor.ts

@@ -0,0 +1,190 @@
+import { createWriteStream, type WriteStream } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import type { Subprocess } from "bun";
+import { nanoid } from "nanoid";
+import stripAnsi from "strip-ansi";
+import { killProcessTree, sanitizeBinaryOutput } from "../../utils/shell";
+import { logger } from "../logger";
+import { DEFAULT_MAX_BYTES, truncateTail } from "../tools/truncate";
+import { ScopeSignal } from "../utils";
+import { buildRemoteCommand, ensureConnection, ensureHostInfo, type SSHConnectionTarget } from "./connection-manager";
+import { hasSshfs, mountRemote } from "./sshfs-mount";
+
+export interface SSHExecutorOptions {
+	/** Timeout in milliseconds */
+	timeout?: number;
+	/** Callback for streaming output chunks (already sanitized) */
+	onChunk?: (chunk: string) => void;
+	/** AbortSignal for cancellation */
+	signal?: AbortSignal;
+	/** Remote path to mount when sshfs is available */
+	remotePath?: string;
+	/** Wrap commands in a POSIX shell for compat mode */
+	compatEnabled?: boolean;
+}
+
+export interface SSHResult {
+	/** Combined stdout + stderr output (sanitized, possibly truncated) */
+	output: string;
+	/** Process exit code (undefined if killed/cancelled) */
+	exitCode: number | undefined;
+	/** Whether the command was cancelled via signal */
+	cancelled: boolean;
+	/** Whether the output was truncated */
+	truncated: boolean;
+	/** Path to temp file containing full output (if output exceeded truncation threshold) */
+	fullOutputPath?: string;
+}
+
+function createSanitizer(): TransformStream<Uint8Array, string> {
+	const decoder = new TextDecoder();
+	return new TransformStream({
+		transform(chunk, controller) {
+			const text = sanitizeBinaryOutput(stripAnsi(decoder.decode(chunk, { stream: true }))).replace(/\r/g, "");
+			controller.enqueue(text);
+		},
+	});
+}
+
+function createOutputSink(
+	spillThreshold: number,
+	maxBuffer: number,
+	onChunk?: (text: string) => void,
+): WritableStream<string> & {
+	dump: (annotation?: string) => { output: string; truncated: boolean; fullOutputPath?: string };
+} {
+	const chunks: string[] = [];
+	let chunkBytes = 0;
+	let totalBytes = 0;
+	let fullOutputPath: string | undefined;
+	let fullOutputStream: WriteStream | undefined;
+
+	const sink = new WritableStream<string>({
+		write(text) {
+			totalBytes += text.length;
+
+			if (totalBytes > spillThreshold && !fullOutputPath) {
+				fullOutputPath = join(tmpdir(), `omp-${nanoid()}.buffer`);
+				const ts = createWriteStream(fullOutputPath);
+				chunks.forEach((c) => {
+					ts.write(c);
+				});
+				fullOutputStream = ts;
+			}
+			fullOutputStream?.write(text);
+
+			chunks.push(text);
+			chunkBytes += text.length;
+			while (chunkBytes > maxBuffer && chunks.length > 1) {
+				chunkBytes -= chunks.shift()!.length;
+			}
+
+			onChunk?.(text);
+		},
+		close() {
+			fullOutputStream?.end();
+		},
+	});
+
+	return Object.assign(sink, {
+		dump(annotation?: string) {
+			if (annotation) {
+				chunks.push(`\n\n${annotation}`);
+			}
+			const full = chunks.join("");
+			const { content, truncated } = truncateTail(full);
+			return { output: truncated ? content : full, truncated, fullOutputPath };
+		},
+	});
+}
+
+function quoteForCompatShell(command: string): string {
+	if (command.length === 0) {
+		return "''";
+	}
+	const escaped = command.replace(/'/g, "'\\''");
+	return `'${escaped}'`;
+}
+
+function buildCompatCommand(shell: "bash" | "sh", command: string): string {
+	return `${shell} -c ${quoteForCompatShell(command)}`;
+}
+
+export async function executeSSH(
+	host: SSHConnectionTarget,
+	command: string,
+	options?: SSHExecutorOptions,
+): Promise<SSHResult> {
+	await ensureConnection(host);
+	if (hasSshfs()) {
+		try {
+			await mountRemote(host, options?.remotePath ?? "/");
+		} catch (err) {
+			logger.warn("SSHFS mount failed", { host: host.name, error: String(err) });
+		}
+	}
+
+	using signal = new ScopeSignal(options);
+
+	let resolvedCommand = command;
+	if (options?.compatEnabled) {
+		const info = await ensureHostInfo(host);
+		if (info.compatShell) {
+			resolvedCommand = buildCompatCommand(info.compatShell, command);
+		} else {
+			logger.warn("SSH compat enabled without detected compat shell", { host: host.name });
+		}
+	}
+	const child: Subprocess = Bun.spawn(["ssh", ...buildRemoteCommand(host, resolvedCommand)], {
+		stdin: "ignore",
+		stdout: "pipe",
+		stderr: "pipe",
+	});
+
+	signal.catch(() => {
+		killProcessTree(child.pid);
+	});
+
+	const sink = createOutputSink(DEFAULT_MAX_BYTES, DEFAULT_MAX_BYTES * 2, options?.onChunk);
+
+	const writer = sink.getWriter();
+	try {
+		async function pumpStream(readable: ReadableStream<Uint8Array>) {
+			const reader = readable.pipeThrough(createSanitizer()).getReader();
+			try {
+				while (true) {
+					const { done, value } = await reader.read();
+					if (done) break;
+					await writer.write(value);
+				}
+			} finally {
+				reader.releaseLock();
+			}
+		}
+		await Promise.all([
+			pumpStream(child.stdout as ReadableStream<Uint8Array>),
+			pumpStream(child.stderr as ReadableStream<Uint8Array>),
+		]);
+	} finally {
+		await writer.close();
+	}
+
+	const exitCode = await child.exited;
+	const cancelled = exitCode === null || (exitCode !== 0 && (options?.signal?.aborted ?? false));
+
+	if (signal.timedOut()) {
+		const secs = Math.round(options!.timeout! / 1000);
+		return {
+			exitCode: undefined,
+			cancelled: true,
+			...sink.dump(`SSH command timed out after ${secs} seconds`),
+		};
+	}
+
+	return {
+		exitCode: cancelled ? undefined : exitCode,
+		cancelled,
+		...sink.dump(),
+	};
+}

package/src/core/ssh/sshfs-mount.ts

@@ -0,0 +1,162 @@
+import { chmodSync, existsSync, mkdirSync } from "node:fs";
+import { homedir } from "node:os";
+import { join } from "node:path";
+import { CONFIG_DIR_NAME } from "../../config";
+import { logger } from "../logger";
+import { getControlDir, getControlPathTemplate, type SSHConnectionTarget } from "./connection-manager";
+
+const REMOTE_DIR = join(homedir(), CONFIG_DIR_NAME, "remote");
+const CONTROL_DIR = getControlDir();
+const CONTROL_PATH = getControlPathTemplate();
+
+const mountedPaths = new Set<string>();
+
+function ensureDir(path: string, mode = 0o700): void {
+	if (!existsSync(path)) {
+		mkdirSync(path, { recursive: true, mode });
+	}
+	try {
+		chmodSync(path, mode);
+	} catch (err) {
+		logger.debug("SSHFS dir chmod failed", { path, error: String(err) });
+	}
+}
+
+function decodeOutput(buffer?: Uint8Array): string {
+	if (!buffer || buffer.length === 0) return "";
+	return new TextDecoder().decode(buffer).trim();
+}
+
+function getMountName(host: SSHConnectionTarget): string {
+	const raw = (host.name ?? host.host).trim();
+	const sanitized = raw.replace(/[^a-zA-Z0-9._-]+/g, "_");
+	return sanitized.length > 0 ? sanitized : "remote";
+}
+
+function getMountPath(host: SSHConnectionTarget): string {
+	return join(REMOTE_DIR, getMountName(host));
+}
+
+function buildSshTarget(host: SSHConnectionTarget): string {
+	return host.username ? `${host.username}@${host.host}` : host.host;
+}
+
+function buildSshfsArgs(host: SSHConnectionTarget): string[] {
+	const args = [
+		"-o",
+		"reconnect",
+		"-o",
+		"ServerAliveInterval=15",
+		"-o",
+		"ServerAliveCountMax=3",
+		"-o",
+		"BatchMode=yes",
+		"-o",
+		"StrictHostKeyChecking=accept-new",
+		"-o",
+		"ControlMaster=auto",
+		"-o",
+		`ControlPath=${CONTROL_PATH}`,
+		"-o",
+		"ControlPersist=3600",
+	];
+
+	if (host.port) {
+		args.push("-p", String(host.port));
+	}
+
+	if (host.keyPath) {
+		args.push("-o", `IdentityFile=${host.keyPath}`);
+	}
+
+	return args;
+}
+
+function unmountPath(path: string): boolean {
+	const fusermount = Bun.which("fusermount") ?? Bun.which("fusermount3");
+	if (fusermount) {
+		const result = Bun.spawnSync([fusermount, "-u", path], {
+			stdin: "ignore",
+			stdout: "ignore",
+			stderr: "pipe",
+		});
+		if (result.exitCode === 0) return true;
+	}
+
+	const umount = Bun.which("umount");
+	if (!umount) return false;
+	const result = Bun.spawnSync([umount, path], {
+		stdin: "ignore",
+		stdout: "ignore",
+		stderr: "pipe",
+	});
+	return result.exitCode === 0;
+}
+
+export function hasSshfs(): boolean {
+	return Bun.which("sshfs") !== null;
+}
+
+export function isMounted(path: string): boolean {
+	const mountpoint = Bun.which("mountpoint");
+	if (!mountpoint) return false;
+	const result = Bun.spawnSync([mountpoint, "-q", path], {
+		stdin: "ignore",
+		stdout: "ignore",
+		stderr: "ignore",
+	});
+	return result.exitCode === 0;
+}
+
+export async function mountRemote(host: SSHConnectionTarget, remotePath = "/"): Promise<string | undefined> {
+	if (!hasSshfs()) return undefined;
+
+	ensureDir(REMOTE_DIR);
+	ensureDir(CONTROL_DIR);
+
+	const mountPath = getMountPath(host);
+	ensureDir(mountPath);
+
+	if (isMounted(mountPath)) {
+		mountedPaths.add(mountPath);
+		return mountPath;
+	}
+
+	const target = `${buildSshTarget(host)}:${remotePath}`;
+	const result = Bun.spawnSync(["sshfs", ...buildSshfsArgs(host), target, mountPath], {
+		stdin: "ignore",
+		stdout: "pipe",
+		stderr: "pipe",
+	});
+
+	if (result.exitCode !== 0) {
+		const detail = decodeOutput(result.stderr);
+		const suffix = detail ? `: ${detail}` : "";
+		throw new Error(`Failed to mount ${target}${suffix}`);
+	}
+
+	mountedPaths.add(mountPath);
+	return mountPath;
+}
+
+export async function unmountRemote(host: SSHConnectionTarget): Promise<boolean> {
+	const mountPath = getMountPath(host);
+	if (!isMounted(mountPath)) {
+		mountedPaths.delete(mountPath);
+		return false;
+	}
+
+	const success = unmountPath(mountPath);
+	if (success) {
+		mountedPaths.delete(mountPath);
+	}
+
+	return success;
+}
+
+export async function unmountAll(): Promise<void> {
+	for (const mountPath of Array.from(mountedPaths)) {
+		unmountPath(mountPath);
+	}
+	mountedPaths.clear();
+}

package/src/core/ssh-executor.ts

@@ -0,0 +1,5 @@
+/**
+ * Backwards-compatible re-export for SSH execution.
+ */
+
+export { executeSSH, type SSHExecutorOptions, type SSHResult } from "./ssh/ssh-executor";