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 ...
2026-06-01 03:14:56 +07:00 · 2026-05-12 08:30:07 -04:00 · 2026-05-12 08:30:07 -04:00 · c9957c136b
commit c9957c136b
parent 7ab0ce6e68
3 changed files with 219 additions and 47 deletions
--- 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::<String>();
+    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::<Vec<_>>(),
+            &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::<String>();
+    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 =
--- a/crates/edit_prediction/src/zeta.rs
+++ b/crates/edit_prediction/src/zeta.rs
@ -495,14 +495,33 @@ fn handle_api_response<T>(
    }
 }

+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<Point>,
+    cursor_row: u32,
    additional_context_token_count: usize,
 ) -> Vec<zeta_prompt::ActiveBufferDiagnostic> {
-    snapshot
+    let mut diagnostics = snapshot
        .diagnostics_in_range::<Point, Point>(diagnostic_search_range, false)
+        .collect::<Vec<_>>();
+    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::<String>();
-            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::<String>();
+                    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,
--- 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*

                <filename>related.rs
                fn helper() {}