@gavin-hu/my-pi-agent 0.0.0-stage → 0.7.0

package/extensions/worktree/index.ts

@@ -0,0 +1,221 @@
+/**
+ * worktree — Isolated git worktrees (EnterWorktree / ExitWorktree) for Pi.
+ *
+ * Mimics Claude Code's worktree isolation:
+ *   - `enter_worktree` creates (or enters) an isolated `git worktree` under
+ *     `.pi/worktrees/<name>` on branch `worktree-<name>`, and rebinds every
+ *     path-taking tool to it.
+ *   - `exit_worktree` returns to the main checkout and cleans up, checking for
+ *     uncommitted work and unpushed commits first.
+ *
+ * Pi has no mutable session cwd and its built-ins capture cwd at construction,
+ * so isolation is implemented by overriding the built-in tools at load time
+ * (`root-tools.ts`) and re-rooting them at call time through the active root.
+ *
+ * This module wires the pieces together: shared state lives in `runtime.ts`,
+ * the lifecycle in `lifecycle.ts`, and the model surface in `tools.ts` and
+ * `commands.ts`.
+ *
+ * This is the `worktree` extension entrypoint: Pi calls the default export,
+ * which registers the extension through {@link registerWorktree}.
+ */
+
+import { existsSync } from "node:fs";
+import { fileURLToPath } from "node:url";
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { isExtensionEnabled } from "../../lib/env.ts";
+import { GLYPHS } from "../../lib/ui.ts";
+import { canonicalize, checkCheckout, repoRoot, unlockWorktree } from "./git.ts";
+import { analyzeBashCommand, guardFileTool } from "./guard.ts";
+import {
+	enterWorktree,
+	findInactiveOverrides,
+	isWorktreeCancelled,
+	refusalMessage,
+	worktreeFacts,
+	worktreeLabel,
+} from "./lifecycle.ts";
+import { registerCommands } from "./commands.ts";
+import { type BuiltinFactory, registerRootTools } from "./root-tools.ts";
+import { loadState, persistState, type WorktreeState } from "./state.ts";
+import { registerTools } from "./tools.ts";
+import {
+	ENV_BRANCH,
+	ENV_MAIN,
+	ENV_ROOT,
+	applyWorktreeEnv,
+	clearConfigCache,
+	configFor,
+	getActive,
+	getInactiveOverrides,
+	getRoot,
+	getSkipOverrides,
+	publishWorktree,
+	setActive,
+	setBaseCwd,
+	setInactiveOverrides,
+	setStatus,
+} from "./runtime.ts";
+
+export interface WorktreeRegistration {
+	/** File Pi records as the extension entry, for inactive-override reporting. */
+	entryPath: string;
+	/** Test seam: override the host built-in factories used by `root-tools.ts`. */
+	createBuiltin?: BuiltinFactory;
+}
+
+/** Prompt section describing the active worktree, shown while the session is isolated. */
+export function formatWorktreeSection(state: WorktreeState): string {
+	return (
+		`You are in an isolated git worktree; relative paths resolve there, not in the main checkout.\n` +
+		`${worktreeFacts(state, "worktree", 10)}\n` +
+		`Do not edit files or run git against the main checkout.` +
+		(state.borrowed ? "\nThis worktree is inherited; do not exit it." : "\nUse exit_worktree to return.")
+	);
+}
+
+export function registerWorktree(pi: ExtensionAPI, { entryPath, createBuiltin }: WorktreeRegistration): void {
+	registerRootTools(pi, { getRoot, getSkipOverrides }, { createBuiltin });
+
+	pi.registerFlag("worktree", {
+		description: "Create or enter a git worktree at session start",
+		type: "string",
+	});
+
+	// --- lifecycle ---------------------------------------------------------
+
+	pi.on("session_start", async (event, ctx) => {
+		setBaseCwd(ctx.cwd);
+		setActive(null);
+		clearConfigCache();
+
+		// A child `pi` process (subagent) inherits an active worktree. Bind to
+		// it without creating, locking, or persisting anything; the parent owns
+		// the lifecycle. Only on a fresh process start: a `reload` re-runs this
+		// factory inside the parent, which must restore its own recorded state.
+		const envRoot = event.reason === "startup" ? process.env[ENV_ROOT] : undefined;
+		if (envRoot) {
+			const borrowedPath = canonicalize(envRoot);
+			if (existsSync(borrowedPath) && (await repoRoot(pi, borrowedPath))) {
+				const borrowedMain = process.env[ENV_MAIN] ? canonicalize(process.env[ENV_MAIN]) : borrowedPath;
+				const borrowed: WorktreeState = {
+					active: true,
+					borrowed: true,
+					path: borrowedPath,
+					branch: process.env[ENV_BRANCH],
+					repoRoot: borrowedMain,
+					baseRef: "HEAD",
+					baseRefMode: "head",
+					createdByUs: false,
+				};
+				setActive(borrowed);
+				setStatus(ctx, `${GLYPHS.worktree} ${worktreeLabel(borrowed)} (inherited)`);
+				publishWorktree(pi, borrowed);
+				ctx.ui.notify(`Using parent worktree ${borrowed.path}`, "info");
+			} else {
+				delete process.env[ENV_ROOT];
+				delete process.env[ENV_BRANCH];
+				delete process.env[ENV_MAIN];
+			}
+		}
+
+		const recorded = loadState(ctx);
+		if (!getActive() && recorded?.active) {
+			const check = await checkCheckout(pi, recorded.path, recorded.repoRoot);
+			if (check.ok) {
+				setActive(recorded);
+				setStatus(ctx, `${GLYPHS.worktree} ${worktreeLabel(recorded)}`);
+				ctx.ui.notify(`Restored worktree ${worktreeLabel(recorded)}`, "info");
+			} else {
+				// "unverified" keeps the recorded binding so a later resume can
+				// retry; the others clear it because the worktree is gone/unsafe.
+				if (check.reason !== "unverified") {
+					persistState((customType, data) => pi.appendEntry(customType, data), { ...recorded, active: false });
+				}
+				setStatus(ctx, undefined);
+				ctx.ui.notify(refusalMessage(check, recorded.path, ctx.cwd), "warning");
+			}
+		} else if (!getActive()) {
+			setStatus(ctx, undefined);
+		}
+
+		setInactiveOverrides(findInactiveOverrides(pi, entryPath));
+		const inactive = getInactiveOverrides();
+		if (inactive.length > 0) {
+			ctx.ui.notify(
+				`Worktree isolation is not active for: ${inactive.join(", ")}. Those tools keep their default behavior; set skipOverrides or remove the conflicting extension.`,
+				"warning",
+			);
+		}
+
+		if (event.reason === "startup" && !getActive()) {
+			const flag = pi.getFlag("worktree");
+			if (typeof flag === "string" && flag.trim()) {
+				const name = flag.trim() === "true" ? undefined : flag.trim();
+				try {
+					const { state } = await enterWorktree(pi, ctx, { name });
+					ctx.ui.notify(`Entered worktree ${worktreeLabel(state)}`, "info");
+				} catch (error) {
+					if (isWorktreeCancelled(error)) {
+						ctx.ui.notify(error instanceof Error ? error.message : String(error), "info");
+						return;
+					}
+					ctx.ui.notify(`Could not enter worktree: ${error instanceof Error ? error.message : String(error)}`, "error");
+				}
+			}
+		}
+
+		// Keep the process environment in step with the restored binding: a peer
+		// extension or a child process reads PI_WORKTREE_* (lib/env.ts), and this
+		// session may have replaced one that exported a different root.
+		applyWorktreeEnv(getActive());
+	});
+
+	pi.on("session_shutdown", async (_event, ctx) => {
+		// Release our lock but never remove the worktree on shutdown; cleanup is
+		// an explicit exit or `/worktree prune`.
+		const state = getActive();
+		if (state && !state.borrowed && state.lockReason) {
+			try {
+				await unlockWorktree(pi, state.repoRoot, state.path);
+			} catch {
+				// Shutdown must not throw.
+			}
+		}
+		// Do not leave this session's root exported for the next one in this process.
+		applyWorktreeEnv(null);
+		setStatus(ctx, undefined);
+	});
+
+	pi.on("before_agent_start", (event) => {
+		const state = getActive();
+		if (!state) return;
+		event.systemPromptOptions.cwd = state.path;
+		event.systemPromptOptions.sections = {
+			...event.systemPromptOptions.sections,
+			worktree: formatWorktreeSection(state),
+		};
+	});
+
+	pi.on("tool_call", (event) => {
+		const state = getActive();
+		if (!state) return;
+		const config = configFor(state.repoRoot);
+		if (event.toolName === "bash") {
+			const command = (event.input as { command?: string }).command ?? "";
+			return analyzeBashCommand(command, state.path, config);
+		}
+		return guardFileTool(event.toolName, event.input, state.path, config);
+	});
+
+	// --- model surface -----------------------------------------------------
+
+	registerTools(pi);
+	registerCommands(pi);
+}
+
+/** Pi extension entrypoint. */
+export default function worktreeExtension(pi: ExtensionAPI): void {
+	if (!isExtensionEnabled("worktree")) return;
+	registerWorktree(pi, { entryPath: canonicalize(fileURLToPath(import.meta.url)) });
+}

package/extensions/worktree/lifecycle.ts

@@ -0,0 +1,14 @@
+/**
+ * Worktree lifecycle, split by phase and re-exported here for callers.
+ *
+ * These are stateless functions over the shared runtime state in `runtime.ts`;
+ * the extension's tools, commands, and event handlers import them from this
+ * barrel.
+ */
+
+export { WorktreeCancelledError, isWorktreeCancelled } from "./errors.ts";
+export { enterWorktree, type EnterOptions } from "./enter.ts";
+export { exitWorktree, type ExitOptions } from "./exit.ts";
+export { pruneWorktrees } from "./prune.ts";
+export { findInactiveOverrides } from "./root-tools.ts";
+export { refusalMessage, stateSummary, worktreeFacts, worktreeLabel, worktreeStatus } from "./status.ts";

package/extensions/worktree/prune.ts

@@ -0,0 +1,99 @@
+/**
+ * Worktree pruning: remove clean, unused managed worktrees past the configured
+ * age, unlocking entries held by dead processes. Age follows the registry's
+ * `lastUsedAt` (falling back to directory mtime for unregistered worktrees).
+ */
+
+import { statSync } from "node:fs";
+import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
+import { loadConfig } from "./config.ts";
+import {
+	defaultBranch,
+	deleteBranch,
+	inspectWork,
+	isStaleLock,
+	listManagedWorktrees,
+	mergeBase,
+	repoRoot,
+	unlockWorktree,
+	worktreeRemove,
+} from "./git.ts";
+import { reconcileRegistry, removeRecord, type WorktreeRecord } from "./registry.ts";
+import { getActive } from "./runtime.ts";
+
+/**
+ * Age in days from the most recent evidence of use: the registry's
+ * `lastUsedAt`, or the directory mtime when there is no record.
+ */
+function ageInDays(path: string, lastUsedAt = 0): number {
+	let newest = lastUsedAt;
+	try {
+		newest = Math.max(newest, statSync(path).mtimeMs);
+	} catch {
+		return 0;
+	}
+	return newest === 0 ? 0 : (Date.now() - newest) / 86_400_000;
+}
+
+/** Remove clean, unused managed worktrees past the prune age. Returns a report. */
+export async function pruneWorktrees(pi: ExtensionAPI, ctx: ExtensionContext): Promise<string> {
+	const root = await repoRoot(pi, ctx.cwd);
+	if (!root) return "Not a git repository.";
+	const config = loadConfig(root);
+	const managed = await listManagedWorktrees(pi, root, config);
+	const { registry } = reconcileRegistry(root, config, managed);
+	const byPath = new Map<string, WorktreeRecord>(registry.worktrees.map((record) => [record.path, record]));
+	const branch = await defaultBranch(pi, root);
+	const current = getActive();
+	const removed: string[] = [];
+	const kept: string[] = [];
+
+	for (const entry of managed) {
+		if (current?.path === entry.path) {
+			kept.push(`${entry.path} (current)`);
+			continue;
+		}
+		if (entry.locked !== undefined) {
+			if (isStaleLock(entry.locked)) {
+				await unlockWorktree(pi, root, entry.path);
+			} else {
+				kept.push(`${entry.path} (locked)`);
+				continue;
+			}
+		}
+		if (!branch) {
+			kept.push(`${entry.path} (no default branch)`);
+			continue;
+		}
+		const fork = await mergeBase(pi, entry.path, `origin/${branch}`);
+		if (!fork) {
+			kept.push(`${entry.path} (cannot verify)`);
+			continue;
+		}
+		const { hasWork, unverifiable } = await inspectWork(pi, entry.path, fork);
+		if (hasWork || unverifiable) {
+			kept.push(`${entry.path} (has work)`);
+			continue;
+		}
+		if (ageInDays(entry.path, byPath.get(entry.path)?.lastUsedAt) < config.pruneAfterDays) {
+			kept.push(`${entry.path} (recent)`);
+			continue;
+		}
+		const result = await worktreeRemove(pi, root, entry.path, false);
+		if (!result.ok) {
+			kept.push(`${entry.path} (${result.error})`);
+			continue;
+		}
+		removed.push(entry.path);
+		removeRecord(root, config, entry.path);
+		if (entry.branch) {
+			const deleted = await deleteBranch(pi, root, entry.branch, false);
+			if (!deleted.ok) kept.push(`branch ${entry.branch} (unmerged, kept)`);
+		}
+	}
+
+	return (
+		[...removed.map((path) => `Removed ${path}`), ...kept.map((path) => `Kept ${path}`)].join("\n") ||
+		"Nothing to prune."
+	);
+}

package/extensions/worktree/registry.ts

@@ -0,0 +1,221 @@
+/**
+ * The managed-worktree registry and derived checkout status.
+ *
+ * Git is authoritative for whether a checkout exists; this registry only records
+ * provenance (`createdByUs`, base, PR, timestamps) and usage (`lastUsedAt`). It
+ * is a disposable cache: a missing or corrupt file loads as empty, writes are
+ * best-effort, and reads reconcile against `git worktree list`.
+ *
+ * `statusOf` is the single derived view used by status, prune, and (later)
+ * switch/exit, so those do not each re-derive dirty/ahead/behind/merged state.
+ */
+
+import { existsSync, mkdirSync, readFileSync, renameSync, writeFileSync } from "node:fs";
+import { dirname, join } from "node:path";
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import type { WorktreeConfig } from "./config.ts";
+import {
+	canonicalize,
+	commitsAhead,
+	commitsBehind,
+	isBranchMerged,
+	isStaleLock,
+	type ManagedWorktree,
+	statusEntries,
+} from "./git.ts";
+
+const REGISTRY_VERSION = 1;
+
+export interface WorktreeRecord {
+	/** Canonical path; the record key. */
+	path: string;
+	name?: string;
+	branch?: string;
+	repoRoot: string;
+	base: { ref: string; commit?: string; mode: "fresh" | "head" };
+	/** Epoch ms. Zero means "unknown", for records synthesized from git. */
+	createdAt: number;
+	/** Epoch ms. Zero means "unknown". */
+	lastUsedAt: number;
+	pr?: { number: number; host: string };
+	createdByUs: boolean;
+}
+
+export interface WorktreeRegistry {
+	version: number;
+	worktrees: WorktreeRecord[];
+}
+
+export interface CheckoutStatus {
+	path: string;
+	branch?: string;
+	head?: string;
+	state: "clean" | "dirty" | "missing";
+	/** Number of changed files (porcelain entries). */
+	changed: number;
+	ahead: number;
+	behind: number;
+	merged: boolean;
+	locked: boolean;
+	staleLock: boolean;
+	lastUsedAt?: number;
+}
+
+export function registryPath(repoRoot: string, config: WorktreeConfig): string {
+	return join(repoRoot, config.dir, "index.json");
+}
+
+function emptyRegistry(): WorktreeRegistry {
+	return { version: REGISTRY_VERSION, worktrees: [] };
+}
+
+function isRecord(value: unknown): value is WorktreeRecord {
+	if (!value || typeof value !== "object") return false;
+	const record = value as Partial<WorktreeRecord>;
+	return typeof record.path === "string" && typeof record.repoRoot === "string";
+}
+
+/** Read the registry. Missing, unreadable, or malformed files load as empty. */
+export function loadRegistry(repoRoot: string, config: WorktreeConfig): WorktreeRegistry {
+	try {
+		const parsed = JSON.parse(readFileSync(registryPath(repoRoot, config), "utf-8")) as Partial<WorktreeRegistry>;
+		if (!parsed || typeof parsed !== "object" || !Array.isArray(parsed.worktrees)) return emptyRegistry();
+		return { version: REGISTRY_VERSION, worktrees: parsed.worktrees.filter(isRecord) };
+	} catch {
+		return emptyRegistry();
+	}
+}
+
+/** Atomically persist the registry. Failures are swallowed: it is a cache. */
+function saveRegistry(repoRoot: string, config: WorktreeConfig, registry: WorktreeRegistry): void {
+	try {
+		const file = registryPath(repoRoot, config);
+		mkdirSync(dirname(file), { recursive: true });
+		const tmp = `${file}.${process.pid}.tmp`;
+		writeFileSync(tmp, `${JSON.stringify(registry, null, 2)}\n`);
+		renameSync(tmp, file);
+	} catch {
+		// Losing the registry only loses provenance, never a worktree.
+	}
+}
+
+/** Insert or replace a record, preserving an existing `createdAt`. */
+export function upsertRecord(repoRoot: string, config: WorktreeConfig, record: WorktreeRecord): void {
+	const registry = loadRegistry(repoRoot, config);
+	const path = canonicalize(record.path);
+	const existing = registry.worktrees.find((candidate) => candidate.path === path);
+	const next: WorktreeRecord = existing
+		? { ...record, path, createdAt: existing.createdAt || record.createdAt }
+		: { ...record, path };
+	registry.worktrees = registry.worktrees.filter((candidate) => candidate.path !== path);
+	registry.worktrees.push(next);
+	registry.worktrees.sort((a, b) => a.path.localeCompare(b.path));
+	saveRegistry(repoRoot, config, registry);
+}
+
+/** Mark a worktree as just used. No-op when it is not registered. */
+export function touchRecord(repoRoot: string, config: WorktreeConfig, path: string, at = Date.now()): void {
+	const registry = loadRegistry(repoRoot, config);
+	const wanted = canonicalize(path);
+	const record = registry.worktrees.find((candidate) => candidate.path === wanted);
+	if (!record) return;
+	record.lastUsedAt = at;
+	saveRegistry(repoRoot, config, registry);
+}
+
+export function removeRecord(repoRoot: string, config: WorktreeConfig, path: string): void {
+	const registry = loadRegistry(repoRoot, config);
+	const wanted = canonicalize(path);
+	const next = registry.worktrees.filter((candidate) => candidate.path !== wanted);
+	if (next.length === registry.worktrees.length) return;
+	registry.worktrees = next;
+	saveRegistry(repoRoot, config, registry);
+}
+
+/**
+ * Reconcile the registry against the worktrees git currently reports: drop
+ * records whose checkout is gone, and synthesize records for managed worktrees
+ * that have none. Persists only when something changed, and only when
+ * `options.persist` is not `false` (the read-only status path passes `false`).
+ */
+export function reconcileRegistry(
+	repoRoot: string,
+	config: WorktreeConfig,
+	managed: ManagedWorktree[],
+	options: { persist?: boolean } = {},
+): { registry: WorktreeRegistry; changed: boolean } {
+	const registry = loadRegistry(repoRoot, config);
+	const live = new Set(managed.map((entry) => canonicalize(entry.path)));
+	const byPath = new Map(registry.worktrees.map((record) => [record.path, record] as const));
+
+	let changed = false;
+	const next: WorktreeRecord[] = [];
+	for (const record of registry.worktrees) {
+		if (live.has(record.path) || existsSync(record.path)) {
+			next.push(record);
+		} else {
+			changed = true;
+		}
+	}
+	for (const entry of managed) {
+		const path = canonicalize(entry.path);
+		if (byPath.has(path)) continue;
+		next.push({
+			path,
+			branch: entry.branch,
+			repoRoot,
+			base: { ref: "HEAD", mode: "head" },
+			createdAt: 0,
+			lastUsedAt: 0,
+			createdByUs: false,
+		});
+		changed = true;
+	}
+
+	registry.worktrees = next.sort((a, b) => a.path.localeCompare(b.path));
+	if (changed && options.persist !== false) saveRegistry(repoRoot, config, registry);
+	return { registry, changed };
+}
+
+export interface StatusOptions {
+	/** Ref to compare against (for example `origin/main`). Omit when unknown. */
+	baseRef?: string;
+	record?: WorktreeRecord;
+}
+
+/** Derive one checkout's status from its git listing and registry record. */
+export async function statusOf(
+	pi: ExtensionAPI,
+	entry: ManagedWorktree,
+	options: StatusOptions = {},
+): Promise<CheckoutStatus> {
+	const path = canonicalize(entry.path);
+	const locked = entry.locked !== undefined;
+	const staleLock = locked && entry.locked ? isStaleLock(entry.locked) : false;
+	const base: CheckoutStatus = {
+		path,
+		branch: entry.branch,
+		head: entry.head,
+		state: "missing",
+		changed: 0,
+		ahead: 0,
+		behind: 0,
+		merged: false,
+		locked,
+		staleLock,
+		lastUsedAt: options.record?.lastUsedAt,
+	};
+
+	if (!existsSync(path)) return base;
+
+	const dirty = await statusEntries(pi, path);
+	const baseRef = options.baseRef;
+	base.state = dirty.length > 0 ? "dirty" : "clean";
+	base.changed = dirty.length;
+	if (baseRef) {
+		base.ahead = await commitsAhead(pi, path, baseRef);
+		base.behind = await commitsBehind(pi, path, baseRef);
+		base.merged = base.ahead === 0 || (await isBranchMerged(pi, path, "HEAD", baseRef));
+	}
+	return base;
+}