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