ssh-controlmaster: reuse system ControlMaster from ~/.ssh/config if present

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
+ * instead of opening a new one each time.
- * pi makes many small file reads in sequence.
+ *
+ * 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<string> {
+  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<CmConfig> {
+  // 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<void> {
   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<void> {
 
 function stopControlMaster(remote: string, socketPath: string): Promise<void> {
   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<Buffer> {
   return new Promise((resolve, reject) => {
     const child = spawn(
       "ssh",
-      [
+      ["-o", "ControlMaster=no", "-o", `ControlPath=${socketPath}`, remote, command],
-        "-o", "ControlMaster=no",   // use existing master, don't create a new one
-        "-o", `ControlPath=${socketPath}`,
-        remote,
-        command,
-      ],
       { stdio: ["ignore", "pipe", "pipe"] },
     );
-    const chunks: Buffer[] = [];
+    const out: Buffer[] = [];
-    const errChunks: Buffer[] = [];
+    const err: Buffer[] = [];
-    child.stdout.on("data", (d: Buffer) => chunks.push(d));
+    child.stdout.on("data", (d: Buffer) => out.push(d));
-    child.stderr.on("data", (d: Buffer) => errChunks.push(d));
+    child.stderr.on("data", (d: Buffer) => err.push(d));
     child.on("error", reject);
     child.on("close", (code) => {
-      if (code !== 0) {
+      if (code !== 0) reject(new Error(`SSH failed (${code}): ${Buffer.concat(err).toString()}`));
-        reject(new Error(`SSH failed (${code}): ${Buffer.concat(errChunks).toString()}`));
+      else resolve(Buffer.concat(out));
-      } else {
-        resolve(Buffer.concat(chunks));
-      }
     });
   });
 }
@@ -108,106 +178,69 @@ function sshExec(remote: string, socketPath: string, command: string): Promise<B
 // ── Remote tool operations ───────────────────────────────────────────────────
 
 function createRemoteReadOps(
-  remote: string,
+  remote: string, socketPath: string, remoteCwd: string, localCwd: string,
-  socketPath: string,
-  remoteCwd: string,
-  localCwd: string,
 ): ReadOperations {
-  const toRemote = (p: string) => 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(
+        const out = await sshExec(remote, socketPath, `file --mime-type -b ${JSON.stringify(r(p))}`);
-          remote,
+        const m = out.toString().trim();
-          socketPath,
-          `file --mime-type -b ${JSON.stringify(toRemote(p))}`,
-        );
-        const m = r.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 {
+      } catch { return null; }
-        return null;
-      }
     },
   };
 }
 
 function createRemoteWriteOps(
-  remote: string,
+  remote: string, socketPath: string, remoteCwd: string, localCwd: 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(
+      await sshExec(remote, socketPath, `echo ${JSON.stringify(b64)} | base64 -d > ${JSON.stringify(r(p))}`);
-        remote,
-        socketPath,
-        `echo ${JSON.stringify(b64)} | base64 -d > ${JSON.stringify(toRemote(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,
+  remote: string, socketPath: string, remoteCwd: string, localCwd: string,
-  socketPath: string,
-  remoteCwd: string,
-  localCwd: string,
 ): EditOperations {
-  const r = createRemoteReadOps(remote, socketPath, remoteCwd, localCwd);
+  const read = createRemoteReadOps(remote, socketPath, remoteCwd, localCwd);
-  const w = createRemoteWriteOps(remote, socketPath, remoteCwd, localCwd);
+  const write = createRemoteWriteOps(remote, socketPath, remoteCwd, localCwd);
-  return { readFile: r.readFile, access: r.access, writeFile: w.writeFile };
+  return { readFile: read.readFile, access: read.access, writeFile: write.writeFile };
 }
 
 function createRemoteBashOps(
-  remote: string,
+  remote: string, socketPath: string, remoteCwd: string, localCwd: 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(() => {
+          ? setTimeout(() => { timedOut = true; child.kill(); }, timeout * 1000)
-              timedOut = true;
-              child.kill();
-            }, timeout * 1000)
           : undefined;
-
         child.stdout.on("data", onData);
         child.stderr.on("data", onData);
-        child.on("error", (e: Error) => {
+        child.on("error", (e: Error) => { if (timer) clearTimeout(timer); reject(e); });
-          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 state: SshState | null = null;
-  let resolvedSsh: {
+  const getState = () => state;
-    remote: string;
-    remoteCwd: string;
-    socketPath: string;
-  } | null = null;
 
-  const getSsh = () => resolvedSsh;
+  // ── Tool overrides ─────────────────────────────────────────────────────────
-
-  // ── Tool overrides ───────────────────────────────────────────────────────
 
   pi.registerTool({
     ...localRead,
-    async execute(id, params, signal, onUpdate, _ctx) {
+    async execute(id, params, signal, onUpdate) {
-      const ssh = getSsh();
+      const s = getState();
-      if (ssh) {
+      if (!s) return localRead.execute(id, params, signal, onUpdate);
-        const tool = createReadTool(localCwd, {
+      return createReadTool(localCwd, {
-          operations: createRemoteReadOps(ssh.remote, ssh.socketPath, ssh.remoteCwd, localCwd),
+        operations: createRemoteReadOps(s.remote, s.socketPath, s.remoteCwd, localCwd),
-        });
+      }).execute(id, params, signal, onUpdate);
-        return tool.execute(id, params, signal, onUpdate);
-      }
-      return localRead.execute(id, params, signal, onUpdate);
     },
   });
 
   pi.registerTool({
     ...localWrite,
-    async execute(id, params, signal, onUpdate, _ctx) {
+    async execute(id, params, signal, onUpdate) {
-      const ssh = getSsh();
+      const s = getState();
-      if (ssh) {
+      if (!s) return localWrite.execute(id, params, signal, onUpdate);
-        const tool = createWriteTool(localCwd, {
+      return createWriteTool(localCwd, {
-          operations: createRemoteWriteOps(ssh.remote, ssh.socketPath, ssh.remoteCwd, localCwd),
+        operations: createRemoteWriteOps(s.remote, s.socketPath, s.remoteCwd, localCwd),
-        });
+      }).execute(id, params, signal, onUpdate);
-        return tool.execute(id, params, signal, onUpdate);
-      }
-      return localWrite.execute(id, params, signal, onUpdate);
     },
   });
 
   pi.registerTool({
     ...localEdit,
-    async execute(id, params, signal, onUpdate, _ctx) {
+    async execute(id, params, signal, onUpdate) {
-      const ssh = getSsh();
+      const s = getState();
-      if (ssh) {
+      if (!s) return localEdit.execute(id, params, signal, onUpdate);
-        const tool = createEditTool(localCwd, {
+      return createEditTool(localCwd, {
-          operations: createRemoteEditOps(ssh.remote, ssh.socketPath, ssh.remoteCwd, localCwd),
+        operations: createRemoteEditOps(s.remote, s.socketPath, s.remoteCwd, localCwd),
-        });
+      }).execute(id, params, signal, onUpdate);
-        return tool.execute(id, params, signal, onUpdate);
-      }
-      return localEdit.execute(id, params, signal, onUpdate);
     },
   });
 
   pi.registerTool({
     ...localBash,
-    async execute(id, params, signal, onUpdate, _ctx) {
+    async execute(id, params, signal, onUpdate) {
-      const ssh = getSsh();
+      const s = getState();
-      if (ssh) {
+      if (!s) return localBash.execute(id, params, signal, onUpdate);
-        const tool = createBashTool(localCwd, {
+      return createBashTool(localCwd, {
-          operations: createRemoteBashOps(ssh.remote, ssh.socketPath, ssh.remoteCwd, localCwd),
+        operations: createRemoteBashOps(s.remote, s.socketPath, s.remoteCwd, localCwd),
-        });
+      }).execute(id, params, signal, onUpdate);
-        return tool.execute(id, params, signal, onUpdate);
-      }
-      return localBash.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
+      remoteCwd = await run(["ssh", remote, "pwd"]).catch((e) => {
-      // ControlMaster is up — this is the only un-multiplexed connection.
+        throw new Error(`Could not resolve remote pwd: ${e.message}`);
-      const result = await new Promise<string>((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 = 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 };
+    state = { remote, remoteCwd, socketPath, ownsmaster };
-    ctx.ui.setStatus("ssh", ctx.ui.theme.fg("accent", `SSH ⚡ ${remote}:${remoteCwd}`));
+
-    ctx.ui.notify(`SSH ControlMaster ready — ${remote}:${remoteCwd}`, "success");
+    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) => {
+  pi.on("session_shutdown", async () => {
-    const ssh = getSsh();
+    const s = getState();
-    if (ssh) {
+    if (s?.ownsmaster) {
-      await stopControlMaster(ssh.remote, ssh.socketPath);
+      await stopControlMaster(s.remote, s.socketPath);
-      resolvedSsh = null;
     }
+    state = null;
   });
 
-  // ── User ! commands via SSH ──────────────────────────────────────────────
+  // ── User ! commands via SSH ────────────────────────────────────────────────
 
-  pi.on("user_bash", (_event) => {
+  pi.on("user_bash", () => {
-    const ssh = getSsh();
+    const s = getState();
-    if (!ssh) return;
+    if (!s) return;
-    return {
+    return { operations: createRemoteBashOps(s.remote, s.socketPath, s.remoteCwd, localCwd) };
-      operations: createRemoteBashOps(ssh.remote, ssh.socketPath, ssh.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();
+    const s = getState();
-    if (ssh) {
+    if (!s) return;
     const modified = event.systemPrompt.replace(
       `Current working directory: ${localCwd}`,
-        `Current working directory: ${ssh.remoteCwd} (via SSH ControlMaster: ${ssh.remote})`,
+      `Current working directory: ${s.remoteCwd} (via SSH ControlMaster: ${s.remote})`,
     );
     return { systemPrompt: modified };
-    }
   });
 }