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