tenant.openagents/omega

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