1  /**
  2   * Roam Extension
  3   *
  4   * Move the current Pi session into a tmux window for remote access
  5   *
  6   * Usage:
  7   *   /roam [window-name]    (default: cwd basename)
  8   *
  9   * Optional config (~/.pi/agent/extensions/roam/config.json):
 10   *   - copy from extensions/roam/config.json.example
 11   *   {
 12   *     "tailscale": {
 13   *       "account": "you@example.com",
 14   *       "binary": "/Applications/Tailscale.app/Contents/MacOS/Tailscale"
 15   *     }
 16   *   }
 17   *
 18   * All Pi sessions share a single tmux session ("pi") on a dedicated socket (-L pi),
 19   * each in its own window. Uses a custom tmux config with dual prefix keys:
 20   *   - Ctrl+S (available on iOS Termius toolbar)
 21   *   - Ctrl+B (tmux default, for local use on Mac)
 22   *
 23   * The dedicated socket ensures the custom config is always applied, regardless
 24   * of other tmux servers that may be running.
 25   *
 26   * From Termius:
 27   *   - Attach: tmux -L pi -f ~/.config/pi-tmux/tmux.conf -u attach -t pi
 28   *   - Window list: Ctrl+S, then w
 29   *   - Next/prev window: Ctrl+S, then n/p
 30   *   - Detach: Ctrl+S, then d
 31   *   - No time limit between prefix and command key
 32   *
 33   * Flow:
 34   *   1. Pre-flight: TTY check, not already in tmux, session exists, tmux installed
 35   *   2. Optionally switch Tailscale account, then ensure Tailscale is up (non-fatal, macOS only)
 36   *   3. Fork the current Pi session to a new file (parentSession cleared)
 37   *   4. Create a tmux window (or session if first time) running the fork
 38   *   5. Tear down parent terminal, attach to tmux
 39   *   6. Trash original session file if in standard sessions dir (no duplicates)
 40   *   7. Parent becomes inert, forwarding exit code
 41   */
 42  
 43  import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
 44  import { SessionManager } from "@mariozechner/pi-coding-agent";
 45  import { spawn } from "node:child_process";
 46  import { basename, join, resolve } from "node:path";
 47  import { homedir } from "node:os";
 48  import {
 49      writeFileSync, existsSync, mkdirSync, realpathSync, readFileSync,
 50      openSync, readSync, writeSync, closeSync, renameSync, unlinkSync,
 51  } from "node:fs";
 52  
 53  const TMUX_SESSION = "pi";
 54  const TMUX_SOCKET = "pi";
 55  const TAILSCALE_BIN = "/Applications/Tailscale.app/Contents/MacOS/Tailscale";
 56  const TAILSCALE_TIMEOUT_MS = 10_000;
 57  const TRASH_TIMEOUT_MS = 5_000;
 58  const HEADER_READ_MAX = 8192;
 59  const COPY_CHUNK_SIZE = 65_536;
 60  
 61  function getBranchSelectionWarning(
 62      sourceSessionFile: string,
 63      currentLeafId: string | null,
 64      commandName: string,
 65      actionName: string,
 66  ): { warning: string | null; hasPersistedEntries: boolean | null } {
 67      try {
 68          const persistedSession = SessionManager.open(sourceSessionFile);
 69          const hasPersistedEntries = persistedSession.getEntries().length > 0;
 70          if (!hasPersistedEntries) {
 71              return { warning: null, hasPersistedEntries };
 72          }
 73  
 74          const persistedLeafId = persistedSession.getLeafId() as string | null;
 75          if (currentLeafId === null) {
 76              return {
 77                  warning:
 78                      `${commandName} will not preserve the current /tree root selection. ` +
 79                      `It reopens at the session file's default branch tip. ` +
 80                      `Consider /fork first or continue from the branch tip before ${actionName}.`,
 81                  hasPersistedEntries,
 82              };
 83          }
 84  
 85          if (currentLeafId !== persistedLeafId) {
 86              return {
 87                  warning:
 88                      `${commandName} will not preserve the current /tree selection. ` +
 89                      `It reopens at the session file's default branch tip. ` +
 90                      `Consider /fork first or continue from the branch tip before ${actionName}.`,
 91                  hasPersistedEntries,
 92              };
 93          }
 94  
 95          return { warning: null, hasPersistedEntries };
 96      } catch {
 97          // Ignore warning-detection failures
 98          return { warning: null, hasPersistedEntries: null };
 99      }
100  }
101  
102  const TMUX_CONFIG_CONTENT = [
103      "# Pi roam config — only used by /roam sessions (dedicated socket: -L pi)",
104      "# Does not affect your global ~/.tmux.conf or other tmux servers",
105      "",
106      "# Dual prefix: Ctrl+S (iOS Termius toolbar) and Ctrl+B (default, for local use)",
107      "set -g prefix C-s",
108      "set -g prefix2 C-b",
109      "bind C-s send-prefix",
110      "bind C-b send-prefix -2",
111      "",
112      "# UTF-8 and modern terminal support",
113      "set -g default-terminal 'screen-256color'",
114      "set -ga terminal-overrides ',xterm-256color:Tc'",
115      "",
116      "# Mouse support (useful for Termius touch scrolling)",
117      "set -g mouse on",
118      "",
119      "# Start window numbering at 1 (easier to reach on mobile)",
120      "set -g base-index 1",
121      "setw -g pane-base-index 1",
122      "",
123      "# Window status shows the name clearly",
124      "set -g status-left '[pi] '",
125      "set -g status-right ''",
126      "",
127  ].join("\n");
128  
129  type RoamConfig = {
130      tailscale?: {
131          account?: string;
132          binary?: string;
133      };
134  };
135  
136  type ResolvedRoamConfig = {
137      tailscaleAccount: string | null;
138      tailscaleBinary: string;
139  };
140  
141  function getAgentDir(): string {
142      return process.env.PI_CODING_AGENT_DIR || join(homedir(), ".pi", "agent");
143  }
144  
145  function getRoamConfigPath(): string {
146      return join(getAgentDir(), "extensions", "roam", "config.json");
147  }
148  
149  /**
150   * Load optional /roam config and validate shape.
151   * Missing config is treated as defaults.
152   */
153  function loadRoamConfig(configPath: string): ResolvedRoamConfig {
154      const defaults: ResolvedRoamConfig = {
155          tailscaleAccount: null,
156          tailscaleBinary: TAILSCALE_BIN,
157      };
158  
159      if (!existsSync(configPath)) {
160          return defaults;
161      }
162  
163      let parsed: RoamConfig;
164      try {
165          parsed = JSON.parse(readFileSync(configPath, "utf-8")) as RoamConfig;
166      } catch (error: any) {
167          throw new Error(`Invalid JSON in ${configPath}: ${error?.message ?? String(error)}`);
168      }
169  
170      const rawAccount = parsed.tailscale?.account;
171      if (rawAccount !== undefined && typeof rawAccount !== "string") {
172          throw new Error(`Invalid tailscale.account in ${configPath}; expected a string`);
173      }
174  
175      const rawBinary = parsed.tailscale?.binary;
176      if (rawBinary !== undefined && typeof rawBinary !== "string") {
177          throw new Error(`Invalid tailscale.binary in ${configPath}; expected a string`);
178      }
179  
180      return {
181          tailscaleAccount: rawAccount?.trim() || null,
182          tailscaleBinary: rawBinary?.trim() || TAILSCALE_BIN,
183      };
184  }
185  
186  /**
187   * Write the tmux config file; throws on FS errors (caller must handle)
188   */
189  function ensureTmuxConfig(): string {
190      const configDir = join(
191          process.env.HOME || process.env.USERPROFILE || "/tmp",
192          ".config", "pi-tmux"
193      );
194      const configPath = join(configDir, "tmux.conf");
195  
196      if (!existsSync(configDir)) {
197          mkdirSync(configDir, { recursive: true });
198      }
199      writeFileSync(configPath, TMUX_CONFIG_CONTENT);
200  
201      return configPath;
202  }
203  
204  /**
205   * Remove the parentSession field from a forked session's JSONL header
206   * without reading the entire file into memory. Reads only the first line
207   * (header), and if modification is needed, rewrites the file via a temp
208   * file using chunked streaming for the remaining content.
209   *
210   * Throws on FS errors (caller must handle).
211   */
212  function clearParentSession(sessionFile: string): void {
213      const fd = openSync(sessionFile, "r");
214      const buf = Buffer.alloc(HEADER_READ_MAX);
215      const bytesRead = readSync(fd, buf, 0, HEADER_READ_MAX, 0);
216      const headerChunk = buf.toString("utf-8", 0, bytesRead);
217      const newlineIdx = headerChunk.indexOf("\n");
218  
219      if (newlineIdx === -1) {
220          closeSync(fd);
221          return;
222      }
223  
224      const header = JSON.parse(headerChunk.slice(0, newlineIdx));
225      if (!header.parentSession) {
226          closeSync(fd);
227          return;
228      }
229  
230      // parentSession exists — stream-rewrite with modified header
231      delete header.parentSession;
232      const newHeaderLine = JSON.stringify(header) + "\n";
233      const originalHeaderBytes = Buffer.byteLength(
234          headerChunk.slice(0, newlineIdx + 1), "utf-8"
235      );
236  
237      const tmpPath = sessionFile + ".roam-tmp";
238      let wfd: number | undefined;
239      try {
240          wfd = openSync(tmpPath, "w");
241          const headerBuf = Buffer.from(newHeaderLine, "utf-8");
242          writeSync(wfd, headerBuf, 0, headerBuf.length);
243  
244          // Copy rest of file in chunks (avoids loading full session into memory)
245          const copyBuf = Buffer.alloc(COPY_CHUNK_SIZE);
246          let pos = originalHeaderBytes;
247          while (true) {
248              const n = readSync(fd, copyBuf, 0, COPY_CHUNK_SIZE, pos);
249              if (n === 0) break;
250              writeSync(wfd, copyBuf, 0, n);
251              pos += n;
252          }
253          closeSync(wfd);
254          wfd = undefined;
255          closeSync(fd);
256          renameSync(tmpPath, sessionFile);
257      } catch (error) {
258          // Clean up temp file on failure
259          if (wfd !== undefined) try { closeSync(wfd); } catch {}
260          closeSync(fd);
261          try { unlinkSync(tmpPath); } catch {}
262          throw error;
263      }
264  }
265  
266  /**
267   * Strip control characters from a window name to prevent tmux breakage
268   * and stdout parsing issues. Returns null if the result is empty.
269   */
270  function sanitizeWindowName(name: string): string | null {
271      const cleaned = name.replace(/[\x00-\x1f\x7f]/g, "").trim();
272      return cleaned || null;
273  }
274  
275  /**
276   * Check if a session file is inside the standard Pi sessions directory (~/.pi/).
277   * Custom --session paths (e.g. /some/custom/path.jsonl) should not be trashed.
278   * Handles symlinks (e.g. ~/.pi/agent -> ~/dot314/agent) via realpathSync.
279   */
280  function isInStandardSessionsDir(sessionFile: string): boolean {
281      const home = process.env.HOME || process.env.USERPROFILE;
282      if (!home) return false;
283      const piDir = join(home, ".pi");
284      try {
285          const resolvedFile = realpathSync(sessionFile);
286          const resolvedPiDir = realpathSync(piDir);
287          return resolvedFile.startsWith(resolvedPiDir + "/");
288      } catch {
289          // realpathSync can fail if file/dir doesn't exist; fall back to resolve()
290          return resolve(sessionFile).startsWith(resolve(piDir) + "/");
291      }
292  }
293  
294  export default function (pi: ExtensionAPI) {
295      const trashFileBestEffort = async (filePath: string): Promise<boolean> => {
296          try {
297              const { code } = await pi.exec("trash", [filePath], { timeout: TRASH_TIMEOUT_MS });
298              return code === 0;
299          } catch {
300              return false;
301          }
302      };
303  
304      pi.registerCommand("roam", {
305          description: "Move session into a tmux window for remote access via Tailscale",
306          handler: async (args, ctx) => {
307              await ctx.waitForIdle();
308  
309              // --- Pre-flight checks ---
310  
311              if (!ctx.hasUI || !process.stdin.isTTY || !process.stdout.isTTY) {
312                  if (ctx.hasUI) {
313                      ctx.ui.notify("/roam requires an interactive terminal", "error");
314                  }
315                  return;
316              }
317  
318              if (process.env.TMUX) {
319                  ctx.ui.notify("Already inside tmux. Use Ctrl+S d (or Ctrl+B d) to detach.", "error");
320                  return;
321              }
322  
323              const tmuxCheck = await pi.exec("which", ["tmux"]);
324              if (tmuxCheck.code !== 0) {
325                  ctx.ui.notify("tmux is not installed", "error");
326                  return;
327              }
328  
329              const sourceSessionFile = ctx.sessionManager.getSessionFile();
330              if (!sourceSessionFile) {
331                  ctx.ui.notify("No persistent session (started with --no-session?)", "error");
332                  return;
333              }
334  
335              const leafId = (ctx.sessionManager.getLeafId?.() as string | null) ?? null;
336              const { warning: branchSelectionWarning, hasPersistedEntries } = getBranchSelectionWarning(
337                  sourceSessionFile,
338                  leafId,
339                  "/roam",
340                  "roaming",
341              );
342              if (leafId === null && hasPersistedEntries === false) {
343                  ctx.ui.notify("No messages yet — nothing to roam", "error");
344                  return;
345              }
346              if (branchSelectionWarning) {
347                  ctx.ui.notify(branchSelectionWarning, "warning");
348              }
349  
350              const cwd = ctx.cwd;
351  
352              // Window name: from args or cwd basename, sanitized for tmux safety
353              const rawName = args.trim() || basename(cwd);
354              const windowName = sanitizeWindowName(rawName);
355              if (!windowName) {
356                  ctx.ui.notify("Invalid window name (empty after sanitization)", "error");
357                  return;
358              }
359  
360              let tailscaleAccount: string | null = null;
361              let tailscaleBinary = TAILSCALE_BIN;
362              const roamConfigPath = getRoamConfigPath();
363              try {
364                  const roamConfig = loadRoamConfig(roamConfigPath);
365                  tailscaleAccount = roamConfig.tailscaleAccount;
366                  tailscaleBinary = roamConfig.tailscaleBinary;
367              } catch (error: any) {
368                  ctx.ui.notify(
369                      `Ignoring invalid /roam config (${roamConfigPath}): ${error?.message ?? String(error)}`,
370                      "warning"
371                  );
372              }
373  
374              // Ensure dedicated tmux config exists and is up to date
375              let tmuxConfig: string;
376              try {
377                  tmuxConfig = ensureTmuxConfig();
378              } catch (error: any) {
379                  ctx.ui.notify(`Failed to write tmux config: ${error?.message ?? String(error)}`, "error");
380                  return;
381              }
382  
383              // Common tmux flags: dedicated socket + config file
384              const tmuxBase = ["-L", TMUX_SOCKET, "-f", tmuxConfig];
385  
386              // Check tmux state on our dedicated socket
387              let sessionExists = false;
388              try {
389                  const { code } = await pi.exec("tmux", [...tmuxBase, "has-session", "-t", TMUX_SESSION]);
390                  sessionExists = code === 0;
391              } catch {
392                  // tmux server not running on this socket — we'll create a new session
393              }
394  
395              // Source the latest config unconditionally — covers the case where the
396              // server is running but our session doesn't exist yet. If the server
397              // isn't running, this fails harmlessly. If the config has errors, we warn.
398              {
399                  const { code: srcCode, stderr: srcStderr } = await pi.exec(
400                      "tmux", [...tmuxBase, "source-file", tmuxConfig]
401                  );
402                  if (srcCode !== 0 && sessionExists) {
403                      // Only warn if we know the server is running (otherwise failure
404                      // just means "no server" which is expected and fine)
405                      ctx.ui.notify(
406                          `tmux config warning: ${srcStderr || "source-file failed"}`,
407                          "warning"
408                      );
409                  }
410              }
411  
412              if (sessionExists) {
413                  // Check for duplicate window name
414                  const { code, stdout } = await pi.exec("tmux", [
415                      ...tmuxBase,
416                      "list-windows", "-t", TMUX_SESSION, "-F", "#{window_name}",
417                  ]);
418                  if (code !== 0) {
419                      ctx.ui.notify("Failed to list tmux windows — tmux may be in an unexpected state", "error");
420                      return;
421                  }
422                  const existingWindows = stdout.trim().split("\n").filter(Boolean);
423                  if (existingWindows.includes(windowName)) {
424                      ctx.ui.notify(
425                          `Window "${windowName}" already exists in tmux session "${TMUX_SESSION}". ` +
426                              `Use: /roam <different-name>`,
427                          "error"
428                      );
429                      return;
430                  }
431              }
432  
433              // --- Tailscale (non-fatal, macOS only) ---
434  
435              if (process.platform === "darwin") {
436                  if (tailscaleAccount) {
437                      ctx.ui.notify(`Switching Tailscale account: ${tailscaleAccount}`, "info");
438                      try {
439                          const { code, stderr } = await pi.exec(
440                              tailscaleBinary,
441                              ["switch", tailscaleAccount],
442                              { timeout: TAILSCALE_TIMEOUT_MS }
443                          );
444                          if (code !== 0) {
445                              ctx.ui.notify(
446                                  `Tailscale switch warning: ${stderr || "switch command failed"}`,
447                                  "warning"
448                              );
449                          }
450                      } catch {
451                          ctx.ui.notify("Tailscale switch unavailable — continuing", "warning");
452                      }
453                  }
454  
455                  ctx.ui.notify("Bringing up Tailscale...", "info");
456                  try {
457                      const { code, stderr } = await pi.exec(tailscaleBinary, ["up"], {
458                          timeout: TAILSCALE_TIMEOUT_MS,
459                      });
460                      if (code !== 0) {
461                          ctx.ui.notify(`Tailscale warning: ${stderr || "failed to start"}`, "warning");
462                      }
463                  } catch {
464                      ctx.ui.notify("Tailscale not available — continuing without it", "warning");
465                  }
466              }
467  
468              // --- Fork session ---
469              // waitForIdle() above ensures the agent has finished streaming. Pi persists
470              // entries via synchronous appendFileSync, so by the time waitForIdle() resolves
471              // and the command handler runs, all entries should be flushed to disk.
472  
473              let destSessionFile: string;
474              try {
475                  const forked = SessionManager.forkFrom(sourceSessionFile, cwd);
476                  const dest = forked.getSessionFile();
477                  if (!dest) {
478                      ctx.ui.notify("Fork produced no session file", "error");
479                      return;
480                  }
481                  destSessionFile = dest;
482              } catch (error: any) {
483                  ctx.ui.notify(`Failed to fork session: ${error?.message ?? String(error)}`, "error");
484                  return;
485              }
486  
487              // Remove parentSession pointer since we intend to trash the original.
488              // A dangling parentSession would break session_lineage and session_ask.
489              try {
490                  clearParentSession(destSessionFile);
491              } catch (error: any) {
492                  // Non-fatal: the session still works, just has a dangling parentSession
493                  ctx.ui.notify(
494                      `Warning: could not clear parent session reference: ${error?.message ?? String(error)}`,
495                      "warning"
496                  );
497              }
498  
499              // --- Create tmux window if session already exists ---
500              // (must happen before terminal teardown so pi.exec still works)
501  
502              let tmuxArgs: string[];
503  
504              if (sessionExists) {
505                  // Add a new window to the existing "pi" session
506                  const { code, stderr, stdout } = await pi.exec("tmux", [
507                      ...tmuxBase,
508                      "new-window", "-t", TMUX_SESSION, "-n", windowName, "-c", cwd,
509                      "pi", "--session", destSessionFile,
510                  ]);
511                  if (code !== 0) {
512                      ctx.ui.notify(
513                          `Failed to create tmux window: ${stderr || stdout || "unknown error"}`,
514                          "error"
515                      );
516                      return;
517                  }
518                  // Attach to the session (new window is now current)
519                  tmuxArgs = [...tmuxBase, "-u", "attach", "-t", TMUX_SESSION];
520              } else {
521                  // Create new session with first window (attaches automatically)
522                  tmuxArgs = [
523                      ...tmuxBase, "-u", "new-session",
524                      "-s", TMUX_SESSION, "-n", windowName, "-c", cwd,
525                      "--", "pi", "--session", destSessionFile,
526                  ];
527              }
528  
529              // --- Tear down parent terminal ---
530  
531              process.stdout.write("\x1b[<u");      // Pop kitty keyboard protocol
532              process.stdout.write("\x1b[?2004l");  // Disable bracketed paste
533              process.stdout.write("\x1b[?25h");    // Show cursor
534              process.stdout.write("\r\n");
535  
536              if (process.stdin.isTTY && process.stdin.setRawMode) {
537                  process.stdin.setRawMode(false);
538              }
539  
540              // --- Spawn tmux ---
541  
542              const child = spawn("tmux", tmuxArgs, { stdio: "inherit" });
543  
544              child.once("spawn", () => {
545                  // Trash the original session file to prevent duplicates in /resume,
546                  // but only if it's in the standard Pi sessions directory. Custom
547                  // --session paths should not be trashed as that would be surprising.
548                  if (isInStandardSessionsDir(sourceSessionFile)) {
549                      void trashFileBestEffort(sourceSessionFile).then((trashed) => {
550                          if (!trashed) {
551                              process.stderr.write(
552                                  `\nNote: Could not trash original session file. Remove manually:\n  ${sourceSessionFile}\n`
553                              );
554                          }
555                      });
556                  } else {
557                      process.stderr.write(
558                          `\nNote: Session file is at a custom path and was not trashed:\n  ${sourceSessionFile}\n` +
559                              `The roamed session in tmux is independent. You may see duplicates in /resume.\n`
560                      );
561                  }
562  
563                  // Stop the parent from stealing keypresses
564                  process.stdin.removeAllListeners();
565                  process.stdin.destroy();
566  
567                  // Parent should not react to signals
568                  process.removeAllListeners("SIGINT");
569                  process.removeAllListeners("SIGTERM");
570                  process.on("SIGINT", () => {});
571                  process.on("SIGTERM", () => {});
572              });
573  
574              child.on("exit", (code) => process.exit(code ?? 0));
575              child.on("error", (err) => {
576                  process.stderr.write(`Failed to launch tmux: ${err.message}\n`);
577                  process.exit(1);
578              });
579          },
580      });
581  }