pi-worktrees/.pi/extensions/pi-worktree.ts

/**
 * pi-worktree Extension
 *
 * Manages git worktrees through tool commands and an interactive `/worktree` command.
 *
 * Tools:
 * - git_worktree_create — Create a linked worktree for a branch
 * - git_worktree_list — List all linked worktrees
 * - git_worktree_switch — Switch working directory to a worktree
 * - git_worktree_remove — Remove a linked worktree
 *
 * Command:
 * - /worktree — Interactive picker to switch between worktrees
 */

import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import { stat } from "node:fs/promises";
import { defineTool } from "@earendil-works/pi-coding-agent";
import { Type } from "@earendil-works/pi-ai";
import { parseWorktreeList, type WorktreeEntry } from "./pi-worktree-parser";
import {
	validateCreateWorktreeInput,
	buildWorktreeAddCommand,
	validatePathDoesNotExist,
	isCurrentBranch,
	type WorktreeError,
} from "./pi-worktree-create";
import {
	validateRemoveWorktreeInput,
	isMainRepository,
	hasUncommittedChanges,
	buildWorktreeRemoveCommand,
} from "./pi-worktree-remove";
import {
	validateSwitchInput,
	resolveWorktreeTarget,
	isAlreadyInWorktree,
} from "./pi-worktree-switch";

// ─── Type Definitions ──────────────────────────────────────────────

/** Result of listing all worktrees. */
interface WorktreeListResult {
	/** All worktrees in the repository. */
	worktrees: WorktreeEntry[];
	/** Number of worktrees found. */
	count: number;
	/** Current working directory used for comparison. */
	currentPath: string;
}

// ─── Tool Parameter Schemas ────────────────────────────────────────

const createWorktreeParams = Type.Object({
	path: Type.String({ description: "Path for the new worktree (relative or absolute)" }),
	branch: Type.Optional(Type.String({ description: "Branch name to check out. Defaults to 'main'." })),
	create_branch: Type.Optional(
		Type.Boolean({
			description: "If true and branch doesn't exist, create it. Defaults to true.",
		}),
	),
});

const listWorktreeParams = Type.Object({});

const switchWorktreeParams = Type.Object({
	target: Type.String({ description: "Path or branch name of the target worktree to switch to." }),
});

const removeWorktreeParams = Type.Object({
	path: Type.String({ description: "Path of the worktree to remove." }),
	force: Type.Optional(Type.Boolean({ description: "Force removal even with uncommitted changes. Defaults to false." })),
});

// ─── Helper Functions ──────────────────────────────────────────────

/**
 * Check if the target path exists using Bun's fs module.
 * @param checkPath - Path to check (resolved against cwd)
 * @param ctx - Extension context for working directory
 * @returns True if the path exists
 */
async function pathExists(checkPath: string, ctx: ExtensionContext): Promise<boolean> {
	const resolvedPath = checkPath.startsWith("/") ? checkPath : `${ctx.cwd}/${checkPath}`;

	try {
		await stat(resolvedPath);
		return true;
	} catch {
		return false;
	}
}

/**
 * Get the current branch name from git.
 * @param pi - Extension API for executing commands
 * @param ctx - Extension context for working directory
 * @returns Current branch name, or "unknown" if detection fails
 */
async function getCurrentBranch(pi: ExtensionAPI, ctx: ExtensionContext): Promise<string> {
	const result = await pi.exec("git", ["rev-parse", "--abbrev-ref", "HEAD"], { cwd: ctx.cwd });

	if (result.exitCode !== 0) {
		return "unknown";
	}

	return result.stdout.trim();
}

/**
 * Format a WorktreeError into the tool response text.
 * @param error - The validation error
 * @returns Formatted error string with suggestion if available
 */
function formatValidationError(error: WorktreeError): string {
	const base = `Error: ${error.message}`;
	return error.suggestion ? `${base} Suggestion: ${error.suggestion}` : base;
}

/**
 * Format a success response for worktree operations.
 * @param action - The action performed (e.g., "Created")
 * @param details - Key-value pairs to include in the message
 * @returns Formatted success string
 */
function formatSuccessResponse(action: string, details: Record<string, string>): string {
	const parts = Object.entries(details).map(([key, value]) => `${key} '${value}'`);
	return `${action} worktree ${parts.join(" on ")}.`;
}

// ─── Tool Factories ────────────────────────────────────────────────

/**
 * Create the git_worktree_create tool definition.
 * @param pi - Extension API for registering and executing tools
 * @returns The configured tool definition
 */
function createListWorktreeTool(pi: ExtensionAPI) {
	return defineTool({
		name: "git_worktree_list",
		label: "Git Worktree List",
		description: "List all linked worktrees in the current repo.",
		parameters: listWorktreeParams,
		async execute(_toolCallId, _params, _signal, _onUpdate, ctx) {
			const result = await pi.exec("git", ["worktree", "list"], { cwd: ctx.cwd });

			if (result.exitCode !== 0) {
				return {
					content: [{ type: "text", text: `Error listing worktrees: ${result.stderr}. Suggestion: Ensure you are inside a git repository.` }],
					isError: true,
				};
			}

			const worktreeEntries = parseWorktreeList(result.stdout, ctx.cwd);

			if (worktreeEntries.length === 0) {
				return {
					content: [{ type: "text", text: "No worktrees found." }],
					details: <WorktreeListResult>{
						worktrees: [],
						count: 0,
						currentPath: ctx.cwd,
						warning: "No worktrees exist. Use git_worktree_create to add one.",
					},
				};
			}

			const displayLines = worktreeEntries.map(
				(entry) => `${entry.branch} - ${entry.path}${entry.isCurrent ? " (current)" : ""}`,
			);

			return {
				content: [{ type: "text", text: displayLines.join("\n") }],
				details: <WorktreeListResult>{
					worktrees: worktreeEntries,
					count: worktreeEntries.length,
					currentPath: ctx.cwd,
				},
			};
		},
	});
}

/**
 * Create the git_worktree_create tool definition.
 * @param pi - Extension API for registering and executing tools
 * @returns The configured tool definition
 */
function createCreateWorktreeTool(pi: ExtensionAPI) {
	return defineTool({
		name: "git_worktree_create",
		label: "Git Worktree Create",
		description: "Create a linked worktree for a branch in the current repo.",
		parameters: createWorktreeParams,
		async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
			const worktreePath = params.path;
			const branchName = params.branch || "main";
			const createBranch = params.create_branch ?? true;

			// 1. Validate input parameters
			const inputError = validateCreateWorktreeInput(worktreePath, branchName);
			if (inputError) {
				return {
					content: [{ type: "text", text: formatValidationError(inputError) }],
					isError: true,
				};
			}

			// 2. Check if path already exists — refuse if so
			const pathExistsResult = await pathExists(worktreePath, ctx);
			const pathError = validatePathDoesNotExist(worktreePath, () => pathExistsResult);
			if (pathError) {
				return {
					content: [{ type: "text", text: formatValidationError(pathError) }],
					isError: true,
				};
			}

			// 3. Check if creating on current branch — refuse if so
			const currentBranch = await getCurrentBranch(pi, ctx);
			if (isCurrentBranch(branchName, currentBranch)) {
				const currentBranchError: WorktreeError = {
					message: "Can't create worktree on the current branch — switch first.",
					suggestion: "Choose a different branch name or switch branches with 'git switch'.",
				};
				return {
					content: [{ type: "text", text: formatValidationError(currentBranchError) }],
					isError: true,
				};
			}

			// 4. Build and run git worktree add command
			const commandArgs = buildWorktreeAddCommand({
				path: worktreePath,
				branch: branchName,
				createBranch,
			});

			const result = await pi.exec("git", commandArgs, { cwd: ctx.cwd });

			if (result.exitCode !== 0) {
				return {
					content: [{
						type: "text",
						text: `Error creating worktree: ${result.stderr}. Suggestion: Check that the branch exists (or use create_branch=true) and that the parent directory is writable.`,
					}],
					isError: true,
				};
			}

			return {
				content: [{
					type: "text",
					text: formatSuccessResponse("Created", { at: worktreePath, branch: branchName }),
				}],
			};
		},
	});
}

/**
 * Create the git_worktree_switch tool definition.
 * @param pi - Extension API for executing commands
 * @returns The configured tool definition
 */
function createSwitchWorktreeTool(pi: ExtensionAPI) {
	return defineTool({
		name: "git_worktree_switch",
		label: "Git Worktree Switch",
		description: "Switch pi's working directory to a linked worktree.",
		parameters: switchWorktreeParams,
		async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
			const target = params.target;

			// 1. Validate input parameters
			const inputError = validateSwitchInput(target);
			if (inputError) {
				return {
					content: [{ type: "text", text: formatValidationError(inputError) }],
					isError: true,
				};
			}

			// 2. List worktrees and resolve target to a path
			const listResult = await pi.exec("git", ["worktree", "list"], { cwd: ctx.cwd });
			if (listResult.exitCode !== 0) {
				return {
					content: [{ type: "text", text: `Error listing worktrees: ${listResult.stderr}. Suggestion: Ensure you are inside a git repository.` }],
					isError: true,
				};
			}

			const worktreeEntries = parseWorktreeList(listResult.stdout, ctx.cwd);
			const resolvedPath = resolveWorktreeTarget(target, worktreeEntries);

			if (!resolvedPath) {
				return {
					content: [{ type: "text", text: `Error: No worktree found at '${target}'. Suggestion: Run 'git_worktree_list' to see available worktrees.` }],
					isError: true,
				};
			}

			// 3. Check if already in that worktree — refuse if so
			if (isAlreadyInWorktree(resolvedPath, ctx.cwd)) {
				return {
					content: [{ type: "text", text: "Already in this worktree." }],
					isError: true,
				};
			}

			// 4. Check for uncommitted changes in source (current worktree)
			const statusResult = await pi.exec("git", ["status", "--porcelain"], { cwd: ctx.cwd });
			if (statusResult.exitCode === 0 && hasUncommittedChanges(statusResult.stdout)) {
				return {
					content: [{ type: "text", text: `Error: Worktree '${ctx.cwd}' has uncommitted changes. Suggestion: Stash them first with 'git stash' before switching.` }],
					isError: true,
				};
			}

			// 5. Change pi's cwd to the target path
			ctx.chdir(resolvedPath);

			// 6. Get branch info for the response
			const targetBranch = worktreeEntries.find((w) => w.path === resolvedPath)?.branch ?? "unknown";

			return {
				content: [{ type: "text", text: `Switched to worktree at '${resolvedPath}' (branch: ${targetBranch}).` }],
				details: {
					previousCwd: ctx.cwd,
					newCwd: resolvedPath,
					branch: targetBranch,
				},
			};
		},
	});
}

/**
 * Create the git_worktree_remove tool definition.
 * @param pi - Extension API for executing commands
 * @returns The configured tool definition
 */
function createRemoveWorktreeTool(pi: ExtensionAPI) {
	return defineTool({
		name: "git_worktree_remove",
		label: "Git Worktree Remove",
		description: "Remove a linked worktree.",
		parameters: removeWorktreeParams,
		async execute(_toolCallId, params, _signal, onUpdate, ctx) {
			const worktreePath = params.path;
			const force = params.force ?? false;

			// 1. Validate input parameters
			const inputError = validateRemoveWorktreeInput(worktreePath);
			if (inputError) {
				return {
					content: [{ type: "text", text: formatValidationError(inputError) }],
					isError: true,
				};
			}

			// 1b. Check that the worktree path exists
			const exists = await pathExists(worktreePath, ctx);
			if (!exists) {
				return {
					content: [{ type: "text", text: `Error: Worktree path '${worktreePath}' does not exist. Suggestion: List worktrees with git_worktree_list to find valid paths.` }],
					isError: true,
				};
			}

			// 2. Get main repo root and check if target is the main repo
			const toplevelResult = await pi.exec("git", ["rev-parse", "--show-toplevel"], { cwd: ctx.cwd });
			if (toplevelResult.exitCode !== 0) {
				return {
					content: [{ type: "text", text: `Error getting repository root: ${toplevelResult.stderr}. Suggestion: Ensure you are inside a git repository.` }],
					isError: true,
				};
			}

			const mainRepoRoot = toplevelResult.stdout.trim();
			if (isMainRepository(worktreePath, mainRepoRoot, ctx.cwd)) {
				return {
					content: [{ type: "text", text: `Error: Can't remove the main repository. Suggestion: Use 'git worktree prune' instead.` }],
					isError: true,
				};
			}

			// 3. Check for uncommitted changes in the target worktree
			const statusResult = await pi.exec("git", ["status", "--porcelain"], { cwd: worktreePath });
			if (statusResult.exitCode !== 0) {
				return {
					content: [{ type: "text", text: `Error checking worktree status: ${statusResult.stderr}. Suggestion: Verify the path exists and is a valid worktree.` }],
					isError: true,
				};
			}
			if (!force && hasUncommittedChanges(statusResult.stdout)) {
				return {
					content: [{ type: "text", text: `Error: Worktree '${worktreePath}' has uncommitted changes. Suggestion: Stash or commit them first, or use force=true to remove anyway.` }],
					isError: true,
				};
			}
			if (force && hasUncommittedChanges(statusResult.stdout)) {
				onUpdate?.({ type: "text", text: `Warning: Force removing worktree '${worktreePath}' despite uncommitted changes.` });
			}

			// 3b. Get the branch name before removal for the success response
			const branchResult = await pi.exec("git", ["-C", worktreePath, "rev-parse", "--abbrev-ref", "HEAD"]);
			const branchName = branchResult.exitCode === 0 ? branchResult.stdout.trim() : "unknown";

			// 4. Build and run git worktree remove command
			const commandArgs = buildWorktreeRemoveCommand(worktreePath, { force });
			const result = await pi.exec("git", commandArgs, { cwd: ctx.cwd });

			if (result.exitCode !== 0) {
				return {
					content: [{ type: "text", text: `Error removing worktree: ${result.stderr}. Suggestion: Check that the path is a valid worktree, or use force=true if the worktree has uncommitted changes.` }],
					isError: true,
				};
			}

			// 5. List remaining worktrees for context
			const listResult = await pi.exec("git", ["worktree", "list"], { cwd: ctx.cwd });
			const remainingWorktrees = listResult.exitCode === 0 ? parseWorktreeList(listResult.stdout, ctx.cwd) : [];

			return {
				content: [{ type: "text", text: formatSuccessResponse("Removed", { at: worktreePath, branch: branchName }) }],
				details: {
					removed: { path: worktreePath, branch: branchName },
					remaining: remainingWorktrees,
					remainingCount: remainingWorktrees.length,
				},
			};
		},
	});
}

// ─── Extension Entry Point ────────────────────────────────────────

export default function (pi: ExtensionAPI) {
	// Register tools via factory functions
	pi.registerTool(createCreateWorktreeTool(pi));
	pi.registerTool(createListWorktreeTool(pi));
	pi.registerTool(createSwitchWorktreeTool(pi));
	pi.registerTool(createRemoveWorktreeTool(pi));

	// ─── Command Registration ──────────────────────────────────────
	pi.registerCommand("worktree", {
		description: "Interactive picker to switch between linked worktrees.",
		handler: async (_args, _ctx) => {
			return {
				content: [{ type: "text", text: "Error: /worktree command is not yet implemented. Suggestion: Use `git worktree list` and navigate manually until this command is complete." }],
				isError: true,
			};
		},
	});
}