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 // Give every DOM boundary its own position, including element boundaries
124 // immediately before text offset 0 and after its final character.
125 const positions = new Map<SelectionNode, number[]>();
126 let position = 0;
127 const index = (node: SelectionNode) => {
128 const offsets = [position++];
129 positions.set(node, offsets);
130 if (node.nodeType === 3) {
131 for (let offset = 0; offset < node.length; offset++) offsets.push(position++);
132 } else {
133 for (const child of node.childNodes) {
134 index(child);
135 offsets.push(position++);
136 }
137 }
138 };
139 index(root);
140 const at = (node: SelectionNode, offset: number) => positions.get(node)![offset]!;
141
142 class SelectionRange {
143 constructor(
144 public startContainer: SelectionNode,
145 public startOffset: number,
146 public endContainer: SelectionNode,
147 public endOffset: number,
148 ) {}
149 get collapsed() {
150 return this.startContainer === this.endContainer && this.startOffset === this.endOffset;
151 }
152 get commonAncestorContainer() {
153 let ancestor = this.startContainer;
154 while (!ancestor.contains(this.endContainer)) ancestor = ancestor.parentElement!;
155 return ancestor;
156 }
157 cloneRange() {
158 return new SelectionRange(
159 this.startContainer,
160 this.startOffset,
161 this.endContainer,
162 this.endOffset,
163 );
164 }
165 setStart(node: SelectionNode, offset: number) {
166 this.startContainer = node;
167 this.startOffset = offset;
168 }
169 setEnd(node: SelectionNode, offset: number) {
170 this.endContainer = node;
171 this.endOffset = offset;
172 }
173 intersectsNode(node: SelectionNode) {
174 const parent = node.parentElement;
175 if (parent === null) return true;
176 const offset = parent.childNodes.indexOf(node);
177 return (
178 at(parent, offset) < at(this.endContainer, this.endOffset) &&
179 at(parent, offset + 1) > at(this.startContainer, this.startOffset)
180 );
181 }
182 }
183
184 const range = new SelectionRange(...start, ...end);
185 const anchor = backwards ? end : start;
186 const focus = backwards ? start : end;
187 return {
188 isCollapsed: range.collapsed,
189 rangeCount: 1,
190 anchorNode: anchor[0],
191 anchorOffset: anchor[1],
192 focusNode: focus[0],
193 focusOffset: focus[1],
194 getRangeAt: () => range,
195 } as unknown as Selection;
196}
197
198function capture(viewport: SelectionNode, selection: Selection) {
199 return captureAssistantTextSelection(viewport as unknown as HTMLElement, selection);
200}
201
202describe("captureAssistantTextSelection", () => {
203 it.each([
204 { boundary: "next text", backwards: false },
205 { boundary: "next block", backwards: false },
206 { boundary: "parent", backwards: false },
207 { boundary: "next text", backwards: true },
208 { boundary: "next block", backwards: true },
209 { boundary: "parent", backwards: true },
210 ])(
211 "captures a paragraph ending at $boundary, backwards=$backwards",
212 ({ boundary, backwards }) => {
213 const quote = textNode("A paragraph šŸ˜€ cafe\u0301.");
214 const source = assistantSource(new SelectionNode("P").append(quote));
215 const nextText = textNode("Another response.");
216 const nextBlock = assistantSource(new SelectionNode("P").append(nextText));
217 const viewport = new SelectionNode("MAIN").append(source, nextBlock);
218 const end: [SelectionNode, number] =
219 boundary === "next text"
220 ? [nextText, 0]
221 : boundary === "next block"
222 ? [nextBlock, 0]
223 : [viewport, 1];
224 const selection = nativeSelection([quote, 0], end, backwards);
225 const captured = capture(viewport, selection);
226
227 expect(captured?.source).toBe(source);
228 expect(captured?.selector).toEqual(selector(quote.data));
229 expect(captured?.range.endContainer).toBe(quote);
230 expect(captured?.range.endOffset).toBe(quote.length);
231 expect(selection.getRangeAt(0).endContainer).toBe(end[0]);
232 expect(selection.getRangeAt(0).endOffset).toBe(end[1]);
233 },
234 );
235
236 it("finds the source when both endpoints are parent boundaries", () => {
237 const quote = textNode("Only this response.");
238 const source = assistantSource(new SelectionNode("P").append(quote));
239 const viewport = new SelectionNode("MAIN").append(source, assistantSource(textNode("Next.")));
240 expect(capture(viewport, nativeSelection([viewport, 0], [viewport, 1]))?.selector).toEqual(
241 selector(quote.data),
242 );
243 });
244
245 it("captures the final paragraph when Chromium extends to the timestamp paragraph boundary", () => {
246 const quote = textNode("The final paragraph.");
247 const source = assistantSource(
248 new SelectionNode("P").append(textNode("Earlier answer.")),
249 new SelectionNode("P").append(quote),
250 );
251 const timestampText = textNode("1:50 AM");
252 const timestamp = new SelectionNode("P").append(timestampText);
253 new SelectionNode("MAIN").append(
254 new SelectionNode("DIV").append(source, new SelectionNode("DIV").append(timestamp)),
255 );
256 const selection = nativeSelection([quote, 0], [timestamp, 0]);
257 selection.toString = () => `${quote.data}\n\n\n`;
258
259 expect(capture(source, selection)?.selector).toEqual(
260 selector(quote.data, {
261 start: "Earlier answer. ".length,
262 end: "Earlier answer. The final paragraph.".length,
263 prefix: "Earlier answer. ",
264 }),
265 );
266 expect(capture(source, nativeSelection([quote, 0], [timestampText, 1]))).toBeNull();
267 });
268
269 it("captures a triple-clicked response that ends before a tool row and the next response", () => {
270 const quote = textNode("Recording your answers in the session note:");
271 const source = assistantSource(new SelectionNode("P").append(quote));
272 const nextParagraph = new SelectionNode("P").append(textNode("1. What it means."));
273 const viewport = new SelectionNode("MAIN").append(
274 source,
275 new SelectionNode("DIV", "", { "data-select-none": "" }).append(textNode("echo ok")),
276 new SelectionNode("H3", "", { "data-select-none": "" }).append(textNode("T3 Code")),
277 assistantSource(nextParagraph),
278 );
279
280 expect(capture(viewport, nativeSelection([quote, 0], [nextParagraph, 0]))?.selector).toEqual(
281 selector(quote.data),
282 );
283 });
284
285 it("does not attribute an empty starting endpoint to the previous response", () => {
286 const previous = textNode("Previous response.");
287 const quote = textNode("Selected response.");
288 const source = assistantSource(quote);
289 const viewport = new SelectionNode("MAIN").append(assistantSource(previous), source);
290 const captured = capture(
291 viewport,
292 nativeSelection([previous, previous.length], [quote, quote.length]),
293 );
294 expect(captured?.source).toBe(source);
295 expect(captured?.selector).toEqual(selector(quote.data));
296 });
297
298 it.each([false, true])(
299 "rejects actual overlap with another response, backwards=%s",
300 (backwards) => {
301 const first = textNode("Repeated quote.");
302 const second = textNode(first.data);
303 const viewport = new SelectionNode("MAIN").append(
304 assistantSource(first),
305 assistantSource(second),
306 );
307 expect(capture(viewport, nativeSelection([first, 0], [second, 1], backwards))).toBeNull();
308 },
309 );
310
311 it("accepts a control endpoint only when none of its text is selected", () => {
312 const quote = textNode("A paragraph.");
313 const control = textNode("Copy");
314 const source = assistantSource(
315 new SelectionNode("P").append(quote),
316 new SelectionNode("BUTTON").append(control),
317 );
318 expect(capture(source, nativeSelection([quote, 0], [control, 0]))?.selector).toEqual(
319 selector(quote.data),
320 );
321 expect(capture(source, nativeSelection([quote, 0], [control, 1]))).toBeNull();
322 });
323
324 it.each([
325 new SelectionNode("BUTTON"),
326 new SelectionNode("SPAN", "", { role: "button" }),
327 new SelectionNode("DIV", "", { contenteditable: "true" }),
328 new SelectionNode("SPAN", "", { hidden: "" }),
329 new SelectionNode("SPAN", "", { "aria-hidden": "true" }),
330 ])("rejects selected text at excluded endpoints: %j", (excluded) => {
331 const quote = textNode("A paragraph.");
332 const control = textNode("Excluded");
333 const source = assistantSource(quote, excluded.append(control));
334 expect(capture(source, nativeSelection([quote, 0], [control, 1]))).toBeNull();
335 expect(capture(source, nativeSelection([control, 0], [control, control.length]))).toBeNull();
336 });
337
338 it("keeps exact multiline Unicode text and omits interior controls and hidden text", () => {
339 const first = textNode("Before šŸ˜€ ");
340 const code = textNode("first line\r\n second line\tšŸš€");
341 const source = assistantSource(
342 new SelectionNode("P").append(first, new SelectionNode("EM").append(textNode("cafe\u0301"))),
343 new SelectionNode("BUTTON").append(textNode("Copy")),
344 new SelectionNode("SPAN", "", { hidden: "" }).append(textNode("Hidden")),
345 new SelectionNode("PRE").append(new SelectionNode("CODE").append(code)),
346 );
347 expect(
348 capture(source, nativeSelection([first, 0], [code, code.length], true))?.selector,
349 ).toEqual(selector("Before šŸ˜€ cafe\u0301\nfirst line\r\n second line\tšŸš€"));
350 });
351
352 it("rejects a source outside the supplied viewport", () => {
353 const quote = textNode("Outside.");
354 assistantSource(quote);
355 expect(
356 capture(new SelectionNode("MAIN"), nativeSelection([quote, 0], [quote, quote.length])),
357 ).toBeNull();
358 });
359});
360
361describe("createAssistantTextSelector", () => {
362 it("keeps exact selected whitespace while normalizing its UTF-16 positions", () => {
363 const text = "Before \n quote\t \n after";
364 expect(
365 createAssistantTextSelector(text, "Before \n".length, "Before \n quote\t".length),
366 ).toEqual({
367 text: " quote\t",
368 start: 6,
369 end: 13,
370 prefix: "Before",
371 suffix: "after",
372 });
373 });
374
375 it.each([
376 0,
377 ASSISTANT_CITATION_CONTEXT_LENGTH - 2,
378 ASSISTANT_CITATION_CONTEXT_LENGTH - 1,
379 ASSISTANT_CITATION_CONTEXT_LENGTH,
380 ])("keeps complete context code points with %i neighboring UTF-16 units", (paddingLength) => {
381 const padding = "x".repeat(paddingLength);
382 const text = `šŸ˜€${padding}quote${padding}šŸš€`;
383 const start = text.indexOf("quote");
384 const captured = createAssistantTextSelector(text, start, start + 5);
385 const keepEmoji = paddingLength + 2 <= ASSISTANT_CITATION_CONTEXT_LENGTH;
386 expect(captured).toEqual({
387 text: "quote",
388 start,
389 end: start + 5,
390 prefix: `${keepEmoji ? "šŸ˜€" : ""}${padding}`,
391 suffix: `${padding}${keepEmoji ? "šŸš€" : ""}`,
392 });
393 expect(roundTripSelector(captured!)).toMatchObject(captured!);
394 });
395
396 it("preserves multiline quotes and their matching location through a citation URL", () => {
397 const quote = "selected\r\n code šŸš€";
398 const prefix = "x".repeat(ASSISTANT_CITATION_CONTEXT_LENGTH - 1);
399 const suffix = "y".repeat(ASSISTANT_CITATION_CONTEXT_LENGTH - 1);
400 const text = `${quote} elsewhere\nšŸ˜€${prefix}${quote}${suffix}šŸš€`;
401 const rawStart = text.lastIndexOf(quote);
402 const captured = createAssistantTextSelector(text, rawStart, rawStart + quote.length);
403 const start = "selected code šŸš€ elsewhere šŸ˜€".length + prefix.length;
404 const end = start + "selected code šŸš€".length;
405 expect(captured).toEqual({ text: quote, start, end, prefix, suffix });
406
407 const parsed = roundTripSelector(captured!);
408 expect(parsed).toMatchObject(captured!);
409 expect(findAssistantCitationText(text, parsed!)).toEqual({ start, end });
410 expect(findAssistantCitationText(`Inserted paragraph.\n${text}`, parsed!)).toEqual({
411 start: "Inserted paragraph. ".length + start,
412 end: "Inserted paragraph. ".length + end,
413 });
414 });
415
416 it.each(["prefix", "suffix"])(
417 "does not relocate to a replacement-character decoy after clipping the %s",
418 (side) => {
419 const prefix = "x".repeat(ASSISTANT_CITATION_CONTEXT_LENGTH - (side === "prefix" ? 1 : 0));
420 const suffix = "y".repeat(ASSISTANT_CITATION_CONTEXT_LENGTH - (side === "suffix" ? 1 : 0));
421 const occurrence = (character: string) =>
422 `${side === "prefix" ? character : ""}${prefix}quote${suffix}${side === "suffix" ? character : ""}`;
423 const text = `${occurrence("šŸ˜€")} / ${occurrence("\uFFFD")}`;
424 const start = text.indexOf("quote");
425 const captured = createAssistantTextSelector(text, start, start + 5);
426 const parsed = roundTripSelector(captured!);
427 expect(parsed).toMatchObject(captured!);
428 expect(findAssistantCitationText(text, captured!)).toBeNull();
429 expect(findAssistantCitationText(text, parsed!)).toBeNull();
430 },
431 );
432
433 it.each(["", " \n\t\r\n"])("rejects empty or whitespace-only selections: %j", (text) => {
434 expect(createAssistantTextSelector(text, 0, text.length)).toBeNull();
435 });
436});
437
438describe("findAssistantCitationText", () => {
439 it("resolves an exact selection with its saved position and context", () => {
440 expect(
441 findAssistantCitationText(
442 "Before the selected text after.",
443 selector("selected text", { start: 11, end: 24, prefix: "Before the ", suffix: " after." }),
444 ),
445 ).toEqual({ start: 11, end: 24 });
446 });
447
448 it("finds a unique quote when insertion before it shifts its offsets", () => {
449 expect(
450 findAssistantCitationText(
451 "An inserted paragraph. Before the selected text after.",
452 selector("selected text", { start: 11, end: 24, prefix: "Before the ", suffix: " after." }),
453 ),
454 ).toEqual({ start: 34, end: 47 });
455 });
456
457 it("still finds a unique quote after its surrounding text changes", () => {
458 expect(
459 findAssistantCitationText(
460 "New selected text nearby.",
461 selector("selected text", { start: 11, end: 24, prefix: "Before the ", suffix: " after." }),
462 ),
463 ).toEqual({ start: 4, end: 17 });
464 });
465
466 it("uses both context sides when each side alone matches several repeated quotes", () => {
467 const text = "left quote one; other quote two; left quote two";
468 expect(
469 findAssistantCitationText(
470 text,
471 selector("quote", { start: 5, end: 10, prefix: "left ", suffix: " two" }),
472 ),
473 ).toEqual({ start: 38, end: 43 });
474 });
475
476 it("does not trust stale offsets that now point at another occurrence", () => {
477 expect(
478 findAssistantCitationText(
479 "wrong quote here; right quote there",
480 selector("quote", { start: 6, end: 11, prefix: "right ", suffix: " there" }),
481 ),
482 ).toEqual({ start: 24, end: 29 });
483 });
484
485 it.each([
486 { prefix: "first ", suffix: "", expected: { start: 6, end: 11 } },
487 { prefix: "", suffix: " last", expected: { start: 14, end: 19 } },
488 ])(
489 "allows a single context side to disambiguate: $prefix / $suffix",
490 ({ expected, ...context }) => {
491 expect(
492 findAssistantCitationText("first quote / quote last", selector("quote", context)),
493 ).toEqual(expected);
494 },
495 );
496
497 it("does not guess between quotes without context, even at the saved position", () => {
498 expect(findAssistantCitationText("quote / quote", selector("quote"))).toBeNull();
499 });
500
501 it("does not use distance or saved offsets to break a context tie", () => {
502 expect(
503 findAssistantCitationText(
504 "same quote end / same quote end",
505 selector("quote", { start: 5, end: 10, prefix: "same ", suffix: " end" }),
506 ),
507 ).toBeNull();
508 });
509
510 it("rejects repeated quotes when neither occurrence matches all supplied context", () => {
511 expect(
512 findAssistantCitationText(
513 "left quote wrong / wrong quote right",
514 selector("quote", { prefix: "left ", suffix: " right" }),
515 ),
516 ).toBeNull();
517 });
518
519 it("counts overlapping occurrences when checking ambiguity", () => {
520 expect(findAssistantCitationText("banana", selector("ana", { start: 1, end: 4 }))).toBeNull();
521 });
522
523 it("matches multiline code after indentation, tabs, and line endings change", () => {
524 const quote = "if (ready) {\r\n run();\r\n}";
525 expect(
526 findAssistantCitationText(
527 "Example:\nif (ready) {\n\trun();\n}\nDone.",
528 selector(quote, { prefix: "Example:\n", suffix: "\nDone." }),
529 ),
530 ).toEqual({ start: 9, end: 30 });
531 });
532
533 it("normalizes nonbreaking spaces and keeps selected boundary whitespace", () => {
534 expect(
535 findAssistantCitationText(
536 "before\u00a0\n quoted\t text \nafter",
537 selector("\nquoted text\t"),
538 ),
539 ).toEqual({ start: 6, end: 19 });
540 });
541
542 it("does not trim the document's leading whitespace out of its offsets", () => {
543 expect(findAssistantCitationText("\n\t quote", selector("quote"))).toEqual({
544 start: 1,
545 end: 6,
546 });
547 });
548
549 it("uses UTF-16 offsets for emoji and combining characters", () => {
550 expect(findAssistantCitationText("šŸ˜€ cafe\u0301 šŸš€ done", selector("cafe\u0301 šŸš€"))).toEqual({
551 start: 3,
552 end: 11,
553 });
554 });
555
556 it("uses up to 32 UTF-16 context units without requiring the entire surrounding text", () => {
557 const prefix = "x".repeat(ASSISTANT_CITATION_CONTEXT_LENGTH - 1) + " ";
558 const suffix = " " + "y".repeat(ASSISTANT_CITATION_CONTEXT_LENGTH - 1);
559 const text = `unrelated quote / ${prefix}quote${suffix} changed further away`;
560 const start = text.lastIndexOf("quote");
561 expect(findAssistantCitationText(text, selector("quote", { prefix, suffix }))).toEqual({
562 start,
563 end: start + 5,
564 });
565 });
566
567 it.each(["", " \n\t\r\n", "absent", "QUOTE", "qu.te"])(
568 "rejects empty or missing literal text: %j",
569 (quote) => {
570 expect(findAssistantCitationText("quote", selector(quote))).toBeNull();
571 },
572 );
573
574 it.each([
575 { start: -1, end: 4 },
576 { start: 2.5, end: 7.5 },
577 { start: Number.NaN, end: Number.NaN },
578 { start: Number.POSITIVE_INFINITY, end: Number.POSITIVE_INFINITY },
579 { start: 8, end: 3 },
580 ])("treats invalid stored offsets as unusable hints: $start / $end", (offsets) => {
581 expect(findAssistantCitationText("a quote", selector("quote", offsets))).toEqual({
582 start: 2,
583 end: 7,
584 });
585 });
586});