c9c13e6c31adf00436dc259584188f85f4595abd

Message

Improve pi-tree-pruning code

Diff

This diff is truncated to protect this page.

  1diff --git a/dot_pi/agent/extensions/pi-tree-pruning/index.ts b/dot_pi/agent/extensions/pi-tree-pruning/index.ts
  2index 4e22f76ea2d2db556de5e55cb89b2f61fc885ca5..ad3aff71dd5fa0e23b191e6b8716c387aac0e5db 100644
  3--- a/dot_pi/agent/extensions/pi-tree-pruning/index.ts
  4+++ b/dot_pi/agent/extensions/pi-tree-pruning/index.ts
  5@@ -1,15 +1,27 @@
  6 import { basename } from "node:path";
  7-import type { ExtensionAPI, ExtensionCommandContext } from "@earendil-works/pi-coding-agent";
  8+import type {
  9+  ExtensionAPI,
 10+  ExtensionCommandContext,
 11+  SessionEntry,
 12+} from "@earendil-works/pi-coding-agent";
 13 import {
 14   captureSessionFile,
 15   countPruneEntries,
 16+  getSessionTreeEntries,
 17   createBackupPath,
 18   listSessionBackups,
 19   pruneSessionFile,
 20   restoreSessionBackup,
 21+  toSessionFilePath,
 22+  type SessionFilePath,
 23 } from "./session-file.ts";
 24 import { selectPruneCandidates } from "./prune-selector.ts";
 25-import { findInactiveBranchCandidates, normalizePruneCandidates } from "./tree-model.ts";
 26+import {
 27+  findInactiveBranchCandidates,
 28+  normalizePruneCandidates,
 29+  toEntryId,
 30+  type EntryId,
 31+} from "./tree-model.ts";
 32 import {
 33   buildPruneTreeItems,
 34   formatEntryCount,
 35@@ -21,9 +33,7 @@ const COMMAND_NAME = "tree-prune";
 36 export default function registerTreePruning(pi: ExtensionAPI): void {
 37   pi.registerCommand(COMMAND_NAME, {
 38     description: "Delete inactive session-tree branches with a backup; use /tree-prune undo to restore",
 39-    handler: async (args, ctx) => {
 40-      await runTreePruneCommand(args, ctx);
 41-    },
 42+    handler: runTreePruneCommand,
 43   });
 44 }
 45 
 46@@ -56,19 +66,20 @@ async function runPrune(ctx: ExtensionCommandContext): Promise<void> {
 47     await ctx.waitForIdle();
 48 
 49     const sessionFile = requireSessionFile(ctx);
 50-    const activeLeafId = ctx.sessionManager.getLeafId();
 51-    if (!activeLeafId) {
 52+    const currentLeafId = ctx.sessionManager.getLeafId();
 53+    if (!currentLeafId) {
 54       ctx.ui.notify("Tree Prune needs a current active leaf", "warning");
 55       return;
 56     }
 57 
 58-    const entries = ctx.sessionManager.getEntries();
 59+    const activeLeafId = toEntryId(currentLeafId);
 60+    const snapshot = await captureSessionFile(sessionFile);
 61+    const entries = getSessionTreeEntries(snapshot);
 62     if (!entries.some((entry) => entry.id === activeLeafId)) {
 63-      ctx.ui.notify("Tree Prune cannot find the active leaf in this session", "error");
 64+      ctx.ui.notify("Tree Prune cannot find the active leaf in this session file", "error");
 65       return;
 66     }
 67 
 68-    const snapshot = await captureSessionFile(sessionFile);
 69     const candidates = findInactiveBranchCandidates(entries, activeLeafId);
 70     if (candidates.length === 0) {
 71       ctx.ui.notify("Tree Prune found no inactive branches", "info");
 72@@ -99,14 +110,12 @@ async function runPrune(ctx: ExtensionCommandContext): Promise<void> {
 73     if (!confirmed) return;
 74 
 75     const result = await pruneSessionFile(snapshot, candidatesToPrune, backupPath);
 76-    const switchResult = await ctx.switchSession(sessionFile, {
 77-      withSession: async (replacement) => {
 78-        replacement.ui.notify(
 79-          `Tree Prune removed ${formatEntryCount(result.removedEntryCount)}. Backup: ${result.backupPath}`,
 80-          "info",
 81-        );
 82-      },
 83-    });
 84+    const switchResult = await reloadSession(
 85+      ctx,
 86+      sessionFile,
 87+      activeLeafId,
 88+      `Tree Prune removed ${formatEntryCount(result.removedEntryCount)}. Backup: ${result.backupPath}`,
 89+    );
 90 
 91     if (switchResult.cancelled) {
 92       ctx.ui.notify(
 93@@ -124,6 +133,8 @@ async function runUndo(ctx: ExtensionCommandContext): Promise<void> {
 94     await ctx.waitForIdle();
 95 
 96     const sessionFile = requireSessionFile(ctx);
 97+    const currentLeafId = ctx.sessionManager.getLeafId();
 98+    const activeLeafId = currentLeafId ? toEntryId(currentLeafId) : undefined;
 99     const currentSnapshot = await captureSessionFile(sessionFile);
100     const backups = await listSessionBackups(sessionFile);
101     if (backups.length === 0) {
102@@ -160,14 +171,12 @@ async function runUndo(ctx: ExtensionCommandContext): Promise<void> {
103       selectedSnapshot,
104       currentBackupPath,
105diff --git a/dot_pi/agent/extensions/pi-tree-pruning/session-file.ts b/dot_pi/agent/extensions/pi-tree-pruning/session-file.ts
106index 99b6ec58ce17a0303d5fb3bb00c935544419c01d..e665d251c1823e08cdffdf57b5b84ced6ec7bbe4 100644
107--- a/dot_pi/agent/extensions/pi-tree-pruning/session-file.ts
108+++ b/dot_pi/agent/extensions/pi-tree-pruning/session-file.ts
109@@ -1,7 +1,17 @@
110 import { randomUUID } from "node:crypto";
111 import { chmod, link, lstat, open, readdir, readFile, rename, unlink } from "node:fs/promises";
112 import { basename, dirname, join, resolve } from "node:path";
113-import type { BranchCandidate, TreeEntry } from "./tree-model.ts";
114+import {
115+  toEntryId,
116+  treeEntryFromUnknown,
117+  type BranchCandidate,
118+  type EntryId,
119+  type TreeEntry,
120+} from "./tree-model.ts";
121+
122+declare const sessionFilePathBrand: unique symbol;
123+
124+export type SessionFilePath = string & { readonly [sessionFilePathBrand]: "SessionFilePath" };
125 
126 export interface FileIdentity {
127   readonly dev: number;
128@@ -12,7 +22,7 @@ export interface FileIdentity {
129 }
130 
131 export interface SessionFileSnapshot {
132-  readonly path: string;
133+  readonly path: SessionFilePath;
134   readonly content: string;
135   readonly identity: FileIdentity;
136 }
137@@ -22,8 +32,12 @@ interface RawLine {
138   readonly record?: Record<string, unknown>;
139 }
140 
141+export function toSessionFilePath(path: string): SessionFilePath {
142+  return resolve(path) as SessionFilePath;
143+}
144+
145 export async function captureSessionFile(sessionFile: string): Promise<SessionFileSnapshot> {
146-  const path = resolve(sessionFile);
147+  const path = toSessionFilePath(sessionFile);
148 
149   for (let attempt = 0; attempt < 2; attempt += 1) {
150     const before = await getFileIdentity(path);
151@@ -45,24 +59,36 @@ export async function assertSessionFileUnchanged(snapshot: SessionFileSnapshot):
152   }
153 }
154 
155+export function getSessionTreeEntries(snapshot: SessionFileSnapshot): TreeEntry[] {
156+  const entries: TreeEntry[] = [];
157+  for (const rawLine of splitRawLines(snapshot.content)) {
158+    const entry = treeEntryFromUnknown(rawLine.record);
159+    if (entry) entries.push(entry);
160+  }
161+  return entries;
162+}
163+
164 export function buildPrunedSessionContent(
165   snapshot: SessionFileSnapshot,
166   candidates: readonly BranchCandidate[],
167 ): string {
168   const expectedEntries = collectPruneEntries(candidates);
169   const rawLines = splitRawLines(snapshot.content);
170-  const linesByEntryId = new Map<string, number[]>();
171+  const linesByEntryId = new Map<EntryId, number[]>();
172 
173   for (let index = 0; index < rawLines.length; index += 1) {
174     const rawLine = rawLines[index];
175     if (!rawLine?.record) continue;
176 
177     const id = rawLine.record.id;
178-    if (typeof id !== "string" || !expectedEntries.has(id)) continue;
179+    if (typeof id !== "string" || id.length === 0) continue;
180 
181-    const indexes = linesByEntryId.get(id) ?? [];
182+    const entryId = toEntryId(id);
183+    if (!expectedEntries.has(entryId)) continue;
184+
185+    const indexes = linesByEntryId.get(entryId) ?? [];
186     indexes.push(index);
187-    linesByEntryId.set(id, indexes);
188+    linesByEntryId.set(entryId, indexes);
189   }
190 
191   const lineIndexesToRemove = new Set<number>();
192@@ -78,7 +104,7 @@ export function buildPrunedSessionContent(
193     }
194 
195     const rawLine = rawLines[lineIndex];
196-    if (!rawLine?.record || !matchesEntry(rawLine.record, entry)) {
197+    if (!rawLine?.record || !hasSameEntryIdentity(entryIdentityFromRecord(rawLine.record), entry)) {
198       throw new Error(`JSONL line for entry ${id} does not match the loaded session`);
199     }
200 
201@@ -91,17 +117,19 @@ export function buildPrunedSessionContent(
202     .join("");
203 }
204 
205-export function createBackupPath(sessionFile: string): string {
206-  const path = resolve(sessionFile);
207+export function createBackupPath(sessionFile: SessionFilePath): SessionFilePath {
208+  const path = toSessionFilePath(sessionFile);
209diff --git a/dot_pi/agent/extensions/pi-tree-pruning/tree-model.ts b/dot_pi/agent/extensions/pi-tree-pruning/tree-model.ts
210index 0259ff06f65fdbe261c8f0112122d3696a4c6fd0..5107e062c9458aba64667dcd072afb848cf9e04d 100644
211--- a/dot_pi/agent/extensions/pi-tree-pruning/tree-model.ts
212+++ b/dot_pi/agent/extensions/pi-tree-pruning/tree-model.ts
213@@ -1,31 +1,42 @@
214-export interface TreeEntry {
215-  readonly id: string;
216-  readonly parentId: string | null;
217+import type { SessionEntry } from "@earendil-works/pi-coding-agent";
218+
219+declare const entryIdBrand: unique symbol;
220+
221+export type EntryId = string & { readonly [entryIdBrand]: "EntryId" };
222+
223+interface UnknownTreeEntry {
224+  readonly id: EntryId;
225+  readonly parentId: EntryId | null;
226   readonly type: string;
227   readonly message?: unknown;
228   readonly summary?: unknown;
229 }
230 
231+export type TreeEntry = (SessionEntry & {
232+  readonly id: EntryId;
233+  readonly parentId: EntryId | null;
234+}) | UnknownTreeEntry;
235+
236 export interface BranchCandidate {
237-  readonly rootId: string;
238-  readonly forkPointId: string;
239+  readonly rootId: EntryId;
240+  readonly forkPointId: EntryId;
241   readonly entries: readonly TreeEntry[];
242-  readonly entryIds: readonly string[];
243+  readonly entryIds: readonly EntryId[];
244   readonly entryCount: number;
245   readonly preview: string;
246 }
247 
248 interface TreeIndex {
249-  readonly entriesById: ReadonlyMap<string, TreeEntry>;
250-  readonly childrenByParentId: ReadonlyMap<string, readonly string[]>;
251-  readonly uncertainChildrenByParentId: ReadonlySet<string>;
252+  readonly entriesById: ReadonlyMap<EntryId, TreeEntry>;
253+  readonly childrenByParentId: ReadonlyMap<EntryId, readonly EntryId[]>;
254+  readonly uncertainChildrenByParentId: ReadonlySet<EntryId>;
255 }
256 
257 const PREVIEW_LIMIT = 72;
258 
259 export function findInactiveBranchCandidates(
260   entries: readonly TreeEntry[],
261-  activeLeafId: string,
262+  activeLeafId: EntryId,
263 ): BranchCandidate[] {
264   const tree = buildTreeIndex(entries);
265   if (!tree.entriesById.has(activeLeafId)) return [];
266@@ -59,7 +70,7 @@ export function findInactiveBranchCandidates(
267 export function normalizePruneCandidates(
268   candidates: readonly BranchCandidate[],
269 ): BranchCandidate[] {
270-  const candidatesByRootId = new Map<string, BranchCandidate>();
271+  const candidatesByRootId = new Map<EntryId, BranchCandidate>();
272   for (const candidate of candidates) {
273     if (!candidatesByRootId.has(candidate.rootId)) {
274       candidatesByRootId.set(candidate.rootId, candidate);
275@@ -84,8 +95,8 @@ export function pruneCandidatesOverlap(left: BranchCandidate, right: BranchCandi
276 
277 function buildTreeIndex(entries: readonly TreeEntry[]): TreeIndex {
278   const validEntries = entries.filter(isTreeEntry);
279-  const occurrences = new Map<string, number>();
280-  const uncertainChildrenByParentId = new Set<string>();
281+  const occurrences = new Map<EntryId, number>();
282+  const uncertainChildrenByParentId = new Set<EntryId>();
283 
284   for (const entry of validEntries) {
285     occurrences.set(entry.id, (occurrences.get(entry.id) ?? 0) + 1);
286@@ -93,15 +104,15 @@ function buildTreeIndex(entries: readonly TreeEntry[]): TreeIndex {
287 
288   for (const entry of entries) {
289     if (isTreeEntry(entry) && occurrences.get(entry.id) === 1) continue;
290-    if (hasStringParentId(entry)) uncertainChildrenByParentId.add(entry.parentId);
291+    if (hasStringParentId(entry)) uncertainChildrenByParentId.add(toEntryId(entry.parentId));
292   }
293 
294-  const entriesById = new Map<string, TreeEntry>();
295+  const entriesById = new Map<EntryId, TreeEntry>();
296   for (const entry of validEntries) {
297     if (occurrences.get(entry.id) === 1) entriesById.set(entry.id, entry);
298   }
299 
300-  const childrenByParentId = new Map<string, string[]>();
301+  const childrenByParentId = new Map<EntryId, EntryId[]>();
302   for (const entry of entriesById.values()) {
303     if (entry.parentId === null || !entriesById.has(entry.parentId)) continue;
304     const children = childrenByParentId.get(entry.parentId) ?? [];
305@@ -112,9 +123,9 @@ function buildTreeIndex(entries: readonly TreeEntry[]): TreeIndex {
306   return { entriesById, childrenByParentId, uncertainChildrenByParentId };
307 }
308 
309-function collectSubtree(rootId: string, tree: TreeIndex): TreeEntry[] | undefined {
310+function collectSubtree(rootId: EntryId, tree: TreeIndex): TreeEntry[] | undefined {
311   const entries: TreeEntry[] = [];
312-  const visitedIds = new Set<string>();