1use crate::udiff::DiffLine;
2use anyhow::{Context as _, Result};
3use serde::{Deserialize, Serialize};
4use std::{borrow::Cow, fmt::Write as _, mem, ops::Range, 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#[derive(Clone, Debug, PartialEq, Hash, Serialize, Deserialize)]
16pub struct ExampleSpec {
17 #[serde(default)]
18 pub name: String,
19 pub repository_url: String,
20 pub revision: String,
21 #[serde(default, skip_serializing_if = "Vec::is_empty")]
22 pub tags: Vec<String>,
23 #[serde(default, skip_serializing_if = "Option::is_none")]
24 pub reasoning: Option<String>,
25 #[serde(default)]
26 pub uncommitted_diff: String,
27 pub cursor_path: Arc<Path>,
28 pub cursor_position: String,
29 pub edit_history: String,
30 pub expected_patches: Vec<String>,
31 #[serde(default, skip_serializing_if = "Option::is_none")]
32 pub rejected_patch: Option<String>,
33 #[serde(default, skip_serializing_if = "Option::is_none")]
34 pub captured_prompt_input: Option<CapturedPromptInput>,
35 #[serde(default, skip_serializing_if = "Option::is_none")]
36 pub telemetry: Option<TelemetrySource>,
37 #[serde(default, skip_serializing_if = "Vec::is_empty")]
38 pub human_feedback: Vec<HumanFeedback>,
39 #[serde(default, skip_serializing_if = "Option::is_none")]
40 pub rating: Option<EditPredictionRating>,
41}
42
43#[derive(Clone, Debug, PartialEq, Hash, Serialize, Deserialize)]
44pub struct HumanFeedback {
45 pub message: String,
46}
47
48/// Metadata for examples sourced from production telemetry (rejected predictions).
49#[derive(Clone, Debug, PartialEq, Hash, Serialize, Deserialize)]
50pub struct TelemetrySource {
51 pub request_id: String,
52 pub device_id: String,
53 pub time: String,
54 pub rejection_reason: String,
55 pub was_shown: bool,
56}
57
58/// All data needed to run format_prompt without loading the project.
59#[derive(Clone, Debug, PartialEq, Hash, Serialize, Deserialize)]
60pub struct CapturedPromptInput {
61 pub cursor_file_content: String,
62 pub cursor_offset: usize,
63 pub cursor_row: u32,
64 pub cursor_column: u32,
65 #[serde(default, skip_serializing_if = "Option::is_none")]
66 pub excerpt_start_row: Option<u32>,
67 pub events: Vec<CapturedEvent>,
68 pub related_files: Vec<CapturedRelatedFile>,
69 #[serde(default)]
70 pub in_open_source_repo: bool,
71}
72
73#[derive(Clone, Debug, PartialEq, Hash, Serialize, Deserialize)]
74pub struct CapturedEvent {
75 pub path: Arc<Path>,
76 pub old_path: Arc<Path>,
77 pub diff: String,
78 pub predicted: bool,
79 #[serde(default)]
80 pub in_open_source_repo: bool,
81}
82
83impl CapturedEvent {
84 pub fn to_event(&self) -> zeta_prompt::Event {
85 zeta_prompt::Event::BufferChange {
86 path: self.path.clone(),
87 old_path: self.old_path.clone(),
88 diff: self.diff.clone(),
89 predicted: self.predicted,
90 in_open_source_repo: self.in_open_source_repo,
91 }
92 }
93}
94
95#[derive(Clone, Debug, PartialEq, Hash, Serialize, Deserialize)]
96pub struct CapturedRelatedFile {
97 pub path: Arc<Path>,
98 pub max_row: u32,
99 pub excerpts: Vec<CapturedRelatedExcerpt>,
100}
101
102impl CapturedRelatedFile {
103 pub fn to_related_file(&self) -> zeta_prompt::RelatedFile {
104 zeta_prompt::RelatedFile {
105 path: self.path.clone(),
106 max_row: self.max_row,
107 in_open_source_repo: false,
108 excerpts: self
109 .excerpts
110 .iter()
111 .map(|e| zeta_prompt::RelatedExcerpt {
112 row_range: e.row_range.clone(),
113 text: e.text.clone().into(),
114 })
115 .collect(),
116 }
117 }
118}
119
120#[derive(Clone, Debug, PartialEq, Hash, Serialize, Deserialize)]
121pub struct CapturedRelatedExcerpt {
122 pub row_range: Range<u32>,
123 pub text: String,
124}
125
126const REASONING_HEADING: &str = "Reasoning";
127const UNCOMMITTED_DIFF_HEADING: &str = "Uncommitted Diff";
128const EDIT_HISTORY_HEADING: &str = "Edit History";
129const CURSOR_POSITION_HEADING: &str = "Cursor Position";
130const EXPECTED_PATCH_HEADING: &str = "Expected Patch";
131const REJECTED_PATCH_HEADING: &str = "Rejected Patch";
132
133#[derive(Serialize, Deserialize)]
134struct FrontMatter<'a> {
135 repository_url: Cow<'a, str>,
136 revision: Cow<'a, str>,
137 #[serde(default, skip_serializing_if = "Vec::is_empty")]
138 tags: Vec<String>,
139}
140
141impl ExampleSpec {
142 /// Generate a sanitized filename for this example.
143 pub fn filename(&self) -> String {
144 self.name
145 .chars()
146 .map(|c| match c {
147 ' ' | ':' | '~' | '^' | '?' | '*' | '[' | '\\' | '@' | '{' | '/' | '<' | '>'
148 | '|' | '"' => '-',
149 c => c,
150 })
151 .collect()
152 }
153
154 /// Format this example spec as markdown.
155 pub fn to_markdown(&self) -> String {
156 use std::fmt::Write as _;
157
158 let front_matter = FrontMatter {
159 repository_url: Cow::Borrowed(&self.repository_url),
160 revision: Cow::Borrowed(&self.revision),
161 tags: self.tags.clone(),
162 };
163 let front_matter_toml =
164 toml::to_string_pretty(&front_matter).unwrap_or_else(|_| String::new());
165
166 let mut markdown = String::new();
167
168 _ = writeln!(markdown, "+++");
169 markdown.push_str(&front_matter_toml);
170 if !markdown.ends_with('\n') {
171 markdown.push('\n');
172 }
173 _ = writeln!(markdown, "+++");
174 markdown.push('\n');
175
176 _ = writeln!(markdown, "# {}", self.name);
177 markdown.push('\n');
178
179 if let Some(reasoning) = &self.reasoning {
180 _ = writeln!(markdown, "## {}", REASONING_HEADING);
181 markdown.push('\n');
182 markdown.push_str(reasoning);
183 if !markdown.ends_with('\n') {
184 markdown.push('\n');
185 }
186 markdown.push('\n');
187 }
188
189 if !self.uncommitted_diff.is_empty() {
190 _ = writeln!(markdown, "## {}", UNCOMMITTED_DIFF_HEADING);
191 _ = writeln!(markdown);
192 _ = writeln!(markdown, "```diff");
193 markdown.push_str(&self.uncommitted_diff);
194 if !markdown.ends_with('\n') {
195 markdown.push('\n');
196 }
197 _ = writeln!(markdown, "```");
198 markdown.push('\n');
199 }
200
201 _ = writeln!(markdown, "## {}", EDIT_HISTORY_HEADING);
202 _ = writeln!(markdown);
203
204 if self.edit_history.is_empty() {
205 _ = writeln!(markdown, "(No edit history)");
206 _ = writeln!(markdown);
207 } else {
208 _ = writeln!(markdown, "```diff");
209 markdown.push_str(&self.edit_history);
210 if !markdown.ends_with('\n') {
211 markdown.push('\n');
212 }
213 _ = writeln!(markdown, "```");
214 markdown.push('\n');
215 }
216
217 _ = writeln!(markdown, "## {}", CURSOR_POSITION_HEADING);
218 _ = writeln!(markdown);
219 _ = writeln!(markdown, "```{}", self.cursor_path.to_string_lossy());
220 markdown.push_str(&self.cursor_position);
221 if !markdown.ends_with('\n') {
222 markdown.push('\n');
223 }
224 _ = writeln!(markdown, "```");
225 markdown.push('\n');
226
227 _ = writeln!(markdown, "## {}", EXPECTED_PATCH_HEADING);
228 markdown.push('\n');
229 for patch in &self.expected_patches {
230 _ = writeln!(markdown, "```diff");
231 markdown.push_str(patch);
232 if !markdown.ends_with('\n') {
233 markdown.push('\n');
234 }
235 _ = writeln!(markdown, "```");
236 markdown.push('\n');
237 }
238
239 if let Some(rejected_patch) = &self.rejected_patch {
240 _ = writeln!(markdown, "## {}", REJECTED_PATCH_HEADING);
241 markdown.push('\n');
242 _ = writeln!(markdown, "```diff");
243 markdown.push_str(rejected_patch);
244 if !markdown.ends_with('\n') {
245 markdown.push('\n');
246 }
247 _ = writeln!(markdown, "```");
248 markdown.push('\n');
249 }
250
251 markdown
252 }
253
254 /// Parse an example spec from markdown.
255 pub fn from_markdown(mut input: &str) -> anyhow::Result<Self> {
256 use pulldown_cmark::{CodeBlockKind, CowStr, Event, HeadingLevel, Parser, Tag, TagEnd};
257
258 let mut spec = ExampleSpec {
259 name: String::new(),
260 repository_url: String::new(),
261 revision: String::new(),
262 tags: Vec::new(),
263 reasoning: None,
264 uncommitted_diff: String::new(),
265 cursor_path: Path::new("").into(),
266 cursor_position: String::new(),
267 edit_history: String::new(),
268 expected_patches: Vec::new(),
269 rejected_patch: None,
270 captured_prompt_input: None,
271 telemetry: None,
272 human_feedback: Vec::new(),
273 rating: None,
274 };
275
276 if let Some(rest) = input.strip_prefix("+++\n")
277 && let Some((front_matter, rest)) = rest.split_once("+++\n")
278 {
279 if let Ok(data) = toml::from_str::<FrontMatter<'_>>(front_matter) {
280 spec.repository_url = data.repository_url.into_owned();
281 spec.revision = data.revision.into_owned();
282 spec.tags = data.tags;
283 }
284 input = rest.trim_start();
285 }
286
287 let parser = Parser::new(input);
288 let mut text = String::new();
289 let mut block_info: CowStr = "".into();
290
291 #[derive(PartialEq)]
292 enum Section {
293 Start,
294 UncommittedDiff,
295 EditHistory,
296 CursorPosition,
297 ExpectedPatch,
298 RejectedPatch,
299 Other,
300 }
301
302 let mut current_section = Section::Start;
303
304 for event in parser {
305 match event {
306 Event::Text(line) => {
307 text.push_str(&line);
308 }
309 Event::End(TagEnd::Heading(HeadingLevel::H1)) => {
310 spec.name = mem::take(&mut text);
311 }
312 Event::End(TagEnd::Heading(HeadingLevel::H2)) => {
313 let title = mem::take(&mut text);
314 current_section = if title.eq_ignore_ascii_case(UNCOMMITTED_DIFF_HEADING) {
315 Section::UncommittedDiff
316 } else if title.eq_ignore_ascii_case(EDIT_HISTORY_HEADING) {
317 Section::EditHistory
318 } else if title.eq_ignore_ascii_case(CURSOR_POSITION_HEADING) {
319 Section::CursorPosition
320 } else if title.eq_ignore_ascii_case(EXPECTED_PATCH_HEADING) {
321 Section::ExpectedPatch
322 } else if title.eq_ignore_ascii_case(REJECTED_PATCH_HEADING) {
323 Section::RejectedPatch
324 } else {
325 Section::Other
326 };
327 }
328 Event::End(TagEnd::Heading(HeadingLevel::H3)) => {
329 mem::take(&mut text);
330 }
331 Event::End(TagEnd::Heading(HeadingLevel::H4)) => {
332 mem::take(&mut text);
333 }
334 Event::End(TagEnd::Heading(level)) => {
335 anyhow::bail!("Unexpected heading level: {level}");
336 }
337 Event::Start(Tag::CodeBlock(kind)) => {
338 match kind {
339 CodeBlockKind::Fenced(info) => {
340 block_info = info;
341 }
342 CodeBlockKind::Indented => {
343 anyhow::bail!("Unexpected indented codeblock");
344 }
345 };
346 }
347 Event::Start(_) => {
348 text.clear();
349 block_info = "".into();
350 }
351 Event::End(TagEnd::CodeBlock) => {
352 let block_info = block_info.trim();
353 match current_section {
354 Section::UncommittedDiff => {
355 spec.uncommitted_diff = mem::take(&mut text);
356 }
357 Section::EditHistory => {
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_editable_region_offset)| {
569 let Some(cursor_offset) = cursor_editable_region_offset else {
570 return patch;
571 };
572
573 let mut result = String::new();
574 let mut line_start_offset = 0usize;
575
576 for line in patch.lines() {
577 if !result.is_empty() {
578 result.push('\n');
579 }
580 result.push_str(line);
581
582 match DiffLine::parse(line) {
583 DiffLine::Addition(content) => {
584 let line_end_offset = line_start_offset + content.len();
585
586 if cursor_offset >= line_start_offset
587 && cursor_offset <= line_end_offset
588 {
589 let cursor_column = cursor_offset - line_start_offset;
590
591 result.push('\n');
592 result.push('#');
593 for _ in 0..cursor_column {
594 result.push(' ');
595 }
596 write!(result, "^{}", CURSOR_POSITION_MARKER).unwrap();
597 }
598
599 line_start_offset = line_end_offset + 1;
600 }
601 DiffLine::Context(content) => {
602 line_start_offset += content.len() + 1;
603 }
604 _ => {}
605 }
606 }
607
608 if patch.ends_with('\n') {
609 result.push('\n');
610 }
611
612 result
613 })
614 .collect();
615 }
616}
617
618#[cfg(test)]
619mod tests {
620 use super::*;
621 use indoc::indoc;
622
623 #[test]
624 fn test_cursor_excerpt_with_caret() {
625 let mut spec = ExampleSpec {
626 name: String::new(),
627 repository_url: String::new(),
628 revision: String::new(),
629 tags: Vec::new(),
630 reasoning: None,
631 uncommitted_diff: String::new(),
632 cursor_path: Path::new("test.rs").into(),
633 cursor_position: String::new(),
634 edit_history: String::new(),
635 expected_patches: Vec::new(),
636 rejected_patch: None,
637 captured_prompt_input: None,
638 telemetry: None,
639 human_feedback: Vec::new(),
640 rating: None,
641 };
642
643 // Cursor before `42`
644 let excerpt = indoc! {"
645 fn main() {
646 let x = 42;
647 println!(\"{}\", x);
648 }"
649 };
650 let offset = excerpt.find("42").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 after `l` in `let`
668 let offset = excerpt.find("et x").unwrap();
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 // Cursor before `let`
686 let offset = excerpt.find("let").unwrap();
687 let position_string = indoc! {"
688 fn main() {
689 let x = 42;
690 // ^[CURSOR_POSITION]
691 println!(\"{}\", x);
692 }"
693 }
694 .to_string();
695
696 spec.set_cursor_excerpt(excerpt, offset, "//");
697 assert_eq!(spec.cursor_position, position_string);
698 assert_eq!(
699 spec.cursor_excerpt().unwrap(),
700 (excerpt.to_string(), offset)
701 );
702
703 // Cursor at beginning of the line with `let`
704 let offset = excerpt.find(" let").unwrap();
705 let position_string = indoc! {"
706 fn main() {
707 let x = 42;
708 // <[CURSOR_POSITION]
709 println!(\"{}\", x);
710 }"
711 }
712 .to_string();
713
714 spec.set_cursor_excerpt(excerpt, offset, "//");
715 assert_eq!(spec.cursor_position, position_string);
716 assert_eq!(
717 spec.cursor_excerpt().unwrap(),
718 (excerpt.to_string(), offset)
719 );
720
721 // Cursor at end of line, after the semicolon
722 let offset = excerpt.find(';').unwrap() + 1;
723 let position_string = indoc! {"
724 fn main() {
725 let x = 42;
726 // ^[CURSOR_POSITION]
727 println!(\"{}\", x);
728 }"
729 }
730 .to_string();
731
732 spec.set_cursor_excerpt(excerpt, offset, "//");
733 assert_eq!(spec.cursor_position, position_string);
734 assert_eq!(
735 spec.cursor_excerpt().unwrap(),
736 (excerpt.to_string(), offset)
737 );
738
739 // Caret at end of file (no trailing newline)
740 let excerpt = indoc! {"
741 fn main() {
742 let x = 42;"
743 };
744 let offset = excerpt.find(';').unwrap() + 1;
745 let position_string = indoc! {"
746 fn main() {
747 let x = 42;
748 // ^[CURSOR_POSITION]"
749 }
750 .to_string();
751
752 spec.set_cursor_excerpt(excerpt, offset, "//");
753 assert_eq!(spec.cursor_position, position_string);
754 assert_eq!(
755 spec.cursor_excerpt().unwrap(),
756 (excerpt.to_string(), offset)
757 );
758 }
759
760 #[test]
761 fn test_cursor_excerpt_with_inline_marker() {
762 let mut spec = ExampleSpec {
763 name: String::new(),
764 repository_url: String::new(),
765 revision: String::new(),
766 tags: Vec::new(),
767 reasoning: None,
768 uncommitted_diff: String::new(),
769 cursor_path: Path::new("test.rs").into(),
770 cursor_position: String::new(),
771 edit_history: String::new(),
772 expected_patches: Vec::new(),
773 rejected_patch: None,
774 captured_prompt_input: None,
775 telemetry: None,
776 human_feedback: Vec::new(),
777 rating: None,
778 };
779
780 // Cursor before `42` using inline marker
781 spec.cursor_position = indoc! {"
782 fn main() {
783 let x = <|user_cursor|>42;
784 println!(\"{}\", x);
785 }"
786 }
787 .to_string();
788
789 let expected_excerpt = indoc! {"
790 fn main() {
791 let x = 42;
792 println!(\"{}\", x);
793 }"
794 };
795 let expected_offset = expected_excerpt.find("42").unwrap();
796
797 assert_eq!(
798 spec.cursor_excerpt().unwrap(),
799 (expected_excerpt.to_string(), expected_offset)
800 );
801
802 // Cursor at beginning of line
803 spec.cursor_position = indoc! {"
804 fn main() {
805 <|user_cursor|> let x = 42;
806 }"
807 }
808 .to_string();
809
810 let expected_excerpt = indoc! {"
811 fn main() {
812 let x = 42;
813 }"
814 };
815 let expected_offset = expected_excerpt.find(" let").unwrap();
816
817 assert_eq!(
818 spec.cursor_excerpt().unwrap(),
819 (expected_excerpt.to_string(), expected_offset)
820 );
821
822 // Cursor at end of file
823 spec.cursor_position = "fn main() {}<|user_cursor|>".to_string();
824 let expected_excerpt = "fn main() {}";
825 let expected_offset = expected_excerpt.len();
826
827 assert_eq!(
828 spec.cursor_excerpt().unwrap(),
829 (expected_excerpt.to_string(), expected_offset)
830 );
831 }
832
833 #[test]
834 fn test_expected_patches_with_cursor_positions() {
835 let mut spec = ExampleSpec {
836 name: String::new(),
837 repository_url: String::new(),
838 revision: String::new(),
839 tags: Vec::new(),
840 reasoning: None,
841 uncommitted_diff: String::new(),
842 cursor_path: Path::new("test.rs").into(),
843 cursor_position: String::new(),
844 edit_history: String::new(),
845 expected_patches: Vec::new(),
846 rejected_patch: None,
847 captured_prompt_input: None,
848 telemetry: None,
849 human_feedback: Vec::new(),
850 rating: None,
851 };
852
853 let new_content = indoc! {r#"
854 // prints a greeting
855 fn main() {
856 println!("hello, {}", );
857 let x = 42;
858 }
859 "#};
860 let cursor_offset = new_content.find(");").unwrap();
861
862 let clean_patch = indoc! {r#"
863 --- a/test.rs
864 +++ b/test.rs
865 @@ -1,3 +1,4 @@
866 +// prints a greeting
867 fn main() {
868 - println!("hi");
869 + println!("hello, {}", );
870 let x = 42;
871 }
872 "#}
873 .to_string();
874
875 let encoded_patch = indoc! {r#"
876 --- a/test.rs
877 +++ b/test.rs
878 @@ -1,3 +1,4 @@
879 +// prints a greeting
880 fn main() {
881 - println!("hi");
882 + println!("hello, {}", );
883 # ^[CURSOR_POSITION]
884 let x = 42;
885 }
886 "#}
887 .to_string();
888
889 spec.set_expected_patches_with_cursor_positions(vec![(
890 clean_patch.clone(),
891 Some(cursor_offset),
892 )]);
893 assert_eq!(spec.expected_patches, vec![encoded_patch]);
894
895 let results = spec.expected_patches_with_cursor_positions();
896 assert_eq!(results, vec![(clean_patch.clone(), Some(cursor_offset))]);
897
898 spec.set_expected_patches_with_cursor_positions(vec![(clean_patch.clone(), None)]);
899 assert_eq!(spec.expected_patches, vec![clean_patch.clone()]);
900
901 let results = spec.expected_patches_with_cursor_positions();
902 assert_eq!(results, vec![(clean_patch, None)]);
903 }
904}