1  /**
  2   * Ephemeral Mode Extension
  3   *
  4   * Toggle session persistence on/off mid-session.
  5   * When ephemeral mode is enabled, the session file is deleted on exit.
  6   *
  7   * Usage:
  8   *   /ephemeral  - Toggle ephemeral mode
  9   *   Alt+E      - Toggle ephemeral mode (shortcut)
 10   */
 11  
 12  import { SessionManager, type ExtensionAPI, type SessionEntry } from "@mariozechner/pi-coding-agent";
 13  import { existsSync, unlinkSync } from "node:fs";
 14  
 15  export default function (pi: ExtensionAPI) {
 16      let ephemeralMode = false;
 17  
 18      const isEphemeralEnabled = (entries: SessionEntry[]) => {
 19          let enabled = false;
 20  
 21          for (const entry of entries) {
 22              if (entry.type === "custom" && entry.customType === "ephemeral-mode") {
 23                  enabled = entry.data?.enabled ?? false;
 24              }
 25          }
 26  
 27          return enabled;
 28      };
 29  
 30      // Reconstruct state from the *current branch* so /tree navigation properly restores branch-specific state
 31      const reconstructStateFromBranch = (ctx: any) => {
 32          ephemeralMode = isEphemeralEnabled(ctx.sessionManager.getBranch());
 33  
 34          if (!ctx.hasUI) {
 35              return;
 36          }
 37  
 38          if (ephemeralMode) {
 39              ctx.ui.setStatus("ephemeral", ctx.ui.theme.fg("warning", "儚"));
 40          } else {
 41              ctx.ui.setStatus("ephemeral", undefined);
 42          }
 43      };
 44  
 45      const toggleEphemeral = async (ctx: any) => {
 46          ephemeralMode = !ephemeralMode;
 47  
 48          // Persist the state in the session branch (for /tree navigation)
 49          pi.appendEntry("ephemeral-mode", { enabled: ephemeralMode });
 50  
 51          if (!ctx.hasUI) {
 52              return;
 53          }
 54  
 55          if (ephemeralMode) {
 56              ctx.ui.setStatus("ephemeral", ctx.ui.theme.fg("warning", "儚"));
 57              ctx.ui.notify("Ephemeral mode ON - session will be deleted on exit", "warning");
 58          } else {
 59              ctx.ui.setStatus("ephemeral", undefined);
 60              ctx.ui.notify("Ephemeral mode OFF - session will be preserved", "info");
 61          }
 62      };
 63  
 64      const clearIndicator = (ctx: any) => {
 65          ephemeralMode = false;
 66          if (ctx?.hasUI) {
 67              ctx.ui.setStatus("ephemeral", undefined);
 68          }
 69      };
 70  
 71      const deleteSessionFileBestEffort = async (sessionFile: string) => {
 72          try {
 73              const { code } = await pi.exec("trash", [sessionFile], { timeout: 5000 });
 74              if (code === 0) {
 75                  return;
 76              }
 77          } catch {
 78              // ignore
 79          }
 80  
 81          // Fallback to direct deletion
 82          try {
 83              if (existsSync(sessionFile)) {
 84                  unlinkSync(sessionFile);
 85              }
 86          } catch {
 87              // ignore
 88          }
 89      };
 90  
 91      const deletePreviousEphemeralSessionIfNeeded = async (previousSessionFile?: string, currentSessionFile?: string) => {
 92          if (!previousSessionFile || previousSessionFile === currentSessionFile) {
 93              return;
 94          }
 95  
 96          try {
 97              const previousSession = SessionManager.open(previousSessionFile);
 98              if (!isEphemeralEnabled(previousSession.getBranch())) {
 99                  return;
100              }
101          } catch {
102              return;
103          }
104  
105          await deleteSessionFileBestEffort(previousSessionFile);
106      };
107  
108      pi.on("session_start", async (event, ctx) => {
109          if (event.reason === "new" || event.reason === "resume" || event.reason === "fork") {
110              await deletePreviousEphemeralSessionIfNeeded(
111                  event.previousSessionFile,
112                  ctx.sessionManager.getSessionFile() ?? undefined,
113              );
114          }
115  
116          reconstructStateFromBranch(ctx);
117      });
118  
119      pi.on("session_before_switch", async (_event, ctx) => {
120          // Clear the indicator immediately so it never visually leaks into the switcher UI or next session
121          clearIndicator(ctx);
122      });
123  
124      pi.on("session_tree", async (_event, ctx) => {
125          reconstructStateFromBranch(ctx);
126      });
127  
128      pi.on("session_before_fork", async (_event, ctx) => {
129          // Clear the indicator immediately so it never visually leaks into the picker UI or next session
130          clearIndicator(ctx);
131      });
132  
133      pi.registerCommand("ephemeral", {
134          description: "Toggle ephemeral mode (delete session file on exit)",
135          handler: async (_args, ctx) => {
136              await toggleEphemeral(ctx);
137          },
138      });
139  
140      pi.registerShortcut("alt+e", {
141          description: "Toggle ephemeral mode",
142          handler: async (ctx) => {
143              await toggleEphemeral(ctx);
144          },
145      });
146  
147      pi.on("session_shutdown", async (_event, ctx) => {
148          if (!ephemeralMode) return;
149  
150          const sessionFile = ctx.sessionManager.getSessionFile();
151          if (!sessionFile) return; // Already in-memory mode
152  
153          await deleteSessionFileBestEffort(sessionFile);
154      });
155  }