tree-model.ts
8981 bytes
1import type { SessionEntry } from "@earendil-works/pi-coding-agent";
2
3declare const entryIdBrand: unique symbol;
4
5export type EntryId = string & { readonly [entryIdBrand]: "EntryId" };
6
7interface UnknownTreeEntry {
8 readonly id: EntryId;
9 readonly parentId: EntryId | null;
10 readonly type: string;
11 readonly message?: unknown;
12 readonly summary?: unknown;
13}
14
15export type TreeEntry = (SessionEntry & {
16 readonly id: EntryId;
17 readonly parentId: EntryId | null;
18}) | UnknownTreeEntry;
19
20export interface BranchCandidate {
21 readonly rootId: EntryId;
22 readonly forkPointId: EntryId;
23 readonly entries: readonly TreeEntry[];
24 readonly entryIds: readonly EntryId[];
25 readonly entryCount: number;
26 readonly preview: string;
27}
28
29interface TreeIndex {
30 readonly entriesById: ReadonlyMap<EntryId, TreeEntry>;
31 readonly childrenByParentId: ReadonlyMap<EntryId | null, readonly EntryId[]>;
32 readonly uncertainChildrenByParentId: ReadonlySet<EntryId>;
33}
34
35const PREVIEW_LIMIT = 72;
36
37export function findInactiveBranchCandidates(
38 entries: readonly TreeEntry[],
39 activeLeafId: EntryId,
40): BranchCandidate[] {
41 const tree = buildTreeIndex(entries);
42 if (!tree.entriesById.has(activeLeafId)) return [];
43
44 const candidates: BranchCandidate[] = [];
45 for (const [forkPointId, childIds] of tree.childrenByParentId) {
46 if (childIds.length < 2 || (forkPointId !== null && tree.uncertainChildrenByParentId.has(forkPointId))) continue;
47
48 for (const rootId of childIds) {
49 const branchEntries = collectSubtree(rootId, tree);
50 if (!branchEntries) continue;
51 if (branchEntries.some((entry) => entry.id === activeLeafId)) continue;
52
53 const root = tree.entriesById.get(rootId);
54 if (!root) continue;
55
56 candidates.push({
57 rootId,
58 forkPointId: forkPointId ?? rootId,
59 entries: branchEntries,
60 entryIds: branchEntries.map((entry) => entry.id),
61 entryCount: branchEntries.length,
62 preview: previewBranch(branchEntries),
63 });
64 }
65 }
66
67 return candidates;
68}
69
70export function normalizePruneCandidates(
71 candidates: readonly BranchCandidate[],
72): BranchCandidate[] {
73 const candidatesByRootId = new Map<EntryId, BranchCandidate>();
74 for (const candidate of candidates) {
75 if (!candidatesByRootId.has(candidate.rootId)) {
76 candidatesByRootId.set(candidate.rootId, candidate);
77 }
78 }
79
80 const uniqueCandidates = [...candidatesByRootId.values()];
81 return uniqueCandidates.filter(
82 (candidate) =>
83 !uniqueCandidates.some(
84 (possibleParent) =>
85 possibleParent.rootId !== candidate.rootId &&
86 possibleParent.entryIds.includes(candidate.rootId),
87 ),
88 );
89}
90
91export function pruneCandidatesOverlap(left: BranchCandidate, right: BranchCandidate): boolean {
92 if (left.rootId === right.rootId) return true;
93 return left.entryIds.some((entryId) => right.entryIds.includes(entryId));
94}
95
96function buildTreeIndex(entries: readonly TreeEntry[]): TreeIndex {
97 const validEntries = entries.filter(isTreeEntry);
98 const occurrences = new Map<EntryId, number>();
99 const uncertainChildrenByParentId = new Set<EntryId>();
100
101 for (const entry of validEntries) {
102 occurrences.set(entry.id, (occurrences.get(entry.id) ?? 0) + 1);
103 }
104
105 for (const entry of entries) {
106 if (isTreeEntry(entry) && occurrences.get(entry.id) === 1) continue;
107 if (hasStringParentId(entry)) uncertainChildrenByParentId.add(toEntryId(entry.parentId));
108 }
109
110 const entriesById = new Map<EntryId, TreeEntry>();
111 for (const entry of validEntries) {
112 if (occurrences.get(entry.id) === 1) entriesById.set(entry.id, entry);
113 }
114
115 const childrenByParentId = new Map<EntryId | null, EntryId[]>();
116 for (const entry of entriesById.values()) {
117 if (!isConversationEntry(entry)) continue;
118
119 const parentId = nearestConversationParent(entry.parentId, entriesById);
120 const children = childrenByParentId.get(parentId ?? null) ?? [];
121 children.push(entry.id);
122 childrenByParentId.set(parentId ?? null, children);
123 }
124
125 return { entriesById, childrenByParentId, uncertainChildrenByParentId };
126}
127
128function nearestConversationParent(
129 parentId: EntryId | null,
130 entriesById: ReadonlyMap<EntryId, TreeEntry>,
131): EntryId | undefined {
132 const visitedIds = new Set<EntryId>();
133 let currentId = parentId;
134
135 while (currentId && !visitedIds.has(currentId)) {
136 visitedIds.add(currentId);
137 const entry = entriesById.get(currentId);
138 if (!entry) return undefined;
139 if (isConversationEntry(entry)) return entry.id;
140 currentId = entry.parentId;
141 }
142
143 return undefined;
144}
145
146function isConversationEntry(entry: TreeEntry): boolean {
147 return (
148 entry.type === "message" ||
149 entry.type === "custom_message" ||
150 entry.type === "branch_summary"
151 );
152}
153
154function collectSubtree(rootId: EntryId, tree: TreeIndex): TreeEntry[] | undefined {
155 const entries: TreeEntry[] = [];
156 const visitedIds = new Set<EntryId>();
157 const pendingIds = [rootId];
158
159 while (pendingIds.length > 0) {
160 const entryId = pendingIds.pop();
161 if (!entryId || visitedIds.has(entryId)) return undefined;
162 if (tree.uncertainChildrenByParentId.has(entryId)) return undefined;
163
164 const entry = tree.entriesById.get(entryId);
165 if (!entry) return undefined;
166
167 visitedIds.add(entryId);
168 entries.push(entry);
169
170 const children = tree.childrenByParentId.get(entryId) ?? [];
171 for (let index = children.length - 1; index >= 0; index -= 1) {
172 const childId = children[index];
173 if (childId) pendingIds.push(childId);
174 }
175 }
176
177 return entries;
178}
179
180export function treeEntryFromUnknown(value: unknown): TreeEntry | undefined {
181 if (!value || typeof value !== "object") return undefined;
182
183 const entry = value as Partial<UnknownTreeEntry>;
184 if (
185 typeof entry.id !== "string" ||
186 entry.id.length === 0 ||
187 (typeof entry.parentId !== "string" && entry.parentId !== null) ||
188 typeof entry.type !== "string" ||
189 entry.type.length === 0
190 ) {
191 return undefined;
192 }
193
194 return entry as TreeEntry;
195}
196
197export function toEntryId(value: string): EntryId {
198 if (value.length === 0) throw new Error("Entry ID must not be empty");
199 return value as EntryId;
200}
201
202function isTreeEntry(value: unknown): value is TreeEntry {
203 return treeEntryFromUnknown(value) !== undefined;
204}
205
206function hasStringParentId(value: unknown): value is { parentId: string } {
207 return (
208 Boolean(value) &&
209 typeof value === "object" &&
210 typeof (value as { parentId?: unknown }).parentId === "string"
211 );
212}
213
214function previewBranch(entries: readonly TreeEntry[]): string {
215 const userMessage = entries.find((entry) => {
216 if (entry.type !== "message") return false;
217 const message = (entry as unknown as Record<string, unknown>).message;
218 return Boolean(message) && typeof message === "object" &&
219 (message as Record<string, unknown>).role === "user";
220 });
221 return previewEntry(userMessage ?? entries[0]!);
222}
223
224function previewEntry(entry: TreeEntry): string {
225 const record = entry as unknown as Record<string, unknown>;
226 const message = record.message;
227
228 if (message && typeof message === "object") {
229 const messageRecord = message as Record<string, unknown>;
230 const role = typeof messageRecord.role === "string" ? messageRecord.role : "message";
231 const text = contentPreview(messageRecord.content);
232
233 if (role === "toolResult") {
234 const toolName = typeof messageRecord.toolName === "string" ? messageRecord.toolName : "result";
235 return truncatePreview(`tool ${toolName}${text ? `: ${text}` : ""}`);
236 }
237
238 return truncatePreview(`${role}${text ? `: ${text}` : ""}`);
239 }
240
241 if (entry.type === "branch_summary" || entry.type === "compaction") {
242 return truncatePreview(`${entry.type.replace("_", " ")}: ${stringField(record, "summary")}`);
243 }
244
245 if (entry.type === "custom") {
246 return truncatePreview(`custom: ${stringField(record, "customType")}`);
247 }
248
249 if (entry.type === "label") {
250 return truncatePreview(`label: ${stringField(record, "label")}`);
251 }
252
253 return entry.type;
254}
255
256function contentPreview(content: unknown): string {
257 if (typeof content === "string") return normalizePreview(content);
258 if (!Array.isArray(content)) return "";
259
260 for (const block of content) {
261 if (!block || typeof block !== "object") continue;
262 const record = block as Record<string, unknown>;
263 if (record.type === "text" && typeof record.text === "string") {
264 return normalizePreview(record.text);
265 }
266 }
267
268 return "";
269}
270
271function stringField(record: Record<string, unknown>, field: string): string {
272 const value = record[field];
273 return typeof value === "string" ? normalizePreview(value) : "";
274}
275
276function normalizePreview(value: string): string {
277 return value.replace(/\s+/g, " ").trim();
278}
279
280function truncatePreview(value: string): string {
281 const normalized = normalizePreview(value);
282 if (normalized.length <= PREVIEW_LIMIT) return normalized || "(empty)";
283 return `${normalized.slice(0, PREVIEW_LIMIT - 1)}…`;
284}