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