edit_prediction: Fix special token check matching `=======` inside comments (#52510)

Shiven Garia and Oleksiy Syvokon created 1 week ago

Closes #52489

The special token check in `prompt_input_contains_special_tokens` used
`String::contains()` to look for `=======\n` in the buffer. This meant
any line containing `=======` (like `// =======` section separators)
would cause edit predictions to bail out entirely.

Fixed by only matching when the token appears at the start of a line,
since the git merge markers are always placed at line boundaries in the
prompt.

Added tests for both the helper function and a regression test for the
reported issue.

---------

Co-authored-by: Oleksiy Syvokon <[e-mail]>

Change summary

crates/zeta_prompt/src/zeta_prompt.rs | 21 ++++++++++++++++++---
1 file changed, 18 insertions(+), 3 deletions(-)

Detailed changes

crates/zeta_prompt/src/zeta_prompt.rs

@@ -211,9 +211,13 @@ pub struct RelatedExcerpt {
 }
 
 pub fn prompt_input_contains_special_tokens(input: &ZetaPromptInput, format: ZetaFormat) -> bool {
-    special_tokens_for_format(format)
-        .iter()
-        .any(|token| input.cursor_excerpt.contains(token))
+    special_tokens_for_format(format).iter().any(|token| {
+        if let Some(line_token) = token.strip_suffix('\n') {
+            input.cursor_excerpt.lines().any(|line| line == line_token)
+        } else {
+            input.cursor_excerpt.contains(token)
+        }
+    })
 }
 
 pub fn format_zeta_prompt(input: &ZetaPromptInput, format: ZetaFormat) -> Option<String> {
@@ -5287,4 +5291,15 @@ mod tests {
         assert_eq!(apply_edit(excerpt, &output1), apply_edit(excerpt, &output2));
         assert_eq!(apply_edit(excerpt, &output1), "new content\n");
     }
+
+    #[test]
+    fn test_special_tokens_not_triggered_by_comment_separator() {
+        // Regression test for https://github.com/zed-industries/zed/issues/52489
+        let excerpt = "fn main() {\n    // =======\n    println!(\"hello\");\n}\n";
+        let input = make_input(excerpt, 0..excerpt.len(), 0, vec![], vec![]);
+        assert!(
+            !prompt_input_contains_special_tokens(&input, ZetaFormat::V0131GitMergeMarkersPrefix),
+            "comment containing ======= should not trigger special token detection"
+        );
+    }
 }