1import { describe, expect, it } from "vite-plus/test";
2import {
3 ASSISTANT_CITATION_CONTEXT_LENGTH,
4 EnvironmentId,
5 MessageId,
6 ThreadId,
7} from "@t3tools/contracts";
8import {
9 formatAssistantCitationHref,
10 parseAssistantCitationHref,
11} from "@t3tools/shared/assistantCitations";
12
13import {
14 type AssistantTextSelector,
15 captureAssistantTextSelection,
16 createAssistantTextSelector,
17 findAssistantCitationText,
18} from "./assistantTextSelection";
19
20function selector(
21 text: string,
22 overrides: Partial<Omit<AssistantTextSelector, "text">> = {},
23): AssistantTextSelector {
24 return {
25 text,
26 start: 0,
27 end: text.replace(/\s+/g, " ").length,
28 prefix: "",
29 suffix: "",
30 ...overrides,
31 };
32}
33
34function roundTripSelector(selector: AssistantTextSelector) {
35 return parseAssistantCitationHref(
36 formatAssistantCitationHref({
37 version: 1,
38 environmentId: EnvironmentId.make("environment"),
39 threadId: ThreadId.make("thread"),
40 messageId: MessageId.make("assistant"),
41 ...selector,
42 }),
43 );
44}
45
46class SelectionNode {
47 parentElement: SelectionNode | null = null;
48 childNodes: SelectionNode[] = [];
49
50 constructor(
51 readonly tagName: string,
52 readonly data = "",
53 readonly attributes: Record<string, string> = {},
54 ) {}
55
56 get nodeType() {
57 return this.tagName === "#text" ? 3 : 1;
58 }
59 get ownerDocument() {
60 // `user-select: none` inherits, like Chromium's computed style.
61 const getComputedStyle = (element: SelectionNode) => ({
62 userSelect: element.closest("[data-select-none]") ? "none" : "auto",
63 });
64 return { defaultView: { getComputedStyle } };
65 }
66 get length() {
67 return this.data.length;
68 }
69 get firstChild(): SelectionNode | null {
70 return this.childNodes[0] ?? null;
71 }
72 get lastChild(): SelectionNode | null {
73 return this.childNodes.at(-1) ?? null;
74 }
75 get previousSibling(): SelectionNode | null {
76 const siblings = this.parentElement?.childNodes;
77 return siblings?.[siblings.indexOf(this) - 1] ?? null;
78 }
79 get nextSibling(): SelectionNode | null {
80 const siblings = this.parentElement?.childNodes;
81 return siblings?.[siblings.indexOf(this) + 1] ?? null;
82 }
83 append(...children: SelectionNode[]) {
84 for (const child of children) child.parentElement = this;
85 this.childNodes.push(...children);
86 return this;
87 }
88 contains(node: SelectionNode): boolean {
89 return node === this || this.childNodes.some((child) => child.contains(node));
90 }
91 matches(selector: string) {
92 return selector.split(", ").some((part) => {
93 if (!part.startsWith("[")) return part === this.tagName.toLowerCase();
94 const [attribute, value] = part.slice(1, -1).split("=");
95 return (
96 Object.hasOwn(this.attributes, attribute!) &&
97 (value === undefined || this.attributes[attribute!] === value)
98 );
99 });
100 }
101 closest(selector: string): SelectionNode | null {
102 return this.matches(selector) ? this : (this.parentElement?.closest(selector) ?? null);
103 }
104}
105
106function textNode(text: string) {
107 return new SelectionNode("#text", text);
108}
109
110function assistantSource(...children: SelectionNode[]) {
111 return new SelectionNode("DIV", "", { "data-assistant-citation-source": "assistant" }).append(
112 ...children,
113 );
114}
115
116function nativeSelection(
117 start: [SelectionNode, number],
118 end: [SelectionNode, number],
119 backwards = false,
120) {
121 let root = start[0];
122 while (root.parentElement !== null) root = root.parentElement;
123 const positions = new Map<SelectionNode, number[]>();
124 let position = 0;
125 const index = (node: SelectionNode) => {
126 const offsets = [position++];
127 positions.set(node, offsets);
128 if (node.nodeType === 3) {
129 for (let offset = 0; offset < node.length; offset++) offsets.push(position++);
130 } else {
131 for (const child of node.childNodes) {
132 index(child);
133 offsets.push(position++);
134 }
135 }
136 };
137 index(root);
138 const at = (node: SelectionNode, offset: number) => positions.get(node)![offset]!;
139
140 class SelectionRange {
141 constructor(
142 public startContainer: SelectionNode,
143 public startOffset: number,
144 public endContainer: SelectionNode,
145 public endOffset: number,
146 ) {}
147 get collapsed() {
148 return this.startContainer === this.endContainer && this.startOffset === this.endOffset;
149 }
150 get commonAncestorContainer() {
151 let ancestor = this.startContainer;
152 while (!ancestor.contains(this.endContainer)) ancestor = ancestor.parentElement!;
153 return ancestor;
154 }
155 cloneRange() {
156 return new SelectionRange(
157 this.startContainer,
158 this.startOffset,
159 this.endContainer,
160 this.endOffset,
161 );
162 }
163 setStart(node: SelectionNode, offset: number) {
164 this.startContainer = node;
165 this.startOffset = offset;
166 }
167 setEnd(node: SelectionNode, offset: number) {
168 this.endContainer = node;
169 this.endOffset = offset;
170 }
171 intersectsNode(node: SelectionNode) {
172 const parent = node.parentElement;
173 if (parent === null) return true;
174 const offset = parent.childNodes.indexOf(node);
175 return (
176 at(parent, offset) < at(this.endContainer, this.endOffset) &&
177 at(parent, offset + 1) > at(this.startContainer, this.startOffset)
178 );
179 }
180 }
181
182 const range = new SelectionRange(...start, ...end);
183 const anchor = backwards ? end : start;
184 const focus = backwards ? start : end;
185 return {
186 isCollapsed: range.collapsed,
187 rangeCount: 1,
188 anchorNode: anchor[0],
189 anchorOffset: anchor[1],
190 focusNode: focus[0],
191 focusOffset: focus[1],
192 getRangeAt: () => range,
193 } as unknown as Selection;
194}
195
196function capture(viewport: SelectionNode, selection: Selection) {
197 return captureAssistantTextSelection(viewport as unknown as HTMLElement, selection);
198}
199
200describe("captureAssistantTextSelection", () => {
201 it.each([
202 { boundary: "next text", backwards: false },
203 { boundary: "next block", backwards: false },
204 { boundary: "parent", backwards: false },
205 { boundary: "next text", backwards: true },
206 { boundary: "next block", backwards: true },
207 { boundary: "parent", backwards: true },
208 ])(
209 "captures a paragraph ending at $boundary, backwards=$backwards",
210 ({ boundary, backwards }) => {
211 const quote = textNode("A paragraph 😀 cafe\u0301.");
212 const source = assistantSource(new SelectionNode("P").append(quote));
213 const nextText = textNode("Another response.");
214 const nextBlock = assistantSource(new SelectionNode("P").append(nextText));
215 const viewport = new SelectionNode("MAIN").append(source, nextBlock);
216 const end: [SelectionNode, number] =
217 boundary === "next text"
218 ? [nextText, 0]
219 : boundary === "next block"
220 ? [nextBlock, 0]
221 : [viewport, 1];
222 const selection = nativeSelection([quote, 0], end, backwards);
223 const captured = capture(viewport, selection);
224
225 expect(captured?.source).toBe(source);
226 expect(captured?.selector).toEqual(selector(quote.data));
227 expect(captured?.range.endContainer).toBe(quote);
228 expect(captured?.range.endOffset).toBe(quote.length);
229 expect(selection.getRangeAt(0).endContainer).toBe(end[0]);
230 expect(selection.getRangeAt(0).endOffset).toBe(end[1]);
231 },
232 );
233
234 it("captures the final paragraph when Chromium extends to the timestamp paragraph boundary", () => {
235 const quote = textNode("The final paragraph.");
236 const source = assistantSource(
237 new SelectionNode("P").append(textNode("Earlier answer.")),
238 new SelectionNode("P").append(quote),
239 );
240 const timestampText = textNode("1:50 AM");
241 const timestamp = new SelectionNode("P").append(timestampText);
242 new SelectionNode("MAIN").append(
243 new SelectionNode("DIV").append(source, new SelectionNode("DIV").append(timestamp)),
244 );
245 const selection = nativeSelection([quote, 0], [timestamp, 0]);
246 selection.toString = () => `${quote.data}\n\n\n`;
247
248 expect(capture(source, selection)?.selector).toEqual(
249 selector(quote.data, {
250 start: "Earlier answer. ".length,
251 end: "Earlier answer. The final paragraph.".length,
252 prefix: "Earlier answer. ",
253 }),
254 );
255 expect(capture(source, nativeSelection([quote, 0], [timestampText, 1]))).toBeNull();
256 });
257
258 it("captures a triple-clicked response that ends before a tool row and the next response", () => {
259 const quote = textNode("Recording your answers in the session note:");
260 const source = assistantSource(new SelectionNode("P").append(quote));
261 const nextParagraph = new SelectionNode("P").append(textNode("1. What it means."));
262 const viewport = new SelectionNode("MAIN").append(
263 source,
264 new SelectionNode("DIV", "", { "data-select-none": "" }).append(textNode("echo ok")),
265 new SelectionNode("H3", "", { "data-select-none": "" }).append(textNode("T3 Code")),
266 assistantSource(nextParagraph),
267 );
268
269 expect(capture(viewport, nativeSelection([quote, 0], [nextParagraph, 0]))?.selector).toEqual(
270 selector(quote.data),
271 );
272 });
273
274 it("does not attribute an empty starting endpoint to the previous response", () => {
275 const previous = textNode("Previous response.");
276 const quote = textNode("Selected response.");
277 const source = assistantSource(quote);
278 const viewport = new SelectionNode("MAIN").append(assistantSource(previous), source);
279 const captured = capture(
280 viewport,
281 nativeSelection([previous, previous.length], [quote, quote.length]),
282 );
283 expect(captured?.source).toBe(source);
284 expect(captured?.selector).toEqual(selector(quote.data));
285 });
286
287 it.each([false, true])(
288 "rejects actual overlap with another response, backwards=%s",
289 (backwards) => {
290 const first = textNode("Repeated quote.");
291 const second = textNode(first.data);
292 const viewport = new SelectionNode("MAIN").append(
293 assistantSource(first),
294 assistantSource(second),
295 );
296 expect(capture(viewport, nativeSelection([first, 0], [second, 1], backwards))).toBeNull();
297 },
298 );
299
300 it("accepts a control endpoint only when none of its text is selected", () => {
301 const quote = textNode("A paragraph.");
302 const control = textNode("Copy");
303 const source = assistantSource(
304 new SelectionNode("P").append(quote),
305 new SelectionNode("BUTTON").append(control),
306 );
307 expect(capture(source, nativeSelection([quote, 0], [control, 0]))?.selector).toEqual(
308 selector(quote.data),
309 );
310 expect(capture(source, nativeSelection([quote, 0], [control, 1]))).toBeNull();
311 });
312
313 it.each([
314 new SelectionNode("BUTTON"),
315 new SelectionNode("SPAN", "", { role: "button" }),
316 new SelectionNode("DIV", "", { contenteditable: "true" }),
317 new SelectionNode("SPAN", "", { hidden: "" }),
318 new SelectionNode("SPAN", "", { "aria-hidden": "true" }),
319 ])("rejects selected text at excluded endpoints: %j", (excluded) => {
320 const quote = textNode("A paragraph.");
321 const control = textNode("Excluded");
322 const source = assistantSource(quote, excluded.append(control));
323 expect(capture(source, nativeSelection([quote, 0], [control, 1]))).toBeNull();
324 expect(capture(source, nativeSelection([control, 0], [control, control.length]))).toBeNull();
325 });
326
327 it("keeps exact multiline Unicode text and omits interior controls and hidden text", () => {
328 const first = textNode("Before 😀 ");
329 const code = textNode("first line\r\n second line\t🚀");
330 const source = assistantSource(
331 new SelectionNode("P").append(first, new SelectionNode("EM").append(textNode("cafe\u0301"))),
332 new SelectionNode("BUTTON").append(textNode("Copy")),
333 new SelectionNode("SPAN", "", { hidden: "" }).append(textNode("Hidden")),
334 new SelectionNode("PRE").append(new SelectionNode("CODE").append(code)),
335 );
336 expect(
337 capture(source, nativeSelection([first, 0], [code, code.length], true))?.selector,
338 ).toEqual(selector("Before 😀 cafe\u0301\nfirst line\r\n second line\t🚀"));
339 });
340});
341
342describe("createAssistantTextSelector", () => {
343 it("keeps exact selected whitespace while normalizing its UTF-16 positions", () => {
344 const text = "Before \n quote\t \n after";
345 expect(
346 createAssistantTextSelector(text, "Before \n".length, "Before \n quote\t".length),
347 ).toEqual({
348 text: " quote\t",
349 start: 6,
350 end: 13,
351 prefix: "Before",
352 suffix: "after",
353 });
354 });
355
356 it.each([
357 0,
358 ASSISTANT_CITATION_CONTEXT_LENGTH - 2,
359 ASSISTANT_CITATION_CONTEXT_LENGTH - 1,
360 ASSISTANT_CITATION_CONTEXT_LENGTH,
361 ])("keeps complete context code points with %i neighboring UTF-16 units", (paddingLength) => {
362 const padding = "x".repeat(paddingLength);
363 const text = `😀${padding}quote${padding}🚀`;
364 const start = text.indexOf("quote");
365 const captured = createAssistantTextSelector(text, start, start + 5);
366 const keepEmoji = paddingLength + 2 <= ASSISTANT_CITATION_CONTEXT_LENGTH;
367 expect(captured).toEqual({
368 text: "quote",
369 start,
370 end: start + 5,
371 prefix: `${keepEmoji ? "😀" : ""}${padding}`,
372 suffix: `${padding}${keepEmoji ? "🚀" : ""}`,
373 });
374 expect(roundTripSelector(captured!)).toMatchObject(captured!);
375 });
376
377 it("preserves multiline quotes and their matching location through a citation URL", () => {
378 const quote = "selected\r\n code 🚀";
379 const prefix = "x".repeat(ASSISTANT_CITATION_CONTEXT_LENGTH - 1);
380 const suffix = "y".repeat(ASSISTANT_CITATION_CONTEXT_LENGTH - 1);
381 const text = `${quote} elsewhere\n😀${prefix}${quote}${suffix}🚀`;
382 const rawStart = text.lastIndexOf(quote);
383 const captured = createAssistantTextSelector(text, rawStart, rawStart + quote.length);
384 const start = "selected code 🚀 elsewhere 😀".length + prefix.length;
385 const end = start + "selected code 🚀".length;
386 expect(captured).toEqual({ text: quote, start, end, prefix, suffix });
387
388 const parsed = roundTripSelector(captured!);
389 expect(parsed).toMatchObject(captured!);
390 expect(findAssistantCitationText(text, parsed!)).toEqual({ start, end });
391 expect(findAssistantCitationText(`Inserted paragraph.\n${text}`, parsed!)).toEqual({
392 start: "Inserted paragraph. ".length + start,
393 end: "Inserted paragraph. ".length + end,
394 });
395 });
396
397 it.each(["prefix", "suffix"])(
398 "does not relocate to a replacement-character decoy after clipping the %s",
399 (side) => {
400 const prefix = "x".repeat(ASSISTANT_CITATION_CONTEXT_LENGTH - (side === "prefix" ? 1 : 0));
401 const suffix = "y".repeat(ASSISTANT_CITATION_CONTEXT_LENGTH - (side === "suffix" ? 1 : 0));
402 const occurrence = (character: string) =>
403 `${side === "prefix" ? character : ""}${prefix}quote${suffix}${side === "suffix" ? character : ""}`;
404 const text = `${occurrence("😀")} / ${occurrence("\uFFFD")}`;
405 const start = text.indexOf("quote");
406 const captured = createAssistantTextSelector(text, start, start + 5);
407 const parsed = roundTripSelector(captured!);
408 expect(parsed).toMatchObject(captured!);
409 expect(findAssistantCitationText(text, captured!)).toBeNull();
410 expect(findAssistantCitationText(text, parsed!)).toBeNull();
411 },
412 );
413});
414
415describe("findAssistantCitationText", () => {
416 it("finds a unique quote when insertion before it shifts its offsets", () => {
417 expect(
418 findAssistantCitationText(
419 "An inserted paragraph. Before the selected text after.",
420 selector("selected text", { start: 11, end: 24, prefix: "Before the ", suffix: " after." }),
421 ),
422 ).toEqual({ start: 34, end: 47 });
423 });
424
425 it("uses both context sides when each side alone matches several repeated quotes", () => {
426 const text = "left quote one; other quote two; left quote two";
427 expect(
428 findAssistantCitationText(
429 text,
430 selector("quote", { start: 5, end: 10, prefix: "left ", suffix: " two" }),
431 ),
432 ).toEqual({ start: 38, end: 43 });
433 });
434
435 it("does not trust stale offsets that now point at another occurrence", () => {
436 expect(
437 findAssistantCitationText(
438 "wrong quote here; right quote there",
439 selector("quote", { start: 6, end: 11, prefix: "right ", suffix: " there" }),
440 ),
441 ).toEqual({ start: 24, end: 29 });
442 });
443
444 it.each([
445 { prefix: "first ", suffix: "", expected: { start: 6, end: 11 } },
446 { prefix: "", suffix: " last", expected: { start: 14, end: 19 } },
447 ])(
448 "allows a single context side to disambiguate: $prefix / $suffix",
449 ({ expected, ...context }) => {
450 expect(
451 findAssistantCitationText("first quote / quote last", selector("quote", context)),
452 ).toEqual(expected);
453 },
454 );
455
456 it("does not guess between quotes without context, even at the saved position", () => {
457 expect(findAssistantCitationText("quote / quote", selector("quote"))).toBeNull();
458 });
459
460 it("does not use distance or saved offsets to break a context tie", () => {
461 expect(
462 findAssistantCitationText(
463 "same quote end / same quote end",
464 selector("quote", { start: 5, end: 10, prefix: "same ", suffix: " end" }),
465 ),
466 ).toBeNull();
467 });
468
469 it("counts overlapping occurrences when checking ambiguity", () => {
470 expect(findAssistantCitationText("banana", selector("ana", { start: 1, end: 4 }))).toBeNull();
471 });
472
473 it("matches multiline code after indentation, tabs, and line endings change", () => {
474 const quote = "if (ready) {\r\n run();\r\n}";
475 expect(
476 findAssistantCitationText(
477 "Example:\nif (ready) {\n\trun();\n}\nDone.",
478 selector(quote, { prefix: "Example:\n", suffix: "\nDone." }),
479 ),
480 ).toEqual({ start: 9, end: 30 });
481 });
482
483 it("normalizes nonbreaking spaces and keeps selected boundary whitespace", () => {
484 expect(
485 findAssistantCitationText(
486 "before\u00a0\n quoted\t text \nafter",
487 selector("\nquoted text\t"),
488 ),
489 ).toEqual({ start: 6, end: 19 });
490 });
491
492 it("uses UTF-16 offsets for emoji and combining characters", () => {
493 expect(findAssistantCitationText("😀 cafe\u0301 🚀 done", selector("cafe\u0301 🚀"))).toEqual({
494 start: 3,
495 end: 11,
496 });
497 });
498
499 it("uses up to 32 UTF-16 context units without requiring the entire surrounding text", () => {
500 const prefix = "x".repeat(ASSISTANT_CITATION_CONTEXT_LENGTH - 1) + " ";
501 const suffix = " " + "y".repeat(ASSISTANT_CITATION_CONTEXT_LENGTH - 1);
502 const text = `unrelated quote / ${prefix}quote${suffix} changed further away`;
503 const start = text.lastIndexOf("quote");
504 expect(findAssistantCitationText(text, selector("quote", { prefix, suffix }))).toEqual({
505 start,
506 end: start + 5,
507 });
508 });
509});