main
1// =============================================================================
2// Git Extension for Pi Coding Agent
3// =============================================================================
4
5import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
6import {
7 listWorktrees,
8 createWorktree,
9 removeWorktree,
10 execInWorktree,
11 pruneWorktrees,
12 getCurrentWorktreeInfo,
13 getWorktreeNote,
14 setWorktreeNote,
15 getCurrentWorktreeContext,
16} from "./worktree.js";
17import { findRepoRoot, hasUncommittedChanges } from "./utils.js";
18
19export default function (pi: ExtensionAPI) {
20 // =========================================================================
21 // Commands: /worktree (alias: /wt)
22 // =========================================================================
23
24 pi.registerCommand("worktree", {
25 description: "Manage git worktrees (list, create, remove, prune)",
26 getArgumentCompletions: (prefix, ctx) => {
27 const parts = prefix.trim().split(/\s+/);
28
29 // First argument: subcommand
30 if (parts.length <= 1) {
31 const subcommands = ["list", "ls", "create", "new", "add", "remove", "rm", "delete", "prune"];
32 const filtered = subcommands.filter((s) => s.startsWith(parts[0] || ""));
33 return filtered.length > 0 ? filtered.map((s) => ({ value: s, label: s })) : null;
34 }
35
36 // Second argument: branch name for remove/delete
37 const subcommand = parts[0];
38 if ((subcommand === "remove" || subcommand === "rm" || subcommand === "delete") && parts.length === 2) {
39 try {
40 const worktrees = listWorktrees(ctx.cwd);
41 // Don't suggest removing the main worktree
42 const branches = worktrees
43 .filter((wt) => !wt.bare)
44 .map((wt) => wt.branch);
45
46 const searchTerm = parts[1] || "";
47 const filtered = branches.filter((b) => b.startsWith(searchTerm));
48 return filtered.length > 0 ? filtered.map((b) => ({ value: b, label: b })) : null;
49 } catch {
50 return null;
51 }
52 }
53
54 return null;
55 },
56 handler: async (args, ctx) => {
57 const parts = args.trim().split(/\s+/).filter(Boolean);
58 const subcommand = parts[0] || "list";
59 const remainingArgs = parts.slice(1);
60
61 try {
62 switch (subcommand) {
63 case "list":
64 case "ls":
65 await handleList(ctx);
66 break;
67
68 case "create":
69 case "new":
70 case "add":
71 await handleCreate(remainingArgs, ctx);
72 break;
73
74 case "remove":
75 case "rm":
76 case "delete":
77 await handleRemove(remainingArgs, ctx);
78 break;
79
80 case "prune":
81 await handlePrune(ctx);
82 break;
83
84 default:
85 ctx.ui.notify(`Unknown worktree command: ${subcommand}`, "error");
86 ctx.ui.notify(
87 "Available: list, create <branch>, remove <branch>, prune",
88 "info"
89 );
90 }
91 } catch (error: any) {
92 ctx.ui.notify(`Error: ${error.message}`, "error");
93 }
94 },
95 });
96
97 // Register /wt alias
98 pi.registerCommand("wt", {
99 description: "Alias for /worktree",
100 getArgumentCompletions: (prefix, ctx) => {
101 const parts = prefix.trim().split(/\s+/);
102
103 // First argument: subcommand
104 if (parts.length <= 1) {
105 const subcommands = ["list", "ls", "create", "new", "add", "remove", "rm", "delete", "prune"];
106 const filtered = subcommands.filter((s) => s.startsWith(parts[0] || ""));
107 return filtered.length > 0 ? filtered.map((s) => ({ value: s, label: s })) : null;
108 }
109
110 // Second argument: branch name for remove/delete
111 const subcommand = parts[0];
112 if ((subcommand === "remove" || subcommand === "rm" || subcommand === "delete") && parts.length === 2) {
113 try {
114 const worktrees = listWorktrees(ctx.cwd);
115 // Don't suggest removing the main worktree
116 const branches = worktrees
117 .filter((wt) => !wt.bare)
118 .map((wt) => wt.branch);
119
120 const searchTerm = parts[1] || "";
121 const filtered = branches.filter((b) => b.startsWith(searchTerm));
122 return filtered.length > 0 ? filtered.map((b) => ({ value: b, label: b })) : null;
123 } catch {
124 return null;
125 }
126 }
127
128 return null;
129 },
130 handler: async (args, ctx) => {
131 const parts = args.trim().split(/\s+/).filter(Boolean);
132 const subcommand = parts[0] || "list";
133 const remainingArgs = parts.slice(1);
134
135 try {
136 switch (subcommand) {
137 case "list":
138 case "ls":
139 await handleList(ctx);
140 break;
141
142 case "create":
143 case "new":
144 case "add":
145 await handleCreate(remainingArgs, ctx);
146 break;
147
148 case "remove":
149 case "rm":
150 case "delete":
151 await handleRemove(remainingArgs, ctx);
152 break;
153
154 case "prune":
155 await handlePrune(ctx);
156 break;
157
158 default:
159 ctx.ui.notify(`Unknown worktree command: ${subcommand}`, "error");
160 ctx.ui.notify(
161 "Available: list, create <branch>, remove <branch>, prune",
162 "info"
163 );
164 }
165 } catch (error: any) {
166 ctx.ui.notify(`Error: ${error.message}`, "error");
167 }
168 },
169 });
170
171 // =========================================================================
172 // Command Handlers
173 // =========================================================================
174
175 async function handleList(ctx: any) {
176 const repoRoot = findRepoRoot(ctx.cwd);
177 if (!repoRoot) {
178 ctx.ui.notify("Not in a git repository", "error");
179 return;
180 }
181
182 const worktrees = listWorktrees(ctx.cwd);
183 const currentInfo = getCurrentWorktreeInfo(ctx.cwd);
184
185 if (worktrees.length === 0) {
186 ctx.ui.notify("No worktrees found", "info");
187 return;
188 }
189
190 // Build output
191 const lines: string[] = [];
192 lines.push("Git Worktrees:");
193 lines.push("");
194
195 for (const wt of worktrees) {
196 const isCurrent = currentInfo && wt.path === currentInfo.path;
197 const marker = isCurrent ? "➜" : " ";
198 const dirty = (wt as any).dirty ?? hasUncommittedChanges(wt.path);
199 const dirtyMark = dirty ? "*" : "";
200
201 lines.push(`${marker} ${wt.branch}${dirtyMark}`);
202 lines.push(` Path: ${wt.path}`);
203 if (wt.commit) lines.push(` Commit: ${wt.commit.substring(0, 8)}`);
204
205 const ahead = (wt as any).ahead;
206 const behind = (wt as any).behind;
207 if (ahead || behind) {
208 lines.push(` ↑${ahead ?? 0} ↓${behind ?? 0}`);
209 }
210
211 const lastActive = (wt as any).lastActive;
212 if (lastActive) lines.push(` Last active: ${lastActive}`);
213
214 if (wt.locked) lines.push(" 🔒 Locked");
215 if (wt.prunable) lines.push(" ♻️ Prunable");
216
217 lines.push("");
218 }
219
220 ctx.ui.notify(lines.join("\n"), "info");
221 }
222
223 async function handleCreate(args: string[], ctx: any) {
224 const branch = args[0];
225 if (!branch) {
226 ctx.ui.notify("Usage: /worktree create <branch> [path]", "error");
227 return;
228 }
229
230 const customPath = args[1];
231
232 try {
233 const worktreePath = createWorktree(
234 {
235 branch,
236 path: customPath,
237 fromRemote: true,
238 },
239 ctx.cwd
240 );
241
242 ctx.ui.notify(`✓ Worktree created: ${worktreePath}`, "success");
243
244 // Show how to switch
245 ctx.ui.notify(`To switch: cd ${worktreePath}`, "info");
246 } catch (error: any) {
247 ctx.ui.notify(`Failed to create worktree: ${error.message}`, "error");
248 }
249 }
250
251 async function handleRemove(args: string[], ctx: any) {
252 const branch = args[0];
253 if (!branch) {
254 ctx.ui.notify("Usage: /worktree remove <branch>", "error");
255 return;
256 }
257
258 // Check for uncommitted changes
259 const worktrees = listWorktrees(ctx.cwd);
260 const worktree = worktrees.find((wt) => wt.branch === branch);
261
262 if (!worktree) {
263 ctx.ui.notify(`Worktree for branch '${branch}' not found`, "error");
264 return;
265 }
266
267 const hasChanges = hasUncommittedChanges(worktree.path);
268
269 if (hasChanges) {
270 const confirmed = await ctx.ui.confirm(
271 `Worktree has uncommitted changes. Remove anyway?`
272 );
273
274 if (!confirmed) {
275 ctx.ui.notify("Cancelled", "info");
276 return;
277 }
278 }
279
280 try {
281 removeWorktree(
282 {
283 branch,
284 force: hasChanges,
285 },
286 ctx.cwd
287 );
288
289 ctx.ui.notify(`✓ Worktree removed: ${branch}`, "success");
290 } catch (error: any) {
291 ctx.ui.notify(`Failed to remove worktree: ${error.message}`, "error");
292 }
293 }
294
295 async function handlePrune(ctx: any) {
296 try {
297 pruneWorktrees(ctx.cwd);
298 ctx.ui.notify("✓ Pruned stale worktree references", "success");
299 } catch (error: any) {
300 ctx.ui.notify(`Failed to prune worktrees: ${error.message}`, "error");
301 }
302 }
303
304 // =========================================================================
305 // Tool: git-worktree (for AI to use)
306 // =========================================================================
307
308 pi.registerTool({
309 name: "git_worktree",
310 label: "Git Worktree",
311 description:
312 "Manage git worktrees in ~/.local/share/worktrees/<org>/<repo>/<branch>. Uses lazyworktree when available for consistent org/repo/name layout, falls back to raw git. Create isolated working directories for different branches without switching in the main repository. USE WHEN starting feature work that needs isolation OR working on multiple branches simultaneously.",
313 parameters: {
314 type: "object",
315 required: ["action"],
316 properties: {
317 action: {
318 anyOf: [
319 { const: "list", type: "string" },
320 { const: "create", type: "string" },
321 { const: "remove", type: "string" },
322 { const: "exec", type: "string" },
323 { const: "note", type: "string" },
324 ],
325 description: "The worktree action to perform",
326 },
327 branch: {
328 type: "string",
329 description: "Branch name (required for create/remove) or worktree name (for exec/note)",
330 },
331 path: {
332 type: "string",
333 description: "Custom path for worktree (defaults to ~/.local/share/worktrees/<org>/<repo>/<branch>)",
334 },
335 force: {
336 type: "boolean",
337 description: "Force removal even with uncommitted changes",
338 },
339 exec: {
340 type: "string",
341 description: "Command to run: post-create hook (with create), command to execute (with exec action), or note content to write (with note action). Omit exec with note action to read the note.",
342 },
343 },
344 },
345 execute: async (_toolCallId, params, _signal, _onUpdate, ctx) => {
346 const { action, branch, path: customPath, force, exec: execCmd } = params;
347
348 try {
349 switch (action) {
350 case "list": {
351 const worktrees = listWorktrees(ctx.cwd);
352 const lines = worktrees.map((wt) => {
353 const dirty = (wt as any).dirty ?? hasUncommittedChanges(wt.path);
354 const dirtyMark = dirty ? "*" : "";
355 const commit = wt.commit ? ` (${wt.commit.substring(0, 8)})` : "";
356 const lastActive = (wt as any).lastActive ? ` [${(wt as any).lastActive}]` : "";
357 const sync = ((wt as any).ahead || (wt as any).behind)
358 ? ` ↑${(wt as any).ahead ?? 0}↓${(wt as any).behind ?? 0}`
359 : "";
360 const wtName = (wt as any).isMain ? wt.branch : (wt.path.split("/").pop() ?? wt.branch);
361 const note = getWorktreeNote(wtName, ctx.cwd);
362 const noteLine = note ? `\n Note: ${note.split("\n")[0].substring(0, 100)}` : "";
363 return `- ${wt.branch}${dirtyMark}: ${wt.path}${commit}${sync}${lastActive}${noteLine}`;
364 });
365
366 return {
367 content: [{
368 type: "text",
369 text: `Worktrees:\n${lines.join("\n")}`
370 }],
371 details: {},
372 };
373 }
374
375 case "create": {
376 if (!branch) {
377 return {
378 content: [{
379 type: "text",
380 text: "Error: Branch name required for create action"
381 }],
382 details: { error: true },
383 };
384 }
385
386 const worktreePath = createWorktree(
387 {
388 branch,
389 path: customPath,
390 fromRemote: true,
391 exec: execCmd,
392 },
393 ctx.cwd
394 );
395
396 return {
397 content: [{
398 type: "text",
399 text: `✓ Worktree created at ${worktreePath}\n\nTo switch: cd ${worktreePath}`
400 }],
401 details: { path: worktreePath, branch },
402 };
403 }
404
405 case "remove": {
406 if (!branch) {
407 return {
408 content: [{
409 type: "text",
410 text: "Error: Branch name required for remove action"
411 }],
412 details: { error: true },
413 };
414 }
415
416 removeWorktree(
417 {
418 branch,
419 force: force ?? false,
420 },
421 ctx.cwd
422 );
423
424 return {
425 content: [{
426 type: "text",
427 text: `✓ Worktree removed for branch ${branch}`
428 }],
429 details: { branch },
430 };
431 }
432
433 case "exec": {
434 if (!branch || !execCmd) {
435 return {
436 content: [{
437 type: "text",
438 text: "Error: Both branch (worktree name/path) and exec (command) are required for exec action"
439 }],
440 details: { error: true },
441 };
442 }
443
444 const output = await execInWorktree(branch, execCmd, ctx.cwd);
445
446 return {
447 content: [{
448 type: "text",
449 text: output || "(no output)"
450 }],
451 details: { branch, command: execCmd },
452 };
453 }
454
455 case "note": {
456 if (!branch) {
457 // No branch specified: read the current worktree's note
458 const wtCtx = getCurrentWorktreeContext(ctx.cwd);
459 if (!wtCtx) {
460 return {
461 content: [{
462 type: "text",
463 text: "Error: Not in a lazyworktree-managed worktree, or worktree name required"
464 }],
465 details: { error: true },
466 };
467 }
468
469 return {
470 content: [{
471 type: "text",
472 text: wtCtx.note
473 ? `Note for "${wtCtx.name}" (branch: ${wtCtx.branch}):\n${wtCtx.note}`
474 : `No note set for "${wtCtx.name}" (branch: ${wtCtx.branch})`
475 }],
476 details: { name: wtCtx.name, branch: wtCtx.branch },
477 };
478 }
479
480 if (execCmd) {
481 // Write note
482 const success = setWorktreeNote(branch, execCmd, ctx.cwd);
483 return {
484 content: [{
485 type: "text",
486 text: success
487 ? `✓ Note updated for "${branch}"`
488 : `Error: Failed to update note for "${branch}"`
489 }],
490 details: { name: branch, written: success },
491 };
492 } else {
493 // Read note
494 const note = getWorktreeNote(branch, ctx.cwd);
495 return {
496 content: [{
497 type: "text",
498 text: note
499 ? `Note for "${branch}":\n${note}`
500 : `No note set for "${branch}"`
501 }],
502 details: { name: branch },
503 };
504 }
505 }
506
507 default:
508 return {
509 content: [{
510 type: "text",
511 text: `Error: Unknown action: ${action}`
512 }],
513 details: { error: true },
514 };
515 }
516 } catch (error: any) {
517 return {
518 content: [{
519 type: "text",
520 text: `Error: ${error.message}`
521 }],
522 details: { error: true },
523 };
524 }
525 },
526 });
527}