Joakim Persson
380 lines
12 KiB
TypeScript
380 lines
12 KiB
TypeScript
/**
|
|
* SSH ControlMaster Remote Execution
|
|
*
|
|
* 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.
|
|
*
|
|
* 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
|
|
*
|
|
* Requirements:
|
|
* - SSH key-based auth (no password prompts)
|
|
* - bash on the remote machine
|
|
*/
|
|
|
|
import { spawn } from "node:child_process";
|
|
import { tmpdir } from "node:os";
|
|
import { join } from "node:path";
|
|
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
|
|
import {
|
|
type BashOperations,
|
|
createBashTool,
|
|
createEditTool,
|
|
createReadTool,
|
|
createWriteTool,
|
|
type EditOperations,
|
|
type ReadOperations,
|
|
type WriteOperations,
|
|
} from "@mariozechner/pi-coding-agent";
|
|
|
|
// ── ControlMaster helpers ────────────────────────────────────────────────────
|
|
|
|
function controlSocketPath(): string {
|
|
// Keep path short — macOS has a ~104-char Unix socket path limit
|
|
return join(tmpdir(), `pi-cm-${process.pid}.sock`);
|
|
}
|
|
|
|
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
|
|
const child = spawn(
|
|
"ssh",
|
|
[
|
|
"-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}`));
|
|
});
|
|
});
|
|
}
|
|
|
|
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
|
|
});
|
|
}
|
|
|
|
// ── SSH exec (multiplexed over ControlMaster socket) ────────────────────────
|
|
|
|
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,
|
|
],
|
|
{ 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));
|
|
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));
|
|
}
|
|
});
|
|
});
|
|
}
|
|
|
|
// ── Remote tool operations ───────────────────────────────────────────────────
|
|
|
|
function createRemoteReadOps(
|
|
remote: string,
|
|
socketPath: string,
|
|
remoteCwd: string,
|
|
localCwd: string,
|
|
): ReadOperations {
|
|
const toRemote = (p: string) => p.replace(localCwd, remoteCwd);
|
|
return {
|
|
readFile: (p) => sshExec(remote, socketPath, `cat ${JSON.stringify(toRemote(p))}`),
|
|
access: (p) =>
|
|
sshExec(remote, socketPath, `test -r ${JSON.stringify(toRemote(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();
|
|
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;
|
|
}
|
|
},
|
|
};
|
|
}
|
|
|
|
function createRemoteWriteOps(
|
|
remote: string,
|
|
socketPath: string,
|
|
remoteCwd: string,
|
|
localCwd: string,
|
|
): WriteOperations {
|
|
const toRemote = (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))}`,
|
|
);
|
|
},
|
|
mkdir: (dir) =>
|
|
sshExec(remote, socketPath, `mkdir -p ${JSON.stringify(toRemote(dir))}`).then(() => {}),
|
|
};
|
|
}
|
|
|
|
function createRemoteEditOps(
|
|
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 };
|
|
}
|
|
|
|
function createRemoteBashOps(
|
|
remote: string,
|
|
socketPath: string,
|
|
remoteCwd: string,
|
|
localCwd: string,
|
|
): BashOperations {
|
|
const toRemote = (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 child = spawn(
|
|
"ssh",
|
|
[
|
|
"-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)
|
|
: undefined;
|
|
|
|
child.stdout.on("data", onData);
|
|
child.stderr.on("data", onData);
|
|
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);
|
|
if (signal?.aborted) reject(new Error("aborted"));
|
|
else if (timedOut) reject(new Error(`timeout:${timeout}`));
|
|
else resolve({ exitCode: code });
|
|
});
|
|
}),
|
|
};
|
|
}
|
|
|
|
// ── Extension ────────────────────────────────────────────────────────────────
|
|
|
|
export default function (pi: ExtensionAPI) {
|
|
pi.registerFlag("ssh", {
|
|
description: "SSH remote: user@host or user@host:/remote/path",
|
|
type: "string",
|
|
});
|
|
|
|
const localCwd = process.cwd();
|
|
const localRead = createReadTool(localCwd);
|
|
const localWrite = createWriteTool(localCwd);
|
|
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;
|
|
|
|
const getSsh = () => resolvedSsh;
|
|
|
|
// ── 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);
|
|
},
|
|
});
|
|
|
|
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);
|
|
},
|
|
});
|
|
|
|
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);
|
|
},
|
|
});
|
|
|
|
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);
|
|
},
|
|
});
|
|
|
|
// ── Session lifecycle ────────────────────────────────────────────────────
|
|
|
|
pi.on("session_start", async (_event, ctx) => {
|
|
const arg = pi.getFlag("ssh") as string | undefined;
|
|
if (!arg) return;
|
|
|
|
let remote: string;
|
|
let remoteCwd: string;
|
|
|
|
if (arg.includes(":")) {
|
|
[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"],
|
|
});
|
|
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…`));
|
|
|
|
try {
|
|
await startControlMaster(remote, socketPath);
|
|
} 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");
|
|
});
|
|
|
|
pi.on("session_shutdown", async (_event, _ctx) => {
|
|
const ssh = getSsh();
|
|
if (ssh) {
|
|
await stopControlMaster(ssh.remote, ssh.socketPath);
|
|
resolvedSsh = null;
|
|
}
|
|
});
|
|
|
|
// ── 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),
|
|
};
|
|
});
|
|
|
|
// ── 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 };
|
|
}
|
|
});
|
|
}
|