pi-extensions/extensions/ext-toggle.ts

/**
 * Ext-Toggle Extension
 *
 * Provides `/ext` to list and toggle pi extensions installed under
 * `~/.pi/agent/extensions/` (the global auto-discovery dir).
 *
 * UX:
 *   ↑/↓     navigate
 *   space   stage a toggle (visual ●/○ flip, not yet applied)
 *   enter   commit all staged changes and reload pi
 *   esc     cancel, no changes applied
 *
 * Toggling works by renaming the file (or symlink) between `name.ts` and
 * `name.ts.off`. Pi discovers `*.ts` only, so a `.ts.off` file is invisible
 * at load time. After commit, `ctx.reload()` is invoked so changes take
 * effect without a restart.
 *
 * Subdirectory-style extensions (`name/index.ts`) are listed read-only —
 * toggling a directory cleanly is more work than the renaming trick is
 * worth. If you need to disable one, move the directory aside manually.
 *
 * DISABLE_GUARDS: per-extension predicates that refuse a stage attempt
 * when toggling would silently break in-flight session state. The space
 * key flip is reverted and a status line explains why.
 */

import * as fs from "node:fs";
import * as path from "node:path";
import * as os from "node:os";
import {
  type ExtensionAPI,
  getSettingsListTheme,
} from "@mariozechner/pi-coding-agent";
import {
  Container,
  Key,
  matchesKey,
  type SettingItem,
  SettingsList,
} from "@mariozechner/pi-tui";

const EXT_DIR = path.join(os.homedir(), ".pi", "agent", "extensions");

// ── Disable guards ────────────────────────────────────────────────────────────

/**
 * Per-extension guards that refuse a disable when toggling would silently
 * break in-flight session state. Returns a string explaining why the
 * disable should be blocked, or null to allow it.
 *
 * Keyed by bare extension name (matches the file basename without `.ts`).
 */
const DISABLE_GUARDS: Record<string, () => string | null> = {
  "ssh-controlmaster": () => {
    const hasSshFlag = process.argv.some(
      (a) => a === "--ssh" || a.startsWith("--ssh="),
    );
    if (!hasSshFlag) return null;
    return "Active --ssh session — disabling would silently revert read/write/edit/bash to local. Exit pi and relaunch without --ssh.";
  },
};

// ── Filesystem helpers ────────────────────────────────────────────────────────

type Entry = {
  name: string;
  fullPath: string;
  enabled: boolean;
  kind: "file" | "dir";
  isSymlink: boolean;
};

function listEntries(): Entry[] {
  if (!fs.existsSync(EXT_DIR)) return [];
  const out: Entry[] = [];
  for (const ent of fs.readdirSync(EXT_DIR, { withFileTypes: true })) {
    const full = path.join(EXT_DIR, ent.name);
    const isSymlink = ent.isSymbolicLink();
    let stat: fs.Stats | undefined;
    try {
      stat = fs.statSync(full);
    } catch {
      continue;
    }

    if (stat.isDirectory()) {
      const indexPath = path.join(full, "index.ts");
      if (fs.existsSync(indexPath)) {
        out.push({
          name: ent.name,
          fullPath: full,
          enabled: true,
          kind: "dir",
          isSymlink,
        });
      }
      continue;
    }

    if (ent.name.endsWith(".ts")) {
      out.push({
        name: ent.name.slice(0, -3),
        fullPath: full,
        enabled: true,
        kind: "file",
        isSymlink,
      });
    } else if (ent.name.endsWith(".ts.off")) {
      out.push({
        name: ent.name.slice(0, -7),
        fullPath: full,
        enabled: false,
        kind: "file",
        isSymlink,
      });
    }
  }
  return out.sort((a, b) => a.name.localeCompare(b.name));
}

/** Rename to match a target enabled state. Returns the new path. */
function applyState(entry: Entry, enabled: boolean): string {
  const dir = path.dirname(entry.fullPath);
  const target = path.join(dir, enabled ? `${entry.name}.ts` : `${entry.name}.ts.off`);
  if (entry.fullPath === target) return target;
  fs.renameSync(entry.fullPath, target);
  return target;
}

// ── Settings list values ──────────────────────────────────────────────────────

const VAL_ON = "● enabled";
const VAL_OFF = "○ disabled";
const VAL_DIR = "[dir] (manage manually)";

function valueFor(enabled: boolean): string {
  return enabled ? VAL_ON : VAL_OFF;
}

// ── Extension ────────────────────────────────────────────────────────────────

export default function extToggleExtension(pi: ExtensionAPI) {
  pi.registerCommand("ext", {
    description: "List and toggle pi extensions in ~/.pi/agent/extensions/",
    handler: async (_args, ctx) => {
      const entries = listEntries();
      if (entries.length === 0) {
        ctx.ui.notify(`No extensions found in ${EXT_DIR}`, "info");
        return;
      }

      // Staged state: name → enabled. Initialized from disk; mutated by space.
      const staged = new Map<string, boolean>();
      for (const e of entries) staged.set(e.name, e.enabled);

      await ctx.ui.custom<void>((tui, theme, _kb, done) => {
        let statusText = "";

        const items: SettingItem[] = entries.map((e) => {
          if (e.kind === "dir") {
            return {
              id: e.name,
              label: e.name,
              currentValue: VAL_DIR,
              description: "Subdirectory extension — toggle by moving the dir manually.",
            };
          }
          const guard = DISABLE_GUARDS[e.name];
          const guardNote = guard?.();
          const desc = [
            e.isSymlink ? "symlink" : "file",
            guardNote ? `guarded: ${guardNote}` : null,
          ]
            .filter(Boolean)
            .join(" · ");
          return {
            id: e.name,
            label: e.name,
            currentValue: valueFor(e.enabled),
            values: [VAL_ON, VAL_OFF],
            description: desc || undefined,
          };
        });

        const container = new Container();

        // Header
        container.addChild({
          render(_w: number) {
            return [
              theme.fg("accent", theme.bold("Extensions — ~/.pi/agent/extensions/")),
              theme.fg(
                "muted",
                "  space: toggle (stage) · enter: apply · esc: cancel",
              ),
              "",
            ];
          },
          invalidate() {},
        });

        const settingsList = new SettingsList(
          items,
          Math.min(items.length + 2, 20),
          getSettingsListTheme(),
          (id, newValue) => {
            // SettingsList just cycled this row. Update staged state, but
            // first check the guard if this is a disable transition.
            const entry = entries.find((e) => e.name === id);
            if (!entry || entry.kind === "dir") {
              // Should not happen — dir items have no `values` and won't cycle.
              settingsList.updateValue(id, VAL_DIR);
              return;
            }

            const stagingEnabled = newValue === VAL_ON;

            // Guard check: disabling a guarded extension is refused.
            if (!stagingEnabled) {
              const reason = DISABLE_GUARDS[entry.name]?.();
              if (reason) {
                settingsList.updateValue(id, VAL_ON);
                statusText = `⊘ ${entry.name}: ${reason}`;
                tui.requestRender();
                return;
              }
            }

            staged.set(id, stagingEnabled);

            // Reflect drift from disk in status line so the user knows what
            // will be applied on Enter.
            const drift: string[] = [];
            for (const e of entries) {
              if (e.kind !== "file") continue;
              const s = staged.get(e.name);
              if (s !== e.enabled) {
                drift.push(`${e.name}→${s ? "on" : "off"}`);
              }
            }
            statusText = drift.length
              ? `pending: ${drift.join(", ")}`
              : "";
            tui.requestRender();
          },
          () => {
            // SettingsList's onCancel fires on escape.
            done(undefined);
          },
        );

        container.addChild(settingsList);

        // Footer status line — shows pending changes or guard rejections.
        container.addChild({
          render(_w: number) {
            if (!statusText) return [""];
            const colored = statusText.startsWith("⊘")
              ? theme.fg("warning", statusText)
              : theme.fg("muted", statusText);
            return ["", colored];
          },
          invalidate() {},
        });

        return {
          render(width: number) {
            return container.render(width);
          },
          invalidate() {
            container.invalidate();
          },
          handleInput(data: string) {
            // Intercept Enter at the wrapper level so SettingsList doesn't
            // consume it for cycling. Everything else (space, arrows, esc)
            // passes through unchanged.
            if (matchesKey(data, Key.enter)) {
              // Commit staged → disk.
              const errors: string[] = [];
              let applied = 0;
              for (const e of entries) {
                if (e.kind !== "file") continue;
                const s = staged.get(e.name);
                if (s === undefined || s === e.enabled) continue;
                try {
                  applyState(e, s);
                  applied++;
                } catch (err) {
                  const msg = err instanceof Error ? err.message : String(err);
                  errors.push(`${e.name}: ${msg}`);
                }
              }
              done(undefined);
              if (errors.length) {
                ctx.ui.notify(`ext-toggle: ${errors.join(" | ")}`, "error");
              } else if (applied > 0) {
                ctx.ui.notify(`ext-toggle: applied ${applied} change(s); reloading…`, "info");
                ctx.reload().catch((err) => {
                  const msg = err instanceof Error ? err.message : String(err);
                  ctx.ui.notify(`reload failed: ${msg}`, "error");
                });
              }
              return;
            }
            settingsList.handleInput?.(data);
            tui.requestRender();
          },
        };
      });
    },
  });
}