1  import { readFileSync } from "node:fs";
  2  import { access } from "node:fs/promises";
  3  import os from "node:os";
  4  import path from "node:path";
  5  
  6  import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
  7  
  8  const AGENT_DIR = path.join(os.homedir(), ".pi", "agent");
  9  
 10  const REWIND_EXTENSION_DIR = path.join(AGENT_DIR, "extensions", "rewind");
 11  const REWIND_EXTENSION_CANDIDATES = [
 12    "index.ts",
 13    "index.js",
 14    path.join("dist", "index.js"),
 15    path.join("build", "index.js"),
 16    "package.json",
 17  ];
 18  
 19  interface RewindForkPendingData {
 20    v: 2;
 21    current: string;
 22    undo?: string;
 23  }
 24  
 25  async function isRewindInstalled(): Promise<boolean> {
 26    try {
 27      await access(REWIND_EXTENSION_DIR);
 28    } catch {
 29      return false;
 30    }
 31  
 32    for (const relPath of REWIND_EXTENSION_CANDIDATES) {
 33      try {
 34        await access(path.join(REWIND_EXTENSION_DIR, relPath));
 35        return true;
 36      } catch {
 37        // keep looking
 38      }
 39    }
 40  
 41    return false;
 42  }
 43  
 44  function extractUserMessageText(content: string | Array<{ type: string; text?: string }>): string {
 45    if (typeof content === "string") {
 46      return content;
 47    }
 48  
 49    return content
 50      .filter((part): part is { type: "text"; text: string } => part.type === "text" && typeof part.text === "string")
 51      .map((part) => part.text)
 52      .join("");
 53  }
 54  
 55  function isRewindForkPendingData(value: unknown): value is RewindForkPendingData {
 56    if (!value || typeof value !== "object") {
 57      return false;
 58    }
 59  
 60    const data = value as Partial<RewindForkPendingData>;
 61    return data.v === 2 && typeof data.current === "string" && data.current.length > 0;
 62  }
 63  
 64  function loadLatestRewindForkPending(sessionFile: string): RewindForkPendingData | null {
 65    const raw = readFileSync(sessionFile, "utf8");
 66    const lines = raw.split("\n");
 67  
 68    for (let index = lines.length - 1; index >= 0; index -= 1) {
 69      const line = lines[index]?.trim();
 70      if (!line) {
 71        continue;
 72      }
 73  
 74      try {
 75        const entry = JSON.parse(line) as {
 76          type?: string;
 77          customType?: string;
 78          data?: unknown;
 79        };
 80  
 81        if (
 82          entry.type === "custom" &&
 83          entry.customType === "rewind-fork-pending" &&
 84          isRewindForkPendingData(entry.data)
 85        ) {
 86          return entry.data;
 87        }
 88      } catch {
 89        // ignore malformed lines and keep scanning backward
 90      }
 91    }
 92  
 93    return null;
 94  }
 95  
 96  function seedChildRewindStateFromParent(
 97    parentSessionFile: string,
 98    sessionManager: { appendCustomEntry(customType: string, data?: unknown): string },
 99  ): void {
100    let pending: RewindForkPendingData | null = null;
101  
102    try {
103      pending = loadLatestRewindForkPending(parentSessionFile);
104    } catch {
105      return;
106    }
107  
108    if (!pending) {
109      return;
110    }
111  
112    const snapshots = [pending.current];
113    const data: { v: 2; snapshots: string[]; current: number; undo?: number } = {
114      v: 2,
115      snapshots,
116      current: 0,
117    };
118  
119    if (typeof pending.undo === "string" && pending.undo.length > 0) {
120      data.snapshots.push(pending.undo);
121      data.undo = 1;
122    }
123  
124    sessionManager.appendCustomEntry("rewind-op", data);
125  }
126  
127  async function requestConversationOnlyForkWhenRewindIsInstalled(pi: ExtensionAPI): Promise<boolean> {
128    if (!(await isRewindInstalled())) {
129      return false;
130    }
131  
132    pi.events.emit("rewind:fork-preference", {
133      mode: "conversation-only",
134      source: "fork-from-first",
135    });
136  
137    return true;
138  }
139  
140  export default function (pi: ExtensionAPI) {
141    pi.registerCommand("fork-from-first", {
142      description: "Fork current session from its first user message",
143      handler: async (_args, ctx) => {
144        await ctx.waitForIdle();
145  
146        const previousSessionFile = ctx.sessionManager.getSessionFile();
147        if (!previousSessionFile) {
148          if (ctx.hasUI) {
149            ctx.ui.notify("/fork-from-first requires a persisted session file", "error");
150          }
151          return;
152        }
153  
154        const firstUserEntry = ctx.sessionManager
155          .getEntries()
156          .find(
157            (entry) =>
158              entry.type === "message" &&
159              entry.message?.role === "user"
160          );
161  
162        if (!firstUserEntry) {
163          if (ctx.hasUI) {
164            ctx.ui.notify("No user message found to fork from", "warning");
165          }
166          return;
167        }
168  
169        const selectedText = extractUserMessageText(firstUserEntry.message.content);
170        const rewindInstalled = await requestConversationOnlyForkWhenRewindIsInstalled(pi);
171        if (ctx.hasUI && rewindInstalled) {
172          ctx.ui.notify("Rewind detected: forcing conversation-only fork", "info");
173        }
174  
175        const result = await ctx.newSession({
176          parentSession: previousSessionFile,
177          setup: async (sessionManager) => {
178            if (!rewindInstalled) {
179              return;
180            }
181  
182            seedChildRewindStateFromParent(previousSessionFile, sessionManager);
183          },
184        });
185  
186        if (result.cancelled) {
187          if (ctx.hasUI) {
188            ctx.ui.notify("Fork cancelled", "warning");
189          }
190          return;
191        }
192  
193        if (ctx.hasUI) {
194          ctx.ui.setEditorText(selectedText);
195          ctx.ui.notify("Forked from first message and switched to new session", "info");
196        }
197      },
198    });
199  }