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

This commit is contained in:
Joakim Persson
2026-05-05 22:58:34 +02:00
parent dee755e291
commit 5aa8036a6f
+179 -175
View File
@@ -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<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", // 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<B
// ── Remote tool operations ───────────────────────────────────────────────────
function createRemoteReadOps(
remote: string,
socketPath: string,
remoteCwd: string,
localCwd: string,
remote: 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(
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<string>((resolve, reject) => {
const child = spawn("ssh", [remote, "pwd"], {
stdio: ["ignore", "pipe", "pipe"],
remoteCwd = await run(["ssh", remote, "pwd"]).catch((e) => {
throw new Error(`Could not resolve remote pwd: ${e.message}`);
});
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 };
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 s = getState();
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 };
}
});
}