vndangkhoa/zed
1
0
Fork 0
mirror of https://github.com/zed-industries/zed.git synced 2026-05-31 19:05:00 +07:00
Code Issues Projects Releases Packages Wiki Activity Actions 6

agent: Batch streaming edit operations (#58037)

Browse source 
## Summary

While profiling agent sessions that make a lot of `edit_file`
operations, I noticed the LSP `textDocument/didChange` handler firing
excessively. Looking into this, I found out that the streaming edit
pipeline was applying each `CharOperation` from `StreamingDiff` as its
own `buffer.edit` transaction, and every transaction emits a
`BufferEvent::Edited` event. Each event can trigger several other
expensive events depending on whether the buffer is being rendered in an
editor or is registered with a language.

For example, there are `didChange` LSP events, the editor's on edit work
(matching brackets, bracket colorization, code actions, outline), and
more. A single `edit_file` could trigger hundreds of these at the higher
end in a single synchronous app update, which would block the foreground
thread for a bit and cause Zed to drop frames.

I fixed this by collecting all of a chunk's `CharOperation`s and
applying them in one `buffer.edit` call, so only a single
`BufferEvent::Edited` event gets emitted. This is safe because
operations are non overlapping by design of streaming diff (the edit
cursor only advances).

## Why this wasn't caught earlier

The cost only fully appears when a buffer is both registered with a
language server and rendered in an editor. Without that, most of the per
transaction observers never run, so the existing `edit_file_tool`
benchmark (which ran the tool against a bare buffer) didn't surface it.
I reworked the benchmark to open the edited buffer in an editor view,
register a fake language server with per edit diagnostics, and lay out a
frame, so it exercises the same cascade as the real editor. I also added
a larger fixture.

## Results

Measured with the `release-fast` profile on the reworked benchmark:

| Fixture | Initial file | Before | After | Improvement |
| --- | --- | --- | --- | --- |
| `tiny_function_rewrite` | 1.4 KB | 31.1 ms | 12.1 ms | −61% |
| `small_function_rewrite` | 3.0 KB | 42.4 ms | 19.3 ms | −55% |
| `medium_many_small_changes` | 4.6 KB | 309.2 ms | 151.5 ms | −51% |
| `medium_insertions` | 4.6 KB | 171.8 ms | 126.1 ms | −27% |
| `large_multi_edit` | 44 KB | 9,549 ms | 919 ms | −90% |

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
- [x] Tests cover the new/changed behavior
- [x] Performance impact has been considered and is acceptable

Release Notes:
- Improved agent's edit file tool performance
This commit is contained in:
Anthony Eid 2026-05-29 05:49:55 -04:00 committed by GitHub
parent ef5da3ccc2
commit 6bca2136a1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 336 additions and 111 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
Cargo.lock generated
View file

@ -234,6 +234,7 @@ dependencies = [
"agent_settings",
"agent_skills",
"anyhow",
"assets",
"async-channel 2.5.0",
"async-io",
"chrono",
@ -290,6 +291,7 @@ dependencies = [
"tempfile",
"text",
"theme",
"theme_settings",
"thiserror 2.0.17",
"ui",
"unindent",

2
crates/agent/Cargo.toml
View file

@ -78,6 +78,7 @@ zed_env_vars.workspace = true
zstd.workspace = true
[dev-dependencies]
assets.workspace = true
async-io.workspace = true
agent_servers = { workspace = true, "features" = ["test-support"] }
client = { workspace = true, "features" = ["test-support"] }
@ -103,6 +104,7 @@ reqwest_client.workspace = true
settings = { workspace = true, "features" = ["test-support"] }
theme = { workspace = true, "features" = ["test-support"] }
theme_settings.workspace = true
unindent = { workspace = true }

414
crates/agent/benches/edit_file_tool.rs
View file

@ -1,4 +1,5 @@
use std::{
any::Any,
future::Future,
path::Path,
sync::Arc,
@ -14,26 +15,40 @@ use agent_settings::{AgentSettings, ToolRules};
use criterion::{
BatchSize, BenchmarkId, Criterion, Throughput, black_box, criterion_group, criterion_main,
};
use futures::{pin_mut, task::noop_waker};
use gpui::{AppContext as _, BackgroundExecutor, Entity, TestAppContext, UpdateGlobal as _};
use editor::{Editor, EditorStyle};
use futures::{StreamExt as _, pin_mut, task::noop_waker};
use gpui::{
AnyWindowHandle, AppContext as _, BackgroundExecutor, Entity, Focusable as _, TestAppContext,
UpdateGlobal as _,
};
use language::{FakeLspAdapter, rust_lang};
use language_model::fake_provider::FakeLanguageModel;
use project::{FakeFs, Project};
use prompt_store::ProjectContext;
use rand::{Rng as _, SeedableRng as _, rngs::StdRng};
use serde_json::{Value, json};
use settings::{Settings as _, SettingsStore};
use ui::IntoElement as _;
const SEED: u64 = 0x5EED_5EED;
const OLD_TEXT_CHUNK_SIZE: usize = 512;
const NEW_TEXT_CHUNK_SIZE: usize = 512;
const FILE_PROJECT_PATH: &str = "root/src/workspace_snapshot.rs";
const FILE_ABS_PATH: &str = "/root/src/workspace_snapshot.rs";
#[derive(Clone)]
struct EditOp {
old_text: String,
new_text: String,
}
#[derive(Clone)]
struct EditFixture {
name: &'static str,
old_file_text: String,
expected_file_text: String,
old_text: String,
new_text: String,
edits: Vec<EditOp>,
}
struct BenchmarkHarness {
@ -43,6 +58,12 @@ struct BenchmarkHarness {
partial_payloads: Vec<Value>,
final_payload: Value,
expected_file_text: String,
editor: Option<Entity<Editor>>,
window: Option<AnyWindowHandle>,
// Keeps the LSP buffer-registration handle and the fake language server alive
// for the lifetime of the benchmark so `didChange`/diagnostics keep flowing
// while edits are applied.
keep_alive: Vec<Box<dyn Any>>,
}
impl Drop for BenchmarkHarness {
@ -50,19 +71,18 @@ impl Drop for BenchmarkHarness {
// Release our handles to the entities first.
self.edit_tool.take();
self.thread.take();
self.editor.take();
self.keep_alive.clear();
if let Some(cx) = self.cx.take() {
// `ActionLog` holds buffers strongly via `tracked_buffers`, and spawns a background
// diff-maintenance task that also captures a strong `Entity<Buffer>`. Releasing the
// last handle to the action log only marks its entity for deferred release; the
// entity's value (and the buffer handles inside) is not actually dropped until
// `flush_effects` runs `release_dropped_entities`. Even then, the cancelled task's
// captured handle does not drop until the executor pumps the cancellation through.
//
// Without this two-step teardown, GPUI's test leak detector panics on
// `TestAppContext` drop because the buffer still appears alive. See
// `ActionLog::track_buffer_internal` and `LeakDetector::drop` in
// `crates/gpui/src/app/entity_map.rs`.
if let Some(mut cx) = self.cx.take() {
// Close the editor window so the editor entity and the buffer handles
// it holds are released, then pump the executor so cancelled editor /
// action-log background tasks drop their captured handles before the
// leak detector runs on `TestAppContext` drop.
if let Some(window) = self.window.take() {
cx.update_window(window, |_, window, _| window.remove_window())
.ok();
}
cx.update(|_| {});
cx.executor().run_until_parked();
cx.quit();
@ -76,9 +96,10 @@ fn edit_file_tool_streaming(c: &mut Criterion) {
group.sample_size(10);
for fixture in fixtures {
group.throughput(Throughput::Bytes(fixture.new_text.len() as u64));
let new_bytes: usize = fixture.edits.iter().map(|edit| edit.new_text.len()).sum();
group.throughput(Throughput::Bytes(new_bytes as u64));
group.bench_with_input(
BenchmarkId::new(fixture.name, fixture.old_text.len()),
BenchmarkId::new(fixture.name, fixture.old_file_text.len()),
&fixture,
|bench, fixture| {
bench.iter_batched(
@ -107,26 +128,168 @@ fn edit_file_tool_streaming(c: &mut Criterion) {
fn setup_harness(fixture: EditFixture) -> BenchmarkHarness {
let mut cx = init_context();
let executor = cx.executor();
let (edit_tool, thread) = block_on_executor(
let parts = block_on_executor(
&executor,
setup_edit_tool(&mut cx, fixture.old_file_text.clone()),
setup_editor_and_tool(&mut cx, fixture.old_file_text.clone()),
);
let partial_payloads = streamed_partial_payloads(&fixture.old_text, &fixture.new_text);
// Let the LSP handshake, initial parse, and first layout settle before timing.
cx.executor().run_until_parked();
let partial_payloads = streamed_partial_payloads(&fixture.edits);
let final_payload = json!({
"path": "root/src/workspace_snapshot.rs",
"edits": [{
"old_text": fixture.old_text,
"new_text": fixture.new_text,
}],
"path": FILE_PROJECT_PATH,
"edits": fixture
.edits
.iter()
.map(|edit| json!({ "old_text": edit.old_text, "new_text": edit.new_text }))
.collect::<Vec<_>>(),
});
BenchmarkHarness {
cx: Some(cx),
edit_tool: Some(edit_tool),
thread: Some(thread),
edit_tool: Some(parts.edit_tool),
thread: Some(parts.thread),
partial_payloads,
final_payload,
expected_file_text: fixture.expected_file_text,
editor: Some(parts.editor),
window: Some(parts.window),
keep_alive: parts.keep_alive,
}
}
struct HarnessParts {
edit_tool: Arc<EditFileTool>,
thread: Entity<Thread>,
editor: Entity<Editor>,
window: AnyWindowHandle,
keep_alive: Vec<Box<dyn Any>>,
}
/// Builds a project + edit tool, opens the target buffer in an editor view inside
/// a window, and attaches a fake Rust language server. This mirrors the real app:
/// the edited file is open in a pane with a language server, so each buffer edit
/// drives the editor's observer cascade (matching brackets, code actions, outline,
/// bracket colorization), a tree-sitter reparse, and an LSP `didChange` +
/// diagnostics round-trip — the costs that dominate a real agent edit.
async fn setup_editor_and_tool(cx: &mut TestAppContext, file_text: String) -> HarnessParts {
let fs = FakeFs::new(cx.executor());
fs.insert_tree(
"/root",
json!({
"src": {
"workspace_snapshot.rs": file_text,
},
}),
)
.await;
let project = Project::test(fs, [Path::new("/root")], cx).await;
let language_registry = project.read_with(cx, |project, _cx| project.languages().clone());
language_registry.add(rust_lang());
let mut fake_servers = language_registry.register_fake_lsp(
"Rust",
FakeLspAdapter {
capabilities: lsp::ServerCapabilities {
text_document_sync: Some(lsp::TextDocumentSyncCapability::Kind(
lsp::TextDocumentSyncKind::INCREMENTAL,
)),
..Default::default()
},
..Default::default()
},
);
let context_server_registry =
cx.new(|cx| ContextServerRegistry::new(project.read(cx).context_server_store(), cx));
let model = Arc::new(FakeLanguageModel::default());
let thread = cx.new(|cx| {
Thread::new(
project.clone(),
cx.new(|_cx| ProjectContext::default()),
context_server_registry,
Templates::new(),
Some(model),
cx,
)
});
let action_log: Entity<ActionLog> =
thread.read_with(cx, |thread, _cx| thread.action_log().clone());
let edit_tool = Arc::new(EditFileTool::new(
project.clone(),
thread.downgrade(),
action_log,
language_registry,
));
// Open the same buffer the tool will edit and register it with the language
// servers so edits produce `didChange` notifications.
let buffer = project
.update(cx, |project, cx| {
project.open_local_buffer(FILE_ABS_PATH, cx)
})
.await
.expect("failed to open buffer");
let lsp_handle = project.update(cx, |project, cx| {
project.register_buffer_with_language_servers(&buffer, cx)
});
let fake_server = fake_servers
.next()
.await
.expect("fake language server should start");
// Publish diagnostics on every edit, mirroring a real server reacting to
// `didChange`, so the editor's diagnostics path runs per edit.
let server = fake_server.clone();
fake_server.handle_notification::<lsp::notification::DidChangeTextDocument, _>(
move |params, _cx| {
server.notify::<lsp::notification::PublishDiagnostics>(lsp::PublishDiagnosticsParams {
uri: params.text_document.uri.clone(),
version: Some(params.text_document.version),
diagnostics: vec![lsp::Diagnostic {
range: lsp::Range::new(lsp::Position::new(0, 0), lsp::Position::new(0, 1)),
severity: Some(lsp::DiagnosticSeverity::WARNING),
message: "bench diagnostic".to_string(),
..Default::default()
}],
});
},
);
// Attach an editor view in a window and lay it out once so the viewport-gated
// observers (bracket colorization, selection highlights) have a visible range.
let window = cx.add_window(|window, cx| {
let mut editor = Editor::for_buffer(buffer.clone(), Some(project.clone()), window, cx);
editor.set_style(EditorStyle::default(), window, cx);
window.focus(&editor.focus_handle(cx), cx);
editor
});
let editor = window.root(cx).expect("window should have an editor root");
let window: AnyWindowHandle = window.into();
// Lay out and paint a real frame so the editor establishes a viewport (this
// is what makes the viewport-gated observers like bracket colorization run).
{
let mut visual_cx = gpui::VisualTestContext::from_window(window, &*cx);
visual_cx.draw(
gpui::point(gpui::px(0.0), gpui::px(0.0)),
gpui::size(gpui::px(1024.0), gpui::px(768.0)),
|_, _| editor.clone().into_any_element(),
);
}
let keep_alive: Vec<Box<dyn Any>> = vec![
Box::new(lsp_handle),
Box::new(fake_server),
Box::new(fake_servers),
Box::new(buffer),
];
HarnessParts {
edit_tool,
thread,
editor,
window,
keep_alive,
}
}
@ -135,6 +298,9 @@ fn init_context() -> TestAppContext {
cx.update(|cx| {
let settings_store = SettingsStore::test(cx);
cx.set_global(settings_store);
assets::Assets.load_test_fonts(cx);
theme_settings::init(theme::LoadThemes::JustBase, cx);
editor::init(cx);
SettingsStore::update_global(cx, |store: &mut SettingsStore, cx| {
store.update_user_settings(cx, |settings| {
settings
@ -142,6 +308,7 @@ fn init_context() -> TestAppContext {
.all_languages
.defaults
.ensure_final_newline_on_save = Some(false);
settings.project.all_languages.defaults.colorize_brackets = Some(true);
});
});
@ -161,48 +328,6 @@ fn init_context() -> TestAppContext {
cx
}
async fn setup_edit_tool(
cx: &mut TestAppContext,
file_text: String,
) -> (Arc<EditFileTool>, Entity<Thread>) {
let fs = FakeFs::new(cx.executor());
fs.insert_tree(
"/root",
json!({
"src": {
"workspace_snapshot.rs": file_text,
},
}),
)
.await;
let project = Project::test(fs, [Path::new("/root")], cx).await;
let language_registry = project.read_with(cx, |project, _cx| project.languages().clone());
let context_server_registry =
cx.new(|cx| ContextServerRegistry::new(project.read(cx).context_server_store(), cx));
let model = Arc::new(FakeLanguageModel::default());
let thread = cx.new(|cx| {
Thread::new(
project.clone(),
cx.new(|_cx| ProjectContext::default()),
context_server_registry,
Templates::new(),
Some(model),
cx,
)
});
let action_log: Entity<ActionLog> =
thread.read_with(cx, |thread, _cx| thread.action_log().clone());
let edit_tool = Arc::new(EditFileTool::new(
project,
thread.downgrade(),
action_log,
language_registry,
));
(edit_tool, thread)
}
fn run_streamed_edit(harness: &mut BenchmarkHarness) -> EditFileToolOutput {
let (mut sender, input): (_, ToolInput<EditFileToolInput>) = ToolInput::test();
for payload in &harness.partial_payloads {
@ -247,33 +372,36 @@ fn block_on_executor<R>(executor: &BackgroundExecutor, future: impl Future<Outpu
panic!("future did not complete while running edit_file_tool benchmark");
}
fn streamed_partial_payloads(old_text: &str, new_text: &str) -> Vec<Value> {
let path = "root/src/workspace_snapshot.rs";
let mut payloads = Vec::new();
/// Builds the streamed partial payloads for a (possibly multi-edit) session,
/// mirroring how the agent reveals one edit at a time: earlier edits stay
/// complete in the array while the current edit streams its `old_text` then its
/// `new_text` in chunks.
fn streamed_partial_payloads(edits: &[EditOp]) -> Vec<Value> {
let path = FILE_PROJECT_PATH;
let mut payloads = vec![json!({ "path": path }), json!({ "path": path })];
payloads.push(json!({ "path": path }));
payloads.push(json!({ "path": path }));
for index in 0..edits.len() {
let completed: Vec<Value> = edits[..index]
.iter()
.map(|edit| json!({ "old_text": edit.old_text, "new_text": edit.new_text }))
.collect();
let edit = &edits[index];
for old_end in chunk_ends(old_text, OLD_TEXT_CHUNK_SIZE) {
payloads.push(json!({
"path": path,
"edits": [{ "old_text": &old_text[..old_end] }],
}));
}
for old_end in chunk_ends(&edit.old_text, OLD_TEXT_CHUNK_SIZE) {
let mut arr = completed.clone();
arr.push(json!({ "old_text": &edit.old_text[..old_end] }));
payloads.push(json!({ "path": path, "edits": arr }));
}
payloads.push(json!({
"path": path,
"edits": [{ "old_text": old_text, "new_text": "" }],
}));
let mut arr = completed.clone();
arr.push(json!({ "old_text": edit.old_text, "new_text": "" }));
payloads.push(json!({ "path": path, "edits": arr }));
for new_end in chunk_ends(new_text, NEW_TEXT_CHUNK_SIZE) {
payloads.push(json!({
"path": path,
"edits": [{
"old_text": old_text,
"new_text": &new_text[..new_end],
}],
}));
for new_end in chunk_ends(&edit.new_text, NEW_TEXT_CHUNK_SIZE) {
let mut arr = completed.clone();
arr.push(json!({ "old_text": edit.old_text, "new_text": &edit.new_text[..new_end] }));
payloads.push(json!({ "path": path, "edits": arr }));
}
}
payloads
@ -326,6 +454,7 @@ fn fixtures() -> Vec<EditFixture> {
EditPattern::InsertHelperBlocks { every_nth_line: 9 },
SEED + 3,
),
make_large_multi_edit_fixture("large_multi_edit", 80, 16, SEED + 4),
]
}
@ -375,11 +504,106 @@ fn make_fixture(
name,
old_file_text,
expected_file_text,
old_text,
new_text,
edits: vec![EditOp { old_text, new_text }],
}
}
fn make_large_multi_edit_fixture(
name: &'static str,
function_count: usize,
edit_count: usize,
seed: u64,
) -> EditFixture {
const HEADER_LINES: usize = 10;
const FUNCTION_LINES: usize = 12;
const FUNCTION_BODY_LINES: usize = 11;
let mut rng = StdRng::seed_from_u64(seed);
let old_lines = random_rust_module(&mut rng, function_count);
let old_file_text = old_lines.join("\n");
let step = (function_count / edit_count).max(1);
let mut picks: Vec<usize> = (0..edit_count)
.map(|k| (k * step).min(function_count - 1))
.collect();
picks.dedup();
let replacements: Vec<(usize, Vec<String>)> = picks
.iter()
.map(|&function_index| {
(
function_index,
large_function_lines(&mut rng, function_index),
)
})
.collect();
let edits = replacements
.iter()
.map(|(function_index, new_function)| {
let start = HEADER_LINES + function_index * FUNCTION_LINES;
let end = start + FUNCTION_BODY_LINES;
EditOp {
old_text: old_lines[start..end].join("\n"),
new_text: new_function.join("\n"),
}
})
.collect();
let mut new_lines = old_lines;
for (function_index, new_function) in replacements.iter().rev() {
let start = HEADER_LINES + function_index * FUNCTION_LINES;
let end = start + FUNCTION_BODY_LINES;
new_lines.splice(start..end, new_function.iter().cloned());
}
let expected_file_text = new_lines.join("\n");
EditFixture {
name,
old_file_text,
expected_file_text,
edits,
}
}
fn large_function_lines(rng: &mut StdRng, index: usize) -> Vec<String> {
let function_name = identifier(rng, index + 40_000);
let argument_name = identifier(rng, index + 41_000);
let mut lines = vec![
format!(
" pub fn {function_name}(&mut self, {argument_name}: usize) -> Result<usize> {{"
),
format!(" let mut accumulator = {argument_name};"),
];
let body_lines = rng.random_range(30..42);
for body_index in 0..body_lines {
let local_name = identifier(rng, index + 50_000 + body_index);
let multiplier = rng.random_range(2..19);
let offset = rng.random_range(1..256);
match body_index % 4 {
0 => lines.push(format!(
" let {local_name} = accumulator.saturating_mul({multiplier}).saturating_add({offset});"
)),
1 => lines.push(format!(
" accumulator = {local_name}.saturating_sub(self.version % {offset}.max(1));"
)),
2 => lines.push(format!(
" if {local_name} % {multiplier} == 0 {{ accumulator = accumulator.saturating_add({local_name}); }}"
)),
_ => lines.push(format!(
" self.buffers.insert(\"{local_name}\".to_string(), accumulator);"
)),
}
}
lines.push(" self.version = self.version.saturating_add(accumulator);".to_string());
lines.push(" Ok(accumulator)".to_string());
lines.push(" }".to_string());
lines
}
fn edit_range(lines: &[String], pattern: &EditPattern) -> std::ops::Range<usize> {
let mut range = match pattern {
EditPattern::LocalizedRewrite {

29
crates/agent/src/tools/edit_session.rs
View file

@ -620,21 +620,14 @@ impl EditPipeline {
log::debug!("new_text_chunk: done=true, final_text='{}'", final_text);
if !final_text.is_empty() {
let char_ops = streaming_diff.push_new(&final_text);
apply_char_operations(
&char_ops,
buffer,
&original_snapshot,
&mut edit_cursor,
&context.action_log,
cx,
);
}
let remaining_ops = streaming_diff.finish();
let mut char_ops = if final_text.is_empty() {
Vec::new()
} else {
streaming_diff.push_new(&final_text)
};
char_ops.extend(streaming_diff.finish());
apply_char_operations(
&remaining_ops,
&char_ops,
buffer,
&original_snapshot,
&mut edit_cursor,
@ -902,16 +895,17 @@ fn apply_char_operations(
action_log: &Entity<ActionLog>,
cx: &mut AsyncApp,
) {
let mut edits: Vec<_> = Vec::new();
for op in ops {
match op {
CharOperation::Insert { text } => {
let anchor = snapshot.anchor_after(*edit_cursor);
agent_edit_buffer(&buffer, [(anchor..anchor, text.as_str())], action_log, cx);
edits.push((anchor..anchor, text.as_str().into()));
}
CharOperation::Delete { bytes } => {
let delete_end = *edit_cursor + bytes;
let anchor_range = snapshot.anchor_range_inside(*edit_cursor..delete_end);
agent_edit_buffer(&buffer, [(anchor_range, "")], action_log, cx);
edits.push((anchor_range, Arc::<str>::from("")));
*edit_cursor = delete_end;
}
CharOperation::Keep { bytes } => {
@ -919,6 +913,9 @@ fn apply_char_operations(
}
}
}
if !edits.is_empty() {
agent_edit_buffer(buffer, edits, action_log, cx);
}
}
fn extract_match(