index.ts
8294 bytes
1import { basename } from "node:path";
2import type {
3 ExtensionAPI,
4 ExtensionCommandContext,
5 SessionEntry,
6} from "@earendil-works/pi-coding-agent";
7import {
8 captureSessionFile,
9 countPruneEntries,
10 getSessionTreeEntries,
11 createBackupPath,
12 listSessionBackups,
13 pruneSessionFile,
14 restoreSessionBackup,
15 toSessionFilePath,
16 type SessionFilePath,
17} from "./session-file.ts";
18import { selectPruneCandidates } from "./prune-selector.ts";
19import {
20 findInactiveBranchCandidates,
21 normalizePruneCandidates,
22 toEntryId,
23 type EntryId,
24} from "./tree-model.ts";
25import {
26 buildPruneTreeItems,
27 formatEntryCount,
28 formatPruneTreeItem,
29} from "./tree-view.ts";
30
31const COMMAND_NAME = "tree-prune";
32
33export default function registerTreePruning(pi: ExtensionAPI): void {
34 pi.registerCommand(COMMAND_NAME, {
35 description: "Delete inactive session-tree branches with a backup; use /tree-prune undo to restore",
36 handler: runTreePruneCommand,
37 });
38}
39
40export async function runTreePruneCommand(
41 args: string,
42 ctx: ExtensionCommandContext,
43): Promise<void> {
44 const mode = args.trim();
45
46 if (!ctx.hasUI) {
47 ctx.ui.notify("/tree-prune requires a user interface", "error");
48 return;
49 }
50
51 if (mode === "") {
52 await runPrune(ctx);
53 return;
54 }
55
56 if (mode === "undo") {
57 await runUndo(ctx);
58 return;
59 }
60
61 ctx.ui.notify("Usage: /tree-prune [undo]", "error");
62}
63
64async function runPrune(ctx: ExtensionCommandContext): Promise<void> {
65 try {
66 await ctx.waitForIdle();
67
68 const sessionFile = requireSessionFile(ctx);
69 const currentLeafId = ctx.sessionManager.getLeafId();
70 if (!currentLeafId) {
71 ctx.ui.notify("Tree Prune needs a current active leaf", "warning");
72 return;
73 }
74
75 const activeLeafId = toEntryId(currentLeafId);
76 const snapshot = await captureSessionFile(sessionFile);
77 const entries = getSessionTreeEntries(snapshot);
78 if (!entries.some((entry) => entry.id === activeLeafId)) {
79 ctx.ui.notify("Tree Prune cannot find the active leaf in this session file", "error");
80 return;
81 }
82
83 const candidates = findInactiveBranchCandidates(entries, activeLeafId);
84 if (candidates.length === 0) {
85 ctx.ui.notify("Tree Prune found no inactive branches", "info");
86 return;
87 }
88
89 const treeItems = buildPruneTreeItems(candidates, (entryId) => ctx.sessionManager.getLabel(entryId));
90 const selectedCandidates = await selectPruneCandidates(ctx, treeItems);
91 if (!selectedCandidates || selectedCandidates.length === 0) return;
92
93 const candidatesToPrune = normalizePruneCandidates(selectedCandidates);
94 const removedEntryCount = countPruneEntries(candidatesToPrune);
95 const branchesToPrune = candidatesToPrune.map((candidate) => {
96 const item = treeItems.find((treeItem) => treeItem.candidate.rootId === candidate.rootId);
97 return `• ${item ? formatPruneTreeItem(item) : candidate.preview}`;
98 });
99 const backupPath = createBackupPath(sessionFile);
100 const confirmed = await ctx.ui.confirm(
101 "Delete inactive branches?",
102 [
103 "Branches to delete:",
104 ...branchesToPrune,
105 `Entries to delete: ${formatEntryCount(removedEntryCount)}`,
106 `Backup: ${backupPath}`,
107 "The branches will be removed from this session file.",
108 ].join("\n"),
109 );
110 if (!confirmed) return;
111
112 const result = await pruneSessionFile(snapshot, candidatesToPrune, backupPath);
113 const switchResult = await reloadSession(
114 ctx,
115 sessionFile,
116 activeLeafId,
117 `Tree Prune removed ${formatEntryCount(result.removedEntryCount)}. Backup: ${result.backupPath}`,
118 );
119
120 if (switchResult.cancelled) {
121 ctx.ui.notify(
122 `Tree Prune changed the session, but reload was cancelled. Resume ${sessionFile} to load it. Backup: ${result.backupPath}`,
123 "error",
124 );
125 }
126 } catch (error) {
127 ctx.ui.notify(`Tree Prune failed: ${errorMessage(error)}`, "error");
128 }
129}
130
131async function runUndo(ctx: ExtensionCommandContext): Promise<void> {
132 try {
133 await ctx.waitForIdle();
134
135 const sessionFile = requireSessionFile(ctx);
136 const currentLeafId = ctx.sessionManager.getLeafId();
137 const activeLeafId = currentLeafId ? toEntryId(currentLeafId) : undefined;
138 const currentSnapshot = await captureSessionFile(sessionFile);
139 const backups = await listSessionBackups(sessionFile);
140 if (backups.length === 0) {
141 ctx.ui.notify("Tree Prune found no backups for this session", "info");
142 return;
143 }
144
145 const selectedBackupPath = await ctx.ui.select(
146 "Restore Tree Prune backup:",
147 backups.map((backup) => basename(backup)),
148 );
149 if (!selectedBackupPath) return;
150
151 const selectedBackup = backups.find((backup) => basename(backup) === selectedBackupPath);
152 if (!selectedBackup) {
153 ctx.ui.notify("Tree Prune could not resolve the selected backup", "error");
154 return;
155 }
156
157 const selectedSnapshot = await captureSessionFile(selectedBackup);
158 const currentBackupPath = createBackupPath(sessionFile);
159 const confirmed = await ctx.ui.confirm(
160 "Restore Tree Prune backup?",
161 [
162 `Restore: ${selectedBackup}`,
163 `Current session backup: ${currentBackupPath}`,
164 "The current session file will be replaced.",
165 ].join("\n"),
166 );
167 if (!confirmed) return;
168
169 const createdBackupPath = await restoreSessionBackup(
170 currentSnapshot,
171 selectedSnapshot,
172 currentBackupPath,
173 );
174 const switchResult = await reloadSession(
175 ctx,
176 sessionFile,
177 activeLeafId,
178 `Tree Prune restored ${selectedBackup}. Previous session backup: ${createdBackupPath}`,
179 );
180
181 if (switchResult.cancelled) {
182 ctx.ui.notify(
183 `Tree Prune restored the backup, but reload was cancelled. Resume ${sessionFile} to load it. Previous session backup: ${createdBackupPath}`,
184 "error",
185 );
186 }
187 } catch (error) {
188 ctx.ui.notify(`Tree Prune undo failed: ${errorMessage(error)}`, "error");
189 }
190}
191
192function requireSessionFile(ctx: ExtensionCommandContext): SessionFilePath {
193 const sessionFile = ctx.sessionManager.getSessionFile();
194 if (!sessionFile) {
195 throw new Error("an on-disk session is required; --no-session is not supported");
196 }
197 return toSessionFilePath(sessionFile);
198}
199
200async function reloadSession(
201 ctx: ExtensionCommandContext,
202 sessionFile: SessionFilePath,
203 activeLeafId: EntryId | undefined,
204 message: string,
205): Promise<{ cancelled: boolean }> {
206 return ctx.switchSession(sessionFile, {
207 withSession: async (replacement) => {
208 const activeBranchRestored = activeLeafId
209 ? await restoreActiveBranch(replacement, activeLeafId)
210 : true;
211 replacement.ui.notify(
212 activeBranchRestored ? message : `${message} Active branch could not be restored.`,
213 activeBranchRestored ? "info" : "warning",
214 );
215 },
216 });
217}
218
219async function restoreActiveBranch(
220 ctx: ExtensionCommandContext,
221 activeLeafId: EntryId,
222): Promise<boolean> {
223 const activeEntry = ctx.sessionManager.getEntry(activeLeafId);
224 if (!activeEntry) return false;
225
226 try {
227 const result = await ctx.navigateTree(activeLeafId, { summarize: false });
228 if (result.cancelled) return false;
229
230 const editorText = editorTextForEntry(activeEntry);
231 if (editorText !== undefined) ctx.ui.setEditorText(editorText);
232 return true;
233 } catch {
234 return false;
235 }
236}
237
238function editorTextForEntry(entry: SessionEntry): string | undefined {
239 if (entry.type === "message" && entry.message.role === "user") {
240 return textContent(entry.message.content);
241 }
242 if (entry.type === "custom_message") return textContent(entry.content);
243 return undefined;
244}
245
246function textContent(content: unknown): string {
247 if (typeof content === "string") return content;
248 if (!Array.isArray(content)) return "";
249
250 return content
251 .filter(
252 (block): block is { type: "text"; text: string } =>
253 Boolean(block) &&
254 typeof block === "object" &&
255 (block as { type?: unknown }).type === "text" &&
256 typeof (block as { text?: unknown }).text === "string",
257 )
258 .map((block) => block.text)
259 .join("");
260}
261
262function errorMessage(error: unknown): string {
263 return error instanceof Error ? error.message : String(error);
264}