From 56104fb17e6c5849900a4c282f42d1bfea294121 Mon Sep 17 00:00:00 2001 From: Shiven Garia Date: Tue, 31 Mar 2026 23:20:49 +0530 Subject: [PATCH] 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 --- crates/zeta_prompt/src/zeta_prompt.rs | 21 ++++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/crates/zeta_prompt/src/zeta_prompt.rs b/crates/zeta_prompt/src/zeta_prompt.rs index 3ec90baf6e7d7781b5ddedb0af3dbdb0994cb3ad..e3aced7ed81d8bf3835a3e711e472651764a314e 100644 --- a/crates/zeta_prompt/src/zeta_prompt.rs +++ b/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 { @@ -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" + ); + } }