Files

370 lines
14 KiB
TypeScript

import { homedir } from "node:os";
import { posix as posixPath } from "node:path";
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import {
createBashTool,
createEditTool,
createReadTool,
createWriteTool,
type EditOperations,
type ReadOperations,
type WriteOperations,
} from "@earendil-works/pi-coding-agent";
import { getPermissionsService, PERMISSIONS_READY_CHANNEL } from "@gotgenes/pi-permission-system";
import {
installSshPermissionIntegration,
type SshPermissionConnection,
} from "./permission-integration.ts";
import {
getConfiguredHost,
parseConnectInput,
SSH_CONNECT_TOOL_METADATA,
type HostSelection,
} from "./src/agent-connection.ts";
import { loadVault } from "./src/vault.ts";
import { Ssh2Transport, type RemoteTransport } from "./src/ssh2-transport.ts";
import type { SshHostConfig } from "./src/config.ts";
import {
runRemoteFind,
runRemoteGrep,
type RemoteFindInput,
type RemoteGrepInput,
} from "./src/remote-search.ts";
import { createRemoteBashOps } from "./src/remote-bash.ts";
import {
changeRemoteCwd,
mapLocalPathToRemote,
SSH_CD_EXECUTION_MODE,
} from "./src/remote-cwd.ts";
import { probeRemotePath } from "./src/remote-probe.ts";
interface SshConnection {
hostId: string;
remote: string;
port: number;
remoteCwd: string;
remoteHome: string;
localCwd: string;
localHome: string;
}
function createRemoteReadOps(connection: SshConnection, transport: RemoteTransport): ReadOperations {
return {
readFile: (absolutePath) => transport.readFile(mapLocalPathToRemote(absolutePath, connection)),
access: (absolutePath) => transport.ensureReadable(mapLocalPathToRemote(absolutePath, connection)),
detectImageMimeType: async (absolutePath) => {
try {
return await transport.detectImageMimeType(mapLocalPathToRemote(absolutePath, connection));
} catch {
return null;
}
},
};
}
function createRemoteWriteOps(connection: SshConnection, transport: RemoteTransport): WriteOperations {
return {
mkdir: (absoluteDir) => transport.mkdir(mapLocalPathToRemote(absoluteDir, connection)),
writeFile: (absolutePath, content) =>
transport.writeFile(mapLocalPathToRemote(absolutePath, connection), Buffer.from(content, "utf8")),
};
}
function createRemoteEditOps(connection: SshConnection, transport: RemoteTransport): EditOperations {
const read = createRemoteReadOps(connection, transport);
const write = createRemoteWriteOps(connection, transport);
return {
readFile: read.readFile,
writeFile: write.writeFile,
access: (absolutePath) => transport.ensureReadableWritable(mapLocalPathToRemote(absolutePath, connection)),
};
}
function resolveRequestedPath(selection: HostSelection, host: SshHostConfig, remoteHome: string, remotePwd: string): string {
const requested = selection.remotePath ?? host.defaultCwd ?? remotePwd;
if (requested === "~") return remoteHome;
if (requested.startsWith("~/")) return posixPath.join(remoteHome, requested.slice(2));
return requested;
}
async function connectSelection(
selection: HostSelection,
localCwd: string,
localHome: string,
signal?: AbortSignal,
): Promise<{ connection: SshConnection; transport: Ssh2Transport }> {
const config = loadVault();
const host = getConfiguredHost(config, selection.hostId);
const transport = new Ssh2Transport(host);
try {
await transport.connect(signal);
const remoteHome = await probeRemotePath(transport, "home", ".", signal);
const remotePwd = await probeRemotePath(transport, "cwd", ".", signal);
const requestedPath = resolveRequestedPath(selection, host, remoteHome, remotePwd);
const remoteCwd = await probeRemotePath(transport, "cwd", requestedPath, signal);
return {
connection: {
hostId: selection.hostId,
remote: `${selection.hostId} [${host.user}@${host.hostName}]`,
port: host.port,
remoteCwd,
remoteHome,
localCwd,
localHome,
},
transport,
};
} catch (error) {
await transport.dispose();
throw error;
}
}
function describeRequestedTarget(input: Record<string, unknown>): SshPermissionConnection | null {
try {
const selection = parseConnectInput(input);
const host = getConfiguredHost(loadVault(), selection.hostId);
return {
remote: `${selection.hostId} [${host.user}@${host.hostName}]`,
port: host.port,
remoteCwd: selection.remotePath ?? host.defaultCwd ?? "<server default>",
};
} catch {
return null;
}
}
export default function piSshExtension(pi: ExtensionAPI): void {
const localCwd = process.cwd();
const localHome = homedir();
const localRead = createReadTool(localCwd);
const localWrite = createWriteTool(localCwd);
const localEdit = createEditTool(localCwd);
const localBash = createBashTool(localCwd, { exposeSessionEnvironment: false });
let connection: SshConnection | null = null;
let transport: Ssh2Transport | null = null;
const getConnection = () => connection;
installSshPermissionIntegration(pi, getConnection, {
getPermissionsService,
permissionsReadyChannel: PERMISSIONS_READY_CHANNEL,
getConnectTarget: describeRequestedTarget,
});
const requireSsh = (toolName: string): { connection: SshConnection; transport: Ssh2Transport } => {
if (!connection || !transport) {
throw new Error(`${toolName} requires an active SSH2 connection. Call ssh_connect with an imported host ID first.`);
}
return { connection, transport };
};
const activateConnection = async (
nextConnection: SshConnection,
nextTransport: Ssh2Transport,
): Promise<void> => {
if (transport) await transport.dispose();
connection = nextConnection;
transport = nextTransport;
console.log(`pi-ssh connected: ${nextConnection.remote}:${nextConnection.remoteCwd} (port ${nextConnection.port})`);
};
pi.registerTool({
...SSH_CONNECT_TOOL_METADATA,
async execute(_id, params, signal) {
const selection = parseConnectInput(params as Record<string, unknown>);
const connected = await connectSelection(selection, localCwd, localHome, signal);
if (signal?.aborted) {
await connected.transport.dispose();
throw new Error("SSH connection aborted");
}
await activateConnection(connected.connection, connected.transport);
const text = `Connected to ${connected.connection.remote}:${connected.connection.remoteCwd} (port ${connected.connection.port}).`;
return {
content: [{ type: "text", text }],
details: {
hostId: connected.connection.hostId,
remote: connected.connection.remote,
port: connected.connection.port,
remoteCwd: connected.connection.remoteCwd,
},
};
},
});
pi.registerTool({
name: "ssh_cd",
label: "ssh_cd",
description: "Change the active SSH workspace directory as a reviewed, persistent state transition. Call ssh_cd separately and wait for it to succeed before issuing ssh_bash or relative remote file/search operations that depend on the new directory.",
parameters: {
type: "object",
properties: {
path: {
type: "string",
minLength: 1,
description: "Remote directory: absolute, relative to the active remote cwd, or ~/ relative to remote HOME",
},
},
required: ["path"],
additionalProperties: false,
},
executionMode: SSH_CD_EXECUTION_MODE,
async execute(_id, params, signal) {
const active = requireSsh("ssh_cd");
const changed = await changeRemoteCwd(
active.connection,
(params as { path: string }).path,
(requestedCwd) => probeRemotePath(active.transport, "cwd", requestedCwd, signal),
);
return {
content: [{ type: "text", text: `Remote cwd changed from ${changed.previousCwd} to ${changed.remoteCwd}. Dependent remote tools may now run.` }],
details: changed,
};
},
});
pi.registerTool({
...localRead,
name: "ssh_read",
label: "ssh_read",
description: `Read a file through the active SSH2 connection. ${localRead.description}`,
async execute(id, params, signal, onUpdate) {
const active = requireSsh("ssh_read");
return createReadTool(localCwd, { operations: createRemoteReadOps(active.connection, active.transport) })
.execute(id, params, signal, onUpdate);
},
});
pi.registerTool({
...localWrite,
name: "ssh_write",
label: "ssh_write",
description: `Write a file through the active SSH2 connection. ${localWrite.description}`,
async execute(id, params, signal, onUpdate) {
const active = requireSsh("ssh_write");
return createWriteTool(localCwd, { operations: createRemoteWriteOps(active.connection, active.transport) })
.execute(id, params, signal, onUpdate);
},
});
pi.registerTool({
...localEdit,
name: "ssh_edit",
label: "ssh_edit",
description: `Edit a file through the active SSH2 connection. ${localEdit.description}`,
async execute(id, params, signal, onUpdate) {
const active = requireSsh("ssh_edit");
return createEditTool(localCwd, { operations: createRemoteEditOps(active.connection, active.transport) })
.execute(id, params, signal, onUpdate);
},
});
pi.registerTool({
name: "ssh_find",
label: "ssh_find",
description: "Find remote files by a fixed filename/path substring using fd, git ls-files, or find. Results are bounded and require an active SSH2 connection.",
parameters: {
type: "object",
properties: {
pattern: { type: "string", description: "Fixed substring to match in remote file paths" },
path: { type: "string", description: "Remote root path; defaults to the active remote cwd" },
caseSensitive: { type: "boolean", description: "Use case-sensitive matching (default: false)" },
includeHidden: { type: "boolean", description: "Include hidden paths (default: false)" },
limit: { type: "integer", minimum: 1, maximum: 200, description: "Maximum results (default: 50)" },
},
required: ["pattern"],
additionalProperties: false,
},
async execute(_id, params, signal) {
const active = requireSsh("ssh_find");
const result = await runRemoteFind(
active.transport,
params as RemoteFindInput,
active.connection.remoteCwd,
active.connection.remoteHome,
signal,
);
return {
content: [{ type: "text", text: result.text }],
details: { backend: result.backend, matchCount: result.matchCount, truncated: result.truncated },
};
},
});
pi.registerTool({
name: "ssh_grep",
label: "ssh_grep",
description: "Search remote file contents using ripgrep, git grep, or grep with consistent hidden-file, basename-glob, and portable POSIX ERE semantics. Literal case-insensitive matching is the default; results are bounded and require an active SSH2 connection.",
parameters: {
type: "object",
properties: {
pattern: { type: "string", description: "Text to search for, or a portable POSIX ERE when literal is false" },
path: { type: "string", description: "Remote root path; defaults to the active remote cwd" },
literal: { type: "boolean", description: "Treat pattern as fixed text; false uses portable POSIX ERE syntax (default: true)" },
caseSensitive: { type: "boolean", description: "Use case-sensitive matching (default: false)" },
include: { type: "string", description: "Optional basename-only glob such as *.ts; / is not allowed" },
includeHidden: { type: "boolean", description: "Search hidden files and directories (default: false)" },
limit: { type: "integer", minimum: 1, maximum: 200, description: "Maximum result lines (default: 50)" },
},
required: ["pattern"],
additionalProperties: false,
},
async execute(_id, params, signal) {
const active = requireSsh("ssh_grep");
const result = await runRemoteGrep(
active.transport,
params as RemoteGrepInput,
active.connection.remoteCwd,
active.connection.remoteHome,
signal,
);
return {
content: [{ type: "text", text: result.text }],
details: { backend: result.backend, matchCount: result.matchCount, truncated: result.truncated },
};
},
});
pi.registerTool({
...localBash,
name: "ssh_bash",
label: "ssh_bash",
description: "Run a non-search shell command on the active SSH server. Every call starts a fresh non-interactive Bash process in the active remote cwd; cd, exported variables, aliases, and other shell state do not persist. For a persistent workspace change, call ssh_cd separately and wait for success before calling ssh_bash. Use command-local cd only for an intentionally temporary directory change. Relative paths are remote. Use ssh_find or ssh_grep instead of find, fd, grep, or rg.",
promptSnippet: undefined,
promptGuidelines: undefined,
async execute(id, params, signal, onUpdate) {
const active = requireSsh("ssh_bash");
return createBashTool(active.connection.remoteCwd, {
operations: createRemoteBashOps(active.connection, active.transport),
exposeSessionEnvironment: false,
}).execute(id, params, signal, onUpdate);
},
});
pi.on("session_shutdown", async () => {
if (transport) await transport.dispose();
transport = null;
connection = null;
});
pi.on("before_agent_start", async (event) => {
if (!connection) return;
const guidance =
`\n\n# Remote SSH connection\n\n` +
`An SSH2 connection to configured host ${connection.remote} on port ${connection.port} is active. ` +
`The default read/write/edit/bash/find/grep tools act on the LOCAL machine. ` +
`Use ssh_cd, ssh_read, ssh_write, ssh_edit, ssh_find, ssh_grep, and ssh_bash for explicit remote operations. ` +
`Use ssh_find before ssh_grep to narrow remote searches; both tools return bounded results and choose the fastest available remote backend. ` +
`Remote operations are rooted at ${connection.remoteCwd}; relative paths resolve against that directory. ` +
`For a persistent workspace change, call ssh_cd as a separate step and wait for its successful result before issuing dependent remote tool calls. ` +
`Each ssh_bash call starts a fresh non-interactive shell, so use command-local cd only for an intentionally temporary change; cd, exports, aliases, and other shell state inside one command do not persist to the next call.`;
return { systemPrompt: `${event.systemPrompt}${guidance}` };
});
}