From d8e04a57e07b9b6490d936eb17aa296f970a012d Mon Sep 17 00:00:00 2001 From: Ben Kunkle Date: Sat, 20 Jun 2026 15:57:43 -0500 Subject: [PATCH] ep: Limit diagnostic message token counts (#59644) # Objective Prevent unbounded ep request sizes ## Solution Use same token limiting logic used for snippets, on diagnostic messages ## Testing - Did you test these changes? If so, how? - Are there any parts that need more testing? - How can other people (reviewers) test your changes? Is there anything specific they need to know? - If relevant, what platforms did you test these changes on, and are there any important ones you can't test? ## 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 adheres to Zed's UI standards ([UX/UI](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist) and [icon](https://github.com/zed-industries/zed/blob/main/crates/icons/README.md) guidelines) - [x] Tests cover the new/changed behavior - [x] Performance impact has been considered and is acceptable ## Showcase > This section is optional. If this PR does not include a visual change or does not add a new user-facing feature, you can delete this section. - Help others understand the result of this PR by showcasing your awesome work! - If this PR includes a visual change, consider adding a screenshot, GIF, or video - A before/after comparison is very useful for changes to existing features! While a showcase should aim to be brief and digestible, you can use a toggleable section to save space on longer showcases:
Click to view showcase My super cool demos here
--- Release Notes: - N/A or Added/Fixed/Improved ... --- crates/edit_prediction/src/zeta.rs | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/crates/edit_prediction/src/zeta.rs b/crates/edit_prediction/src/zeta.rs index 172726a1d5..013929b391 100644 --- a/crates/edit_prediction/src/zeta.rs +++ b/crates/edit_prediction/src/zeta.rs @@ -567,6 +567,7 @@ 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_MESSAGE_TOKENS_TO_COLLECT: usize = 512; const MAX_ACTIVE_BUFFER_DIAGNOSTIC_SNIPPET_TOKENS_TO_COLLECT: usize = 512; pub(crate) fn active_buffer_diagnostics( @@ -601,7 +602,11 @@ pub(crate) fn active_buffer_diagnostics( }; ( severity, - entry.diagnostic.message.clone(), + zeta_prompt::clamp_text_to_token_count( + &entry.diagnostic.message, + MAX_ACTIVE_BUFFER_DIAGNOSTIC_MESSAGE_TOKENS_TO_COLLECT, + ) + .to_string(), diagnostic_point_range, snippet_point_range, )