Something went wrong. Try again.
more up to date at codeberg https://codeberg.org/cowie/tau.git
Something went wrong. Try again.
7.3 kB · 220 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221/** * SSH Remote Execution Example * * Demonstrates delegating tool operations to a remote machine via SSH. * When --ssh is provided, read/write/edit/bash run on the remote. * * Usage: * pi -e ./ssh.ts --ssh user@host * pi -e ./ssh.ts --ssh user@host:/remote/path * * Requirements: * - SSH key-based auth (no password prompts) * - bash on remote */
import { spawn } from "node:child_process";import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";import { type BashOperations, createBashTool, createEditTool, createReadTool, createWriteTool, type EditOperations, type ReadOperations, type WriteOperations,} from "@earendil-works/pi-coding-agent";
function sshExec(remote: string, command: string): Promise<Buffer> { return new Promise((resolve, reject) => { const child = spawn("ssh", [remote, command], { stdio: ["ignore", "pipe", "pipe"] }); const chunks: Buffer[] = []; const errChunks: Buffer[] = []; child.stdout.on("data", (data) => chunks.push(data)); child.stderr.on("data", (data) => errChunks.push(data)); 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)); } }); });}
function createRemoteReadOps(remote: string, remoteCwd: string, localCwd: string): ReadOperations { const toRemote = (p: string) => p.replace(localCwd, remoteCwd); return { readFile: (p) => sshExec(remote, `cat ${JSON.stringify(toRemote(p))}`), access: (p) => sshExec(remote, `test -r ${JSON.stringify(toRemote(p))}`).then(() => {}), detectImageMimeType: async (p) => { try { const r = await sshExec(remote, `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 : null; } catch { return null; } }, };}
function createRemoteWriteOps(remote: 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, `echo ${JSON.stringify(b64)} | base64 -d > ${JSON.stringify(toRemote(p))}`); }, mkdir: (dir) => sshExec(remote, `mkdir -p ${JSON.stringify(toRemote(dir))}`).then(() => {}), };}
function createRemoteEditOps(remote: string, remoteCwd: string, localCwd: string): EditOperations { const r = createRemoteReadOps(remote, remoteCwd, localCwd); const w = createRemoteWriteOps(remote, remoteCwd, localCwd); return { readFile: r.readFile, access: r.access, writeFile: w.writeFile };}
function createRemoteBashOps(remote: 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", [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) => { if (timer) clearTimeout(timer); reject(e); }); const onAbort = () => child.kill(); signal?.addEventListener("abort", onAbort, { once: true }); child.on("close", (code) => { 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 }); }); }), };}
export default function (pi: ExtensionAPI) { pi.registerFlag("ssh", { description: "SSH remote: user@host or user@host:/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 on session_start (CLI flags not available during factory) let resolvedSsh: { remote: string; remoteCwd: string } | null = null;
const getSsh = () => resolvedSsh;
pi.registerTool({ ...localRead, async execute(id, params, signal, onUpdate, _ctx) { const ssh = getSsh(); if (ssh) { const tool = createReadTool(localCwd, { operations: createRemoteReadOps(ssh.remote, 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.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.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.remoteCwd, localCwd), }); return tool.execute(id, params, signal, onUpdate); } return localBash.execute(id, params, signal, onUpdate); }, });
pi.on("session_start", async (_event, ctx) => { // Resolve SSH config now that CLI flags are available const arg = pi.getFlag("ssh") as string | undefined; if (arg) { if (arg.includes(":")) { const [remote, path] = arg.split(":"); resolvedSsh = { remote, remoteCwd: path }; } else { // No path given, evaluate pwd on remote const remote = arg; const pwd = (await sshExec(remote, "pwd")).toString().trim(); resolvedSsh = { remote, remoteCwd: pwd }; } ctx.ui.setStatus("ssh", ctx.ui.theme.fg("accent", `SSH: ${resolvedSsh.remote}:${resolvedSsh.remoteCwd}`)); ctx.ui.notify(`SSH mode: ${resolvedSsh.remote}:${resolvedSsh.remoteCwd}`, "info"); } });
// Handle user ! commands via SSH pi.on("user_bash", (_event) => { const ssh = getSsh(); if (!ssh) return; // No SSH, use local execution return { operations: createRemoteBashOps(ssh.remote, ssh.remoteCwd, localCwd) }; });
// Replace local cwd with remote cwd in system prompt 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: ${ssh.remote})`, ); return { systemPrompt: modified }; } });}