1import { ASSISTANT_CITATION_CONTEXT_LENGTH, type AssistantCitation } from "@t3tools/contracts";
2
3import { mathWrapperOf } from "../markdown-clipboard";
4
5export type AssistantTextSelector = {
6 readonly text: string;
7 readonly start: number;
8 readonly end: number;
9 readonly prefix: string;
10 readonly suffix: string;
11};
12
13export type AssistantCitationSourceAnchor = {
14 source: HTMLElement;
15 range: Range;
16 viewport: HTMLElement;
17};
18
19export function findAssistantCitationSourceAnchor(
20 document: Document,
21 citation: AssistantCitation,
22): AssistantCitationSourceAnchor | null {
23 const source = [
24 ...document.querySelectorAll<HTMLElement>("[data-assistant-citation-source]"),
25 ].find(
26 (element) =>
27 element.dataset.assistantCitationSource === citation.messageId &&
28 element.dataset.assistantCitationEnvironment === citation.environmentId &&
29 element.dataset.assistantCitationThread === citation.threadId,
30 );
31 const viewport = source?.closest<HTMLElement>("[data-assistant-citation-viewport]");
32 if (!source || !viewport) return null;
33 const range = resolveAssistantCitationRange(source, citation);
34 return range ? { source, range, viewport } : null;
35}
36
37const CONTROL_SELECTOR = "button, input, textarea, select, [role=button], [contenteditable]";
38const EXCLUDED_SELECTOR = `${CONTROL_SELECTOR}, [hidden], [aria-hidden=true], script, style, template, noscript, svg`;
39const BLOCK_SELECTOR =
40 "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";
41
42function normalizeWhitespace(text: string): string {
43 return text.replace(/\s+/g, " ");
44}
45
46function splitsSurrogatePair(text: string, offset: number): boolean {
47 const before = text.charCodeAt(offset - 1);
48 const after = text.charCodeAt(offset);
49 return before >= 0xd800 && before <= 0xdbff && after >= 0xdc00 && after <= 0xdfff;
50}
51
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 if (rawStart > 0 && /\s/.test(text[rawStart - 1]!) && /\s/.test(text[rawStart]!)) {
63 start -= 1;
64 }
65 const end = normalizeWhitespace(text.slice(0, rawEnd)).length;
66 let prefixStart = Math.max(0, start - ASSISTANT_CITATION_CONTEXT_LENGTH);
67 let suffixEnd = Math.min(normalized.length, end + ASSISTANT_CITATION_CONTEXT_LENGTH);
68 // A split pair becomes a replacement character when the context enters a URL.
69 if (splitsSurrogatePair(normalized, prefixStart)) prefixStart += 1;
70 if (splitsSurrogatePair(normalized, suffixEnd)) suffixEnd -= 1;
71 return {
72 text: quote,
73 start,
74 end,
75 prefix: normalized.slice(prefixStart, start),
76 suffix: normalized.slice(end, suffixEnd),
77 };
78}
79
80/** Offsets use normalized UTF-16 units. Saved offsets may drift, so they cannot break context ties. */
81export function findAssistantCitationText(
82 text: string,
83 selector: AssistantTextSelector,
84): { start: number; end: number } | null {
85 const normalized = normalizeWhitespace(text);
86 const quote = normalizeWhitespace(selector.text);
87 if (quote.trim().length === 0) return null;
88
89 const prefix = normalizeWhitespace(selector.prefix);
90 const suffix = normalizeWhitespace(selector.suffix);
91 const matchesContext = (start: number, end: number) =>
92 normalized.slice(Math.max(0, start - prefix.length), start) === prefix &&
93 normalized.slice(end, end + suffix.length) === suffix;
94
95 let match =
96 Number.isSafeInteger(selector.start) &&
97 Number.isSafeInteger(selector.end) &&
98 selector.start >= 0 &&
99 selector.end - selector.start === quote.length &&
100 normalized.slice(selector.start, selector.end) === quote &&
101 matchesContext(selector.start, selector.end)
102 ? { start: selector.start, end: selector.end }
103 : null;
104 let onlyQuote: { start: number; end: number } | null = null;
105 let quoteCount = 0;
106
107 for (
108 let start = normalized.indexOf(quote);
109 start !== -1;
110 start = normalized.indexOf(quote, start + 1)
111 ) {
112 const end = start + quote.length;
113 quoteCount += 1;
114 onlyQuote = { start, end };
115 if (!matchesContext(start, end)) continue;
116 if (match !== null && match.start !== start) return null;
117 match = { start, end };
118 }
119
120 return match ?? (quoteCount === 1 ? onlyQuote : null);
121}
122
123/** A text node, or a whole rendered formula standing for its TeX. */
124type TextChunk = { node: Text | Element; start: number; end: number };
125
126const isText = (node: Node): node is Text => node.nodeType === 3;
127
128function readAssistantText(root: HTMLElement) {
129 const parts: string[] = [];
130 const chunks: TextChunk[] = [];
131 let length = 0;
132 let separator = false;
133
134 const append = (node: Text | Element, text: string) => {
135 if (separator && length > 0) {
136 parts.push("\n");
137 length += 1;
138 }
139 separator = false;
140 chunks.push({ node, start: length, end: length + text.length });
141 parts.push(text);
142 length += text.length;
143 };
144
145 const visit = (node: Node) => {
146 if (node.nodeType === 3) {
147 const text = node as Text;
148 if (text.length > 0) append(text, text.data);
149 return;
150 }
151 if (node.nodeType !== 1) return;
152 const element = node as Element;
153 if (element.matches(EXCLUDED_SELECTOR)) return;
154 const block = element.matches(BLOCK_SELECTOR);
155 if (block || element.tagName === "BR") separator = true;
156 if (element.matches("[data-markdown-math]")) {
157 append(element, element.getAttribute("data-markdown-copy")?.trim() ?? "");
158 } else {
159 for (const child of element.childNodes) visit(child);
160 }
161 if (block) separator = true;
162 };
163
164 visit(root);
165 return { text: parts.join(""), chunks };
166}
167
168function excludedAncestor(node: Node): Element | null {
169 const element = node.nodeType === 1 ? (node as Element) : node.parentElement;
170 return element?.closest(EXCLUDED_SELECTOR) ?? null;
171}
172
173function isUsableRange(root: HTMLElement, range: Range): boolean {
174 if (
175 range.collapsed ||
176 !root.contains(range.startContainer) ||
177 !root.contains(range.endContainer) ||
178 excludedAncestor(range.startContainer) !== null ||
179 excludedAncestor(range.endContainer) !== null
180 ) {
181 return false;
182 }
183 return true;
184}
185
186function isUnselectable(text: Text): boolean {
187 const element = mathWrapperOf(text.parentElement) ?? text.parentElement;
188 const view = element?.ownerDocument?.defaultView;
189 return !!element && !!view && view.getComputedStyle(element).userSelect === "none";
190}
191
192function selectedTextBoundary(range: Range, node: Node, last: boolean): Text | null {
193 if (!range.intersectsNode(node)) return null;
194 if (node.nodeType === 3) {
195 const text = node as Text;
196 const start = node === range.startContainer ? range.startOffset : 0;
197 const end = node === range.endContainer ? range.endOffset : text.length;
198 // Triple-click ranges include unselected text from unselectable rows and the next response.
199 return start < end && !isUnselectable(text) ? text : null;
200 }
201 for (
202 let child = last ? node.lastChild : node.firstChild;
203 child !== null;
204 child = last ? child.previousSibling : child.nextSibling
205 ) {
206 const boundary = selectedTextBoundary(range, child, last);
207 if (boundary !== null) return boundary;
208 }
209 return null;
210}
211
212export function captureAssistantTextSelection(
213 viewport: HTMLElement,
214 selection: Selection | null,
215): { source: HTMLElement; selector: AssistantTextSelector; range: Range } | null {
216 if (selection === null || selection.isCollapsed || selection.rangeCount !== 1) return null;
217 const range = selection.getRangeAt(0).cloneRange();
218 const first = selectedTextBoundary(range, range.commonAncestorContainer, false);
219 const last = selectedTextBoundary(range, range.commonAncestorContainer, true);
220 if (first === null || last === null) return null;
221 const source = first.parentElement?.closest<HTMLElement>("[data-assistant-citation-source]");
222 if (!source || !viewport.contains(source)) return null;
223
224 // Paragraph selection can end at the next block's offset 0 or a parent
225 // boundary. Validate the text actually selected, not that empty endpoint.
226 // An endpoint inside a formula's glyphs takes in the whole formula.
227 const startFormula = mathWrapperOf(first.parentElement);
228 const endFormula = mathWrapperOf(last.parentElement);
229 if (startFormula) range.setStartBefore(startFormula);
230 else range.setStart(first, first === range.startContainer ? range.startOffset : 0);
231 if (endFormula) range.setEndAfter(endFormula);
232 else range.setEnd(last, last === range.endContainer ? range.endOffset : last.length);
233 if (!isUsableRange(source, range)) return null;
234
235 const stream = readAssistantText(source);
236 let rawStart: number | null = null;
237 let rawEnd = 0;
238 for (const chunk of stream.chunks) {
239 if (!range.intersectsNode(chunk.node)) continue;
240 const start = range.startContainer === chunk.node ? range.startOffset : 0;
241 const end = range.endContainer === chunk.node ? range.endOffset : chunk.end - chunk.start;
242 if (start === end) continue;
243 rawStart ??= chunk.start + start;
244 rawEnd = chunk.start + end;
245 }
246 if (rawStart === null) return null;
247 const selector = createAssistantTextSelector(stream.text, rawStart, rawEnd);
248 return selector === null ? null : { source, selector, range };
249}
250
251/** How much of a response a reply carries: enough to recognise it, not to resend it. */
252const RESPONSE_OPENING_LENGTH = 300;
253
254/** The opening of a response. Falls back to the raw text when the source is virtualized away or
255 * renders no text, as in an image-only response. */
256export function captureAssistantMessageOpening(
257 source: HTMLElement | null,
258 fallbackText: string,
259): AssistantTextSelector | null {
260 const rendered = source ? readAssistantText(source).text : "";
261 const text = rendered.trim() ? rendered : fallbackText;
262 const rawStart = text.length - text.trimStart().length;
263 let rawEnd = Math.min(text.trimEnd().length, rawStart + RESPONSE_OPENING_LENGTH);
264 if (splitsSurrogatePair(text, rawEnd)) rawEnd -= 1;
265 return createAssistantTextSelector(text, rawStart, rawEnd);
266}
267
268function rawTextOffset(text: string, normalizedOffset: number): number {
269 let offset = 0;
270 for (const match of text.matchAll(/\s+|\S+/g)) {
271 const whitespace = /\s/.test(match[0][0]!);
272 const length = whitespace ? 1 : match[0].length;
273 if (normalizedOffset <= offset + length) {
274 return (
275 match.index +
276 (whitespace && normalizedOffset > offset ? match[0].length : normalizedOffset - offset)
277 );
278 }
279 offset += length;
280 }
281 return text.length;
282}
283
284export function resolveAssistantCitationRange(
285 root: HTMLElement,
286 selector: AssistantTextSelector,
287): Range | null {
288 if (excludedAncestor(root) !== null) return null;
289 const stream = readAssistantText(root);
290 const match = findAssistantCitationText(stream.text, selector);
291 if (match === null) return null;
292
293 const start = rawTextOffset(stream.text, match.start);
294 const end = rawTextOffset(stream.text, match.end);
295 const first = stream.chunks.find((chunk) => chunk.end > start);
296 const last = stream.chunks.findLast((chunk) => chunk.start < end);
297 if (first === undefined || last === undefined) return null;
298
299 const range = root.ownerDocument.createRange();
300 // A match that touches a formula covers the formula as rendered.
301 if (isText(first.node)) range.setStart(first.node, Math.max(0, start - first.start));
302 else range.setStartBefore(first.node);
303 if (isText(last.node)) range.setEnd(last.node, Math.min(last.node.length, end - last.start));
304 else range.setEndAfter(last.node);
305 return isUsableRange(root, range) ? range : null;
306}