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

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 <oleksiy.syvokon@gmail.com>
This commit is contained in:
Shiven Garia 2026-03-31 23:20:49 +05:30 committed by GitHub
parent 3ed687d147
commit 56104fb17e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -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"
);
}
}