/** * 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 { 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 { 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 { 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: { 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: { 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, }; }, }); }