my-pi/pi-ssh/index.ts

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;
    // Contribute structurally via a named section: the bundle's single
    // system-prompt owner (extensions/system-prompt.ts) composes the final
    // prompt; extensions must not return `systemPrompt` themselves.
    const guidance =
      `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.`;
    const promptOptions = event.systemPromptOptions as unknown as
      | { sections?: Record<string, string> }
      | undefined;
    if (!promptOptions) return; // hosts without structured options just skip the section
    (promptOptions.sections ??= {}).remote_ssh = guidance;
  });
}