1import { ASSISTANT_CITATION_CONTEXT_LENGTH, type AssistantCitation } from "@t3tools/contracts";
2
3export type AssistantTextSelector = {
4 readonly text: string;
5 readonly start: number;
6 readonly end: number;
7 readonly prefix: string;
8 readonly suffix: string;
9};
10
11/** Live DOM state for an open comment, never part of a saved citation. */
12export type AssistantCitationSourceAnchor = {
13 source: HTMLElement;
14 range: Range;
15 viewport: HTMLElement;
16};
17
18export function findAssistantCitationSourceAnchor(
19 document: Document,
20 citation: AssistantCitation,
21): AssistantCitationSourceAnchor | null {
22 const source = [
23 ...document.querySelectorAll<HTMLElement>("[data-assistant-citation-source]"),
24 ].find(
25 (element) =>
26 element.dataset.assistantCitationSource === citation.messageId &&
27 element.dataset.assistantCitationEnvironment === citation.environmentId &&
28 element.dataset.assistantCitationThread === citation.threadId,
29 );
30 const viewport = source?.closest<HTMLElement>("[data-assistant-citation-viewport]");
31 if (!source || !viewport) return null;
32 const range = resolveAssistantCitationRange(source, citation);
33 return range ? { source, range, viewport } : null;
34}
35
36const CONTROL_SELECTOR = "button, input, textarea, select, [role=button], [contenteditable]";
37const EXCLUDED_SELECTOR = `${CONTROL_SELECTOR}, [hidden], [aria-hidden=true], script, style, template, noscript, svg`;
38const BLOCK_SELECTOR =
39 "address, article, aside, blockquote, dd, div, dl, dt, figcaption, figure, footer, h1, h2, h3, h4, h5, h6, header, hr, li, main, nav, ol, p, pre, section, table, td, th, tr, ul";
40
41function normalizeWhitespace(text: string): string {
42 return text.replace(/\s+/g, " ");
43}
44
45function splitsSurrogatePair(text: string, offset: number): boolean {
46 const before = text.charCodeAt(offset - 1);
47 const after = text.charCodeAt(offset);
48 return before >= 0xd800 && before <= 0xdbff && after >= 0xdc00 && after <= 0xdfff;
49}
50
51/** Keeps the exact captured text while storing normalized UTF-16 positions and context. */
52export function createAssistantTextSelector(
53 text: string,
54 rawStart: number,
55 rawEnd: number,
56): AssistantTextSelector | null {
57 const quote = text.slice(rawStart, rawEnd);
58 if (quote.trim().length === 0) return null;
59
60 const normalized = normalizeWhitespace(text);
61 let start = normalizeWhitespace(text.slice(0, rawStart)).length;
62 // A selection starting inside a whitespace run includes its normalized space.
63 if (rawStart > 0 && /\s/.test(text[rawStart - 1]!) && /\s/.test(text[rawStart]!)) {
64 start -= 1;
65 }
66 const end = normalizeWhitespace(text.slice(0, rawEnd)).length;
67 let prefixStart = Math.max(0, start - ASSISTANT_CITATION_CONTEXT_LENGTH);
68 let suffixEnd = Math.min(normalized.length, end + ASSISTANT_CITATION_CONTEXT_LENGTH);
69 // A split pair becomes a replacement character when the context enters a URL.
70 if (splitsSurrogatePair(normalized, prefixStart)) prefixStart += 1;
71 if (splitsSurrogatePair(normalized, suffixEnd)) suffixEnd -= 1;
72 return {
73 text: quote,
74 start,
75 end,
76 prefix: normalized.slice(prefixStart, start),
77 suffix: normalized.slice(end, suffixEnd),
78 };
79}
80
81/**
82 * Matches case-sensitive text after collapsing each JS whitespace run to one
83 * space, without trimming. All positions and context lengths use UTF-16 units
84 * in that normalized stream, not markdown offsets or Unicode code points.
85 * `selector.text` can retain the original line breaks and code indentation.
86 * Repeated quotes require one match for all supplied context. Even matching
87 * saved offsets cannot break a context tie, since those offsets may have drifted.
88 */
89export function findAssistantCitationText(
90 text: string,
91 selector: AssistantTextSelector,
92): { start: number; end: number } | null {
93 const normalized = normalizeWhitespace(text);
94 const quote = normalizeWhitespace(selector.text);
95 if (quote.trim().length === 0) return null;
96
97 const prefix = normalizeWhitespace(selector.prefix);
98 const suffix = normalizeWhitespace(selector.suffix);
99 const matchesContext = (start: number, end: number) =>
100 normalized.slice(Math.max(0, start - prefix.length), start) === prefix &&
101 normalized.slice(end, end + suffix.length) === suffix;
102
103 let match =
104 Number.isSafeInteger(selector.start) &&
105 Number.isSafeInteger(selector.end) &&
106 selector.start >= 0 &&
107 selector.end - selector.start === quote.length &&
108 normalized.slice(selector.start, selector.end) === quote &&
109 matchesContext(selector.start, selector.end)
110 ? { start: selector.start, end: selector.end }
111 : null;
112 let onlyQuote: { start: number; end: number } | null = null;
113 let quoteCount = 0;
114
115 for (
116 let start = normalized.indexOf(quote);
117 start !== -1;
118 start = normalized.indexOf(quote, start + 1)
119 ) {
120 const end = start + quote.length;
121 quoteCount += 1;
122 onlyQuote = { start, end };
123 if (!matchesContext(start, end)) continue;
124 if (match !== null && match.start !== start) return null;
125 match = { start, end };
126 }
127
128 return match ?? (quoteCount === 1 ? onlyQuote : null);
129}
130
131type TextChunk = { node: Text; start: number; end: number };
132
133/**
134 * Uses DOM text order, with a line break between HTML blocks and at <br>.
135 * Inline markup, including code and links, contributes its displayed text.
136 * Controls and subtrees marked hidden/aria-hidden do not contribute. No layout
137 * reads, CSS-generated content, or soft-wrap line breaks enter the stream, so
138 * reflow cannot move it.
139 */
140function readAssistantText(root: HTMLElement) {
141 const parts: string[] = [];
142 const chunks: TextChunk[] = [];
143 let length = 0;
144 let separator = false;
145
146 const visit = (node: Node) => {
147 if (node.nodeType === 3) {
148 const text = node as Text;
149 if (text.length === 0) return;
150 if (separator && length > 0) {
151 parts.push("\n");
152 length += 1;
153 }
154 separator = false;
155 chunks.push({ node: text, start: length, end: length + text.length });
156 parts.push(text.data);
157 length += text.length;
158 return;
159 }
160 if (node.nodeType !== 1) return;
161 const element = node as Element;
162 if (element.matches(EXCLUDED_SELECTOR)) return;
163 const block = element.matches(BLOCK_SELECTOR);
164 if (block || element.tagName === "BR") separator = true;
165 for (const child of element.childNodes) visit(child);
166 if (block) separator = true;
167 };
168
169 visit(root);
170 return { text: parts.join(""), chunks };
171}
172
173function excludedAncestor(node: Node): Element | null {
174 const element = node.nodeType === 1 ? (node as Element) : node.parentElement;
175 return element?.closest(EXCLUDED_SELECTOR) ?? null;
176}
177
178function isUsableRange(root: HTMLElement, range: Range): boolean {
179 if (
180 range.collapsed ||
181 !root.contains(range.startContainer) ||
182 !root.contains(range.endContainer) ||
183 excludedAncestor(range.startContainer) !== null ||
184 excludedAncestor(range.endContainer) !== null
185 ) {
186 return false;
187 }
188 // Interior controls are allowed; readAssistantText omits them from the stream.
189 return true;
190}
191
192function isUnselectable(text: Text): boolean {
193 const element = text.parentElement;
194 const view = element?.ownerDocument?.defaultView;
195 return !!element && !!view && view.getComputedStyle(element).userSelect === "none";
196}
197
198function selectedTextBoundary(range: Range, node: Node, last: boolean): Text | null {
199 if (!range.intersectsNode(node)) return null;
200 if (node.nodeType === 3) {
201 const text = node as Text;
202 const start = node === range.startContainer ? range.startOffset : 0;
203 const end = node === range.endContainer ? range.endOffset : text.length;
204 // Triple-click ranges run through unselectable rows (tool calls, timestamps,
205 // author headings) into the next response; that text is not selected.
206 return start < end && !isUnselectable(text) ? text : null;
207 }
208 for (
209 let child = last ? node.lastChild : node.firstChild;
210 child !== null;
211 child = last ? child.previousSibling : child.nextSibling
212 ) {
213 const boundary = selectedTextBoundary(range, child, last);
214 if (boundary !== null) return boundary;
215 }
216 return null;
217}
218
219/** Captures the ordered native range, including selections dragged backwards. */
220export function captureAssistantTextSelection(
221 viewport: HTMLElement,
222 selection: Selection | null,
223): { source: HTMLElement; selector: AssistantTextSelector; range: Range } | null {
224 if (selection === null || selection.isCollapsed || selection.rangeCount !== 1) return null;
225 const range = selection.getRangeAt(0).cloneRange();
226 const first = selectedTextBoundary(range, range.commonAncestorContainer, false);
227 const last = selectedTextBoundary(range, range.commonAncestorContainer, true);
228 if (first === null || last === null) return null;
229 const source = first.parentElement?.closest<HTMLElement>("[data-assistant-citation-source]");
230 if (!source || !viewport.contains(source)) return null;
231
232 // Paragraph selection can end at the next block's offset 0 or a parent
233 // boundary. Validate the text actually selected, not that empty endpoint.
234 range.setStart(first, first === range.startContainer ? range.startOffset : 0);
235 range.setEnd(last, last === range.endContainer ? range.endOffset : last.length);
236 if (!isUsableRange(source, range)) return null;
237
238 const stream = readAssistantText(source);
239 let rawStart: number | null = null;
240 let rawEnd = 0;
241 for (const chunk of stream.chunks) {
242 if (!range.intersectsNode(chunk.node)) continue;
243 const start = range.startContainer === chunk.node ? range.startOffset : 0;
244 const end = range.endContainer === chunk.node ? range.endOffset : chunk.node.length;
245 if (start === end) continue;
246 rawStart ??= chunk.start + start;
247 rawEnd = chunk.start + end;
248 }
249 if (rawStart === null) return null;
250 const selector = createAssistantTextSelector(stream.text, rawStart, rawEnd);
251 return selector === null ? null : { source, selector, range };
252}
253
254function rawTextOffset(text: string, normalizedOffset: number): number {
255 let offset = 0;
256 for (const match of text.matchAll(/\s+|\S+/g)) {
257 const whitespace = /\s/.test(match[0][0]!);
258 const length = whitespace ? 1 : match[0].length;
259 if (normalizedOffset <= offset + length) {
260 return (
261 match.index +
262 (whitespace && normalizedOffset > offset ? match[0].length : normalizedOffset - offset)
263 );
264 }
265 offset += length;
266 }
267 return text.length;
268}
269
270/** Resolves against the current DOM without changing the user's selection. */
271export function resolveAssistantCitationRange(
272 root: HTMLElement,
273 selector: AssistantTextSelector,
274): Range | null {
275 if (excludedAncestor(root) !== null) return null;
276 const stream = readAssistantText(root);
277 const match = findAssistantCitationText(stream.text, selector);
278 if (match === null) return null;
279
280 const start = rawTextOffset(stream.text, match.start);
281 const end = rawTextOffset(stream.text, match.end);
282 const first = stream.chunks.find((chunk) => chunk.end > start);
283 const last = stream.chunks.findLast((chunk) => chunk.start < end);
284 if (first === undefined || last === undefined) return null;
285
286 const range = root.ownerDocument.createRange();
287 range.setStart(first.node, Math.max(0, start - first.start));
288 range.setEnd(last.node, Math.min(last.node.length, end - last.start));
289 return isUsableRange(root, range) ? range : null;
290}