tenant.openagents/omega

1use std::{
2    any::Any,
3    future::Future,
4    path::Path,
5    sync::Arc,
6    task::{Context, Poll},
7};
8
9use action_log::ActionLog;
10use agent::{
11    AgentTool, ContextServerRegistry, EditFileTool, EditFileToolInput, EditFileToolOutput,
12    Templates, Thread, ToolCallEventStream, ToolInput,
13};
14use agent_settings::{AgentSettings, ToolRules};
15use benchmarks::bench_utils::{
16    RUST_FUNCTION_BODY_LINES, RUST_FUNCTION_LINES, RUST_MODULE_HEADER_LINES, random_rust_file,
17    rust_file_line_count, rust_identifier as identifier,
18};
19use criterion::{
20    BatchSize, BenchmarkId, Criterion, Throughput, black_box, criterion_group, criterion_main,
21};
22use editor::{Editor, EditorStyle};
23use futures::{StreamExt as _, pin_mut, task::noop_waker};
24use gpui::{
25    AnyWindowHandle, AppContext as _, BackgroundExecutor, Entity, Focusable as _, TestAppContext,
26    UpdateGlobal as _,
27};
28use language::{FakeLspAdapter, rust_lang};
29use language_model::fake_provider::FakeLanguageModel;
30use project::{FakeFs, Project};
31use prompt_store::ProjectContext;
32use rand::{Rng as _, SeedableRng as _, rngs::StdRng};
33use serde_json::{Value, json};
34use settings::{Settings as _, SettingsStore};
35use ui::IntoElement as _;
36
37const SEED: u64 = 0x5EED_5EED;
38const OLD_TEXT_CHUNK_SIZE: usize = 512;
39const NEW_TEXT_CHUNK_SIZE: usize = 512;
40
41const FILE_PROJECT_PATH: &str = "root/src/workspace_snapshot.rs";
42const FILE_ABS_PATH: &str = "/root/src/workspace_snapshot.rs";
43
44#[derive(Clone)]
45struct EditOp {
46    old_text: String,
47    new_text: String,
48}
49
50#[derive(Clone)]
51struct EditFixture {
52    name: &'static str,
53    old_file_text: String,
54    expected_file_text: String,
55    edits: Vec<EditOp>,
56}
57
58struct BenchmarkHarness {
59    cx: Option<TestAppContext>,
60    edit_tool: Option<Arc<EditFileTool>>,
61    thread: Option<Entity<Thread>>,
62    partial_payloads: Vec<Value>,
63    final_payload: Value,
64    expected_file_text: String,
65    editor: Option<Entity<Editor>>,
66    window: Option<AnyWindowHandle>,
67    // Keeps the LSP buffer-registration handle and the fake language server alive
68    // for the lifetime of the benchmark so `didChange`/diagnostics keep flowing
69    // while edits are applied.
70    keep_alive: Vec<Box<dyn Any>>,
71}
72
73impl Drop for BenchmarkHarness {
74    fn drop(&mut self) {
75        // Release our handles to the entities first.
76        self.edit_tool.take();
77        self.thread.take();
78        self.editor.take();
79        self.keep_alive.clear();
80
81        if let Some(mut cx) = self.cx.take() {
82            // Close the editor window so the editor entity and the buffer handles
83            // it holds are released, then pump the executor so cancelled editor /
84            // action-log background tasks drop their captured handles before the
85            // leak detector runs on `TestAppContext` drop.
86            if let Some(window) = self.window.take() {
87                cx.update_window(window, |_, window, _| window.remove_window())
88                    .ok();
89            }
90            cx.update(|_| {});
91            cx.executor().run_until_parked();
92            cx.quit();
93        }
94    }
95}
96
97fn edit_file_tool_streaming(c: &mut Criterion) {
98    let fixtures = fixtures();
99    let mut group = c.benchmark_group("edit_file_tool_streaming");
100    group.sample_size(10);
101
102    for fixture in fixtures {
103        let new_bytes: usize = fixture.edits.iter().map(|edit| edit.new_text.len()).sum();
104        group.throughput(Throughput::Bytes(new_bytes as u64));
105        group.bench_with_input(
106            BenchmarkId::new(fixture.name, fixture.old_file_text.len()),
107            &fixture,
108            |bench, fixture| {
109                bench.iter_batched(
110                    || setup_harness(fixture.clone()),
111                    |mut harness| {
112                        let output = run_streamed_edit(&mut harness);
113                        let EditFileToolOutput::Success { new_text, .. } = &output else {
114                            panic!("expected edit_file tool to succeed");
115                        };
116                        assert_eq!(new_text, &harness.expected_file_text);
117                        // Return the harness as part of the output so its teardown (which has
118                        // to pump the executor to release `Entity<Buffer>` handles captured by
119                        // background tasks) runs in criterion's drop phase after the timer has
120                        // stopped, rather than inside the timed region.
121                        (black_box(output), harness)
122                    },
123                    BatchSize::SmallInput,
124                );
125            },
126        );
127    }
128
129    group.finish();
130}
131
132fn setup_harness(fixture: EditFixture) -> BenchmarkHarness {
133    let mut cx = init_context();
134    let executor = cx.executor();
135    let parts = block_on_executor(
136        &executor,
137        setup_editor_and_tool(&mut cx, fixture.old_file_text.clone()),
138    );
139    // Let the LSP handshake, initial parse, and first layout settle before timing.
140    cx.executor().run_until_parked();
141
142    let partial_payloads = streamed_partial_payloads(&fixture.edits);
143    let final_payload = json!({
144        "path": FILE_PROJECT_PATH,
145        "edits": fixture
146            .edits
147            .iter()
148            .map(|edit| json!({ "old_text": edit.old_text, "new_text": edit.new_text }))
149            .collect::<Vec<_>>(),
150    });
151
152    BenchmarkHarness {
153        cx: Some(cx),
154        edit_tool: Some(parts.edit_tool),
155        thread: Some(parts.thread),
156        partial_payloads,
157        final_payload,
158        expected_file_text: fixture.expected_file_text,
159        editor: Some(parts.editor),
160        window: Some(parts.window),
161        keep_alive: parts.keep_alive,
162    }
163}
164
165struct HarnessParts {
166    edit_tool: Arc<EditFileTool>,
167    thread: Entity<Thread>,
168    editor: Entity<Editor>,
169    window: AnyWindowHandle,
170    keep_alive: Vec<Box<dyn Any>>,
171}
172
173/// Builds a project + edit tool, opens the target buffer in an editor view inside
174/// a window, and attaches a fake Rust language server. This mirrors the real app:
175/// the edited file is open in a pane with a language server, so each buffer edit
176/// drives the editor's observer cascade (matching brackets, code actions, outline,
177/// bracket colorization), a tree-sitter reparse, and an LSP `didChange` +
178/// diagnostics round-trip — the costs that dominate a real agent edit.
179async fn setup_editor_and_tool(cx: &mut TestAppContext, file_text: String) -> HarnessParts {
180    let fs = FakeFs::new(cx.executor());
181    fs.insert_tree(
182        "/root",
183        json!({
184            "src": {
185                "workspace_snapshot.rs": file_text,
186            },
187        }),
188    )
189    .await;
190
191    let project = Project::test(fs, [Path::new("/root")], cx).await;
192    let language_registry = project.read_with(cx, |project, _cx| project.languages().clone());
193    language_registry.add(rust_lang());
194    let mut fake_servers = language_registry.register_fake_lsp(
195        "Rust",
196        FakeLspAdapter {
197            capabilities: lsp::ServerCapabilities {
198                text_document_sync: Some(lsp::TextDocumentSyncCapability::Kind(
199                    lsp::TextDocumentSyncKind::INCREMENTAL,
200                )),
201                ..Default::default()
202            },
203            ..Default::default()
204        },
205    );
206
207    let context_server_registry =
208        cx.new(|cx| ContextServerRegistry::new(project.read(cx).context_server_store(), cx));
209    let model = Arc::new(FakeLanguageModel::default());
210    let thread = cx.new(|cx| {
211        Thread::new(
212            project.clone(),
213            cx.new(|_cx| ProjectContext::default()),
214            context_server_registry,
215            Templates::new(),
216            Some(model),
217            cx,
218        )
219    });
220    let action_log: Entity<ActionLog> =
221        thread.read_with(cx, |thread, _cx| thread.action_log().clone());
222    let edit_tool = Arc::new(EditFileTool::new(
223        project.clone(),
224        thread.downgrade(),
225        action_log,
226        language_registry,
227    ));
228
229    // Open the same buffer the tool will edit and register it with the language
230    // servers so edits produce `didChange` notifications.
231    let buffer = project
232        .update(cx, |project, cx| {
233            project.open_local_buffer(FILE_ABS_PATH, cx)
234        })
235        .await
236        .expect("failed to open buffer");
237    let lsp_handle = project.update(cx, |project, cx| {
238        project.register_buffer_with_language_servers(&buffer, cx)
239    });
240
241    let fake_server = fake_servers
242        .next()
243        .await
244        .expect("fake language server should start");
245    // Publish diagnostics on every edit, mirroring a real server reacting to
246    // `didChange`, so the editor's diagnostics path runs per edit.
247    let server = fake_server.clone();
248    fake_server.handle_notification::<lsp::notification::DidChangeTextDocument, _>(
249        move |params, _cx| {
250            server.notify::<lsp::notification::PublishDiagnostics>(lsp::PublishDiagnosticsParams {
251                uri: params.text_document.uri.clone(),
252                version: Some(params.text_document.version),
253                diagnostics: vec![lsp::Diagnostic {
254                    range: lsp::Range::new(lsp::Position::new(0, 0), lsp::Position::new(0, 1)),
255                    severity: Some(lsp::DiagnosticSeverity::WARNING),
256                    message: "bench diagnostic".to_string(),
257                    ..Default::default()
258                }],
259            });
260        },
261    );
262
263    // Attach an editor view in a window and lay it out once so the viewport-gated
264    // observers (bracket colorization, selection highlights) have a visible range.
265    let window = cx.add_window(|window, cx| {
266        let mut editor = Editor::for_buffer(buffer.clone(), Some(project.clone()), window, cx);
267        editor.set_style(EditorStyle::default(), window, cx);
268        window.focus(&editor.focus_handle(cx), cx);
269        editor
270    });
271    let editor = window.root(cx).expect("window should have an editor root");
272    let window: AnyWindowHandle = window.into();
273    // Lay out and paint a real frame so the editor establishes a viewport (this
274    // is what makes the viewport-gated observers like bracket colorization run).
275    {
276        let mut visual_cx = gpui::VisualTestContext::from_window(window, &*cx);
277        visual_cx.draw(
278            gpui::point(gpui::px(0.0), gpui::px(0.0)),
279            gpui::size(gpui::px(1024.0), gpui::px(768.0)),
280            |_, _| editor.clone().into_any_element(),
281        );
282    }
283
284    let keep_alive: Vec<Box<dyn Any>> = vec![
285        Box::new(lsp_handle),
286        Box::new(fake_server),
287        Box::new(fake_servers),
288        Box::new(buffer),
289    ];
290
291    HarnessParts {
292        edit_tool,
293        thread,
294        editor,
295        window,
296        keep_alive,
297    }
298}
299
300fn init_context() -> TestAppContext {
301    let cx = TestAppContext::single();
302    cx.update(|cx| {
303        let settings_store = SettingsStore::test(cx);
304        cx.set_global(settings_store);
305        assets::Assets.load_test_fonts(cx);
306        theme_settings::init(theme::LoadThemes::JustBase, cx);
307        editor::init(cx);
308        SettingsStore::update_global(cx, |store: &mut SettingsStore, cx| {
309            store.update_user_settings(cx, |settings| {
310                settings
311                    .project
312                    .all_languages
313                    .defaults
314                    .ensure_final_newline_on_save = Some(false);
315                settings.project.all_languages.defaults.colorize_brackets = Some(true);
316            });
317        });
318
319        let mut agent_settings = AgentSettings::get_global(cx).clone();
320        agent_settings.tool_permissions.tools.insert(
321            EditFileTool::NAME.into(),
322            ToolRules {
323                default: Some(settings::ToolPermissionMode::Allow),
324                always_allow: vec![],
325                always_deny: vec![],
326                always_confirm: vec![],
327                invalid_patterns: vec![],
328            },
329        );
330        AgentSettings::override_global(agent_settings, cx);
331    });
332    cx
333}
334
335fn run_streamed_edit(harness: &mut BenchmarkHarness) -> EditFileToolOutput {
336    let (mut sender, input): (_, ToolInput<EditFileToolInput>) = ToolInput::test();
337    for payload in &harness.partial_payloads {
338        sender.send_partial(payload.clone());
339    }
340    sender.send_full(harness.final_payload.clone());
341
342    let (event_stream, _event_rx) = ToolCallEventStream::test();
343    let cx = harness
344        .cx
345        .as_ref()
346        .expect("benchmark harness should have a cx");
347    let task = cx.update(|cx| {
348        harness
349            .edit_tool
350            .as_ref()
351            .expect("benchmark harness should have an edit tool")
352            .clone()
353            .run(input, event_stream, cx)
354    });
355
356    let executor = harness
357        .cx
358        .as_ref()
359        .expect("benchmark harness should have a cx")
360        .executor();
361    block_on_executor(&executor, task).unwrap()
362}
363
364fn block_on_executor<R>(executor: &BackgroundExecutor, future: impl Future<Output = R>) -> R {
365    pin_mut!(future);
366    let waker = noop_waker();
367    let mut task_context = Context::from_waker(&waker);
368
369    for _ in 0..10_000 {
370        if let Poll::Ready(output) = future.as_mut().poll(&mut task_context) {
371            return output;
372        }
373        executor.run_until_parked();
374    }
375
376    panic!("future did not complete while running edit_file_tool benchmark");
377}
378
379/// Builds the streamed partial payloads for a (possibly multi-edit) session,
380/// mirroring how the agent reveals one edit at a time: earlier edits stay
381/// complete in the array while the current edit streams its `old_text` then its
382/// `new_text` in chunks.
383fn streamed_partial_payloads(edits: &[EditOp]) -> Vec<Value> {
384    let path = FILE_PROJECT_PATH;
385    let mut payloads = vec![json!({ "path": path }), json!({ "path": path })];
386
387    for index in 0..edits.len() {
388        let completed: Vec<Value> = edits[..index]
389            .iter()
390            .map(|edit| json!({ "old_text": edit.old_text, "new_text": edit.new_text }))
391            .collect();
392        let edit = &edits[index];
393
394        for old_end in chunk_ends(&edit.old_text, OLD_TEXT_CHUNK_SIZE) {
395            let mut arr = completed.clone();
396            arr.push(json!({ "old_text": &edit.old_text[..old_end] }));
397            payloads.push(json!({ "path": path, "edits": arr }));
398        }
399
400        let mut arr = completed.clone();
401        arr.push(json!({ "old_text": edit.old_text, "new_text": "" }));
402        payloads.push(json!({ "path": path, "edits": arr }));
403
404        for new_end in chunk_ends(&edit.new_text, NEW_TEXT_CHUNK_SIZE) {
405            let mut arr = completed.clone();
406            arr.push(json!({ "old_text": edit.old_text, "new_text": &edit.new_text[..new_end] }));
407            payloads.push(json!({ "path": path, "edits": arr }));
408        }
409    }
410
411    payloads
412}
413
414fn chunk_ends(text: &str, chunk_size: usize) -> impl Iterator<Item = usize> + '_ {
415    let mut end = 0;
416    std::iter::from_fn(move || {
417        if end == text.len() {
418            return None;
419        }
420
421        end = (end + chunk_size).min(text.len());
422        while !text.is_char_boundary(end) {
423            end -= 1;
424        }
425        Some(end)
426    })
427}
428
429fn fixtures() -> Vec<EditFixture> {
430    vec![
431        make_fixture(
432            "tiny_function_rewrite",
433            2,
434            EditPattern::LocalizedRewrite {
435                start_line: 12,
436                line_count: 6,
437            },
438            SEED,
439        ),
440        make_fixture(
441            "small_function_rewrite",
442            5,
443            EditPattern::LocalizedRewrite {
444                start_line: 22,
445                line_count: 12,
446            },
447            SEED + 1,
448        ),
449        make_fixture(
450            "medium_many_small_changes",
451            8,
452            EditPattern::ManySmallChanges { every_nth_line: 7 },
453            SEED + 2,
454        ),
455        make_fixture(
456            "medium_insertions",
457            8,
458            EditPattern::InsertHelperBlocks { every_nth_line: 9 },
459            SEED + 3,
460        ),
461        make_large_multi_edit_fixture("large_multi_edit", 80, 16, SEED + 4),
462    ]
463}
464
465enum EditPattern {
466    LocalizedRewrite {
467        start_line: usize,
468        line_count: usize,
469    },
470    ManySmallChanges {
471        every_nth_line: usize,
472    },
473    InsertHelperBlocks {
474        every_nth_line: usize,
475    },
476}
477
478fn make_fixture(
479    name: &'static str,
480    function_count: usize,
481    pattern: EditPattern,
482    seed: u64,
483) -> EditFixture {
484    let mut rng = StdRng::seed_from_u64(seed);
485    let old_lines = random_rust_file(&mut rng, rust_file_line_count(function_count));
486    let edit_range = edit_range(&old_lines, &pattern);
487    let old_text = old_lines[edit_range.clone()].join("\n");
488    let mut new_lines = old_lines.clone();
489
490    match pattern {
491        EditPattern::LocalizedRewrite { .. } => {
492            rewrite_local_block(&mut new_lines[edit_range.clone()], &mut rng)
493        }
494        EditPattern::ManySmallChanges { every_nth_line } => {
495            rewrite_many_small_lines(&mut new_lines[edit_range.clone()], every_nth_line, &mut rng)
496        }
497        EditPattern::InsertHelperBlocks { every_nth_line } => {
498            insert_helper_blocks(&mut new_lines, edit_range.clone(), every_nth_line, &mut rng)
499        }
500    }
501
502    let new_text_end = edit_range.end + new_lines.len().saturating_sub(old_lines.len());
503    let old_file_text = old_lines.join("\n");
504    let expected_file_text = new_lines.join("\n");
505    let new_text = new_lines[edit_range.start..new_text_end].join("\n");
506
507    EditFixture {
508        name,
509        old_file_text,
510        expected_file_text,
511        edits: vec![EditOp { old_text, new_text }],
512    }
513}
514
515fn make_large_multi_edit_fixture(
516    name: &'static str,
517    function_count: usize,
518    edit_count: usize,
519    seed: u64,
520) -> EditFixture {
521    let mut rng = StdRng::seed_from_u64(seed);
522    let old_lines = random_rust_file(&mut rng, rust_file_line_count(function_count));
523    let old_file_text = old_lines.join("\n");
524
525    let step = (function_count / edit_count).max(1);
526    let mut picks: Vec<usize> = (0..edit_count)
527        .map(|k| (k * step).min(function_count - 1))
528        .collect();
529    picks.dedup();
530
531    let replacements: Vec<(usize, Vec<String>)> = picks
532        .iter()
533        .map(|&function_index| {
534            (
535                function_index,
536                large_function_lines(&mut rng, function_index),
537            )
538        })
539        .collect();
540
541    let edits = replacements
542        .iter()
543        .map(|(function_index, new_function)| {
544            let start = RUST_MODULE_HEADER_LINES + function_index * RUST_FUNCTION_LINES;
545            let end = start + RUST_FUNCTION_BODY_LINES;
546            EditOp {
547                old_text: old_lines[start..end].join("\n"),
548                new_text: new_function.join("\n"),
549            }
550        })
551        .collect();
552
553    let mut new_lines = old_lines;
554    for (function_index, new_function) in replacements.iter().rev() {
555        let start = RUST_MODULE_HEADER_LINES + function_index * RUST_FUNCTION_LINES;
556        let end = start + RUST_FUNCTION_BODY_LINES;
557        new_lines.splice(start..end, new_function.iter().cloned());
558    }
559    let expected_file_text = new_lines.join("\n");
560
561    EditFixture {
562        name,
563        old_file_text,
564        expected_file_text,
565        edits,
566    }
567}
568
569fn large_function_lines(rng: &mut StdRng, index: usize) -> Vec<String> {
570    let function_name = identifier(rng, index + 40_000);
571    let argument_name = identifier(rng, index + 41_000);
572
573    let mut lines = vec![
574        format!(
575            "    pub fn {function_name}(&mut self, {argument_name}: usize) -> Result<usize> {{"
576        ),
577        format!("        let mut accumulator = {argument_name};"),
578    ];
579
580    let body_lines = rng.random_range(30..42);
581    for body_index in 0..body_lines {
582        let local_name = identifier(rng, index + 50_000 + body_index);
583        let multiplier = rng.random_range(2..19);
584        let offset = rng.random_range(1..256);
585        match body_index % 4 {
586            0 => lines.push(format!(
587                "        let {local_name} = accumulator.saturating_mul({multiplier}).saturating_add({offset});"
588            )),
589            1 => lines.push(format!(
590                "        accumulator = {local_name}.saturating_sub(self.version % {offset}.max(1));"
591            )),
592            2 => lines.push(format!(
593                "        if {local_name} % {multiplier} == 0 {{ accumulator = accumulator.saturating_add({local_name}); }}"
594            )),
595            _ => lines.push(format!(
596                "        self.buffers.insert(\"{local_name}\".to_string(), accumulator);"
597            )),
598        }
599    }
600
601    lines.push("        self.version = self.version.saturating_add(accumulator);".to_string());
602    lines.push("        Ok(accumulator)".to_string());
603    lines.push("    }".to_string());
604    lines
605}
606
607fn edit_range(lines: &[String], pattern: &EditPattern) -> std::ops::Range<usize> {
608    let mut range = match pattern {
609        EditPattern::LocalizedRewrite {
610            start_line,
611            line_count,
612        } => *start_line..(*start_line + *line_count).min(lines.len()),
613        EditPattern::ManySmallChanges { .. } | EditPattern::InsertHelperBlocks { .. } => {
614            10..lines.len().saturating_sub(5)
615        }
616    };
617
618    while range.end > range.start && lines[range.end - 1].is_empty() {
619        range.end -= 1;
620    }
621
622    range
623}
624
625fn rewrite_local_block(lines: &mut [String], rng: &mut StdRng) {
626    for (line_index, line) in lines.iter_mut().enumerate() {
627        let suffix = identifier(rng, line_index + 10_000);
628        if line.contains("saturating_add") {
629            *line = format!(
630                "        let {suffix} = self.version.checked_add({line_index}).context(\"version overflow\")?;"
631            );
632        } else if line.contains("saturating_sub") {
633            *line = format!(
634                "            {suffix}.saturating_sub({});",
635                rng.random_range(8..256)
636            );
637        } else if line.trim().is_empty() {
638            *line =
639                format!("        tracing::trace!(target: \"agent_bench\", value = {line_index});");
640        } else {
641            *line = format!("{line} // updated {suffix}");
642        }
643    }
644}
645
646fn rewrite_many_small_lines(lines: &mut [String], every_nth_line: usize, rng: &mut StdRng) {
647    for (line_index, line) in lines.iter_mut().enumerate() {
648        if line_index.is_multiple_of(every_nth_line) || line.trim().is_empty() {
649            continue;
650        }
651
652        let suffix = identifier(rng, line_index + 20_000);
653        *line = format!("{line} // audited {suffix}");
654    }
655}
656
657fn insert_helper_blocks(
658    lines: &mut Vec<String>,
659    range: std::ops::Range<usize>,
660    every_nth_line: usize,
661    rng: &mut StdRng,
662) {
663    let mut line_index = range.start;
664    while line_index < range.end.min(lines.len()) {
665        if line_index.is_multiple_of(every_nth_line) && !lines[line_index].trim().is_empty() {
666            let suffix = identifier(rng, line_index + 30_000);
667            lines.splice(
668                line_index..line_index,
669                [
670                    format!("        let {suffix}_before = self.version;"),
671                    format!("        tracing::debug!(version = {suffix}_before);"),
672                ],
673            );
674            line_index += 2;
675        }
676        line_index += 1;
677    }
678}
679
680criterion_group!(benches, edit_file_tool_streaming);
681criterion_main!(benches);
682