1  /**
  2   * Session Move Extension
  3   *
  4   * Move the current session into another cwd bucket and relaunch pi in that directory
  5   *
  6   * Implementation strategy:
  7   *  1) Fork the current session file into the target cwd bucket using SessionManager.forkFrom()
  8   *  2) Clear the fork header's parentSession pointer
  9   *  3) Tear down the parent's terminal usage (pop kitty protocol, reset modes)
 10   *  4) Spawn a new pi process in the target cwd with inherited stdio
 11   *  5) Once the child has spawned, trash the old session file
 12   *  6) Once the child has spawned, destroy the parent's stdin so it cannot steal key presses
 13   *  7) Parent stays alive as an inert wrapper, forwarding the child's exit code
 14   *
 15   * Usage:
 16   *   /move-session <targetCwd>
 17   */
 18  
 19  import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
 20  import { SessionManager } from "@mariozechner/pi-coding-agent";
 21  import { spawn } from "node:child_process";
 22  import {
 23      statSync,
 24      openSync,
 25      readSync,
 26      writeSync,
 27      closeSync,
 28      renameSync,
 29      unlinkSync,
 30  } from "node:fs";
 31  
 32  import { normalizeTargetCwd } from "./_shared/normalize-target-cwd";
 33  
 34  const TRASH_TIMEOUT_MS = 5000;
 35  const HEADER_READ_MAX = 8192;
 36  const COPY_CHUNK_SIZE = 65_536;
 37  
 38  function getBranchSelectionWarning(
 39      sourceSessionFile: string,
 40      currentLeafId: string | null,
 41      commandName: string,
 42      actionName: string,
 43  ): string | null {
 44      try {
 45          const persistedSession = SessionManager.open(sourceSessionFile);
 46          const hasPersistedEntries = persistedSession.getEntries().length > 0;
 47          if (!hasPersistedEntries) {
 48              return null;
 49          }
 50  
 51          const persistedLeafId = persistedSession.getLeafId() as string | null;
 52          if (currentLeafId === null) {
 53              return `${commandName} will not preserve the current /tree root selection. It reopens at the session file's default branch tip. Consider /fork first or continue from the branch tip before ${actionName}.`;
 54          }
 55  
 56          if (currentLeafId !== persistedLeafId) {
 57              return `${commandName} will not preserve the current /tree selection. It reopens at the session file's default branch tip. Consider /fork first or continue from the branch tip before ${actionName}.`;
 58          }
 59      } catch {
 60          // Ignore warning-detection failures
 61      }
 62  
 63      return null;
 64  }
 65  
 66  /**
 67   * Remove parentSession from the first JSONL header line without loading
 68   * the entire file into memory
 69   */
 70  function clearParentSession(sessionFile: string): void {
 71      const fd = openSync(sessionFile, "r");
 72      const headerBuffer = Buffer.alloc(HEADER_READ_MAX);
 73      const bytesRead = readSync(fd, headerBuffer, 0, HEADER_READ_MAX, 0);
 74      const headerChunk = headerBuffer.toString("utf-8", 0, bytesRead);
 75      const newlineIndex = headerChunk.indexOf("\n");
 76  
 77      if (newlineIndex === -1) {
 78          closeSync(fd);
 79          return;
 80      }
 81  
 82      const header = JSON.parse(headerChunk.slice(0, newlineIndex));
 83      if (!header.parentSession) {
 84          closeSync(fd);
 85          return;
 86      }
 87  
 88      delete header.parentSession;
 89      const newHeaderLine = JSON.stringify(header) + "\n";
 90      const originalHeaderBytes = Buffer.byteLength(headerChunk.slice(0, newlineIndex + 1), "utf-8");
 91  
 92      const temporaryPath = sessionFile + ".move-session-tmp";
 93      let writeFd: number | undefined;
 94      try {
 95          writeFd = openSync(temporaryPath, "w");
 96          const newHeaderBuffer = Buffer.from(newHeaderLine, "utf-8");
 97          writeSync(writeFd, newHeaderBuffer, 0, newHeaderBuffer.length);
 98  
 99          const copyBuffer = Buffer.alloc(COPY_CHUNK_SIZE);
100          let position = originalHeaderBytes;
101          while (true) {
102              const readCount = readSync(fd, copyBuffer, 0, COPY_CHUNK_SIZE, position);
103              if (readCount === 0) break;
104              writeSync(writeFd, copyBuffer, 0, readCount);
105              position += readCount;
106          }
107  
108          closeSync(writeFd);
109          writeFd = undefined;
110          closeSync(fd);
111          renameSync(temporaryPath, sessionFile);
112      } catch (error) {
113          if (writeFd !== undefined) {
114              try {
115                  closeSync(writeFd);
116              } catch {
117                  // ignore cleanup close errors
118              }
119          }
120  
121          closeSync(fd);
122          try {
123              unlinkSync(temporaryPath);
124          } catch {
125              // ignore cleanup unlink errors
126          }
127          throw error;
128      }
129  }
130  
131  export default function (pi: ExtensionAPI) {
132      const trashFileBestEffort = async (filePath: string) => {
133          try {
134              const { code } = await pi.exec("trash", [filePath], { timeout: TRASH_TIMEOUT_MS });
135              if (code === 0) {
136                  return;
137              }
138          } catch {
139              // ignore
140          }
141  
142          // If "trash" isn't available, do not fall back to unlink.
143          // This extension should never permanently delete session files.
144      };
145  
146      pi.registerCommand("move-session", {
147          description: "Move session to another directory and relaunch pi there",
148          handler: async (args, ctx) => {
149              await ctx.waitForIdle();
150  
151              const rawTargetCwd = args.trim();
152              if (!rawTargetCwd) {
153                  ctx.ui.notify("Usage: /move-session <targetCwd>", "error");
154                  return;
155              }
156  
157              let targetCwd: string;
158              try {
159                  targetCwd = normalizeTargetCwd(rawTargetCwd);
160              } catch (error: any) {
161                  ctx.ui.notify(error?.message ?? String(error), "error");
162                  return;
163              }
164  
165              let targetCwdStat;
166              try {
167                  targetCwdStat = statSync(targetCwd);
168              } catch (error: any) {
169                  const code = error?.code;
170                  if (code === "ENOENT") {
171                      ctx.ui.notify(`Path does not exist: ${targetCwd}`, "error");
172                  } else {
173                      ctx.ui.notify(`Cannot access path: ${targetCwd}`, "error");
174                  }
175                  return;
176              }
177  
178              if (!targetCwdStat.isDirectory()) {
179                  ctx.ui.notify(`Not a directory: ${targetCwd}`, "error");
180                  return;
181              }
182  
183              const sourceSessionFile = ctx.sessionManager.getSessionFile();
184              if (!sourceSessionFile) {
185                  ctx.ui.notify("No persistent session file (maybe started with --no-session)", "error");
186                  return;
187              }
188  
189              const branchSelectionWarning = getBranchSelectionWarning(
190                  sourceSessionFile,
191                  (ctx.sessionManager.getLeafId?.() as string | null) ?? null,
192                  "/move-session",
193                  "moving",
194              );
195              if (branchSelectionWarning) {
196                  ctx.ui.notify(branchSelectionWarning, "warning");
197              }
198  
199              try {
200                  const forked = SessionManager.forkFrom(sourceSessionFile, targetCwd);
201                  const destSessionFile = forked.getSessionFile();
202  
203                  if (!destSessionFile) {
204                      ctx.ui.notify("Internal error: forkFrom() produced no session file", "error");
205                      return;
206                  }
207  
208                  // We intend to move/replace the original session, so avoid leaving
209                  // a parentSession pointer that may dangle after trashing the source.
210                  try {
211                      clearParentSession(destSessionFile);
212                  } catch (error: any) {
213                      ctx.ui.notify(
214                          `Warning: could not clear parent session reference: ${error?.message ?? String(error)}`,
215                          "warning"
216                      );
217                  }
218  
219                  // --- Tear down the parent's terminal usage ---
220                  // We do this BEFORE spawning, to avoid nesting Kitty protocol flags.
221                  process.stdout.write("\x1b[<u");      // Pop kitty keyboard protocol
222                  process.stdout.write("\x1b[?2004l");  // Disable bracketed paste
223                  process.stdout.write("\x1b[?25h");    // Show cursor
224                  process.stdout.write("\r\n");         // Ensure child starts on a clean line
225  
226                  if (process.stdin.isTTY && process.stdin.setRawMode) {
227                      process.stdin.setRawMode(false);
228                  }
229  
230                  // Spawn new pi in the target directory
231                  const child = spawn("pi", ["--session", destSessionFile], {
232                      cwd: targetCwd,
233                      stdio: "inherit",
234                  });
235  
236                  child.once("spawn", () => {
237                      // Trash the old session file *after* the new process is actually running
238                      void trashFileBestEffort(sourceSessionFile);
239  
240                      // Stop the parent from stealing keypresses.
241                      // destroy() is important; pause() was not sufficient.
242                      process.stdin.removeAllListeners();
243                      process.stdin.destroy();
244  
245                      // Avoid the parent reacting to Ctrl+C / termination signals.
246                      // (The child is the process the user is interacting with.)
247                      process.removeAllListeners("SIGINT");
248                      process.removeAllListeners("SIGTERM");
249                      process.on("SIGINT", () => {});
250                      process.on("SIGTERM", () => {});
251                  });
252  
253                  child.on("exit", (code) => process.exit(code ?? 0));
254                  child.on("error", (err) => {
255                      process.stderr.write(`Failed to launch pi: ${err.message}\n`);
256                      process.exit(1);
257                  });
258              } catch (error: any) {
259                  ctx.ui.notify(`Failed to move session: ${error?.message ?? String(error)}`, "error");
260              }
261          },
262      });
263  }