Skip to content

Commit fa9fbbf

Browse files
FranciscoMorettiFranciscoMoretti
authored
Threads 2: Add the canonical message tree (#235)
* feat(thread): add canonical message tree * refactor(thread): serialize ordered message nodes * fix(thread): remove root sentinel collision * refactor(thread): simplify message tree path updates * refactor(thread): separate path updates from selection * refactor(thread): reuse ordered tree traversal
1 parent e3189c5 commit fa9fbbf

5 files changed

Lines changed: 477 additions & 2 deletions

File tree

‎packages/thread/package.json‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -40,8 +40,8 @@
4040
},
4141
"scripts": {
4242
"build": "bun -e \"import { rmSync } from 'node:fs'; rmSync('dist', { recursive: true, force: true })\" && tsc -p tsconfig.build.json && bun build ./src/index.ts --outfile ./dist/index.js --target=browser --format=esm --packages external",
43-
"format": "bunx @biomejs/biome@2.4.10 check --write src *.md package.json tsconfig.json tsconfig.build.json biome.jsonc",
44-
"lint": "bunx @biomejs/biome@2.4.10 check src ARCHITECTURE.md package.json tsconfig.json tsconfig.build.json biome.jsonc",
43+
"format": "bunx @biomejs/biome@2.4.10 check --write src test *.md package.json tsconfig.json tsconfig.build.json biome.jsonc",
44+
"lint": "bunx @biomejs/biome@2.4.10 check src test ARCHITECTURE.md package.json tsconfig.json tsconfig.build.json biome.jsonc",
4545
"prepublishOnly": "bun run build",
4646
"test": "bun test --pass-with-no-tests",
4747
"test:unit": "bun test --pass-with-no-tests",

‎packages/thread/src/index.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
export { getMessageText } from "./message-utils";
2+
13
export type {
24
MessageTreeNode,
35
MessageTreeSnapshot,

‎packages/thread/src/message-tree.ts‎

Lines changed: 283 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,283 @@
1+
import type { UIMessage } from "ai";
2+
import type { MessageTreeSnapshot } from "./types";
3+
4+
function clone<T>(value: T): T {
5+
return structuredClone(value);
6+
}
7+
8+
export class MessageTree<TMessage extends UIMessage = UIMessage> {
9+
readonly #childrenByParentId = new Map<string | null, string[]>();
10+
readonly #messagesById = new Map<string, TMessage>();
11+
readonly #parentById = new Map<string, string | null>();
12+
#cursorId: string | null = null;
13+
14+
constructor(
15+
options: {
16+
messages?: TMessage[];
17+
snapshot?: MessageTreeSnapshot<TMessage>;
18+
} = {},
19+
) {
20+
if (options.snapshot) {
21+
this.restore(options.snapshot);
22+
} else if (options.messages) {
23+
this.setPath(options.messages);
24+
}
25+
}
26+
27+
get cursorId() {
28+
return this.#cursorId;
29+
}
30+
31+
has(messageId: string) {
32+
return this.#messagesById.has(messageId);
33+
}
34+
35+
getMessage(messageId: string) {
36+
const message = this.#messagesById.get(messageId);
37+
return message ? clone(message) : undefined;
38+
}
39+
40+
getParentId(messageId: string) {
41+
return this.#parentById.get(messageId);
42+
}
43+
44+
getParent(messageId: string) {
45+
const parentId = this.#parentById.get(messageId);
46+
return parentId ? this.getMessage(parentId) : undefined;
47+
}
48+
49+
getChildren(messageId: string | null) {
50+
return (this.#childrenByParentId.get(messageId) ?? [])
51+
.map((id) => this.#messagesById.get(id))
52+
.filter((message): message is TMessage => Boolean(message))
53+
.map(clone);
54+
}
55+
56+
getSiblings(messageId: string) {
57+
if (!this.#messagesById.has(messageId)) {
58+
return [];
59+
}
60+
return this.getChildren(this.#parentById.get(messageId) ?? null);
61+
}
62+
63+
getLeaves(messageId: string | null = null) {
64+
const leaves: TMessage[] = [];
65+
66+
for (const id of this.walkDescendantIds(messageId)) {
67+
const children = this.#childrenByParentId.get(id) ?? [];
68+
if (children.length === 0) {
69+
const message = this.#messagesById.get(id);
70+
if (message) {
71+
leaves.push(clone(message));
72+
}
73+
}
74+
}
75+
76+
return leaves;
77+
}
78+
79+
getPathIds(messageId: string | null | undefined = this.#cursorId) {
80+
if (!messageId) {
81+
return [];
82+
}
83+
const ids: string[] = [];
84+
let currentId: string | null = messageId;
85+
while (currentId) {
86+
if (!this.#messagesById.has(currentId)) {
87+
break;
88+
}
89+
ids.unshift(currentId);
90+
currentId = this.#parentById.get(currentId) ?? null;
91+
}
92+
return ids;
93+
}
94+
95+
getPath(messageId: string | null | undefined = this.#cursorId) {
96+
return this.getPathIds(messageId)
97+
.map((id) => this.#messagesById.get(id))
98+
.filter((message): message is TMessage => Boolean(message))
99+
.map(clone);
100+
}
101+
102+
getSnapshot(): MessageTreeSnapshot<TMessage> {
103+
const nodes: MessageTreeSnapshot<TMessage>["nodes"] = [];
104+
105+
for (const messageId of this.walkDescendantIds(null)) {
106+
const message = this.#messagesById.get(messageId);
107+
if (message) {
108+
nodes.push({
109+
message: clone(message),
110+
parentId: this.#parentById.get(messageId) ?? null,
111+
});
112+
}
113+
}
114+
115+
return {
116+
cursorId: this.#cursorId,
117+
nodes,
118+
version: 1,
119+
};
120+
}
121+
122+
getIndexes() {
123+
return {
124+
childrenByParentId: Object.fromEntries(
125+
Array.from(this.#childrenByParentId.entries())
126+
.filter((entry): entry is [string, string[]] => entry[0] !== null)
127+
.map(([id, children]) => [id, [...children]]),
128+
),
129+
messagesById: Object.fromEntries(
130+
Array.from(this.#messagesById.entries(), ([id, message]) => [
131+
id,
132+
clone(message),
133+
]),
134+
),
135+
parentById: Object.fromEntries(this.#parentById.entries()),
136+
rootIds: [...(this.#childrenByParentId.get(null) ?? [])],
137+
};
138+
}
139+
140+
setCursor(messageId: string | null) {
141+
if (messageId !== null && !this.#messagesById.has(messageId)) {
142+
throw new Error(`Unknown message ${messageId}`);
143+
}
144+
this.#cursorId = messageId;
145+
}
146+
147+
setCursorToParentOf(messageId: string) {
148+
if (!this.#messagesById.has(messageId)) {
149+
throw new Error(`Unknown message ${messageId}`);
150+
}
151+
this.setCursor(this.#parentById.get(messageId) ?? null);
152+
}
153+
154+
upsertMessage(
155+
message: TMessage,
156+
parentId: string | null,
157+
options: { index?: number } = {},
158+
) {
159+
if (parentId !== null && !this.#messagesById.has(parentId)) {
160+
throw new Error(`Unknown parent message ${parentId}`);
161+
}
162+
let ancestorId = parentId;
163+
while (ancestorId !== null) {
164+
if (ancestorId === message.id) {
165+
throw new Error(`Cannot create a cycle involving ${message.id}`);
166+
}
167+
ancestorId = this.#parentById.get(ancestorId) ?? null;
168+
}
169+
170+
const existingParentId = this.#parentById.get(message.id);
171+
if (existingParentId !== undefined && existingParentId !== parentId) {
172+
throw new Error(
173+
`Cannot move message ${message.id} from ${existingParentId ?? "root"} to ${parentId ?? "root"}`,
174+
);
175+
}
176+
177+
this.#messagesById.set(message.id, clone(message));
178+
this.#parentById.set(message.id, parentId);
179+
const children = this.#childrenByParentId.get(parentId) ?? [];
180+
if (!children.includes(message.id)) {
181+
const index = Math.min(options.index ?? children.length, children.length);
182+
this.#childrenByParentId.set(parentId, [
183+
...children.slice(0, index),
184+
message.id,
185+
...children.slice(index),
186+
]);
187+
}
188+
}
189+
190+
removeLeaf(messageId: string) {
191+
if (!this.#messagesById.has(messageId)) {
192+
return;
193+
}
194+
const children = this.#childrenByParentId.get(messageId) ?? [];
195+
if (children.length > 0) {
196+
throw new Error(`Cannot remove non-leaf message ${messageId}`);
197+
}
198+
const parentId = this.#parentById.get(messageId) ?? null;
199+
this.#messagesById.delete(messageId);
200+
this.#parentById.delete(messageId);
201+
this.#childrenByParentId.delete(messageId);
202+
this.#childrenByParentId.set(
203+
parentId,
204+
(this.#childrenByParentId.get(parentId) ?? []).filter(
205+
(id) => id !== messageId,
206+
),
207+
);
208+
if (this.#cursorId === messageId) {
209+
this.#cursorId = parentId;
210+
}
211+
}
212+
213+
setPath(messages: TMessage[]) {
214+
this.updatePath(messages);
215+
this.#cursorId = messages.at(-1)?.id ?? null;
216+
}
217+
218+
updatePath(messages: TMessage[]) {
219+
this.validatePath(messages, true);
220+
let parentId: string | null = null;
221+
for (const message of messages) {
222+
this.upsertMessage(message, parentId);
223+
parentId = message.id;
224+
}
225+
}
226+
227+
restore(snapshot: MessageTreeSnapshot<TMessage>) {
228+
const restored = new MessageTree<TMessage>();
229+
for (const { message, parentId } of snapshot.nodes) {
230+
if (restored.has(message.id)) {
231+
throw new Error(`Duplicate message id ${message.id} in snapshot`);
232+
}
233+
restored.upsertMessage(message, parentId);
234+
}
235+
restored.setCursor(snapshot.cursorId);
236+
237+
this.clear();
238+
for (const [id, message] of restored.#messagesById) {
239+
this.#messagesById.set(id, message);
240+
}
241+
for (const [id, parentId] of restored.#parentById) {
242+
this.#parentById.set(id, parentId);
243+
}
244+
for (const [id, children] of restored.#childrenByParentId) {
245+
this.#childrenByParentId.set(id, children);
246+
}
247+
this.#cursorId = restored.#cursorId;
248+
}
249+
250+
clear() {
251+
this.#childrenByParentId.clear();
252+
this.#messagesById.clear();
253+
this.#parentById.clear();
254+
this.#cursorId = null;
255+
}
256+
257+
private *walkDescendantIds(parentId: string | null): Generator<string> {
258+
for (const childId of this.#childrenByParentId.get(parentId) ?? []) {
259+
yield childId;
260+
yield* this.walkDescendantIds(childId);
261+
}
262+
}
263+
264+
private validatePath(messages: TMessage[], validateExistingParents = false) {
265+
const ids = new Set<string>();
266+
let parentId: string | null = null;
267+
for (const message of messages) {
268+
if (ids.has(message.id)) {
269+
throw new Error(`Duplicate message id ${message.id} in path`);
270+
}
271+
ids.add(message.id);
272+
if (validateExistingParents) {
273+
const existingParentId = this.#parentById.get(message.id);
274+
if (existingParentId !== undefined && existingParentId !== parentId) {
275+
throw new Error(
276+
`Cannot move message ${message.id} from ${existingParentId ?? "root"} to ${parentId ?? "root"}`,
277+
);
278+
}
279+
}
280+
parentId = message.id;
281+
}
282+
}
283+
}

‎packages/thread/src/message-utils.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
import type { UIMessage } from "ai";
2+
3+
export function getMessageText(message: UIMessage) {
4+
return message.parts
5+
.map((part) => (part.type === "text" ? part.text : ""))
6+
.join("");
7+
}

0 commit comments

Comments
 (0)