From d405c509d2b12678c3ba4b15dd9e154246840ead Mon Sep 17 00:00:00 2001 From: Ben Kunkle Date: Tue, 12 May 2026 08:30:07 -0400 Subject: [PATCH] ep: Limit collected diagnostics (#56431) Self-Review Checklist: - [x] I've reviewed my own diff for quality, security, and reliability - [x] Unsafe blocks (if any) have justifying comments - [x] The content is consistent with the [UI/UX checklist](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist) - [x] Tests cover the new/changed behavior - [x] Performance impact has been considered and is acceptable Closes #ISSUE Release Notes: - N/A or Added/Fixed/Improved ... --- .../src/edit_prediction_tests.rs | 126 ++++++++++++++++-- crates/edit_prediction/src/zeta.rs | 92 ++++++++++--- crates/zeta_prompt/src/zeta_prompt.rs | 48 ++++--- 3 files changed, 219 insertions(+), 47 deletions(-) diff --git a/crates/edit_prediction/src/edit_prediction_tests.rs b/crates/edit_prediction/src/edit_prediction_tests.rs index 00d52023f4..d404be7920 100644 --- a/crates/edit_prediction/src/edit_prediction_tests.rs +++ b/crates/edit_prediction/src/edit_prediction_tests.rs @@ -2241,19 +2241,48 @@ fn test_active_buffer_diagnostics_fetching(cx: &mut TestAppContext) { let search_range = snapshot.offset_to_point(search_ranges[0].start) ..snapshot.offset_to_point(search_ranges[0].end); - let active_buffer_diagnostics = zeta::active_buffer_diagnostics(&snapshot, search_range, 100); + let active_buffer_diagnostics = zeta::active_buffer_diagnostics(&snapshot, search_range, 5, 0); assert_eq!( active_buffer_diagnostics, vec![zeta_prompt::ActiveBufferDiagnostic { severity: Some(1), message: "second error".to_string(), - snippet: text, + snippet: " let second_value = 2;".to_string(), snippet_buffer_row_range: 5..5, - diagnostic_range_in_snippet: 61..73, + diagnostic_range_in_snippet: 8..20, }] ); + let active_buffer_diagnostics = + zeta::active_buffer_diagnostics(&snapshot, Point::new(0, 0)..snapshot.max_point(), 5, 100); + assert_eq!( + active_buffer_diagnostics, + vec![ + zeta_prompt::ActiveBufferDiagnostic { + severity: Some(1), + message: "second error".to_string(), + snippet: String::new(), + snippet_buffer_row_range: 5..5, + diagnostic_range_in_snippet: 0..0, + }, + zeta_prompt::ActiveBufferDiagnostic { + severity: Some(2), + message: "first warning".to_string(), + snippet: String::new(), + snippet_buffer_row_range: 1..1, + diagnostic_range_in_snippet: 0..0, + }, + zeta_prompt::ActiveBufferDiagnostic { + severity: Some(4), + message: "third hint".to_string(), + snippet: String::new(), + snippet_buffer_row_range: 10..10, + diagnostic_range_in_snippet: 0..0, + }, + ] + ); + let buffer = cx.new(|cx| { Buffer::local( indoc! {" @@ -2313,7 +2342,7 @@ fn test_active_buffer_diagnostics_fetching(cx: &mut TestAppContext) { let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot()); let active_buffer_diagnostics = - zeta::active_buffer_diagnostics(&snapshot, Point::new(2, 0)..Point::new(4, 0), 100); + zeta::active_buffer_diagnostics(&snapshot, Point::new(2, 0)..Point::new(4, 0), 3, 0); assert_eq!( active_buffer_diagnostics @@ -2330,21 +2359,102 @@ fn test_active_buffer_diagnostics_fetching(cx: &mut TestAppContext) { ( Some(2), "row two".to_string(), - "one\ntwo\nthree\nfour\nfive\n".to_string(), + "three".to_string(), 2..2, - 8..13, + 0..5, ), ( Some(3), "row four".to_string(), - "one\ntwo\nthree\nfour\nfive\n".to_string(), + "five".to_string(), 4..4, - 19..23, + 0..4, ), ] ); } +#[gpui::test] +fn test_active_buffer_diagnostics_collection_limits(cx: &mut TestAppContext) { + let text = (0..25) + .map(|row| format!("line {row}\n")) + .collect::(); + let buffer = cx.new(|cx| Buffer::local(&text, cx)); + + buffer.update(cx, |buffer, cx| { + let snapshot = buffer.snapshot(); + let diagnostics = DiagnosticSet::new( + (0..25) + .map(|row| DiagnosticEntry { + range: text::PointUtf16::new(row, 0)..text::PointUtf16::new(row, 4), + diagnostic: Diagnostic { + severity: DiagnosticSeverity::ERROR, + message: format!("row {row}"), + group_id: row as usize, + is_primary: true, + source_kind: language::DiagnosticSourceKind::Pushed, + ..Diagnostic::default() + }, + }) + .collect::>(), + &snapshot, + ); + buffer.update_diagnostics(LanguageServerId(0), diagnostics, cx); + }); + + let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot()); + let active_buffer_diagnostics = + zeta::active_buffer_diagnostics(&snapshot, Point::new(0, 0)..Point::new(25, 0), 12, 0); + + assert_eq!(active_buffer_diagnostics.len(), 20); + assert!( + active_buffer_diagnostics + .iter() + .any(|diagnostic| diagnostic.message == "row 12") + ); + assert!( + active_buffer_diagnostics + .iter() + .all(|diagnostic| diagnostic.message != "row 0" && diagnostic.message != "row 24") + ); + + let text = (0..300) + .map(|row| format!("line {row} has some diagnostic context\n")) + .collect::(); + let buffer = cx.new(|cx| Buffer::local(&text, cx)); + + buffer.update(cx, |buffer, cx| { + let snapshot = buffer.snapshot(); + let diagnostics = DiagnosticSet::new( + vec![DiagnosticEntry { + range: text::PointUtf16::new(150, 0)..text::PointUtf16::new(150, 4), + diagnostic: Diagnostic { + severity: DiagnosticSeverity::ERROR, + message: "long snippet".to_string(), + group_id: 1, + is_primary: true, + source_kind: language::DiagnosticSourceKind::Pushed, + ..Diagnostic::default() + }, + }], + &snapshot, + ); + buffer.update_diagnostics(LanguageServerId(0), diagnostics, cx); + }); + + let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot()); + let active_buffer_diagnostics = zeta::active_buffer_diagnostics( + &snapshot, + Point::new(100, 0)..Point::new(200, 0), + 150, + 2000, + ); + + assert_eq!(active_buffer_diagnostics.len(), 1); + assert!(active_buffer_diagnostics[0].snippet.len() <= 512 * 3 + 2); + assert!(active_buffer_diagnostics[0].snippet.len() < text.len()); +} + // Generate a model response that would apply the given diff to the active file. fn model_response(request: &PredictEditsV3Request, diff_to_apply: &str) -> PredictEditsV3Response { let editable_range = diff --git a/crates/edit_prediction/src/zeta.rs b/crates/edit_prediction/src/zeta.rs index a5637ca3ce..39dd897502 100644 --- a/crates/edit_prediction/src/zeta.rs +++ b/crates/edit_prediction/src/zeta.rs @@ -495,14 +495,33 @@ fn handle_api_response( } } +const ACTIVE_BUFFER_DIAGNOSTIC_ADDITIONAL_CONTEXT_TOKEN_COUNT: usize = 100; +const MAX_ACTIVE_BUFFER_DIAGNOSTICS_TO_COLLECT: usize = 20; +const MAX_ACTIVE_BUFFER_DIAGNOSTIC_SNIPPET_TOKENS_TO_COLLECT: usize = 512; + pub(crate) fn active_buffer_diagnostics( snapshot: &language::BufferSnapshot, diagnostic_search_range: Range, + cursor_row: u32, additional_context_token_count: usize, ) -> Vec { - snapshot + let mut diagnostics = snapshot .diagnostics_in_range::(diagnostic_search_range, false) + .collect::>(); + diagnostics.sort_by_key(|entry| { + cursor_row.abs_diff(entry.range.start.row) + cursor_row.abs_diff(entry.range.end.row) + }); + + diagnostics + .into_iter() .map(|entry| { + let diagnostic_point_range = entry.range.clone(); + let snippet_point_range = cursor_excerpt::expand_context_syntactically_then_linewise( + snapshot, + diagnostic_point_range.clone(), + additional_context_token_count, + ); + let severity = match entry.diagnostic.severity { DiagnosticSeverity::ERROR => Some(1), DiagnosticSeverity::WARNING => Some(2), @@ -510,27 +529,52 @@ pub(crate) fn active_buffer_diagnostics( DiagnosticSeverity::HINT => Some(4), _ => None, }; - let diagnostic_point_range = entry.range.clone(); - let snippet_point_range = cursor_excerpt::expand_context_syntactically_then_linewise( - snapshot, - diagnostic_point_range.clone(), - additional_context_token_count, - ); - let snippet = snapshot - .text_for_range(snippet_point_range.clone()) - .collect::(); - let snippet_start_offset = snippet_point_range.start.to_offset(snapshot); - let diagnostic_offset_range = diagnostic_point_range.to_offset(snapshot); - zeta_prompt::ActiveBufferDiagnostic { + ( severity, - message: entry.diagnostic.message.clone(), - snippet, - snippet_buffer_row_range: diagnostic_point_range.start.row - ..diagnostic_point_range.end.row, - diagnostic_range_in_snippet: diagnostic_offset_range.start - snippet_start_offset - ..diagnostic_offset_range.end - snippet_start_offset, - } + entry.diagnostic.message.clone(), + diagnostic_point_range, + snippet_point_range, + ) }) + .take(MAX_ACTIVE_BUFFER_DIAGNOSTICS_TO_COLLECT) + .map( + |(severity, message, diagnostic_point_range, snippet_point_range)| { + let (snippet, diagnostic_range_in_snippet) = if snippet_point_range.start + == Point::new(0, 0) + && snippet_point_range.end == snapshot.max_point() + { + (String::new(), 0..0) + } else { + let snippet = snapshot + .text_for_range(snippet_point_range.clone()) + .collect::(); + let snippet = zeta_prompt::clamp_text_to_token_count( + &snippet, + MAX_ACTIVE_BUFFER_DIAGNOSTIC_SNIPPET_TOKENS_TO_COLLECT, + ) + .to_string(); + let snippet_start_offset = snippet_point_range.start.to_offset(snapshot); + let diagnostic_offset_range = diagnostic_point_range.to_offset(snapshot); + let diagnostic_range_start = diagnostic_offset_range + .start + .saturating_sub(snippet_start_offset) + .min(snippet.len()); + let diagnostic_range_end = diagnostic_offset_range + .end + .saturating_sub(snippet_start_offset) + .min(snippet.len()); + (snippet, diagnostic_range_start..diagnostic_range_end) + }; + zeta_prompt::ActiveBufferDiagnostic { + severity, + message, + snippet, + snippet_buffer_row_range: diagnostic_point_range.start.row + ..diagnostic_point_range.end.row, + diagnostic_range_in_snippet, + } + }, + ) .collect() } @@ -559,8 +603,12 @@ pub fn zeta2_prompt_input( &syntax_ranges, ); - let active_buffer_diagnostics = - active_buffer_diagnostics(snapshot, diagnostic_search_range, 100); + let active_buffer_diagnostics = active_buffer_diagnostics( + snapshot, + diagnostic_search_range, + snapshot.offset_to_point(cursor_offset).row, + ACTIVE_BUFFER_DIAGNOSTIC_ADDITIONAL_CONTEXT_TOKEN_COUNT, + ); let prompt_input = zeta_prompt::ZetaPromptInput { cursor_path: excerpt_path, diff --git a/crates/zeta_prompt/src/zeta_prompt.rs b/crates/zeta_prompt/src/zeta_prompt.rs index 7bee9fc9b0..747aa4e0de 100644 --- a/crates/zeta_prompt/src/zeta_prompt.rs +++ b/crates/zeta_prompt/src/zeta_prompt.rs @@ -870,16 +870,20 @@ fn format_active_buffer_diagnostics_with_budget( let diagnostic = &diagnostics[diagnostic_index]; let snippet = clamp_text_to_token_count(&diagnostic.snippet, 256); - let diagnostic_section = format!( - "*{}*:\n```\n{}{}\n```\n", - diagnostic.message, - snippet, - if snippet.len() < diagnostic.snippet.len() { - "..." - } else { - "" - } - ); + let diagnostic_section = if snippet.is_empty() { + format!("*{}*\n", diagnostic.message) + } else { + format!( + "*{}*:\n```\n{}{}\n```\n", + diagnostic.message, + snippet, + if snippet.len() < diagnostic.snippet.len() { + "..." + } else { + "" + } + ) + }; let diagnostic_tokens = estimate_tokens(diagnostic_section.len()); if used_tokens + diagnostic_tokens > budget { break; @@ -5299,13 +5303,22 @@ mod tests { vec![], vec![make_related_file("related.rs", "fn helper() {}\n")], ); - input.active_buffer_diagnostics = vec![ActiveBufferDiagnostic { - severity: Some(1), - message: "missing semicolon".to_string(), - snippet: "let value = 1".to_string(), - snippet_buffer_row_range: 1..2, - diagnostic_range_in_snippet: 12..13, - }]; + input.active_buffer_diagnostics = vec![ + ActiveBufferDiagnostic { + severity: Some(1), + message: "missing semicolon".to_string(), + snippet: "let value = 1".to_string(), + snippet_buffer_row_range: 1..2, + diagnostic_range_in_snippet: 12..13, + }, + ActiveBufferDiagnostic { + severity: Some(2), + message: "file-level warning".to_string(), + snippet: String::new(), + snippet_buffer_row_range: 0..0, + diagnostic_range_in_snippet: 0..0, + }, + ]; let prompt = format_prompt_with_budget_for_format(&input, ZetaFormat::V0420Diagnostics, 10000) @@ -5321,6 +5334,7 @@ mod tests { ``` let value = 1 ``` + *file-level warning* related.rs fn helper() {}