my-pi/pi-ssh/permission-integration.ts

import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import type { PermissionsService } from "@gotgenes/pi-permission-system";

export interface SshPermissionConnection {
  remote: string;
  port?: number;
  remoteCwd: string;
}

export interface SshPermissionIntegrationDependencies {
  getPermissionsService: () => PermissionsService | undefined;
  permissionsReadyChannel: string;
  getConnectTarget?: (input: Record<string, unknown>) => SshPermissionConnection | null;
  warn?: (message: string) => void;
}

type PermissionIntegrationApi = Pick<ExtensionAPI, "events" | "on">;
type ToolInput = Record<string, unknown>;

const REMOTE_FILE_TOOLS = ["ssh_read", "ssh_write", "ssh_edit", "ssh_find", "ssh_grep"] as const;
const REMOTE_TOOLS = ["ssh_connect", ...REMOTE_FILE_TOOLS, "ssh_bash"] as const;

function inline(value: string, limit = 240): string {
  const normalized = value.replace(/\s+/g, " ").trim();
  return normalized.length > limit ? `${normalized.slice(0, limit)}…` : normalized;
}

function stringField(input: ToolInput, key: string): string | undefined {
  const value = input[key];
  return typeof value === "string" && value.length > 0 ? value : undefined;
}

function countLines(value: string): number {
  return value.length === 0 ? 0 : value.split(/\r\n|\r|\n/).length;
}

function formatTarget(connection: SshPermissionConnection | null): string {
  if (connection === null) return "an inactive SSH connection";
  const port = connection.port === undefined ? "" : `:${connection.port}`;
  return `SSH target '${inline(connection.remote)}${port}' in remote cwd '${inline(connection.remoteCwd)}'`;
}

export function formatSshPermissionInput(
  toolName: string,
  input: ToolInput,
  connection: SshPermissionConnection | null,
): string {
  const target = formatTarget(connection);
  const path = stringField(input, "path");

  if (toolName === "ssh_connect") {
    if (connection !== null) return `${target}; establish a persistent SSH2 connection`;
    const hostId = inline(stringField(input, "hostId") ?? "<unspecified>");
    const remotePath = stringField(input, "remotePath");
    return `requested imported SSH host '${hostId}'${remotePath ? ` in remote cwd '${inline(remotePath)}'` : ""}; establish a persistent SSH2 connection`;
  }

  if (toolName === "ssh_read") {
    const details = path ? [`remote path '${inline(path)}'`] : ["an unspecified remote path"];
    if (typeof input.offset === "number") details.push(`offset ${input.offset}`);
    if (typeof input.limit === "number") details.push(`limit ${input.limit}`);
    return `${target}; read ${details.join(", ")}`;
  }

  if (toolName === "ssh_write") {
    const content = stringField(input, "content") ?? "";
    return `${target}; write remote path '${inline(path ?? "<unspecified>")}' (${countLines(content)} lines, ${content.length} characters)`;
  }

  if (toolName === "ssh_edit") {
    const oldText = stringField(input, "oldText") ?? "";
    const newText = stringField(input, "newText") ?? "";
    return `${target}; edit remote path '${inline(path ?? "<unspecified>")}' (replace ${countLines(oldText)} lines with ${countLines(newText)} lines)`;
  }

  if (toolName === "ssh_find" || toolName === "ssh_grep") {
    const pattern = inline(stringField(input, "pattern") ?? "<unspecified>");
    const operation = toolName === "ssh_find" ? "find remote files" : "search remote file contents";
    const details = [
      `under '${inline(path ?? ".")}'`,
      `for '${pattern}'`,
      `limit ${typeof input.limit === "number" ? input.limit : 50}`,
    ];
    const include = stringField(input, "include");
    if (include) details.push(`file glob '${inline(include)}'`);
    return `${target}; ${operation} ${details.join(", ")}`;
  }

  if (toolName === "ssh_bash") {
    return `${target}; execute the separately displayed remote shell command`;
  }

  return target;
}

/**
 * Register pi-ssh's permission previews and remote-path semantics with the
 * bundle's published permission service.
 *
 * Remote paths must not enter the local `path` / `external_directory` gates:
 * those gates resolve against the local cwd and filesystem. Returning
 * `undefined` from an explicit extractor disables the default `input.path`
 * convention, leaving the dedicated ssh_* policy surfaces authoritative.
 */
export function installSshPermissionIntegration(
  pi: PermissionIntegrationApi,
  getConnection: () => SshPermissionConnection | null,
  dependencies: SshPermissionIntegrationDependencies,
): () => void {
  const getPermissionsService = dependencies.getPermissionsService;
  const warn = dependencies.warn ?? ((message: string) => console.warn(`[pi-ssh] ${message}`));
  let registeredService: PermissionsService | undefined;
  let disposers: Array<() => void> = [];

  const unregister = (): void => {
    for (const dispose of disposers.splice(0).reverse()) {
      try {
        dispose();
      } catch (error) {
        warn(`failed to unregister permission integration: ${error instanceof Error ? error.message : String(error)}`);
      }
    }
    registeredService = undefined;
  };

  const tryRegister = (): void => {
    const service = getPermissionsService();
    if (service === undefined || (service === registeredService && disposers.length > 0)) return;
    unregister();

    const pending: Array<() => void> = [];
    try {
      for (const toolName of REMOTE_TOOLS) {
        pending.push(
          service.registerToolInputFormatter(toolName, (input) => {
            const target = toolName === "ssh_connect"
              ? dependencies.getConnectTarget?.(input) ?? null
              : getConnection();
            return formatSshPermissionInput(toolName, input, target);
          }),
        );
      }
      for (const toolName of REMOTE_FILE_TOOLS) {
        pending.push(service.registerToolAccessExtractor(toolName, () => undefined));
      }
      disposers = pending;
      registeredService = service;
    } catch (error) {
      for (const dispose of pending.reverse()) {
        try {
          dispose();
        } catch {
          // Best-effort rollback; the permission gate remains conservative.
        }
      }
      warn(`failed to register permission integration: ${error instanceof Error ? error.message : String(error)}`);
    }
  };

  pi.on("session_start", tryRegister);
  pi.events.on(dependencies.permissionsReadyChannel, tryRegister);
  pi.on("session_shutdown", unregister);

  // Handles extension load after the permission service has already published.
  tryRegister();
  return unregister;
}