1use crate::udiff::DiffLine;
2use anyhow::{Context as _, Result};
3use serde::{Deserialize, Serialize};
4use std::{borrow::Cow, fmt::Write as _, mem, path::Path, sync::Arc};
5use telemetry_events::EditPredictionRating;
6
7pub const CURSOR_POSITION_MARKER: &str = "[CURSOR_POSITION]";
8pub const INLINE_CURSOR_MARKER: &str = "<|user_cursor|>";
9
10/// Maximum cursor file size to capture (64KB).
11/// Files larger than this will not have their content captured,
12/// falling back to git-based loading.
13pub const MAX_CURSOR_FILE_SIZE: usize = 64 * 1024;
14
15/// Encodes a cursor position into a diff patch by adding a comment line with a caret
16/// pointing to the cursor column.
17///
18/// The cursor offset is relative to the start of the new text content (additions and context lines).
19/// Returns the patch with cursor marker comment lines inserted after the relevant addition line.
20pub fn encode_cursor_in_patch(patch: &str, cursor_offset: Option<usize>) -> String {
21 let Some(cursor_offset) = cursor_offset else {
22 return patch.to_string();
23 };
24
25 let mut result = String::new();
26 let mut line_start_offset = 0usize;
27
28 for line in patch.lines() {
29 if matches!(
30 DiffLine::parse(line),
31 DiffLine::Garbage(content)
32 if content.starts_with('#') && content.contains(CURSOR_POSITION_MARKER)
33 ) {
34 continue;
35 }
36
37 if !result.is_empty() {
38 result.push('\n');
39 }
40 result.push_str(line);
41
42 match DiffLine::parse(line) {
43 DiffLine::Addition(content) => {
44 let line_end_offset = line_start_offset + content.len();
45
46 if cursor_offset >= line_start_offset && cursor_offset <= line_end_offset {
47 let cursor_column = cursor_offset - line_start_offset;
48
49 result.push('\n');
50 result.push('#');
51 for _ in 0..cursor_column {
52 result.push(' ');
53 }
54 write!(result, "^{}", CURSOR_POSITION_MARKER).unwrap();
55 }
56
57 line_start_offset = line_end_offset + 1;
58 }
59 DiffLine::Context(content) => {
60 line_start_offset += content.len() + 1;
61 }
62 _ => {}
63 }
64 }
65
66 if patch.ends_with('\n') {
67 result.push('\n');
68 }
69
70 result
71}
72
73#[derive(Clone, Debug, PartialEq, Hash, Serialize, Deserialize)]
74pub struct ExampleSpec {
75 #[serde(default)]
76 pub name: String,
77 pub repository_url: String,
78 pub revision: String,
79 #[serde(default, skip_serializing_if = "Vec::is_empty")]
80 pub tags: Vec<String>,
81 #[serde(default, skip_serializing_if = "Option::is_none")]
82 pub reasoning: Option<String>,
83 #[serde(default)]
84 pub uncommitted_diff: String,
85 pub cursor_path: Arc<Path>,
86 pub cursor_position: String,
87 pub edit_history: String,
88 pub expected_patches: Vec<String>,
89 #[serde(default, skip_serializing_if = "Option::is_none")]
90 pub rejected_patch: Option<String>,
91 #[serde(default, skip_serializing_if = "Option::is_none")]
92 pub telemetry: Option<TelemetrySource>,
93 #[serde(default, skip_serializing_if = "Vec::is_empty")]
94 pub human_feedback: Vec<HumanFeedback>,
95 #[serde(default, skip_serializing_if = "Option::is_none")]
96 pub rating: Option<EditPredictionRating>,
97}
98
99#[derive(Clone, Debug, PartialEq, Hash, Serialize, Deserialize)]
100pub struct HumanFeedback {
101 pub message: String,
102}
103
104/// Metadata for examples sourced from production telemetry (rejected predictions).
105#[derive(Clone, Debug, PartialEq, Hash, Serialize, Deserialize)]
106pub struct TelemetrySource {
107 pub request_id: String,
108 pub device_id: String,
109 pub time: String,
110 pub rejection_reason: String,
111 pub was_shown: bool,
112}
113
114const REASONING_HEADING: &str = "Reasoning";
115const UNCOMMITTED_DIFF_HEADING: &str = "Uncommitted Diff";
116const EDIT_HISTORY_HEADING: &str = "Edit History";
117const CURSOR_POSITION_HEADING: &str = "Cursor Position";
118const EXPECTED_PATCH_HEADING: &str = "Expected Patch";
119const REJECTED_PATCH_HEADING: &str = "Rejected Patch";
120const ACCEPTED_PREDICTION_MARKER: &str = "// User accepted prediction:";
121
122#[derive(Serialize, Deserialize)]
123struct FrontMatter<'a> {
124 repository_url: Cow<'a, str>,
125 revision: Cow<'a, str>,
126 #[serde(default, skip_serializing_if = "Vec::is_empty")]
127 tags: Vec<String>,
128}
129
130impl ExampleSpec {
131 /// Generate a sanitized filename for this example.
132 pub fn filename(&self) -> String {
133 self.name
134 .chars()
135 .map(|c| match c {
136 ' ' | ':' | '~' | '^' | '?' | '*' | '[' | '\\' | '@' | '{' | '/' | '<' | '>'
137 | '|' | '"' => '-',
138 c => c,
139 })
140 .collect()
141 }
142
143 /// Format this example spec as markdown.
144 pub fn to_markdown(&self) -> String {
145 use std::fmt::Write as _;
146
147 let front_matter = FrontMatter {
148 repository_url: Cow::Borrowed(&self.repository_url),
149 revision: Cow::Borrowed(&self.revision),
150 tags: self.tags.clone(),
151 };
152 let front_matter_toml =
153 toml::to_string_pretty(&front_matter).unwrap_or_else(|_| String::new());
154
155 let mut markdown = String::new();
156
157 _ = writeln!(markdown, "+++");
158 markdown.push_str(&front_matter_toml);
159 if !markdown.ends_with('\n') {
160 markdown.push('\n');
161 }
162 _ = writeln!(markdown, "+++");
163 markdown.push('\n');
164
165 _ = writeln!(markdown, "# {}", self.name);
166 markdown.push('\n');
167
168 if let Some(reasoning) = &self.reasoning {
169 _ = writeln!(markdown, "## {}", REASONING_HEADING);
170 markdown.push('\n');
171 markdown.push_str(reasoning);
172 if !markdown.ends_with('\n') {
173 markdown.push('\n');
174 }
175 markdown.push('\n');
176 }
177
178 if !self.uncommitted_diff.is_empty() {
179 _ = writeln!(markdown, "## {}", UNCOMMITTED_DIFF_HEADING);
180 _ = writeln!(markdown);
181 _ = writeln!(markdown, "```diff");
182 markdown.push_str(&self.uncommitted_diff);
183 if !markdown.ends_with('\n') {
184 markdown.push('\n');
185 }
186 _ = writeln!(markdown, "```");
187 markdown.push('\n');
188 }
189
190 _ = writeln!(markdown, "## {}", EDIT_HISTORY_HEADING);
191 _ = writeln!(markdown);
192
193 if self.edit_history.is_empty() {
194 _ = writeln!(markdown, "(No edit history)");
195 _ = writeln!(markdown);
196 } else {
197 _ = writeln!(markdown, "```diff");
198 markdown.push_str(&self.edit_history);
199 if !markdown.ends_with('\n') {
200 markdown.push('\n');
201 }
202 _ = writeln!(markdown, "```");
203 markdown.push('\n');
204 }
205
206 _ = writeln!(markdown, "## {}", CURSOR_POSITION_HEADING);
207 _ = writeln!(markdown);
208 _ = writeln!(markdown, "```{}", self.cursor_path.to_string_lossy());
209 markdown.push_str(&self.cursor_position);
210 if !markdown.ends_with('\n') {
211 markdown.push('\n');
212 }
213 _ = writeln!(markdown, "```");
214 markdown.push('\n');
215
216 _ = writeln!(markdown, "## {}", EXPECTED_PATCH_HEADING);
217 markdown.push('\n');
218 for patch in &self.expected_patches {
219 _ = writeln!(markdown, "```diff");
220 markdown.push_str(patch);
221 if !markdown.ends_with('\n') {
222 markdown.push('\n');
223 }
224 _ = writeln!(markdown, "```");
225 markdown.push('\n');
226 }
227
228 if let Some(rejected_patch) = &self.rejected_patch {
229 _ = writeln!(markdown, "## {}", REJECTED_PATCH_HEADING);
230 markdown.push('\n');
231 _ = writeln!(markdown, "```diff");
232 markdown.push_str(rejected_patch);
233 if !markdown.ends_with('\n') {
234 markdown.push('\n');
235 }
236 _ = writeln!(markdown, "```");
237 markdown.push('\n');
238 }
239
240 markdown
241 }
242
243 /// Parse an example spec from markdown.
244 pub fn from_markdown(mut input: &str) -> anyhow::Result<Self> {
245 use pulldown_cmark::{CodeBlockKind, CowStr, Event, HeadingLevel, Parser, Tag, TagEnd};
246
247 let mut spec = ExampleSpec {
248 name: String::new(),
249 repository_url: String::new(),
250 revision: String::new(),
251 tags: Vec::new(),
252 reasoning: None,
253 uncommitted_diff: String::new(),
254 cursor_path: Path::new("").into(),
255 cursor_position: String::new(),
256 edit_history: String::new(),
257 expected_patches: Vec::new(),
258 rejected_patch: None,
259 telemetry: None,
260 human_feedback: Vec::new(),
261 rating: None,
262 };
263
264 if let Some(rest) = input.strip_prefix("+++\n")
265 && let Some((front_matter, rest)) = rest.split_once("+++\n")
266 {
267 if let Ok(data) = toml::from_str::<FrontMatter<'_>>(front_matter) {
268 spec.repository_url = data.repository_url.into_owned();
269 spec.revision = data.revision.into_owned();
270 spec.tags = data.tags;
271 }
272 input = rest.trim_start();
273 }
274
275 let parser = Parser::new(input);
276 let mut text = String::new();
277 let mut block_info: CowStr = "".into();
278
279 #[derive(PartialEq)]
280 enum Section {
281 Start,
282 UncommittedDiff,
283 EditHistory,
284 CursorPosition,
285 ExpectedPatch,
286 RejectedPatch,
287 Other,
288 }
289
290 let mut current_section = Section::Start;
291 let mut next_edit_predicted = false;
292
293 for event in parser {
294 match event {
295 Event::Text(line) => {
296 text.push_str(&line);
297 }
298 Event::End(TagEnd::Heading(HeadingLevel::H1)) => {
299 spec.name = mem::take(&mut text);
300 }
301 Event::End(TagEnd::Heading(HeadingLevel::H2)) => {
302 let title = mem::take(&mut text);
303 current_section = if title.eq_ignore_ascii_case(UNCOMMITTED_DIFF_HEADING) {
304 Section::UncommittedDiff
305 } else if title.eq_ignore_ascii_case(EDIT_HISTORY_HEADING) {
306 Section::EditHistory
307 } else if title.eq_ignore_ascii_case(CURSOR_POSITION_HEADING) {
308 Section::CursorPosition
309 } else if title.eq_ignore_ascii_case(EXPECTED_PATCH_HEADING) {
310 Section::ExpectedPatch
311 } else if title.eq_ignore_ascii_case(REJECTED_PATCH_HEADING) {
312 Section::RejectedPatch
313 } else {
314 Section::Other
315 };
316 }
317 Event::End(TagEnd::Heading(HeadingLevel::H3)) => {
318 mem::take(&mut text);
319 }
320 Event::End(TagEnd::Heading(HeadingLevel::H4)) => {
321 mem::take(&mut text);
322 }
323 Event::End(TagEnd::Heading(level)) => {
324 anyhow::bail!("Unexpected heading level: {level}");
325 }
326 Event::Start(Tag::CodeBlock(kind)) => {
327 if current_section == Section::EditHistory
328 && text.trim() == ACCEPTED_PREDICTION_MARKER
329 {
330 next_edit_predicted = true;
331 }
332 text.clear();
333 match kind {
334 CodeBlockKind::Fenced(info) => {
335 block_info = info;
336 }
337 CodeBlockKind::Indented => {
338 anyhow::bail!("Unexpected indented codeblock");
339 }
340 };
341 }
342 Event::Start(_) => {
343 text.clear();
344 block_info = "".into();
345 }
346 Event::End(TagEnd::CodeBlock) => {
347 let block_info = block_info.trim();
348 match current_section {
349 Section::UncommittedDiff => {
350 spec.uncommitted_diff = mem::take(&mut text);
351 }
352 Section::EditHistory => {
353 if next_edit_predicted {
354 spec.edit_history
355 .push_str(&format!("{}\n", ACCEPTED_PREDICTION_MARKER));
356 next_edit_predicted = false;
357 }
358 spec.edit_history.push_str(&mem::take(&mut text));
359 }
360 Section::CursorPosition => {
361 spec.cursor_path = Path::new(block_info).into();
362 spec.cursor_position = mem::take(&mut text);
363 }
364 Section::ExpectedPatch => {
365 spec.expected_patches.push(mem::take(&mut text));
366 }
367 Section::RejectedPatch => {
368 spec.rejected_patch = Some(mem::take(&mut text));
369 }
370 Section::Start | Section::Other => {}
371 }
372 }
373 _ => {}
374 }
375 }
376
377 if spec.cursor_path.as_ref() == Path::new("") || spec.cursor_position.is_empty() {
378 anyhow::bail!("Missing cursor position codeblock");
379 }
380
381 Ok(spec)
382 }
383
384 /// Returns the excerpt of text around the cursor, and the offset of the cursor within that
385 /// excerpt.
386 ///
387 /// The cursor's position is marked with a special comment that appears
388 /// below the cursor line, which contains the string `[CURSOR_POSITION]`,
389 /// preceded by an arrow marking the cursor's column. The arrow can be
390 /// either:
391 /// - `^` - The cursor column is at the position of the `^` character (pointing up to the cursor)
392 /// - `<` - The cursor column is at the first non-whitespace character on that line.
393 pub fn cursor_excerpt(&self) -> Result<(String, usize)> {
394 let input = &self.cursor_position;
395
396 // Check for inline cursor marker first
397 if let Some(inline_offset) = input.find(INLINE_CURSOR_MARKER) {
398 let excerpt = input[..inline_offset].to_string()
399 + &input[inline_offset + INLINE_CURSOR_MARKER.len()..];
400 return Ok((excerpt, inline_offset));
401 }
402
403 let marker_offset = input
404 .find(CURSOR_POSITION_MARKER)
405 .context("missing [CURSOR_POSITION] marker")?;
406 let marker_line_start = input[..marker_offset]
407 .rfind('\n')
408 .map(|pos| pos + 1)
409 .unwrap_or(0);
410 let marker_line_end = input[marker_line_start..]
411 .find('\n')
412 .map(|pos| marker_line_start + pos + 1)
413 .unwrap_or(input.len());
414 let marker_line = &input[marker_line_start..marker_line_end].trim_end_matches('\n');
415
416 let cursor_column = if let Some(cursor_offset) = marker_line.find('^') {
417 cursor_offset
418 } else if let Some(less_than_pos) = marker_line.find('<') {
419 marker_line
420 .find(|c: char| !c.is_whitespace())
421 .unwrap_or(less_than_pos)
422 } else {
423 anyhow::bail!(
424 "cursor position marker line must contain '^' or '<' before [CURSOR_POSITION]"
425 );
426 };
427
428 let mut excerpt = input[..marker_line_start].to_string() + &input[marker_line_end..];
429 excerpt.truncate(excerpt.trim_end_matches('\n').len());
430
431 // The cursor is on the line above the marker line.
432 let cursor_line_end = marker_line_start.saturating_sub(1);
433 let cursor_line_start = excerpt[..cursor_line_end]
434 .rfind('\n')
435 .map(|pos| pos + 1)
436 .unwrap_or(0);
437 let cursor_offset = cursor_line_start + cursor_column;
438
439 Ok((excerpt, cursor_offset))
440 }
441
442 /// Sets the cursor position excerpt from a plain excerpt and cursor byte offset.
443 ///
444 /// The `line_comment_prefix` is used to format the marker line as a comment.
445 /// If the cursor column is less than the comment prefix length, the `<` format is used.
446 /// Otherwise, the `^` format is used.
447 pub fn set_cursor_excerpt(
448 &mut self,
449 excerpt: &str,
450 cursor_offset: usize,
451 line_comment_prefix: &str,
452 ) {
453 // Find which line the cursor is on and its column
454 let cursor_line_start = excerpt[..cursor_offset]
455 .rfind('\n')
456 .map(|pos| pos + 1)
457 .unwrap_or(0);
458 let cursor_line_end = excerpt[cursor_line_start..]
459 .find('\n')
460 .map(|pos| cursor_line_start + pos + 1)
461 .unwrap_or(excerpt.len());
462 let cursor_line = &excerpt[cursor_line_start..cursor_line_end];
463 let cursor_line_indent = &cursor_line[..cursor_line.len() - cursor_line.trim_start().len()];
464 let cursor_column = cursor_offset - cursor_line_start;
465
466 // Build the marker line
467 let mut marker_line = String::new();
468 if cursor_column < line_comment_prefix.len() {
469 for _ in 0..cursor_column {
470 marker_line.push(' ');
471 }
472 marker_line.push_str(line_comment_prefix);
473 write!(marker_line, " <{}", CURSOR_POSITION_MARKER).unwrap();
474 } else {
475 if cursor_column >= cursor_line_indent.len() + line_comment_prefix.len() {
476 marker_line.push_str(cursor_line_indent);
477 }
478 marker_line.push_str(line_comment_prefix);
479 while marker_line.len() < cursor_column {
480 marker_line.push(' ');
481 }
482 write!(marker_line, "^{}", CURSOR_POSITION_MARKER).unwrap();
483 }
484
485 // Build the final cursor_position string
486 let mut result = String::with_capacity(excerpt.len() + marker_line.len() + 2);
487 result.push_str(&excerpt[..cursor_line_end]);
488 if !result.ends_with('\n') {
489 result.push('\n');
490 }
491 result.push_str(&marker_line);
492 if cursor_line_end < excerpt.len() {
493 result.push('\n');
494 result.push_str(&excerpt[cursor_line_end..]);
495 }
496
497 self.cursor_position = result;
498 }
499
500 /// Returns all of the possible expected patches for this example, each with an optional
501 /// cursor offset.
502 ///
503 /// The cursor offset is an offset within the new text (after applying the patch), relative
504 /// to the start of the hunk.
505 ///
506 /// In the serialized representation of this example, the cursor position is represented
507 /// using a comment line in the diff, beginning with `#`, and containing a `[CURSOR_POSITION]`
508 /// marker with the same format as the [`Self::cursor_excerpt`].
509 pub fn expected_patches_with_cursor_positions(&self) -> Vec<(String, Option<usize>)> {
510 self.expected_patches
511 .iter()
512 .map(|patch| {
513 let mut clean_patch = String::new();
514 let mut cursor_offset: Option<usize> = None;
515 let mut line_start_offset = 0usize;
516 let mut prev_line_start_offset = 0usize;
517
518 for line in patch.lines() {
519 let diff_line = DiffLine::parse(line);
520
521 match &diff_line {
522 DiffLine::Garbage(content)
523 if content.starts_with('#')
524 && content.contains(CURSOR_POSITION_MARKER) =>
525 {
526 let caret_column = if let Some(caret_pos) = content.find('^') {
527 caret_pos
528 } else if let Some(_) = content.find('<') {
529 0
530 } else {
531 continue;
532 };
533 let cursor_column = caret_column.saturating_sub('#'.len_utf8());
534 cursor_offset = Some(prev_line_start_offset + cursor_column);
535 }
536 _ => {
537 if !clean_patch.is_empty() {
538 clean_patch.push('\n');
539 }
540 clean_patch.push_str(line);
541
542 match diff_line {
543 DiffLine::Addition(content) | DiffLine::Context(content) => {
544 prev_line_start_offset = line_start_offset;
545 line_start_offset += content.len() + 1;
546 }
547 _ => {}
548 }
549 }
550 }
551 }
552
553 if patch.ends_with('\n') && !clean_patch.is_empty() {
554 clean_patch.push('\n');
555 }
556
557 (clean_patch, cursor_offset)
558 })
559 .collect()
560 }
561
562 pub fn set_expected_patches_with_cursor_positions(
563 &mut self,
564 patches: Vec<(String, Option<usize>)>,
565 ) {
566 self.expected_patches = patches
567 .into_iter()
568 .map(|(patch, cursor_offset)| encode_cursor_in_patch(&patch, cursor_offset))
569 .collect();
570 }
571}
572
573#[cfg(test)]
574mod tests {
575 use super::*;
576 use indoc::indoc;
577
578 #[test]
579 fn test_cursor_excerpt_with_caret() {
580 let mut spec = ExampleSpec {
581 name: String::new(),
582 repository_url: String::new(),
583 revision: String::new(),
584 tags: Vec::new(),
585 reasoning: None,
586 uncommitted_diff: String::new(),
587 cursor_path: Path::new("test.rs").into(),
588 cursor_position: String::new(),
589 edit_history: String::new(),
590 expected_patches: Vec::new(),
591 rejected_patch: None,
592 telemetry: None,
593 human_feedback: Vec::new(),
594 rating: None,
595 };
596
597 // Cursor before `42`
598 let excerpt = indoc! {"
599 fn main() {
600 let x = 42;
601 println!(\"{}\", x);
602 }"
603 };
604 let offset = excerpt.find("42").unwrap();
605 let position_string = indoc! {"
606 fn main() {
607 let x = 42;
608 // ^[CURSOR_POSITION]
609 println!(\"{}\", x);
610 }"
611 }
612 .to_string();
613
614 spec.set_cursor_excerpt(excerpt, offset, "//");
615 assert_eq!(spec.cursor_position, position_string);
616 assert_eq!(
617 spec.cursor_excerpt().unwrap(),
618 (excerpt.to_string(), offset)
619 );
620
621 // Cursor after `l` in `let`
622 let offset = excerpt.find("et x").unwrap();
623 let position_string = indoc! {"
624 fn main() {
625 let x = 42;
626 // ^[CURSOR_POSITION]
627 println!(\"{}\", x);
628 }"
629 }
630 .to_string();
631
632 spec.set_cursor_excerpt(excerpt, offset, "//");
633 assert_eq!(spec.cursor_position, position_string);
634 assert_eq!(
635 spec.cursor_excerpt().unwrap(),
636 (excerpt.to_string(), offset)
637 );
638
639 // Cursor before `let`
640 let offset = excerpt.find("let").unwrap();
641 let position_string = indoc! {"
642 fn main() {
643 let x = 42;
644 // ^[CURSOR_POSITION]
645 println!(\"{}\", x);
646 }"
647 }
648 .to_string();
649
650 spec.set_cursor_excerpt(excerpt, offset, "//");
651 assert_eq!(spec.cursor_position, position_string);
652 assert_eq!(
653 spec.cursor_excerpt().unwrap(),
654 (excerpt.to_string(), offset)
655 );
656
657 // Cursor at beginning of the line with `let`
658 let offset = excerpt.find(" let").unwrap();
659 let position_string = indoc! {"
660 fn main() {
661 let x = 42;
662 // <[CURSOR_POSITION]
663 println!(\"{}\", x);
664 }"
665 }
666 .to_string();
667
668 spec.set_cursor_excerpt(excerpt, offset, "//");
669 assert_eq!(spec.cursor_position, position_string);
670 assert_eq!(
671 spec.cursor_excerpt().unwrap(),
672 (excerpt.to_string(), offset)
673 );
674
675 // Cursor at end of line, after the semicolon
676 let offset = excerpt.find(';').unwrap() + 1;
677 let position_string = indoc! {"
678 fn main() {
679 let x = 42;
680 // ^[CURSOR_POSITION]
681 println!(\"{}\", x);
682 }"
683 }
684 .to_string();
685
686 spec.set_cursor_excerpt(excerpt, offset, "//");
687 assert_eq!(spec.cursor_position, position_string);
688 assert_eq!(
689 spec.cursor_excerpt().unwrap(),
690 (excerpt.to_string(), offset)
691 );
692
693 // Caret at end of file (no trailing newline)
694 let excerpt = indoc! {"
695 fn main() {
696 let x = 42;"
697 };
698 let offset = excerpt.find(';').unwrap() + 1;
699 let position_string = indoc! {"
700 fn main() {
701 let x = 42;
702 // ^[CURSOR_POSITION]"
703 }
704 .to_string();
705
706 spec.set_cursor_excerpt(excerpt, offset, "//");
707 assert_eq!(spec.cursor_position, position_string);
708 assert_eq!(
709 spec.cursor_excerpt().unwrap(),
710 (excerpt.to_string(), offset)
711 );
712 }
713
714 #[test]
715 fn test_cursor_excerpt_with_inline_marker() {
716 let mut spec = ExampleSpec {
717 name: String::new(),
718 repository_url: String::new(),
719 revision: String::new(),
720 tags: Vec::new(),
721 reasoning: None,
722 uncommitted_diff: String::new(),
723 cursor_path: Path::new("test.rs").into(),
724 cursor_position: String::new(),
725 edit_history: String::new(),
726 expected_patches: Vec::new(),
727 rejected_patch: None,
728 telemetry: None,
729 human_feedback: Vec::new(),
730 rating: None,
731 };
732
733 // Cursor before `42` using inline marker
734 spec.cursor_position = indoc! {"
735 fn main() {
736 let x = <|user_cursor|>42;
737 println!(\"{}\", x);
738 }"
739 }
740 .to_string();
741
742 let expected_excerpt = indoc! {"
743 fn main() {
744 let x = 42;
745 println!(\"{}\", x);
746 }"
747 };
748 let expected_offset = expected_excerpt.find("42").unwrap();
749
750 assert_eq!(
751 spec.cursor_excerpt().unwrap(),
752 (expected_excerpt.to_string(), expected_offset)
753 );
754
755 // Cursor at beginning of line
756 spec.cursor_position = indoc! {"
757 fn main() {
758 <|user_cursor|> let x = 42;
759 }"
760 }
761 .to_string();
762
763 let expected_excerpt = indoc! {"
764 fn main() {
765 let x = 42;
766 }"
767 };
768 let expected_offset = expected_excerpt.find(" let").unwrap();
769
770 assert_eq!(
771 spec.cursor_excerpt().unwrap(),
772 (expected_excerpt.to_string(), expected_offset)
773 );
774
775 // Cursor at end of file
776 spec.cursor_position = "fn main() {}<|user_cursor|>".to_string();
777 let expected_excerpt = "fn main() {}";
778 let expected_offset = expected_excerpt.len();
779
780 assert_eq!(
781 spec.cursor_excerpt().unwrap(),
782 (expected_excerpt.to_string(), expected_offset)
783 );
784 }
785
786 #[test]
787 fn test_expected_patches_with_cursor_positions() {
788 let mut spec = ExampleSpec {
789 name: String::new(),
790 repository_url: String::new(),
791 revision: String::new(),
792 tags: Vec::new(),
793 reasoning: None,
794 uncommitted_diff: String::new(),
795 cursor_path: Path::new("test.rs").into(),
796 cursor_position: String::new(),
797 edit_history: String::new(),
798 expected_patches: Vec::new(),
799 rejected_patch: None,
800 telemetry: None,
801 human_feedback: Vec::new(),
802 rating: None,
803 };
804
805 let new_content = indoc! {r#"
806 // prints a greeting
807 fn main() {
808 println!("hello, {}", );
809 let x = 42;
810 }
811 "#};
812 let cursor_offset = new_content.find(");").unwrap();
813
814 let clean_patch = indoc! {r#"
815 --- a/test.rs
816 +++ b/test.rs
817 @@ -1,3 +1,4 @@
818 +// prints a greeting
819 fn main() {
820 - println!("hi");
821 + println!("hello, {}", );
822 let x = 42;
823 }
824 "#}
825 .to_string();
826
827 let encoded_patch = indoc! {r#"
828 --- a/test.rs
829 +++ b/test.rs
830 @@ -1,3 +1,4 @@
831 +// prints a greeting
832 fn main() {
833 - println!("hi");
834 + println!("hello, {}", );
835 # ^[CURSOR_POSITION]
836 let x = 42;
837 }
838 "#}
839 .to_string();
840
841 spec.set_expected_patches_with_cursor_positions(vec![(
842 clean_patch.clone(),
843 Some(cursor_offset),
844 )]);
845 assert_eq!(spec.expected_patches, vec![encoded_patch]);
846
847 let results = spec.expected_patches_with_cursor_positions();
848 assert_eq!(results, vec![(clean_patch.clone(), Some(cursor_offset))]);
849
850 spec.set_expected_patches_with_cursor_positions(vec![(clean_patch.clone(), None)]);
851 assert_eq!(spec.expected_patches, vec![clean_patch.clone()]);
852
853 let results = spec.expected_patches_with_cursor_positions();
854 assert_eq!(results, vec![(clean_patch, None)]);
855 }
856
857 #[test]
858 fn test_encode_cursor_in_patch_is_idempotent() {
859 let patch = indoc! {r#"
860 --- a/test.rs
861 +++ b/test.rs
862 @@ -1,2 +1,2 @@
863 -fn old() {}
864 +fn new_name() {}
865 # ^[CURSOR_POSITION]
866 "#};
867
868 let cursor_offset = "fn new_name() {}".find("name").unwrap();
869 let encoded_once = encode_cursor_in_patch(patch, Some(cursor_offset));
870 let encoded_twice = encode_cursor_in_patch(&encoded_once, Some(cursor_offset));
871
872 assert_eq!(encoded_once, encoded_twice);
873 assert_eq!(
874 encoded_once
875 .lines()
876 .filter(|line| line.contains(CURSOR_POSITION_MARKER))
877 .count(),
878 1
879 );
880 }
881
882 #[test]
883 fn test_from_markdown_accepted_prediction_marker() {
884 let markdown = indoc! {r#"
885 +++
886 repository_url = "https://github.com/example/repo"
887 revision = "abc123"
888 +++
889
890 ## Edit History
891
892 ```diff
893 --- a/src/main.rs
894 +++ b/src/main.rs
895 @@ -1,3 +1,3 @@
896 -fn hello() {}
897 +fn hello_world() {}
898 ```
899
900 // User accepted prediction:
901 ```diff
902 --- a/src/main.rs
903 +++ b/src/main.rs
904 @@ -1,3 +1,3 @@
905 -fn hello_world() {}
906 +fn hello_world() { println!("hi"); }
907 ```
908
909 ```diff
910 --- a/src/main.rs
911 +++ b/src/main.rs
912 @@ -1,3 +1,3 @@
913 -fn hello_world() { println!("hi"); }
914 +fn hello_world() { println!("hello"); }
915 ```
916
917 ## Cursor Position
918
919 ```src/main.rs
920 fn hello_world() { println!("hello"); }
921 # ^[CURSOR_POSITION]
922 ```
923
924 ## Expected Patch
925
926 ```diff
927 --- a/src/main.rs
928 +++ b/src/main.rs
929 @@ -1,3 +1,3 @@
930 -fn hello_world() { println!("hello"); }
931 +fn hello_world() { println!("hello, world!"); }
932 ```
933 "#};
934
935 let spec = ExampleSpec::from_markdown(markdown).unwrap();
936
937 // The first diff should NOT have the marker
938 assert!(spec.edit_history.starts_with("--- a/src/main.rs"));
939
940 // The second diff should be preceded by the accepted prediction marker
941 assert!(
942 spec.edit_history
943 .contains("// User accepted prediction:\n--- a/src/main.rs")
944 );
945
946 // Count occurrences of the marker - should be exactly one
947 let marker_count = spec
948 .edit_history
949 .matches("// User accepted prediction:")
950 .count();
951 assert_eq!(marker_count, 1);
952
953 // The third diff should NOT have the marker
954 // Verify all three diffs are present
955 let diff_count = spec.edit_history.matches("--- a/src/main.rs").count();
956 assert_eq!(diff_count, 3);
957 }
958}