gsd-pi 0.1.0

package/src/resources/extensions/worktree/index.ts

@@ -0,0 +1,420 @@
+/**
+ * Worktree Extension — /wt
+ *
+ * Manage git worktrees for parallel extension development.
+ * Each worktree gets its own branch, isolated from main.
+ *
+ * Commands:
+ *   /wt new <name>   — create worktree + branch, switch agent cwd there
+ *   /wt ls           — list all worktrees with status
+ *   /wt switch <name>— switch agent cwd to a worktree (or main)
+ *   /wt merge <name> — squash merge into main, delete worktree + branch
+ *   /wt rm <name>    — delete worktree + branch (discard changes)
+ *   /wt status       — show current worktree and git status summary
+ */
+
+import { execSync } from "node:child_process";
+import * as path from "node:path";
+import type { ExtensionAPI, ExtensionCommandContext, ExtensionContext } from "@mariozechner/pi-coding-agent";
+
+// ─── Helpers ──────────────────────────────────────────────────────────────────
+
+const REPO_ROOT = "/Users/lexchristopherson/.pi";
+const WORKTREES_DIR = path.join(REPO_ROOT, "worktrees");
+
+function run(cmd: string, cwd = REPO_ROOT): string {
+	try {
+		return execSync(cmd, { cwd, encoding: "utf8", stdio: ["pipe", "pipe", "pipe"] }).trim();
+	} catch (e: unknown) {
+		const err = e as { stderr?: string; message?: string };
+		throw new Error(err.stderr?.trim() || err.message || String(e));
+	}
+}
+
+interface WorktreeInfo {
+	name: string;
+	path: string;
+	branch: string;
+	isMain: boolean;
+	commitCount: number;
+	dirty: boolean;
+}
+
+function listWorktrees(): WorktreeInfo[] {
+	const raw = run("git worktree list --porcelain");
+	const entries = raw.split("\n\n").filter(Boolean);
+
+	return entries.map((entry) => {
+		const lines = entry.split("\n");
+		const wtPath = lines.find((l) => l.startsWith("worktree "))?.replace("worktree ", "") ?? "";
+		const branch = lines.find((l) => l.startsWith("branch "))?.replace("branch refs/heads/", "") ?? "(detached)";
+		const isMain = wtPath === REPO_ROOT;
+
+		let commitCount = 0;
+		let dirty = false;
+
+		if (!isMain) {
+			try {
+				const count = run(`git rev-list --count main..${branch}`, wtPath);
+				commitCount = parseInt(count, 10) || 0;
+				const status = run("git status --porcelain", wtPath);
+				dirty = status.length > 0;
+			} catch {
+				// branch might not exist yet
+			}
+		}
+
+		const name = isMain ? "main" : path.basename(wtPath);
+		return { name, path: wtPath, branch, isMain, commitCount, dirty };
+	});
+}
+
+function getCurrentWorktreeName(cwd: string): string {
+	if (cwd === REPO_ROOT) return "main";
+	const worktrees = listWorktrees();
+	const match = worktrees.find((w) => w.path === cwd);
+	return match?.name ?? "main";
+}
+
+function formatWorktreeList(worktrees: WorktreeInfo[], currentPath: string): string[] {
+	return worktrees.map((w) => {
+		const isCurrent = w.path === currentPath;
+		const marker = isCurrent ? "→ " : "  ";
+		const dirty = w.dirty ? " *" : "";
+		const commits = w.isMain ? "" : ` (${w.commitCount} commit${w.commitCount !== 1 ? "s" : ""} ahead)`;
+		return `${marker}${w.name}${dirty}${commits}  [${w.branch}]`;
+	});
+}
+
+function updateStatus(ctx: ExtensionContext, currentWorktree: string): void {
+	if (currentWorktree === "main") {
+		ctx.ui.setStatus("worktree", undefined);
+	} else {
+		ctx.ui.setStatus("worktree", ctx.ui.theme.fg("accent", `⑂ ${currentWorktree}`));
+	}
+}
+
+function getWorktreeCdCommand(wtPath: string): string {
+	return `cd ${JSON.stringify(wtPath)}`;
+}
+
+function getWorktreeExtensionLaunchCommand(wtPath: string): string {
+	return `pi -e ${JSON.stringify(wtPath)}`;
+}
+
+function formatWorktreeHandoff(wtPath: string): string[] {
+	return [
+		"Human terminal handoff:",
+		`  ${getWorktreeCdCommand(wtPath)}`,
+		"",
+		"To run pi against the worktree copy instead of main:",
+		`  ${getWorktreeExtensionLaunchCommand(wtPath)}`,
+	];
+}
+
+// ─── Extension ────────────────────────────────────────────────────────────────
+
+export default function worktreeExtension(pi: ExtensionAPI): void {
+	// Track current cwd — starts at repo root
+	let currentCwd = REPO_ROOT;
+
+	// Inject cwd into every bash call so the agent operates in the right worktree
+	pi.on("tool_call", async (event) => {
+		if (event.toolName !== "bash" || currentCwd === REPO_ROOT) return;
+		// Prepend cd so all bash commands run in the active worktree
+		const cmd = event.input.command as string;
+		if (!cmd.startsWith(`cd ${currentCwd}`)) {
+			return {
+				input: {
+					...event.input,
+					command: `cd ${currentCwd} && ${cmd}`,
+				},
+			};
+		}
+	});
+
+	// Restore cwd on session resume
+	pi.on("session_start", async (_event, ctx) => {
+		const entries = ctx.sessionManager.getEntries();
+		const last = [...entries]
+			.reverse()
+			.find(
+				(e: { type: string; customType?: string }) => e.type === "custom" && e.customType === "worktree-cwd",
+			) as { data?: { cwd: string } } | undefined;
+		if (last?.data?.cwd) {
+			currentCwd = last.data?.cwd;
+		}
+		updateStatus(ctx, getCurrentWorktreeName(currentCwd));
+	});
+
+	// ── /wt command ─────────────────────────────────────────────────────────────
+
+	pi.registerCommand("wt", {
+		description: "Manage git worktrees: /wt new|ls|switch|merge|rm|status [name]",
+
+		getArgumentCompletions: (prefix: string) => {
+			const subcommands = ["new", "ls", "switch", "merge", "rm", "status"];
+			const parts = prefix.trim().split(/\s+/);
+
+			if (parts.length <= 1) {
+				return subcommands
+					.filter((s) => s.startsWith(parts[0] ?? ""))
+					.map((s) => ({ value: s, label: s }));
+			}
+
+			const sub = parts[0];
+			if (["switch", "merge", "rm"].includes(sub)) {
+				try {
+					const worktrees = listWorktrees();
+					const namePrefix = parts[1] ?? "";
+					return worktrees
+						.filter((w) => !w.isMain && w.name.startsWith(namePrefix))
+						.map((w) => ({ value: `${sub} ${w.name}`, label: w.name }));
+				} catch {
+					return [];
+				}
+			}
+
+			return [];
+		},
+
+		handler: async (args: string, ctx: ExtensionCommandContext) => {
+			const parts = args.trim().split(/\s+/);
+			const sub = parts[0];
+			const name = parts[1];
+
+			await ctx.waitForIdle();
+
+			// ── ls ────────────────────────────────────────────────────────────────
+			if (!sub || sub === "ls") {
+				try {
+					const worktrees = listWorktrees();
+					const lines = formatWorktreeList(worktrees, currentCwd);
+					const header = `Worktrees (${worktrees.length}):`;
+					ctx.ui.notify([header, "", ...lines].join("\n"), "info");
+				} catch (e) {
+					ctx.ui.notify(`Error: ${(e as Error).message}`, "error");
+				}
+				return;
+			}
+
+			// ── status ────────────────────────────────────────────────────────────
+			if (sub === "status") {
+				try {
+					const wtName = getCurrentWorktreeName(currentCwd);
+					const status = run("git status --short", currentCwd);
+					const branch = run("git branch --show-current", currentCwd);
+					const lines = [
+						`Current worktree: ${wtName}`,
+						`Branch: ${branch}`,
+						`Path: ${currentCwd}`,
+						"",
+						...(wtName === "main" ? [] : [...formatWorktreeHandoff(currentCwd), ""]),
+						status || "(clean)",
+					];
+					ctx.ui.notify(lines.join("\n"), "info");
+				} catch (e) {
+					ctx.ui.notify(`Error: ${(e as Error).message}`, "error");
+				}
+				return;
+			}
+
+			// ── new ───────────────────────────────────────────────────────────────
+			if (sub === "new") {
+				let wtName = name;
+				if (!wtName) {
+					wtName = (await ctx.ui.input("Worktree name:", "my-experiment")) ?? undefined;
+					if (!wtName) return;
+				}
+				wtName = wtName.replace(/[^a-z0-9_-]/gi, "-").toLowerCase();
+
+				const wtPath = path.join(WORKTREES_DIR, wtName);
+				try {
+					// Check if already exists
+					const existing = listWorktrees();
+					if (existing.some((w) => w.name === wtName)) {
+						ctx.ui.notify(`Worktree '${wtName}' already exists. Use /wt switch ${wtName}`, "error");
+						return;
+					}
+
+					// Create from current HEAD of main
+					run(`git worktree add ${wtPath} -b ${wtName}`, REPO_ROOT);
+
+					// Switch agent cwd
+					currentCwd = wtPath;
+					pi.appendEntry("worktree-cwd", { cwd: currentCwd });
+					updateStatus(ctx, wtName);
+
+					ctx.ui.notify(
+						[
+							`✓ Created worktree '${wtName}'`,
+							`  Branch: ${wtName}`,
+							`  Path: ${wtPath}`,
+							`  Agent is now working in this worktree.`,
+							"",
+							...formatWorktreeHandoff(wtPath),
+							"",
+							`  Use /wt merge ${wtName} when done.`,
+						].join("\n"),
+						"info",
+					);
+				} catch (e) {
+					ctx.ui.notify(`Failed to create worktree: ${(e as Error).message}`, "error");
+				}
+				return;
+			}
+
+			// ── switch ────────────────────────────────────────────────────────────
+			if (sub === "switch") {
+				if (!name) {
+					ctx.ui.notify("Usage: /wt switch <name>  (use 'main' to go back)", "error");
+					return;
+				}
+
+				if (name === "main") {
+					currentCwd = REPO_ROOT;
+					pi.appendEntry("worktree-cwd", { cwd: currentCwd });
+					updateStatus(ctx, "main");
+					ctx.ui.notify("Switched to main repo.", "info");
+					return;
+				}
+
+				const worktrees = listWorktrees();
+				const target = worktrees.find((w) => w.name === name);
+				if (!target) {
+					ctx.ui.notify(`No worktree named '${name}'. Use /wt ls to see all.`, "error");
+					return;
+				}
+
+				currentCwd = target.path;
+				pi.appendEntry("worktree-cwd", { cwd: currentCwd });
+				updateStatus(ctx, name);
+				ctx.ui.notify(
+					[
+						`Switched to worktree '${name}' (${target.branch})`,
+						`  Path: ${target.path}`,
+						"",
+						...formatWorktreeHandoff(target.path),
+					].join("\n"),
+					"info",
+				);
+				return;
+			}
+
+			// ── merge ─────────────────────────────────────────────────────────────
+			if (sub === "merge") {
+				const wtName = name ?? getCurrentWorktreeName(currentCwd);
+				if (wtName === "main") {
+					ctx.ui.notify("Can't merge main into itself.", "error");
+					return;
+				}
+
+				const worktrees = listWorktrees();
+				const target = worktrees.find((w) => w.name === wtName);
+				if (!target) {
+					ctx.ui.notify(`No worktree named '${wtName}'.`, "error");
+					return;
+				}
+
+				// Check for uncommitted changes
+				const dirty = run("git status --porcelain", target.path);
+				if (dirty) {
+					const ok = await ctx.ui.confirm(
+						`'${wtName}' has uncommitted changes`,
+						"These will be lost on merge. Continue?",
+					);
+					if (!ok) return;
+				}
+
+				// Get commit log for the editor prefill
+				let commitInfo = "";
+				try {
+					const count = run(`git rev-list --count main..${wtName}`, REPO_ROOT);
+					const log = run(`git log --oneline main..${wtName}`, REPO_ROOT);
+					commitInfo = `\n\n# Squashing ${count} commit(s):\n# ${log.split("\n").join("\n# ")}`;
+				} catch {}
+
+				// Ask for squash commit message
+				const msg = await ctx.ui.editor(
+					`Squash merge '${wtName}' → main\n(Edit the commit message)`,
+					`feat(${wtName}): ${commitInfo}`,
+				);
+				if (!msg?.trim()) {
+					ctx.ui.notify("Merge cancelled — no commit message.", "info");
+					return;
+				}
+
+				try {
+					// If currently in this worktree, switch to main first
+					if (currentCwd === target.path) {
+						currentCwd = REPO_ROOT;
+						pi.appendEntry("worktree-cwd", { cwd: currentCwd });
+						updateStatus(ctx, "main");
+					}
+
+					// Squash merge
+					run(`git merge --squash ${wtName}`, REPO_ROOT);
+					run(`git commit -m ${JSON.stringify(msg.trim())}`, REPO_ROOT);
+
+					// Clean up worktree and branch
+					run(`git worktree remove --force ${target.path}`, REPO_ROOT);
+					run(`git branch -D ${wtName}`, REPO_ROOT);
+
+					ctx.ui.notify(
+						[`✓ Squash merged '${wtName}' → main`, `  Branch and worktree deleted.`, `  Now on: main`].join("\n"),
+						"info",
+					);
+				} catch (e) {
+					ctx.ui.notify(`Merge failed: ${(e as Error).message}`, "error");
+				}
+				return;
+			}
+
+			// ── rm ────────────────────────────────────────────────────────────────
+			if (sub === "rm") {
+				const wtName = name ?? getCurrentWorktreeName(currentCwd);
+				if (wtName === "main") {
+					ctx.ui.notify("Can't remove main.", "error");
+					return;
+				}
+
+				const worktrees = listWorktrees();
+				const target = worktrees.find((w) => w.name === wtName);
+				if (!target) {
+					ctx.ui.notify(`No worktree named '${wtName}'.`, "error");
+					return;
+				}
+
+				let warn = `Delete worktree '${wtName}' and branch '${target.branch}'?`;
+				if (target.commitCount > 0) {
+					warn += `\n\n⚠️  ${target.commitCount} commit${target.commitCount !== 1 ? "s" : ""} will be lost.`;
+				}
+
+				const ok = await ctx.ui.confirm(warn, "This cannot be undone.");
+				if (!ok) return;
+
+				try {
+					if (currentCwd === target.path) {
+						currentCwd = REPO_ROOT;
+						pi.appendEntry("worktree-cwd", { cwd: currentCwd });
+						updateStatus(ctx, "main");
+					}
+
+					run(`git worktree remove --force ${target.path}`, REPO_ROOT);
+					run(`git branch -D ${target.branch}`, REPO_ROOT);
+
+					ctx.ui.notify(`✓ Deleted worktree '${wtName}' and branch '${target.branch}'.`, "info");
+				} catch (e) {
+					ctx.ui.notify(`Failed: ${(e as Error).message}`, "error");
+				}
+				return;
+			}
+
+			// ── fallthrough ───────────────────────────────────────────────────────
+			ctx.ui.notify(
+				["Usage: /wt <subcommand> [name]", "", "  new [name]    create worktree + branch, switch there", "  ls            list all worktrees", "  switch <name> switch agent cwd (use 'main' to go back)", "  merge [name]  squash merge → main, delete worktree + branch", "  rm [name]     discard worktree + branch", "  status        show current worktree git status"].join("\n"),
+				"info",
+			);
+		},
+	});
+}