index.ts
5847 bytes
1import type {
2 ExtensionAPI,
3 ExtensionCommandContext,
4 SessionEntry,
5} from "@earendil-works/pi-coding-agent";
6import {
7 captureSessionFile,
8 countPruneEntries,
9 getSessionTreeEntries,
10 pruneSessionFile,
11 toSessionFilePath,
12 type SessionFilePath,
13} from "./session-file.ts";
14import { selectPruneCandidates } from "./prune-selector.ts";
15import {
16 findInactiveBranchCandidates,
17 normalizePruneCandidates,
18 toEntryId,
19 type EntryId,
20} from "./tree-model.ts";
21import {
22 buildPruneTreeItems,
23 formatEntryCount,
24 formatPruneTreeItem,
25} from "./tree-view.ts";
26
27const COMMAND_NAME = "tree-prune";
28
29export default function registerTreePruning(pi: ExtensionAPI): void {
30 pi.registerCommand(COMMAND_NAME, {
31 description: "Delete inactive session-tree branches",
32 handler: runTreePruneCommand,
33 });
34}
35
36export async function runTreePruneCommand(
37 args: string,
38 ctx: ExtensionCommandContext,
39): Promise<void> {
40 const mode = args.trim();
41
42 if (!ctx.hasUI) {
43 ctx.ui.notify("/tree-prune requires a user interface", "error");
44 return;
45 }
46
47 if (mode === "") {
48 await runPrune(ctx);
49 return;
50 }
51
52 ctx.ui.notify("Usage: /tree-prune", "error");
53}
54
55async function runPrune(ctx: ExtensionCommandContext): Promise<void> {
56 try {
57 await ctx.waitForIdle();
58
59 const sessionFile = requireSessionFile(ctx);
60 const currentLeafId = ctx.sessionManager.getLeafId();
61 if (!currentLeafId) {
62 ctx.ui.notify("Tree Prune needs a current active leaf", "warning");
63 return;
64 }
65
66 const activeLeafId = toEntryId(currentLeafId);
67 const snapshot = await captureSessionFile(sessionFile);
68 const entries = getSessionTreeEntries(snapshot);
69 if (!entries.some((entry) => entry.id === activeLeafId)) {
70 ctx.ui.notify("Tree Prune cannot find the active leaf in this session file", "error");
71 return;
72 }
73
74 const candidates = findInactiveBranchCandidates(entries, activeLeafId);
75 if (candidates.length === 0) {
76 ctx.ui.notify("Tree Prune found no inactive branches", "info");
77 return;
78 }
79
80 const treeItems = buildPruneTreeItems(candidates, (entryId) => ctx.sessionManager.getLabel(entryId));
81 const selectedCandidates = await selectPruneCandidates(ctx, treeItems);
82 if (!selectedCandidates || selectedCandidates.length === 0) return;
83
84 const candidatesToPrune = normalizePruneCandidates(selectedCandidates);
85 const removedEntryCount = countPruneEntries(candidatesToPrune);
86 const branchesToPrune = candidatesToPrune.map((candidate) => {
87 const item = treeItems.find((treeItem) => treeItem.candidate.rootId === candidate.rootId);
88 return `• ${item ? formatPruneTreeItem(item) : candidate.preview}`;
89 });
90 const confirmed = await ctx.ui.confirm(
91 "Delete inactive branches?",
92 [
93 "Branches to delete:",
94 ...branchesToPrune,
95 `Entries to delete: ${formatEntryCount(removedEntryCount)}`,
96 "The branches will be permanently removed from this session file.",
97 ].join("\n"),
98 );
99 if (!confirmed) return;
100
101 await pruneSessionFile(snapshot, candidatesToPrune);
102 const switchResult = await reloadSession(
103 ctx,
104 sessionFile,
105 activeLeafId,
106 `Tree Prune removed ${formatEntryCount(removedEntryCount)}.`,
107 );
108
109 if (switchResult.cancelled) {
110 ctx.ui.notify(
111 `Tree Prune changed the session, but reload was cancelled. Resume ${sessionFile} to load it.`,
112 "error",
113 );
114 }
115 } catch (error) {
116 ctx.ui.notify(`Tree Prune failed: ${errorMessage(error)}`, "error");
117 }
118}
119
120function requireSessionFile(ctx: ExtensionCommandContext): SessionFilePath {
121 const sessionFile = ctx.sessionManager.getSessionFile();
122 if (!sessionFile) {
123 throw new Error("an on-disk session is required; --no-session is not supported");
124 }
125 return toSessionFilePath(sessionFile);
126}
127
128async function reloadSession(
129 ctx: ExtensionCommandContext,
130 sessionFile: SessionFilePath,
131 activeLeafId: EntryId | undefined,
132 message: string,
133): Promise<{ cancelled: boolean }> {
134 return ctx.switchSession(sessionFile, {
135 withSession: async (replacement) => {
136 const activeBranchRestored = activeLeafId
137 ? await restoreActiveBranch(replacement, activeLeafId)
138 : true;
139 replacement.ui.notify(
140 activeBranchRestored ? message : `${message} Active branch could not be restored.`,
141 activeBranchRestored ? "info" : "warning",
142 );
143 },
144 });
145}
146
147async function restoreActiveBranch(
148 ctx: ExtensionCommandContext,
149 activeLeafId: EntryId,
150): Promise<boolean> {
151 const activeEntry = ctx.sessionManager.getEntry(activeLeafId);
152 if (!activeEntry) return false;
153
154 try {
155 const result = await ctx.navigateTree(activeLeafId, { summarize: false });
156 if (result.cancelled) return false;
157
158 const editorText = editorTextForEntry(activeEntry);
159 if (editorText !== undefined) ctx.ui.setEditorText(editorText);
160 return true;
161 } catch {
162 return false;
163 }
164}
165
166function editorTextForEntry(entry: SessionEntry): string | undefined {
167 if (entry.type === "message" && entry.message.role === "user") {
168 return textContent(entry.message.content);
169 }
170 if (entry.type === "custom_message") return textContent(entry.content);
171 return undefined;
172}
173
174function textContent(content: unknown): string {
175 if (typeof content === "string") return content;
176 if (!Array.isArray(content)) return "";
177
178 return content
179 .filter(
180 (block): block is { type: "text"; text: string } =>
181 Boolean(block) &&
182 typeof block === "object" &&
183 (block as { type?: unknown }).type === "text" &&
184 typeof (block as { text?: unknown }).text === "string",
185 )
186 .map((block) => block.text)
187 .join("");
188}
189
190function errorMessage(error: unknown): string {
191 return error instanceof Error ? error.message : String(error);
192}