feat(pi-ssh): add reviewed agent connection flow

pi-ssh/index.ts

@@ -1,6 +1,6 @@
 import { homedir } from "node:os";
 import { posix as posixPath } from "node:path";
-import type { CustomEntry, ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
 import {
   createBashTool,
   createEditTool,
@@ -12,7 +12,15 @@ import {
   type WriteOperations,
 } from "@earendil-works/pi-coding-agent";
 import { getPermissionsService, PERMISSIONS_READY_CHANNEL } from "@gotgenes/pi-permission-system";
-import { installSshPermissionIntegration } from "./permission-integration.ts";
+import {
+  installSshPermissionIntegration,
+  type SshPermissionConnection,
+} from "./permission-integration.ts";
+import {
+  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 { PiSshConfig, SshHostConfig } from "./src/config.ts";
@@ -23,11 +31,6 @@ import {
   type RemoteGrepInput,
 } from "./src/remote-search.ts";
 
-interface SshStoredConfig {
-  hostId: string;
-  remoteCwd: string;
-  remoteHome: string;
-}
 
 interface SshConnection {
   hostId: string;
@@ -39,24 +42,6 @@ interface SshConnection {
   localHome: string;
 }
 
-interface HostSelection {
-  hostId: string;
-  remotePath?: string;
-}
-
-function parseHostSelection(raw: string): HostSelection {
-  const value = raw.trim();
-  if (!value) throw new Error("SSH host id is required");
-  const colon = value.indexOf(":");
-  const hostId = (colon < 0 ? value : value.slice(0, colon)).trim();
-  if (!/^[A-Za-z0-9][A-Za-z0-9._-]*$/.test(hostId)) throw new Error(`invalid pi-ssh host id: ${hostId}`);
-  if (colon < 0) return { hostId };
-  const remotePath = value.slice(colon + 1).trim();
-  if (!remotePath || !(remotePath === "~" || remotePath.startsWith("~/") || remotePath.startsWith("/"))) {
-    throw new Error("remote path must be absolute or start with ~/");
-  }
-  return { hostId, remotePath };
-}
 
 function mapLocalPathToRemote(path: string, connection: SshConnection): string {
   if (path === connection.localCwd) return connection.remoteCwd;
@@ -163,22 +148,21 @@ async function connectSelection(
   }
 }
 
-function pickerEntries(config: PiSshConfig): Array<{ value: string; hostId: string }> {
-  const groupByHost = new Map<string, string>();
-  for (const group of Object.values(config.groups ?? {})) {
-    for (const hostId of group.hosts) if (!groupByHost.has(hostId)) groupByHost.set(hostId, group.label);
+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;
   }
-  return Object.entries(config.hosts).map(([hostId, host]) => ({
-    hostId,
-    value: `${groupByHost.get(hostId) ? `${groupByHost.get(hostId)} / ` : ""}${host.label ?? hostId} [${hostId}]`,
-  }));
 }
 
 export default function piSshExtension(pi: ExtensionAPI): void {
-  pi.registerFlag("ssh", {
-    description: "Configured pi-ssh host id, optionally followed by :/absolute/remote/path",
-    type: "string",
-  });
 
   const localCwd = process.cwd();
   const localHome = homedir();
@@ -193,11 +177,12 @@ export default function piSshExtension(pi: ExtensionAPI): void {
   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. Run /ssh and select an imported host.`);
+      throw new Error(`${toolName} requires an active SSH2 connection. Call ssh_connect with an imported host ID first.`);
     }
     return { connection, transport };
   };
@@ -205,33 +190,31 @@ export default function piSshExtension(pi: ExtensionAPI): void {
   const activateConnection = async (
     nextConnection: SshConnection,
     nextTransport: Ssh2Transport,
-    ctx: ExtensionContext,
-    options: { persist: boolean; verb: string },
   ): Promise<void> => {
     if (transport) await transport.dispose();
     connection = nextConnection;
     transport = nextTransport;
-    if (options.persist) {
-      pi.appendEntry("pi-ssh-config", {
-        hostId: nextConnection.hostId,
-        remoteCwd: nextConnection.remoteCwd,
-        remoteHome: nextConnection.remoteHome,
-      } satisfies SshStoredConfig);
-    }
-    const message = `pi-ssh ${options.verb}: ${nextConnection.remote}:${nextConnection.remoteCwd} (port ${nextConnection.port})`;
-    console.log(message);
-    if (ctx.hasUI) {
-      ctx.ui.setStatus("pi-ssh", ctx.ui.theme.fg("accent", `SSH ${nextConnection.hostId}:${nextConnection.remoteCwd}`));
-      ctx.ui.notify(message, "info");
-    }
+    console.log(`pi-ssh connected: ${nextConnection.remote}:${nextConnection.remoteCwd} (port ${nextConnection.port})`);
   };
 
-  const deactivateConnection = async (ctx: ExtensionContext): Promise<void> => {
-    if (transport) await transport.dispose();
-    transport = null;
-    connection = null;
-    if (ctx.hasUI) ctx.ui.setStatus("pi-ssh", undefined);
-  };
+  pi.registerTool({
+    ...SSH_CONNECT_TOOL_METADATA,
+    async execute(_id, params) {
+      const selection = parseConnectInput(params as Record<string, unknown>);
+      const connected = await connectSelection(selection, localCwd, localHome);
+      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({
     ...localRead,
@@ -346,115 +329,12 @@ export default function piSshExtension(pi: ExtensionAPI): void {
     },
   });
 
-  pi.on("session_start", async (event, ctx) => {
-    const flag = pi.getFlag("ssh") as string | undefined;
-    if (flag) {
-      try {
-        const connected = await connectSelection(parseHostSelection(flag), localCwd, localHome);
-        await activateConnection(connected.connection, connected.transport, ctx, { persist: true, verb: "enabled" });
-        return;
-      } catch (error) {
-        const message = error instanceof Error ? error.message : String(error);
-        await deactivateConnection(ctx);
-        console.error(`pi-ssh failed to connect: ${message}`);
-        if (ctx.hasUI) ctx.ui.notify(`pi-ssh failed to connect: ${message}`, "error");
-        throw error;
-      }
-    }
-
-    if (event.reason !== "startup" && event.reason !== "resume") return;
-    const entries = ctx.sessionManager.getEntries();
-    let stored: SshStoredConfig | undefined;
-    for (let index = entries.length - 1; index >= 0; index -= 1) {
-      const entry = entries[index];
-      if (entry.type === "custom" && (entry as CustomEntry<unknown>).customType === "pi-ssh-config") {
-        stored = (entry as CustomEntry<SshStoredConfig>).data;
-        if (stored) break;
-      }
-    }
-    if (!stored) return;
-    try {
-      const connected = await connectSelection(
-        { hostId: stored.hostId, remotePath: stored.remoteCwd },
-        localCwd,
-        localHome,
-      );
-      await activateConnection(connected.connection, connected.transport, ctx, { persist: false, verb: "resumed" });
-    } catch (error) {
-      const message = error instanceof Error ? error.message : String(error);
-      await deactivateConnection(ctx);
-      console.error(`pi-ssh resume failed: ${message}`);
-      if (ctx.hasUI) ctx.ui.notify(`pi-ssh resume failed: ${message}`, "warning");
-    }
-  });
-
-  pi.registerCommand("ssh", {
-    description: "Connect configured SSH2 hosts: /ssh [host-id[:/path]], /ssh status, /ssh off",
-    getArgumentCompletions: (prefix) => {
-      try {
-        const options = ["off", "status", ...Object.keys(loadVault().hosts)];
-        const filtered = options.filter((option) => option.startsWith(prefix));
-        return filtered.length > 0 ? filtered.map((option) => ({ value: option, label: option })) : null;
-      } catch {
-        return ["off", "status"].filter((option) => option.startsWith(prefix)).map((option) => ({ value: option, label: option }));
-      }
-    },
-    handler: async (args, ctx) => {
-      const input = args.trim();
-      if (input === "status") {
-        ctx.ui.notify(connection
-          ? `pi-ssh: ${connection.remote}:${connection.remoteCwd} (port ${connection.port})`
-          : "pi-ssh: not connected (local tools active)", "info");
-        return;
-      }
-      if (input === "off") {
-        await deactivateConnection(ctx);
-        ctx.ui.notify("pi-ssh: disconnected", "info");
-        return;
-      }
-
-      let target = input;
-      if (!target) {
-        let config: PiSshConfig;
-        try {
-          config = loadVault();
-        } catch (error) {
-          ctx.ui.notify(error instanceof Error ? error.message : String(error), "warning");
-          return;
-        }
-        const entries = pickerEntries(config);
-        if (entries.length === 0) {
-          ctx.ui.notify("No pi-ssh hosts configured. Run ssh_config.sh import.", "warning");
-          return;
-        }
-        const values = [...(connection ? ["Disconnect [off]"] : []), ...entries.map((entry) => entry.value)];
-        const selected = await ctx.ui.select("SSH2 host", values);
-        if (!selected) return;
-        if (selected === "Disconnect [off]") {
-          await deactivateConnection(ctx);
-          ctx.ui.notify("pi-ssh: disconnected", "info");
-          return;
-        }
-        target = entries.find((entry) => entry.value === selected)?.hostId ?? "";
-      }
-
-      try {
-        const connected = await connectSelection(parseHostSelection(target), localCwd, localHome);
-        await activateConnection(connected.connection, connected.transport, ctx, { persist: true, verb: "connected" });
-      } catch (error) {
-        const message = error instanceof Error ? error.message : String(error);
-        ctx.ui.notify(`pi-ssh: failed to connect: ${message}`, "error");
-      }
-    },
-  });
-
   pi.on("session_shutdown", async () => {
     if (transport) await transport.dispose();
     transport = null;
+    connection = null;
   });
 
-  pi.on("user_bash", () => transport ? { operations: createRemoteBashOps(transport) } : undefined);
-
   pi.on("before_agent_start", async (event) => {
     if (!connection) return;
     const guidance =
@@ -463,7 +343,6 @@ export default function piSshExtension(pi: ExtensionAPI): void {
       `The default read/write/edit/bash/find/grep tools act on the LOCAL machine. ` +
       `Use 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. ` +
-      `(User \`!\` commands run remotely.) ` +
       `Remote operations are rooted at ${connection.remoteCwd}; relative paths resolve against that directory.`;
     return { systemPrompt: `${event.systemPrompt}${guidance}` };
   });