pi-extensions/extensions/ssh-controlmaster.ts

/**
 * 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.
 *
 * 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 --ssh user@host
 *   pi --ssh user@host:/remote/path
 *
 * Requirements:
 *   - SSH key-based auth (preferred) or password auth (see --ssh-ask-pass)
 *   - bash on the remote machine
 *
 * Password authentication:
 *   pi --ssh user@host --ssh-ask-pass
 *
 *   Prompts for the password via pi's input dialog before connecting.
 *   The password is passed to SSH via SSH_ASKPASS (a temp script written to
 *   /tmp, chmod 700, deleted immediately after the master is established).
 *   Note: input is NOT masked — the password is visible while typing.
 */

import { spawn } from "node:child_process";
import { writeFile, unlink } from "node:fs/promises";
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";

// ── Types ────────────────────────────────────────────────────────────────────

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`);
}

function askpassScriptPath(): string {
  return join(tmpdir(), `pi-askpass-${process.pid}.sh`);
}

/**
 * Write a temporary SSH_ASKPASS script that prints the given password, then
 * start the ControlMaster with that script set as the askpass helper.
 * The script is deleted immediately after the master process exits.
 */
async function startControlMasterWithPassword(
  remote: string,
  socketPath: string,
  password: string,
): Promise<void> {
  const scriptPath = askpassScriptPath();
  // Write script before spawning so SSH finds it when it calls askpass
  await writeFile(scriptPath, `#!/bin/sh\necho ${JSON.stringify(password)}\n`, { mode: 0o700 });
  try {
    await new Promise<void>((resolve, reject) => {
      const child = spawn(
        "ssh",
        ["-fN", "-o", "ControlMaster=yes", "-o", `ControlPath=${socketPath}`, "-o", "ControlPersist=yes", remote],
        {
          stdio: "ignore",
          env: {
            ...process.env,
            SSH_ASKPASS: scriptPath,
            SSH_ASKPASS_REQUIRE: "force", // OpenSSH 8.4+ — no DISPLAY needed
            DISPLAY: "dummy",             // fallback for older SSH
          },
        },
      );
      child.on("error", reject);
      child.on("close", (code) => {
        if (code === 0) resolve();
        else reject(new Error(`ControlMaster exited with code ${code}`));
      });
    });
  } finally {
    await unlink(scriptPath).catch(() => {}); // best-effort delete
  }
}

/** 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 };
}

async function negotiateMasterWithPassword(
  remote: string,
  password: string,
): Promise<{ socketPath: string; ownsmaster: boolean }> {
  const cfg = await readSshConfig(remote);
  const systemHasMaster = cfg.master === "auto" || cfg.master === "yes";

  if (systemHasMaster && cfg.path) {
    // System master handles auth on its own — password flag is ignored
    return { socketPath: cfg.path, ownsmaster: false };
  }

  const socketPath = ownSocketPath();
  await startControlMasterWithPassword(remote, socketPath, password);
  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 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) => {
      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) => {
    const child = spawn(
      "ssh",
      ["-O", "exit", "-o", `ControlPath=${socketPath}`, remote],
      { stdio: "ignore" },
    );
    child.on("close", () => resolve()); // best-effort; ignore errors
  });
}

// ── 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],
      { 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) reject(new Error(`SSH failed (${code}): ${Buffer.concat(err).toString()}`));
      else resolve(Buffer.concat(out));
    });
  });
}

// ── Remote tool operations ───────────────────────────────────────────────────

function createRemoteReadOps(
  remote: string, socketPath: string, remoteCwd: string, localCwd: string,
): ReadOperations {
  const r = (p: string) => p.replace(localCwd, remoteCwd);
  return {
    readFile: (p) => sshExec(remote, socketPath, `cat ${JSON.stringify(r(p))}`),
    access: (p) =>
      sshExec(remote, socketPath, `test -r ${JSON.stringify(r(p))}`).then(() => {}),
    detectImageMimeType: async (p) => {
      try {
        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; }
    },
  };
}

function createRemoteWriteOps(
  remote: string, socketPath: string, remoteCwd: string, localCwd: string,
): WriteOperations {
  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(r(p))}`);
    },
    mkdir: (dir) =>
      sshExec(remote, socketPath, `mkdir -p ${JSON.stringify(r(dir))}`).then(() => {}),
  };
}

function createRemoteEditOps(
  remote: string, socketPath: string, remoteCwd: string, localCwd: string,
): EditOperations {
  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,
): BashOperations {
  const r = (p: string) => p.replace(localCwd, remoteCwd);
  return {
    exec: (command, cwd, { onData, signal, timeout }) =>
      new Promise((resolve, reject) => {
        const cmd = `cd ${JSON.stringify(r(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",
  });

  pi.registerFlag("ssh-ask-pass", {
    description: "Prompt for SSH password instead of using key-based auth",
    type: "boolean",
    default: false,
  });

  const localCwd = process.cwd();
  const localRead = createReadTool(localCwd);
  const localWrite = createWriteTool(localCwd);
  const localEdit = createEditTool(localCwd);
  const localBash = createBashTool(localCwd);

  let state: SshState | null = null;
  const getState = () => state;

  // ── Tool overrides ─────────────────────────────────────────────────────────

  pi.registerTool({
    ...localRead,
    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) {
      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) {
      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) {
      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 ──────────────────────────────────────────────────────

  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;
      remoteCwd = await run(["ssh", remote, "pwd"]).catch((e) => {
        throw new Error(`Could not resolve remote pwd: ${e.message}`);
      });
    }

    ctx.ui.setStatus("ssh", ctx.ui.theme.fg("accent", `SSH: ${remote} ⟳ connecting…`));

    const askPass = pi.getFlag("ssh-ask-pass") as boolean;
    let password: string | undefined;

    if (askPass) {
      const input = await ctx.ui.input("SSH password:", "");
      if (input === null) {
        ctx.ui.setStatus("ssh", "");
        ctx.ui.notify("SSH connection cancelled", "info");
        return;
      }
      password = input;
    }

    let socketPath: string;
    let ownsmaster: boolean;

    try {
      ({ socketPath, ownsmaster } =
        password !== undefined
          ? await negotiateMasterWithPassword(remote, password)
          : 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;
    }

    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 () => {
    const s = getState();
    if (s?.ownsmaster) {
      await stopControlMaster(s.remote, s.socketPath);
    }
    state = null;
  });

  // ── User ! commands via SSH ────────────────────────────────────────────────

  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 ─────────────────────────────────────

  pi.on("before_agent_start", async (event) => {
    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 };
  });
}