From 5aa8036a6fc6376c0e1020fa73369d997ccc5655 Mon Sep 17 00:00:00 2001 From: Joakim Persson Date: Tue, 5 May 2026 22:58:34 +0200 Subject: [PATCH] ssh-controlmaster: reuse system ControlMaster from ~/.ssh/config if present --- extensions/ssh-controlmaster.ts | 362 ++++++++++++++++---------------- 1 file changed, 183 insertions(+), 179 deletions(-) diff --git a/extensions/ssh-controlmaster.ts b/extensions/ssh-controlmaster.ts index 0769e5d..3eb196c 100644 --- a/extensions/ssh-controlmaster.ts +++ b/extensions/ssh-controlmaster.ts @@ -3,15 +3,16 @@ * * Like the ssh.ts example but uses a persistent ControlMaster socket so all * tool calls (read, write, edit, bash) multiplex over a single SSH connection - * instead of opening a new one each time. Much faster on slow links or when - * pi makes many small file reads in sequence. + * instead of opening a new one each time. + * + * If ~/.ssh/config already configures ControlMaster (auto or yes) for the + * target host, the extension reuses the system socket rather than creating a + * second parallel connection. In that case pi does NOT tear down the master + * on exit — it was the system's to manage before pi arrived. * * Usage: - * pi -e ~/.pi/agent/extensions/ssh-controlmaster.ts --ssh user@host - * pi -e ~/.pi/agent/extensions/ssh-controlmaster.ts --ssh user@host:/remote/path - * - * Or, once installed globally, just: * pi --ssh user@host + * pi --ssh user@host:/remote/path * * Requirements: * - SSH key-based auth (no password prompts) @@ -33,31 +34,104 @@ import { type WriteOperations, } from "@mariozechner/pi-coding-agent"; -// ── ControlMaster helpers ──────────────────────────────────────────────────── +// ── Types ──────────────────────────────────────────────────────────────────── -function controlSocketPath(): string { +interface SshState { + remote: string; + remoteCwd: string; + socketPath: string; + ownsmaster: boolean; // true = we started the master, we must stop it on exit +} + +// ── Helpers ────────────────────────────────────────────────────────────────── + +function ownSocketPath(): string { // Keep path short — macOS has a ~104-char Unix socket path limit return join(tmpdir(), `pi-cm-${process.pid}.sock`); } +/** Run a one-shot command, return stdout as string. Rejects on non-zero exit. */ +function run(args: string[]): Promise { + return new Promise((resolve, reject) => { + const child = spawn(args[0], args.slice(1), { stdio: ["ignore", "pipe", "pipe"] }); + const out: Buffer[] = []; + const err: Buffer[] = []; + child.stdout.on("data", (d: Buffer) => out.push(d)); + child.stderr.on("data", (d: Buffer) => err.push(d)); + child.on("error", reject); + child.on("close", (code) => { + if (code === 0) resolve(Buffer.concat(out).toString().trim()); + else reject(new Error(`${args.join(" ")} exited ${code}: ${Buffer.concat(err).toString().trim()}`)); + }); + }); +} + +// ── ControlMaster negotiation ──────────────────────────────────────────────── + +interface CmConfig { + master: string; // "auto" | "yes" | "no" | "ask" | "autoask" + path: string; // controlpath value (may be "" if unset) +} + +/** + * Read the effective SSH config for a host via `ssh -G`. + * Returns the controlmaster and controlpath values SSH would actually use. + */ +async function readSshConfig(remote: string): Promise { + // Strip any user@ prefix for -G (ssh -G takes a hostname or alias, not user@host) + const host = remote.includes("@") ? remote.split("@")[1] : remote; + try { + const output = await run(["ssh", "-G", host]); + const get = (key: string): string => { + const m = output.match(new RegExp(`^${key}\\s+(.+)$`, "im")); + return m ? m[1].trim() : ""; + }; + return { master: get("controlmaster").toLowerCase(), path: get("controlpath") }; + } catch { + return { master: "no", path: "" }; + } +} + +/** + * Decide whether to reuse the system ControlMaster or create our own. + * + * "System master" path: ~/.ssh/config has ControlMaster auto/yes for this host. + * - The first plain `ssh` call will establish the master automatically (auto), + * or it may already be running. Either way SSH handles it. + * - We do NOT own the master and must not shut it down on exit. + * - socketPath is the system-configured path (used only for status display). + * + * "Owned master" path: no ControlMaster configured (or explicitly disabled). + * - We start ssh -fN ControlMaster=yes at our own socket. + * - We shut it down on exit. + */ +async function negotiateMaster( + remote: string, +): Promise<{ socketPath: string; ownsmaster: boolean }> { + const cfg = await readSshConfig(remote); + const systemHasMaster = cfg.master === "auto" || cfg.master === "yes"; + + if (systemHasMaster && cfg.path) { + return { socketPath: cfg.path, ownsmaster: false }; + } + + // No system master — create our own + const socketPath = ownSocketPath(); + await startControlMaster(remote, socketPath); + return { socketPath, ownsmaster: true }; +} + function startControlMaster(remote: string, socketPath: string): Promise { return new Promise((resolve, reject) => { // -fN: fork to background after auth, don't run a remote command - // ControlPersist=yes: keep the master alive in the background indefinitely + // ControlPersist=yes: keep master alive in the background indefinitely const child = spawn( "ssh", - [ - "-fN", - "-o", "ControlMaster=yes", - "-o", `ControlPath=${socketPath}`, - "-o", "ControlPersist=yes", - remote, - ], + ["-fN", "-o", "ControlMaster=yes", "-o", `ControlPath=${socketPath}`, "-o", "ControlPersist=yes", remote], { stdio: "ignore" }, ); child.on("error", reject); child.on("close", (code) => { - // -f forks to background; parent exits 0 once master is ready if (code === 0) resolve(); else reject(new Error(`ControlMaster exited with code ${code}`)); }); @@ -66,41 +140,37 @@ function startControlMaster(remote: string, socketPath: string): Promise { function stopControlMaster(remote: string, socketPath: string): Promise { return new Promise((resolve) => { - // -O exit sends the exit signal to the master process via the socket const child = spawn( "ssh", ["-O", "exit", "-o", `ControlPath=${socketPath}`, remote], { stdio: "ignore" }, ); - child.on("close", () => resolve()); // best-effort; ignore errors on cleanup + child.on("close", () => resolve()); // best-effort; ignore errors }); } -// ── SSH exec (multiplexed over ControlMaster socket) ──────────────────────── +// ── SSH exec (multiplexed) ─────────────────────────────────────────────────── +/** + * Run a remote command over SSH. + * When socketPath is set we pin to that socket; otherwise we let SSH use + * whatever is configured in ~/.ssh/config (which includes auto-multiplexing). + */ function sshExec(remote: string, socketPath: string, command: string): Promise { return new Promise((resolve, reject) => { const child = spawn( "ssh", - [ - "-o", "ControlMaster=no", // use existing master, don't create a new one - "-o", `ControlPath=${socketPath}`, - remote, - command, - ], + ["-o", "ControlMaster=no", "-o", `ControlPath=${socketPath}`, remote, command], { stdio: ["ignore", "pipe", "pipe"] }, ); - const chunks: Buffer[] = []; - const errChunks: Buffer[] = []; - child.stdout.on("data", (d: Buffer) => chunks.push(d)); - child.stderr.on("data", (d: Buffer) => errChunks.push(d)); + const out: Buffer[] = []; + const err: Buffer[] = []; + child.stdout.on("data", (d: Buffer) => out.push(d)); + child.stderr.on("data", (d: Buffer) => err.push(d)); 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)); - } + if (code !== 0) reject(new Error(`SSH failed (${code}): ${Buffer.concat(err).toString()}`)); + else resolve(Buffer.concat(out)); }); }); } @@ -108,106 +178,69 @@ function sshExec(remote: string, socketPath: string, command: string): Promise p.replace(localCwd, remoteCwd); + const r = (p: string) => p.replace(localCwd, remoteCwd); return { - readFile: (p) => sshExec(remote, socketPath, `cat ${JSON.stringify(toRemote(p))}`), + readFile: (p) => sshExec(remote, socketPath, `cat ${JSON.stringify(r(p))}`), access: (p) => - sshExec(remote, socketPath, `test -r ${JSON.stringify(toRemote(p))}`).then(() => {}), + sshExec(remote, socketPath, `test -r ${JSON.stringify(r(p))}`).then(() => {}), detectImageMimeType: async (p) => { try { - const r = await sshExec( - remote, - socketPath, - `file --mime-type -b ${JSON.stringify(toRemote(p))}`, - ); - const m = r.toString().trim(); + const out = await sshExec(remote, socketPath, `file --mime-type -b ${JSON.stringify(r(p))}`); + const m = out.toString().trim(); return ["image/jpeg", "image/png", "image/gif", "image/webp"].includes(m) ? (m as "image/jpeg" | "image/png" | "image/gif" | "image/webp") : null; - } catch { - return null; - } + } catch { return null; } }, }; } function createRemoteWriteOps( - remote: string, - socketPath: string, - remoteCwd: string, - localCwd: string, + remote: string, socketPath: string, remoteCwd: string, localCwd: string, ): WriteOperations { - const toRemote = (p: string) => p.replace(localCwd, remoteCwd); + const r = (p: string) => p.replace(localCwd, remoteCwd); return { writeFile: async (p, content) => { const b64 = Buffer.from(content).toString("base64"); - await sshExec( - remote, - socketPath, - `echo ${JSON.stringify(b64)} | base64 -d > ${JSON.stringify(toRemote(p))}`, - ); + await sshExec(remote, socketPath, `echo ${JSON.stringify(b64)} | base64 -d > ${JSON.stringify(r(p))}`); }, mkdir: (dir) => - sshExec(remote, socketPath, `mkdir -p ${JSON.stringify(toRemote(dir))}`).then(() => {}), + sshExec(remote, socketPath, `mkdir -p ${JSON.stringify(r(dir))}`).then(() => {}), }; } function createRemoteEditOps( - remote: string, - socketPath: string, - remoteCwd: string, - localCwd: string, + remote: string, socketPath: string, remoteCwd: string, localCwd: string, ): EditOperations { - const r = createRemoteReadOps(remote, socketPath, remoteCwd, localCwd); - const w = createRemoteWriteOps(remote, socketPath, remoteCwd, localCwd); - return { readFile: r.readFile, access: r.access, writeFile: w.writeFile }; + const read = createRemoteReadOps(remote, socketPath, remoteCwd, localCwd); + const write = createRemoteWriteOps(remote, socketPath, remoteCwd, localCwd); + return { readFile: read.readFile, access: read.access, writeFile: write.writeFile }; } function createRemoteBashOps( - remote: string, - socketPath: string, - remoteCwd: string, - localCwd: string, + remote: string, socketPath: string, remoteCwd: string, localCwd: string, ): BashOperations { - const toRemote = (p: string) => p.replace(localCwd, remoteCwd); + const r = (p: string) => p.replace(localCwd, remoteCwd); return { exec: (command, cwd, { onData, signal, timeout }) => new Promise((resolve, reject) => { - const cmd = `cd ${JSON.stringify(toRemote(cwd))} && ${command}`; + const cmd = `cd ${JSON.stringify(r(cwd))} && ${command}`; const child = spawn( "ssh", - [ - "-o", "ControlMaster=no", - "-o", `ControlPath=${socketPath}`, - remote, - cmd, - ], + ["-o", "ControlMaster=no", "-o", `ControlPath=${socketPath}`, remote, cmd], { stdio: ["ignore", "pipe", "pipe"] }, ); - let timedOut = false; const timer = timeout - ? setTimeout(() => { - timedOut = true; - child.kill(); - }, timeout * 1000) + ? setTimeout(() => { timedOut = true; child.kill(); }, timeout * 1000) : undefined; - child.stdout.on("data", onData); child.stderr.on("data", onData); - child.on("error", (e: Error) => { - if (timer) clearTimeout(timer); - reject(e); - }); - + child.on("error", (e: Error) => { if (timer) clearTimeout(timer); reject(e); }); const onAbort = () => child.kill(); signal?.addEventListener("abort", onAbort, { once: true }); - child.on("close", (code: number) => { if (timer) clearTimeout(timer); signal?.removeEventListener("abort", onAbort); @@ -233,74 +266,56 @@ export default function (pi: ExtensionAPI) { const localEdit = createEditTool(localCwd); const localBash = createBashTool(localCwd); - // Resolved lazily in session_start once CLI flags are available - let resolvedSsh: { - remote: string; - remoteCwd: string; - socketPath: string; - } | null = null; + let state: SshState | null = null; + const getState = () => state; - const getSsh = () => resolvedSsh; - - // ── Tool overrides ─────────────────────────────────────────────────────── + // ── Tool overrides ───────────────────────────────────────────────────────── pi.registerTool({ ...localRead, - async execute(id, params, signal, onUpdate, _ctx) { - const ssh = getSsh(); - if (ssh) { - const tool = createReadTool(localCwd, { - operations: createRemoteReadOps(ssh.remote, ssh.socketPath, ssh.remoteCwd, localCwd), - }); - return tool.execute(id, params, signal, onUpdate); - } - return localRead.execute(id, params, signal, onUpdate); + async execute(id, params, signal, onUpdate) { + const s = getState(); + if (!s) return localRead.execute(id, params, signal, onUpdate); + return createReadTool(localCwd, { + operations: createRemoteReadOps(s.remote, s.socketPath, s.remoteCwd, localCwd), + }).execute(id, params, signal, onUpdate); }, }); pi.registerTool({ ...localWrite, - async execute(id, params, signal, onUpdate, _ctx) { - const ssh = getSsh(); - if (ssh) { - const tool = createWriteTool(localCwd, { - operations: createRemoteWriteOps(ssh.remote, ssh.socketPath, ssh.remoteCwd, localCwd), - }); - return tool.execute(id, params, signal, onUpdate); - } - return localWrite.execute(id, params, signal, onUpdate); + async execute(id, params, signal, onUpdate) { + const s = getState(); + if (!s) return localWrite.execute(id, params, signal, onUpdate); + return createWriteTool(localCwd, { + operations: createRemoteWriteOps(s.remote, s.socketPath, s.remoteCwd, localCwd), + }).execute(id, params, signal, onUpdate); }, }); pi.registerTool({ ...localEdit, - async execute(id, params, signal, onUpdate, _ctx) { - const ssh = getSsh(); - if (ssh) { - const tool = createEditTool(localCwd, { - operations: createRemoteEditOps(ssh.remote, ssh.socketPath, ssh.remoteCwd, localCwd), - }); - return tool.execute(id, params, signal, onUpdate); - } - return localEdit.execute(id, params, signal, onUpdate); + async execute(id, params, signal, onUpdate) { + const s = getState(); + if (!s) return localEdit.execute(id, params, signal, onUpdate); + return createEditTool(localCwd, { + operations: createRemoteEditOps(s.remote, s.socketPath, s.remoteCwd, localCwd), + }).execute(id, params, signal, onUpdate); }, }); pi.registerTool({ ...localBash, - async execute(id, params, signal, onUpdate, _ctx) { - const ssh = getSsh(); - if (ssh) { - const tool = createBashTool(localCwd, { - operations: createRemoteBashOps(ssh.remote, ssh.socketPath, ssh.remoteCwd, localCwd), - }); - return tool.execute(id, params, signal, onUpdate); - } - return localBash.execute(id, params, signal, onUpdate); + async execute(id, params, signal, onUpdate) { + const s = getState(); + if (!s) return localBash.execute(id, params, signal, onUpdate); + return createBashTool(localCwd, { + operations: createRemoteBashOps(s.remote, s.socketPath, s.remoteCwd, localCwd), + }).execute(id, params, signal, onUpdate); }, }); - // ── Session lifecycle ──────────────────────────────────────────────────── + // ── Session lifecycle ────────────────────────────────────────────────────── pi.on("session_start", async (_event, ctx) => { const arg = pi.getFlag("ssh") as string | undefined; @@ -313,67 +328,56 @@ export default function (pi: ExtensionAPI) { [remote, remoteCwd] = arg.split(":"); } else { remote = arg; - // Resolve remote home dir via a plain one-shot SSH call before the - // ControlMaster is up — this is the only un-multiplexed connection. - const result = await new Promise((resolve, reject) => { - const child = spawn("ssh", [remote, "pwd"], { - stdio: ["ignore", "pipe", "pipe"], - }); - const chunks: Buffer[] = []; - child.stdout.on("data", (d: Buffer) => chunks.push(d)); - child.on("error", reject); - child.on("close", (code) => { - if (code === 0) resolve(Buffer.concat(chunks).toString().trim()); - else reject(new Error(`Could not resolve remote pwd (exit ${code})`)); - }); + remoteCwd = await run(["ssh", remote, "pwd"]).catch((e) => { + throw new Error(`Could not resolve remote pwd: ${e.message}`); }); - remoteCwd = result; } - const socketPath = controlSocketPath(); ctx.ui.setStatus("ssh", ctx.ui.theme.fg("accent", `SSH: ${remote} ⟳ connecting…`)); + let socketPath: string; + let ownsmaster: boolean; + try { - await startControlMaster(remote, socketPath); + ({ socketPath, ownsmaster } = await negotiateMaster(remote)); } catch (err) { ctx.ui.setStatus("ssh", ctx.ui.theme.fg("error", `SSH: ${remote} ✗ failed`)); ctx.ui.notify(`SSH ControlMaster failed: ${err}`, "error"); return; } - resolvedSsh = { remote, remoteCwd, socketPath }; - ctx.ui.setStatus("ssh", ctx.ui.theme.fg("accent", `SSH ⚡ ${remote}:${remoteCwd}`)); - ctx.ui.notify(`SSH ControlMaster ready — ${remote}:${remoteCwd}`, "success"); + state = { remote, remoteCwd, socketPath, ownsmaster }; + + const tag = ownsmaster ? "⚡ own master" : "⚡ system master"; + ctx.ui.setStatus("ssh", ctx.ui.theme.fg("accent", `SSH ${tag} ${remote}:${remoteCwd}`)); + ctx.ui.notify(`SSH ready (${ownsmaster ? "own" : "system"} master) — ${remote}:${remoteCwd}`, "success"); }); - pi.on("session_shutdown", async (_event, _ctx) => { - const ssh = getSsh(); - if (ssh) { - await stopControlMaster(ssh.remote, ssh.socketPath); - resolvedSsh = null; + pi.on("session_shutdown", async () => { + const s = getState(); + if (s?.ownsmaster) { + await stopControlMaster(s.remote, s.socketPath); } + state = null; }); - // ── User ! commands via SSH ────────────────────────────────────────────── + // ── User ! commands via SSH ──────────────────────────────────────────────── - pi.on("user_bash", (_event) => { - const ssh = getSsh(); - if (!ssh) return; - return { - operations: createRemoteBashOps(ssh.remote, ssh.socketPath, ssh.remoteCwd, localCwd), - }; + pi.on("user_bash", () => { + const s = getState(); + if (!s) return; + return { operations: createRemoteBashOps(s.remote, s.socketPath, s.remoteCwd, localCwd) }; }); - // ── Tell the LLM where it's operating ─────────────────────────────────── + // ── Tell the LLM where it's operating ───────────────────────────────────── pi.on("before_agent_start", async (event) => { - const ssh = getSsh(); - if (ssh) { - const modified = event.systemPrompt.replace( - `Current working directory: ${localCwd}`, - `Current working directory: ${ssh.remoteCwd} (via SSH ControlMaster: ${ssh.remote})`, - ); - return { systemPrompt: modified }; - } + const s = getState(); + if (!s) return; + const modified = event.systemPrompt.replace( + `Current working directory: ${localCwd}`, + `Current working directory: ${s.remoteCwd} (via SSH ControlMaster: ${s.remote})`, + ); + return { systemPrompt: modified }; }); }