yoyo-pi 0.1.4

package/extensions/clear-context.ts

@@ -0,0 +1,164 @@
+/**
+ * Clear/restore context snapshots.
+ *
+ * /clear
+ *   Saves the current branch context to .tmp/<random>.jsonl, then starts a new empty session.
+ *
+ * /restore <name>
+ *   Restores .tmp/<name>.jsonl into a fresh session file and switches to it.
+ *   /resume remains pi's built-in session picker.
+ */
+
+import { randomBytes, randomUUID } from "node:crypto";
+import * as fsp from "node:fs/promises";
+import * as path from "node:path";
+import type { ExtensionAPI, ExtensionCommandContext, SessionHeader } from "@earendil-works/pi-coding-agent";
+import { CURRENT_SESSION_VERSION } from "@earendil-works/pi-coding-agent";
+
+const SNAPSHOT_DIR = ".tmp";
+const SNAPSHOT_EXT = ".jsonl";
+const SNAPSHOT_NAME_RE = /^[A-Za-z0-9_-]+$/;
+
+function snapshotDir(cwd: string): string {
+	return path.resolve(cwd, SNAPSHOT_DIR);
+}
+
+function snapshotPath(cwd: string, name: string): string {
+	return path.join(snapshotDir(cwd), `${name}${SNAPSHOT_EXT}`);
+}
+
+function randomName(): string {
+	return randomBytes(5).toString("hex");
+}
+
+function normalizeSnapshotName(raw: string): string | undefined {
+	const firstArg = raw.trim().split(/\s+/)[0] ?? "";
+	const name = firstArg.endsWith(SNAPSHOT_EXT) ? firstArg.slice(0, -SNAPSHOT_EXT.length) : firstArg;
+	if (!name || !SNAPSHOT_NAME_RE.test(name)) return undefined;
+	return name;
+}
+
+async function createUniqueSnapshotPath(cwd: string): Promise<{ name: string; filePath: string }> {
+	await fsp.mkdir(snapshotDir(cwd), { recursive: true });
+
+	for (let i = 0; i < 100; i++) {
+		const name = randomName();
+		const filePath = snapshotPath(cwd, name);
+		try {
+			await fsp.access(filePath);
+		} catch {
+			return { name, filePath };
+		}
+	}
+
+	throw new Error("Failed to generate a unique snapshot name");
+}
+
+async function saveCurrentBranch(ctx: ExtensionCommandContext): Promise<{ name: string; filePath: string; entryCount: number }> {
+	const { name, filePath } = await createUniqueSnapshotPath(ctx.cwd);
+	const now = new Date().toISOString();
+	const currentFile = ctx.sessionManager.getSessionFile();
+	const currentHeader = ctx.sessionManager.getHeader();
+
+	const header: SessionHeader = {
+		type: "session",
+		version: CURRENT_SESSION_VERSION,
+		id: randomUUID(),
+		timestamp: now,
+		cwd: currentHeader?.cwd ?? ctx.cwd,
+		...(currentFile ? { parentSession: currentFile } : {}),
+	};
+
+	const entries = ctx.sessionManager.getBranch();
+	const content = [header, ...entries].map((entry) => JSON.stringify(entry)).join("\n") + "\n";
+	await fsp.writeFile(filePath, content, { encoding: "utf8", flag: "wx" });
+
+	return { name, filePath, entryCount: entries.length };
+}
+
+async function materializeSnapshot(snapshotFile: string, ctx: ExtensionCommandContext): Promise<string> {
+	const raw = await fsp.readFile(snapshotFile, "utf8");
+	const lines = raw.split(/\r?\n/).filter((line) => line.trim().length > 0);
+	if (lines.length === 0) {
+		throw new Error(`Snapshot is empty: ${snapshotFile}`);
+	}
+
+	const first = JSON.parse(lines[0]!) as SessionHeader;
+	if (first.type !== "session") {
+		throw new Error(`Snapshot is not a pi session file: ${snapshotFile}`);
+	}
+
+	const id = randomUUID();
+	const now = new Date().toISOString();
+	const header: SessionHeader = {
+		...first,
+		version: first.version ?? CURRENT_SESSION_VERSION,
+		id,
+		timestamp: now,
+		cwd: ctx.cwd,
+		parentSession: snapshotFile,
+	};
+	lines[0] = JSON.stringify(header);
+
+	const sessionDir = ctx.sessionManager.getSessionDir() || snapshotDir(ctx.cwd);
+	await fsp.mkdir(sessionDir, { recursive: true });
+	const destPath = path.join(sessionDir, `${now.replace(/[:.]/g, "-")}_${id}${SNAPSHOT_EXT}`);
+	await fsp.writeFile(destPath, lines.join("\n") + "\n", { encoding: "utf8", flag: "wx" });
+	return destPath;
+}
+
+export default function clearContextExtension(pi: ExtensionAPI) {
+	pi.registerCommand("clear", {
+		description: "Save current context to .tmp/<name>.jsonl and start a clean session",
+		handler: async (_args, ctx) => {
+			await ctx.waitForIdle();
+
+			const snapshot = await saveCurrentBranch(ctx);
+			const shortPath = path.join(SNAPSHOT_DIR, `${snapshot.name}${SNAPSHOT_EXT}`);
+
+			const result = await ctx.newSession({
+				parentSession: snapshot.filePath,
+				withSession: async (newCtx) => {
+					newCtx.ui.notify(`Context saved to ${shortPath}. Restore with /restore ${snapshot.name}`, "info");
+				},
+			});
+
+			if (result.cancelled) {
+				ctx.ui.notify(`Context saved to ${shortPath}, but clearing was cancelled.`, "warning");
+			}
+		},
+	});
+
+	pi.registerCommand("restore", {
+		description: "Restore a /clear snapshot by name: /restore <name>",
+		handler: async (args, ctx) => {
+			const name = normalizeSnapshotName(args);
+			if (!name) {
+				ctx.ui.notify("Usage: /restore <snapshot-name>", "warning");
+				return;
+			}
+
+			await ctx.waitForIdle();
+
+			const sourcePath = snapshotPath(ctx.cwd, name);
+			try {
+				await fsp.access(sourcePath);
+			} catch {
+				ctx.ui.notify(`No snapshot found at ${path.join(SNAPSHOT_DIR, `${name}${SNAPSHOT_EXT}`)}`, "error");
+				return;
+			}
+
+			const restorePath = await materializeSnapshot(sourcePath, ctx);
+			const result = await ctx.switchSession(restorePath, {
+				withSession: async (newCtx) => {
+					newCtx.ui.notify(`Restored context from ${path.join(SNAPSHOT_DIR, `${name}${SNAPSHOT_EXT}`)}`, "info");
+				},
+			});
+
+			if (result.cancelled) {
+				await fsp.rm(restorePath, { force: true });
+				ctx.ui.notify("Restore cancelled.", "warning");
+			}
+		},
+	});
+}