tenant.openagents/omega

1use std::{collections::BTreeMap, path::PathBuf, sync::Arc};
2
3use criterion::{BatchSize, BenchmarkId, Criterion, criterion_group, criterion_main};
4use extension::{
5    ExtensionCapability, ExtensionHostProxy, ExtensionLibraryKind, ExtensionManifest,
6    LanguageServerManifestEntry, LibManifestEntry, SchemaVersion,
7    extension_builder::{CompilationConcurrency, CompileExtensionOptions, ExtensionBuilder},
8};
9use extension_host::wasm_host::WasmHost;
10use fs::{Fs, RealFs};
11use gpui::{TestAppContext, TestDispatcher};
12use http_client::{FakeHttpClient, Response};
13use node_runtime::NodeRuntime;
14
15use reqwest_client::ReqwestClient;
16use serde_json::json;
17use settings::SettingsStore;
18use util::test::TempTree;
19
20fn extension_benchmarks(c: &mut Criterion) {
21    let cx = init();
22    cx.update(gpui_tokio::init);
23
24    let mut group = c.benchmark_group("load");
25
26    let mut manifest = manifest();
27    let wasm_bytes = wasm_bytes(
28        &cx,
29        &mut manifest,
30        Arc::new(RealFs::new(None, cx.executor())),
31    );
32    let manifest = Arc::new(manifest);
33    let extensions_dir = TempTree::new(json!({
34        "installed": {},
35        "work": {}
36    }));
37    let wasm_host = wasm_host(&cx, &extensions_dir);
38
39    group.bench_function(BenchmarkId::from_parameter(1), |b| {
40        b.iter_batched(
41            || wasm_bytes.clone(),
42            |wasm_bytes| {
43                let _extension = cx
44                    .foreground_executor()
45                    .block_on(wasm_host.load_extension(wasm_bytes, &manifest, &cx.to_async()))
46                    .unwrap();
47            },
48            BatchSize::SmallInput,
49        );
50    });
51}
52
53fn init() -> TestAppContext {
54    const SEED: u64 = 9999;
55    let dispatcher = TestDispatcher::new(SEED);
56    let cx = TestAppContext::build(dispatcher, None);
57    cx.executor().allow_parking();
58    cx.update(|cx| {
59        let store = SettingsStore::test(cx);
60        cx.set_global(store);
61        release_channel::init(semver::Version::new(0, 0, 0), cx);
62    });
63
64    cx
65}
66
67fn wasm_bytes(cx: &TestAppContext, manifest: &mut ExtensionManifest, fs: Arc<dyn Fs>) -> Vec<u8> {
68    let extension_builder = extension_builder();
69    let path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
70        .parent()
71        .unwrap()
72        .parent()
73        .unwrap()
74        .join("extensions/test-extension");
75    cx.foreground_executor()
76        .block_on(extension_builder.compile_extension(
77            &path,
78            manifest,
79            CompileExtensionOptions {
80                release: true,
81                max_concurrency: CompilationConcurrency::Unbounded,
82            },
83            fs,
84        ))
85        .unwrap();
86    std::fs::read(path.join("extension.wasm")).unwrap()
87}
88
89fn extension_builder() -> ExtensionBuilder {
90    let user_agent = format!(
91        "Zed Extension CLI/{} ({}; {})",
92        env!("CARGO_PKG_VERSION"),
93        std::env::consts::OS,
94        std::env::consts::ARCH
95    );
96    let http_client = Arc::new(ReqwestClient::user_agent(&user_agent).unwrap());
97    // Local dir so that we don't have to download it on every run
98    let build_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("benches/.build");
99    ExtensionBuilder::new(http_client, build_dir)
100}
101
102fn wasm_host(cx: &TestAppContext, extensions_dir: &TempTree) -> Arc<WasmHost> {
103    let http_client = FakeHttpClient::create(async |_| {
104        Ok(Response::builder().status(404).body("not found".into())?)
105    });
106    let extensions_dir = extensions_dir.path().canonicalize().unwrap();
107    let work_dir = extensions_dir.join("work");
108    let fs = Arc::new(RealFs::new(None, cx.executor()));
109
110    cx.update(|cx| {
111        WasmHost::new(
112            fs,
113            http_client,
114            NodeRuntime::unavailable(),
115            Arc::new(ExtensionHostProxy::new()),
116            work_dir,
117            cx,
118        )
119    })
120}
121
122fn manifest() -> ExtensionManifest {
123    ExtensionManifest {
124        id: "test-extension".into(),
125        name: "Test Extension".into(),
126        version: "0.1.0".into(),
127        schema_version: SchemaVersion(1),
128        description: Some("An extension for use in tests.".into()),
129        authors: Vec::new(),
130        repository: None,
131        themes: Default::default(),
132        icon_themes: Vec::new(),
133        lib: LibManifestEntry {
134            kind: Some(ExtensionLibraryKind::Rust),
135            version: Some(semver::Version::new(0, 1, 0)),
136        },
137        languages: Vec::new(),
138        grammars: BTreeMap::default(),
139        language_servers: [("gleam".into(), LanguageServerManifestEntry::default())]
140            .into_iter()
141            .collect(),
142        context_servers: BTreeMap::default(),
143        slash_commands: BTreeMap::default(),
144        snippets: None,
145        capabilities: vec![ExtensionCapability::ProcessExec(
146            extension::ProcessExecCapability {
147                command: "echo".into(),
148                args: vec!["hello!".into()],
149            },
150        )],
151        debug_adapters: Default::default(),
152        debug_locators: Default::default(),
153        language_model_providers: BTreeMap::default(),
154    }
155}
156
157criterion_group!(benches, extension_benchmarks);
158criterion_main!(benches);
159