1use serde::Serialize;
2
3#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
4#[serde(rename_all = "snake_case")]
5pub enum VisualTokenKind {
6 ComponentOpen,
7 ComponentClose,
8 ComponentBody,
9 PatchOpen,
10 PatchClose,
11 Boundary,
12 ScratchComment,
13 ScratchCommentBody,
14 Prompt,
15 ResponseHeading,
16 TrackedId,
17 LabelTag,
18}
19
20#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
21pub struct VisualToken {
22 pub kind: VisualTokenKind,
23 pub start: usize,
24 pub end: usize,
25}
26
27pub fn collect_visual_tokens(doc: &str) -> Vec<VisualToken> {
28 let code_ranges = crate::component::find_code_ranges(doc);
29 let mut tokens = Vec::new();
30
31 collect_comment_tokens(doc, &code_ranges, &mut tokens);
32 collect_component_body_tokens(doc, &code_ranges, &mut tokens);
33 collect_scratch_comment_body_tokens(doc, &code_ranges, &mut tokens);
34 collect_line_tokens(doc, &code_ranges, &mut tokens);
35 collect_label_tags(doc, &code_ranges, &mut tokens);
36 collect_tracked_ids(doc, &code_ranges, &mut tokens);
37
38 tokens.sort_by_key(|token| (token.start, token.end));
39 tokens
40}
41
42enum CommentKind<'a> {
43 ComponentOpen(&'a str),
44 ComponentClose(&'a str),
45 PatchOpen,
46 PatchClose,
47 Boundary,
48 ScratchComment,
49}
50
51fn classify_comment_inner(inner: &str) -> CommentKind<'_> {
52 if inner.starts_with("agent:boundary:") {
53 CommentKind::Boundary
54 } else if let Some(rest) = inner.strip_prefix("agent:") {
55 let name = rest.split_whitespace().next().unwrap_or_default();
56 if name.is_empty() {
57 CommentKind::ScratchComment
58 } else {
59 CommentKind::ComponentOpen(name)
60 }
61 } else if let Some(rest) = inner.strip_prefix("/agent:") {
62 let name = rest.split_whitespace().next().unwrap_or_default();
63 if name.is_empty() {
64 CommentKind::ScratchComment
65 } else {
66 CommentKind::ComponentClose(name)
67 }
68 } else if inner.starts_with("patch:") {
69 CommentKind::PatchOpen
70 } else if inner.starts_with("/patch:") {
71 CommentKind::PatchClose
72 } else {
73 CommentKind::ScratchComment
74 }
75}
76
77fn collect_comment_tokens(doc: &str, code_ranges: &[(usize, usize)], out: &mut Vec<VisualToken>) {
78 let mut search_from = 0usize;
79 while let Some(rel) = doc[search_from..].find("<!--") {
80 let start = search_from + rel;
81 let Some(close_rel) = doc[start + 4..].find("-->") else {
82 break;
83 };
84 let end = start + 4 + close_rel + 3;
85 search_from = end;
86
87 if overlaps_code(start, end, code_ranges) {
88 continue;
89 }
90
91 let inner = doc[start + 4..end - 3].trim();
92 let kind = match classify_comment_inner(inner) {
93 CommentKind::ComponentOpen(_) => VisualTokenKind::ComponentOpen,
94 CommentKind::ComponentClose(_) => VisualTokenKind::ComponentClose,
95 CommentKind::PatchOpen => VisualTokenKind::PatchOpen,
96 CommentKind::PatchClose => VisualTokenKind::PatchClose,
97 CommentKind::Boundary => VisualTokenKind::Boundary,
98 CommentKind::ScratchComment => VisualTokenKind::ScratchComment,
99 };
100 out.push(VisualToken { kind, start, end });
101 }
102}
103
104fn collect_component_body_tokens(
105 doc: &str,
106 code_ranges: &[(usize, usize)],
107 out: &mut Vec<VisualToken>,
108) {
109 let mut stack: Vec<(String, usize)> = Vec::new();
110 let mut search_from = 0usize;
111 while let Some(rel) = doc[search_from..].find("<!--") {
112 let start = search_from + rel;
113 let Some(close_rel) = doc[start + 4..].find("-->") else {
114 break;
115 };
116 let end = start + 4 + close_rel + 3;
117 search_from = end;
118
119 if overlaps_code(start, end, code_ranges) {
120 continue;
121 }
122
123 let inner = doc[start + 4..end - 3].trim();
124 match classify_comment_inner(inner) {
125 CommentKind::ComponentOpen(name) => {
126 stack.push((name.to_string(), end));
127 }
128 CommentKind::ComponentClose(name) => {
129 if let Some(idx) = stack.iter().rposition(|(open_name, _)| open_name == name) {
130 let (_, body_start) = stack.remove(idx);
131 if body_start < start && !doc[body_start..start].trim().is_empty() {
132 out.push(VisualToken {
133 kind: VisualTokenKind::ComponentBody,
134 start: body_start,
135 end: start,
136 });
137 }
138 }
139 }
140 CommentKind::PatchOpen
141 | CommentKind::PatchClose
142 | CommentKind::Boundary
143 | CommentKind::ScratchComment => {}
144 }
145 }
146}
147
148fn collect_scratch_comment_body_tokens(
149 doc: &str,
150 code_ranges: &[(usize, usize)],
151 out: &mut Vec<VisualToken>,
152) {
153 let mut search_from = 0usize;
154 while let Some(rel) = doc[search_from..].find("<!--") {
155 let start = search_from + rel;
156 let Some(close_rel) = doc[start + 4..].find("-->") else {
157 break;
158 };
159 let end = start + 4 + close_rel + 3;
160 search_from = end;
161
162 if overlaps_code(start, end, code_ranges) {
163 continue;
164 }
165
166 let inner = doc[start + 4..end - 3].trim();
167 if !matches!(classify_comment_inner(inner), CommentKind::ScratchComment) {
168 continue;
169 }
170
171 let body_start = start + 4;
172 let body_end = end - 3;
173 if body_start < body_end && !doc[body_start..body_end].trim().is_empty() {
174 out.push(VisualToken {
175 kind: VisualTokenKind::ScratchCommentBody,
176 start: body_start,
177 end: body_end,
178 });
179 }
180 }
181}
182
183fn collect_line_tokens(doc: &str, code_ranges: &[(usize, usize)], out: &mut Vec<VisualToken>) {
184 let mut offset = 0usize;
185 for chunk in doc.split_inclusive('\n') {
186 let line = chunk.strip_suffix('\n').unwrap_or(chunk);
187 let line_end = offset + line.len();
188 if !overlaps_code(offset, line_end, code_ranges) {
189 let leading_ws = line.len() - line.trim_start_matches([' ', '\t']).len();
190 let trimmed = &line[leading_ws..];
191 if trimmed.starts_with("❯ ") {
192 out.push(VisualToken {
193 kind: VisualTokenKind::Prompt,
194 start: offset + leading_ws,
195 end: line_end,
196 });
197 }
198 if trimmed.starts_with("### Re:") {
199 out.push(VisualToken {
200 kind: VisualTokenKind::ResponseHeading,
201 start: offset + leading_ws,
202 end: line_end,
203 });
204 }
205 }
206 offset += chunk.len();
207 }
208}
209
210fn collect_tracked_ids(doc: &str, code_ranges: &[(usize, usize)], out: &mut Vec<VisualToken>) {
211 let mut search_from = 0usize;
212 while let Some(rel) = doc[search_from..].find("[#") {
213 let start = search_from + rel;
214 let rest = &doc[start + 2..];
215 let Some(close_rel) = rest.find(']') else {
216 break;
217 };
218 let end = start + 2 + close_rel + 1;
219 search_from = end;
220
221 if overlaps_code(start, end, code_ranges) {
222 continue;
223 }
224
225 let inner = &rest[..close_rel];
226 if !inner.is_empty()
227 && inner
228 .chars()
229 .all(|ch| ch.is_ascii_alphanumeric() || ch == '-' || ch == '_')
230 {
231 out.push(VisualToken {
232 kind: VisualTokenKind::TrackedId,
233 start,
234 end,
235 });
236 }
237 }
238}
239
240fn collect_label_tags(doc: &str, code_ranges: &[(usize, usize)], out: &mut Vec<VisualToken>) {
241 let mut search_from = 0usize;
242 while let Some(rel) = doc[search_from..].find('[') {
243 let start = search_from + rel;
244 let rest = &doc[start + 1..];
245 let Some(close_rel) = rest.find(']') else {
246 break;
247 };
248 let end = start + 1 + close_rel + 1;
249 search_from = end;
250
251 if overlaps_code(start, end, code_ranges) {
252 continue;
253 }
254 if start > 0 && doc.as_bytes()[start - 1] == b'!' {
255 continue;
256 }
257 if doc[end..].starts_with('(') {
258 continue;
259 }
260
261 let inner = &rest[..close_rel];
262 if inner.starts_with('#') || inner == "x" {
263 continue;
264 }
265 if inner.len() <= 1 {
266 continue;
267 }
268 if !inner
269 .chars()
270 .all(|ch| ch.is_ascii_lowercase() || ch.is_ascii_digit() || ch == '-' || ch == '_')
271 {
272 continue;
273 }
274 if !inner.chars().any(|ch| ch.is_ascii_alphabetic()) {
275 continue;
276 }
277
278 out.push(VisualToken {
279 kind: VisualTokenKind::LabelTag,
280 start,
281 end,
282 });
283 }
284}
285
286fn overlaps_code(start: usize, end: usize, code_ranges: &[(usize, usize)]) -> bool {
287 code_ranges
288 .iter()
289 .any(|&(code_start, code_end)| start < code_end && end > code_start)
290}
291
292#[cfg(test)]
293mod tests {
294 use super::{VisualTokenKind, collect_visual_tokens};
295
296 #[test]
297 fn collects_agent_and_scratch_comments() {
298 let doc = "\
299<!-- agent:exchange patch=append -->
300hello
301<!-- /agent:exchange -->
302<!-- patch:exchange -->
303<!-- /patch:exchange -->
304<!-- agent:boundary:abc12345 -->
305<!-- scratch note -->
306";
307 let tokens = collect_visual_tokens(doc);
308 let kinds: Vec<VisualTokenKind> = tokens.iter().map(|token| token.kind).collect();
309 assert_eq!(
310 kinds,
311 vec![
312 VisualTokenKind::ComponentOpen,
313 VisualTokenKind::ComponentBody,
314 VisualTokenKind::ComponentClose,
315 VisualTokenKind::PatchOpen,
316 VisualTokenKind::PatchClose,
317 VisualTokenKind::Boundary,
318 VisualTokenKind::ScratchComment,
319 VisualTokenKind::ScratchCommentBody,
320 ]
321 );
322 }
323
324 #[test]
325 fn ignores_visual_markers_inside_fenced_code() {
326 let doc = "\
327```md
328<!-- agent:exchange -->
329❯ hidden prompt
330### Re: hidden
331[#hide]
332```
333";
334 assert!(collect_visual_tokens(doc).is_empty());
335 }
336
337 #[test]
338 fn collects_prompts_response_headings_and_tracked_ids() {
339 let doc = "\
340❯ do #qey0. spec-test-build-install-commit-push
341### Re: #qey0 editor highlighting — gpt-5
342- [ ] [#qey0] Track this work item
343- [ ] [#multi-user] Alias-style id should still highlight
344- [ ] [recommended] Treat labels as metadata, not broken markdown refs
345";
346 let tokens = collect_visual_tokens(doc);
347 assert!(
348 tokens
349 .iter()
350 .any(|token| token.kind == VisualTokenKind::Prompt)
351 );
352 assert!(
353 tokens
354 .iter()
355 .any(|token| token.kind == VisualTokenKind::ResponseHeading)
356 );
357 assert_eq!(
358 tokens
359 .iter()
360 .filter(|token| token.kind == VisualTokenKind::TrackedId)
361 .count(),
362 2
363 );
364 assert!(
365 tokens
366 .iter()
367 .any(|token| token.kind == VisualTokenKind::LabelTag)
368 );
369 }
370
371 #[test]
372 fn ignores_image_alt_text_and_checklists_for_label_tags() {
373 let doc = "\
374- [ ] keep checklist
375
376- [recommended] keep label
377";
378 let tokens = collect_visual_tokens(doc);
379 let label_tags: Vec<&super::VisualToken> = tokens
380 .iter()
381 .filter(|token| token.kind == VisualTokenKind::LabelTag)
382 .collect();
383 assert_eq!(label_tags.len(), 1);
384 assert_eq!(
385 &doc[label_tags[0].start..label_tags[0].end],
386 "[recommended]"
387 );
388 }
389
390 #[test]
391 fn collects_multiline_scratch_comment_body() {
392 let doc = "\
393<!--
394scratch note
395```md
396
397```
398-->
399";
400 let tokens = collect_visual_tokens(doc);
401 let body = tokens
402 .iter()
403 .find(|token| token.kind == VisualTokenKind::ScratchCommentBody)
404 .expect("expected scratch comment body token");
405 let body_text = &doc[body.start..body.end];
406 assert!(body_text.contains("scratch note"));
407 assert!(body_text.contains(""));
408 }
409}