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