tenant.openagents/omega

1pub mod wit;
2
3use crate::capability_granter::CapabilityGranter;
4use crate::{ExtensionManifest, ExtensionSettings};
5use anyhow::{Context as _, Result, anyhow, bail};
6use async_trait::async_trait;
7use dap::{DebugRequest, StartDebuggingRequestArgumentsRequest};
8use extension::{
9    CodeLabel, Command, Completion, ContextServerConfiguration, DebugAdapterBinary,
10    DebugTaskDefinition, ExtensionCapability, ExtensionHostProxy, KeyValueStoreDelegate,
11    ProjectDelegate, SlashCommand, SlashCommandArgumentCompletion, SlashCommandOutput, Symbol,
12    WorktreeDelegate,
13};
14use fs::Fs;
15use futures::future::LocalBoxFuture;
16use futures::{
17    Future, FutureExt, StreamExt as _,
18    channel::{
19        mpsc::{self, UnboundedSender},
20        oneshot,
21    },
22    future::BoxFuture,
23};
24use gpui::{App, AsyncApp, BackgroundExecutor, Task};
25use http_client::HttpClient;
26use language::LanguageName;
27use lsp::LanguageServerName;
28use moka::sync::Cache;
29use node_runtime::NodeRuntime;
30use release_channel::ReleaseChannel;
31use semver::Version;
32use settings::Settings;
33use std::{
34    borrow::Cow,
35    path::{Path, PathBuf},
36    sync::{Arc, LazyLock, OnceLock},
37    time::Duration,
38};
39use task::{DebugScenario, SpawnInTerminal, TaskTemplate, ZedDebugConfig};
40use util::paths::SanitizedPath;
41use wasmtime::{
42    CacheStore, Engine, Store,
43    component::{Component, ResourceTable},
44};
45use wasmtime_wasi::{WasiCtx, WasiCtxBuilder, WasiCtxView, WasiView};
46use wit::Extension;
47
48pub struct WasmHost {
49    engine: Engine,
50    release_channel: ReleaseChannel,
51    http_client: Arc<dyn HttpClient>,
52    node_runtime: NodeRuntime,
53    pub(crate) proxy: Arc<ExtensionHostProxy>,
54    fs: Arc<dyn Fs>,
55    pub work_dir: PathBuf,
56    /// The capabilities granted to extensions running on the host.
57    pub(crate) granted_capabilities: Vec<ExtensionCapability>,
58    _main_thread_message_task: Task<()>,
59    main_thread_message_tx: mpsc::UnboundedSender<MainThreadCall>,
60}
61
62#[derive(Clone, Debug)]
63pub struct WasmExtension {
64    tx: UnboundedSender<ExtensionCall>,
65    pub manifest: Arc<ExtensionManifest>,
66    pub work_dir: Arc<Path>,
67    #[allow(unused)]
68    pub zed_api_version: Version,
69    _task: Arc<Task<Result<(), gpui_tokio::JoinError>>>,
70}
71
72impl Drop for WasmExtension {
73    fn drop(&mut self) {
74        self.tx.close_channel();
75    }
76}
77
78#[async_trait]
79impl extension::Extension for WasmExtension {
80    fn manifest(&self) -> Arc<ExtensionManifest> {
81        self.manifest.clone()
82    }
83
84    fn work_dir(&self) -> Arc<Path> {
85        self.work_dir.clone()
86    }
87
88    async fn language_server_command(
89        &self,
90        language_server_id: LanguageServerName,
91        language_name: LanguageName,
92        worktree: Arc<dyn WorktreeDelegate>,
93    ) -> Result<Command> {
94        self.call(|extension, store| {
95            async move {
96                let resource = store.data_mut().table.push(worktree)?;
97                let command = extension
98                    .call_language_server_command(
99                        store,
100                        &language_server_id,
101                        &language_name,
102                        resource,
103                    )
104                    .await?
105                    .map_err(|err| store.data().extension_error(err))?;
106
107                Ok(command.into())
108            }
109            .boxed()
110        })
111        .await?
112    }
113
114    async fn language_server_initialization_options(
115        &self,
116        language_server_id: LanguageServerName,
117        language_name: LanguageName,
118        worktree: Arc<dyn WorktreeDelegate>,
119    ) -> Result<Option<String>> {
120        self.call(|extension, store| {
121            async move {
122                let resource = store.data_mut().table.push(worktree)?;
123                let options = extension
124                    .call_language_server_initialization_options(
125                        store,
126                        &language_server_id,
127                        &language_name,
128                        resource,
129                    )
130                    .await?
131                    .map_err(|err| store.data().extension_error(err))?;
132                anyhow::Ok(options)
133            }
134            .boxed()
135        })
136        .await?
137    }
138
139    async fn language_server_workspace_configuration(
140        &self,
141        language_server_id: LanguageServerName,
142        worktree: Arc<dyn WorktreeDelegate>,
143    ) -> Result<Option<String>> {
144        self.call(|extension, store| {
145            async move {
146                let resource = store.data_mut().table.push(worktree)?;
147                let options = extension
148                    .call_language_server_workspace_configuration(
149                        store,
150                        &language_server_id,
151                        resource,
152                    )
153                    .await?
154                    .map_err(|err| store.data().extension_error(err))?;
155                anyhow::Ok(options)
156            }
157            .boxed()
158        })
159        .await?
160    }
161
162    async fn language_server_initialization_options_schema(
163        &self,
164        language_server_id: LanguageServerName,
165        worktree: Arc<dyn WorktreeDelegate>,
166    ) -> Result<Option<String>> {
167        self.call(|extension, store| {
168            async move {
169                let resource = store.data_mut().table.push(worktree)?;
170                extension
171                    .call_language_server_initialization_options_schema(
172                        store,
173                        &language_server_id,
174                        resource,
175                    )
176                    .await
177            }
178            .boxed()
179        })
180        .await?
181    }
182
183    async fn language_server_workspace_configuration_schema(
184        &self,
185        language_server_id: LanguageServerName,
186        worktree: Arc<dyn WorktreeDelegate>,
187    ) -> Result<Option<String>> {
188        self.call(|extension, store| {
189            async move {
190                let resource = store.data_mut().table.push(worktree)?;
191                extension
192                    .call_language_server_workspace_configuration_schema(
193                        store,
194                        &language_server_id,
195                        resource,
196                    )
197                    .await
198            }
199            .boxed()
200        })
201        .await?
202    }
203
204    async fn language_server_additional_initialization_options(
205        &self,
206        language_server_id: LanguageServerName,
207        target_language_server_id: LanguageServerName,
208        worktree: Arc<dyn WorktreeDelegate>,
209    ) -> Result<Option<String>> {
210        self.call(|extension, store| {
211            async move {
212                let resource = store.data_mut().table.push(worktree)?;
213                let options = extension
214                    .call_language_server_additional_initialization_options(
215                        store,
216                        &language_server_id,
217                        &target_language_server_id,
218                        resource,
219                    )
220                    .await?
221                    .map_err(|err| store.data().extension_error(err))?;
222                anyhow::Ok(options)
223            }
224            .boxed()
225        })
226        .await?
227    }
228
229    async fn language_server_additional_workspace_configuration(
230        &self,
231        language_server_id: LanguageServerName,
232        target_language_server_id: LanguageServerName,
233        worktree: Arc<dyn WorktreeDelegate>,
234    ) -> Result<Option<String>> {
235        self.call(|extension, store| {
236            async move {
237                let resource = store.data_mut().table.push(worktree)?;
238                let options = extension
239                    .call_language_server_additional_workspace_configuration(
240                        store,
241                        &language_server_id,
242                        &target_language_server_id,
243                        resource,
244                    )
245                    .await?
246                    .map_err(|err| store.data().extension_error(err))?;
247                anyhow::Ok(options)
248            }
249            .boxed()
250        })
251        .await?
252    }
253
254    async fn labels_for_completions(
255        &self,
256        language_server_id: LanguageServerName,
257        completions: Vec<Completion>,
258    ) -> Result<Vec<Option<CodeLabel>>> {
259        self.call(|extension, store| {
260            async move {
261                let labels = extension
262                    .call_labels_for_completions(
263                        store,
264                        &language_server_id,
265                        completions.into_iter().map(Into::into).collect(),
266                    )
267                    .await?
268                    .map_err(|err| store.data().extension_error(err))?;
269
270                Ok(labels
271                    .into_iter()
272                    .map(|label| label.map(Into::into))
273                    .collect())
274            }
275            .boxed()
276        })
277        .await?
278    }
279
280    async fn labels_for_symbols(
281        &self,
282        language_server_id: LanguageServerName,
283        symbols: Vec<Symbol>,
284    ) -> Result<Vec<Option<CodeLabel>>> {
285        self.call(|extension, store| {
286            async move {
287                let labels = extension
288                    .call_labels_for_symbols(
289                        store,
290                        &language_server_id,
291                        symbols.into_iter().map(Into::into).collect(),
292                    )
293                    .await?
294                    .map_err(|err| store.data().extension_error(err))?;
295
296                Ok(labels
297                    .into_iter()
298                    .map(|label| label.map(Into::into))
299                    .collect())
300            }
301            .boxed()
302        })
303        .await?
304    }
305
306    async fn complete_slash_command_argument(
307        &self,
308        command: SlashCommand,
309        arguments: Vec<String>,
310    ) -> Result<Vec<SlashCommandArgumentCompletion>> {
311        self.call(|extension, store| {
312            async move {
313                let completions = extension
314                    .call_complete_slash_command_argument(store, &command.into(), &arguments)
315                    .await?
316                    .map_err(|err| store.data().extension_error(err))?;
317
318                Ok(completions.into_iter().map(Into::into).collect())
319            }
320            .boxed()
321        })
322        .await?
323    }
324
325    async fn run_slash_command(
326        &self,
327        command: SlashCommand,
328        arguments: Vec<String>,
329        delegate: Option<Arc<dyn WorktreeDelegate>>,
330    ) -> Result<SlashCommandOutput> {
331        self.call(|extension, store| {
332            async move {
333                let resource = if let Some(delegate) = delegate {
334                    Some(store.data_mut().table.push(delegate)?)
335                } else {
336                    None
337                };
338
339                let output = extension
340                    .call_run_slash_command(store, &command.into(), &arguments, resource)
341                    .await?
342                    .map_err(|err| store.data().extension_error(err))?;
343
344                Ok(output.into())
345            }
346            .boxed()
347        })
348        .await?
349    }
350
351    async fn context_server_command(
352        &self,
353        context_server_id: Arc<str>,
354        project: Arc<dyn ProjectDelegate>,
355    ) -> Result<Command> {
356        self.call(|extension, store| {
357            async move {
358                let project_resource = store.data_mut().table.push(project)?;
359                let command = extension
360                    .call_context_server_command(store, context_server_id.clone(), project_resource)
361                    .await?
362                    .map_err(|err| store.data().extension_error(err))?;
363                anyhow::Ok(command.into())
364            }
365            .boxed()
366        })
367        .await?
368    }
369
370    async fn context_server_configuration(
371        &self,
372        context_server_id: Arc<str>,
373        project: Arc<dyn ProjectDelegate>,
374    ) -> Result<Option<ContextServerConfiguration>> {
375        self.call(|extension, store| {
376            async move {
377                let project_resource = store.data_mut().table.push(project)?;
378                let Some(configuration) = extension
379                    .call_context_server_configuration(
380                        store,
381                        context_server_id.clone(),
382                        project_resource,
383                    )
384                    .await?
385                    .map_err(|err| store.data().extension_error(err))?
386                else {
387                    return Ok(None);
388                };
389
390                Ok(Some(configuration.try_into()?))
391            }
392            .boxed()
393        })
394        .await?
395    }
396
397    async fn suggest_docs_packages(&self, provider: Arc<str>) -> Result<Vec<String>> {
398        self.call(|extension, store| {
399            async move {
400                let packages = extension
401                    .call_suggest_docs_packages(store, provider.as_ref())
402                    .await?
403                    .map_err(|err| store.data().extension_error(err))?;
404
405                Ok(packages)
406            }
407            .boxed()
408        })
409        .await?
410    }
411
412    async fn index_docs(
413        &self,
414        provider: Arc<str>,
415        package_name: Arc<str>,
416        kv_store: Arc<dyn KeyValueStoreDelegate>,
417    ) -> Result<()> {
418        self.call(|extension, store| {
419            async move {
420                let kv_store_resource = store.data_mut().table.push(kv_store)?;
421                extension
422                    .call_index_docs(
423                        store,
424                        provider.as_ref(),
425                        package_name.as_ref(),
426                        kv_store_resource,
427                    )
428                    .await?
429                    .map_err(|err| store.data().extension_error(err))?;
430
431                anyhow::Ok(())
432            }
433            .boxed()
434        })
435        .await?
436    }
437
438    async fn get_dap_binary(
439        &self,
440        dap_name: Arc<str>,
441        config: DebugTaskDefinition,
442        user_installed_path: Option<PathBuf>,
443        worktree: Arc<dyn WorktreeDelegate>,
444    ) -> Result<DebugAdapterBinary> {
445        self.call(|extension, store| {
446            async move {
447                let resource = store.data_mut().table.push(worktree)?;
448                let dap_binary = extension
449                    .call_get_dap_binary(store, dap_name, config, user_installed_path, resource)
450                    .await?
451                    .map_err(|err| store.data().extension_error(err))?;
452                let dap_binary = dap_binary.try_into()?;
453                Ok(dap_binary)
454            }
455            .boxed()
456        })
457        .await?
458    }
459    async fn dap_request_kind(
460        &self,
461        dap_name: Arc<str>,
462        config: serde_json::Value,
463    ) -> Result<StartDebuggingRequestArgumentsRequest> {
464        self.call(|extension, store| {
465            async move {
466                let kind = extension
467                    .call_dap_request_kind(store, dap_name, config)
468                    .await?
469                    .map_err(|err| store.data().extension_error(err))?;
470                Ok(kind.into())
471            }
472            .boxed()
473        })
474        .await?
475    }
476
477    async fn dap_config_to_scenario(&self, config: ZedDebugConfig) -> Result<DebugScenario> {
478        self.call(|extension, store| {
479            async move {
480                let kind = extension
481                    .call_dap_config_to_scenario(store, config)
482                    .await?
483                    .map_err(|err| store.data().extension_error(err))?;
484                Ok(kind)
485            }
486            .boxed()
487        })
488        .await?
489    }
490
491    async fn dap_locator_create_scenario(
492        &self,
493        locator_name: String,
494        build_config_template: TaskTemplate,
495        resolved_label: String,
496        debug_adapter_name: String,
497    ) -> Result<Option<DebugScenario>> {
498        self.call(|extension, store| {
499            async move {
500                extension
501                    .call_dap_locator_create_scenario(
502                        store,
503                        locator_name,
504                        build_config_template,
505                        resolved_label,
506                        debug_adapter_name,
507                    )
508                    .await
509            }
510            .boxed()
511        })
512        .await?
513    }
514    async fn run_dap_locator(
515        &self,
516        locator_name: String,
517        config: SpawnInTerminal,
518    ) -> Result<DebugRequest> {
519        self.call(|extension, store| {
520            async move {
521                extension
522                    .call_run_dap_locator(store, locator_name, config)
523                    .await?
524                    .map_err(|err| store.data().extension_error(err))
525            }
526            .boxed()
527        })
528        .await?
529    }
530}
531
532pub struct WasmState {
533    manifest: Arc<ExtensionManifest>,
534    pub table: ResourceTable,
535    ctx: WasiCtx,
536    pub host: Arc<WasmHost>,
537    pub(crate) capability_granter: CapabilityGranter,
538}
539
540type MainThreadCall = Box<dyn Send + for<'a> FnOnce(&'a mut AsyncApp) -> LocalBoxFuture<'a, ()>>;
541
542type ExtensionCall = Box<
543    dyn Send + for<'a> FnOnce(&'a mut Extension, &'a mut Store<WasmState>) -> BoxFuture<'a, ()>,
544>;
545
546fn wasm_engine(executor: &BackgroundExecutor) -> wasmtime::Engine {
547    static WASM_ENGINE: OnceLock<wasmtime::Engine> = OnceLock::new();
548    WASM_ENGINE
549        .get_or_init(|| {
550            let mut config = wasmtime::Config::new();
551            config.wasm_component_model(true);
552            config.async_support(true);
553            config
554                .enable_incremental_compilation(cache_store())
555                .unwrap();
556            // Async support introduces the issue that extension execution happens during `Future::poll`,
557            // which could block an async thread.
558            // https://docs.rs/wasmtime/latest/wasmtime/struct.Config.html#execution-in-poll
559            //
560            // Epoch interruption is a lightweight mechanism to allow the extensions to yield control
561            // back to the executor at regular intervals.
562            config.epoch_interruption(true);
563
564            let engine = wasmtime::Engine::new(&config).unwrap();
565
566            // It might be safer to do this on a non-async thread to make sure it makes progress
567            // regardless of if extensions are blocking.
568            // However, due to our current setup, this isn't a likely occurrence and we'd rather
569            // not have a dedicated thread just for this. If it becomes an issue, we can consider
570            // creating a separate thread for epoch interruption.
571            let engine_ref = engine.weak();
572            let executor2 = executor.clone();
573            executor
574                .spawn(async move {
575                    // Somewhat arbitrary interval, as it isn't a guaranteed interval.
576                    // But this is a rough upper bound for how long the extension execution can block on
577                    // `Future::poll`.
578                    const EPOCH_INTERVAL: Duration = Duration::from_millis(100);
579                    loop {
580                        executor2.timer(EPOCH_INTERVAL).await;
581                        // Exit the loop and thread once the engine is dropped.
582                        let Some(engine) = engine_ref.upgrade() else {
583                            break;
584                        };
585                        engine.increment_epoch();
586                    }
587                })
588                .detach();
589
590            engine
591        })
592        .clone()
593}
594
595fn cache_store() -> Arc<IncrementalCompilationCache> {
596    static CACHE_STORE: LazyLock<Arc<IncrementalCompilationCache>> =
597        LazyLock::new(|| Arc::new(IncrementalCompilationCache::new()));
598    CACHE_STORE.clone()
599}
600
601impl WasmHost {
602    pub fn new(
603        fs: Arc<dyn Fs>,
604        http_client: Arc<dyn HttpClient>,
605        node_runtime: NodeRuntime,
606        proxy: Arc<ExtensionHostProxy>,
607        work_dir: PathBuf,
608        cx: &mut App,
609    ) -> Arc<Self> {
610        let (tx, mut rx) = mpsc::unbounded::<MainThreadCall>();
611        let task = cx.spawn(async move |cx| {
612            while let Some(message) = rx.next().await {
613                message(cx).await;
614            }
615        });
616
617        let extension_settings = ExtensionSettings::get_global(cx);
618
619        Arc::new(Self {
620            engine: wasm_engine(cx.background_executor()),
621            fs,
622            work_dir,
623            http_client,
624            node_runtime,
625            proxy,
626            release_channel: ReleaseChannel::global(cx),
627            granted_capabilities: extension_settings.granted_capabilities.clone(),
628            _main_thread_message_task: task,
629            main_thread_message_tx: tx,
630        })
631    }
632
633    pub fn load_extension(
634        self: &Arc<Self>,
635        wasm_bytes: Vec<u8>,
636        manifest: &Arc<ExtensionManifest>,
637        cx: &AsyncApp,
638    ) -> Task<Result<WasmExtension>> {
639        let this = self.clone();
640        let manifest = manifest.clone();
641        let executor = cx.background_executor().clone();
642
643        // Parse version and compile component on gpui's background executor.
644        // These are cpu-bound operations that don't require a tokio runtime.
645        let compile_task = {
646            let manifest_id = manifest.id.clone();
647            let engine = this.engine.clone();
648
649            executor.spawn(async move {
650                let zed_api_version = parse_wasm_extension_version(&manifest_id, &wasm_bytes)?;
651                let component = Component::from_binary(&engine, &wasm_bytes)
652                    .context("failed to compile wasm component")?;
653
654                anyhow::Ok((zed_api_version, component))
655            })
656        };
657
658        let load_extension = |zed_api_version: Version, component| async move {
659            let wasi_ctx = this.build_wasi_ctx(&manifest).await?;
660            let mut store = wasmtime::Store::new(
661                &this.engine,
662                WasmState {
663                    ctx: wasi_ctx,
664                    manifest: manifest.clone(),
665                    table: ResourceTable::new(),
666                    host: this.clone(),
667                    capability_granter: CapabilityGranter::new(
668                        this.granted_capabilities.clone(),
669                        manifest.clone(),
670                    ),
671                },
672            );
673            // Store will yield after 1 tick, and get a new deadline of 1 tick after each yield.
674            store.set_epoch_deadline(1);
675            store.epoch_deadline_async_yield_and_update(1);
676
677            let mut extension = Extension::instantiate_async(
678                &executor,
679                &mut store,
680                this.release_channel,
681                zed_api_version.clone(),
682                &component,
683            )
684            .await?;
685
686            extension
687                .call_init_extension(&mut store)
688                .await
689                .context("failed to initialize wasm extension")?;
690
691            let (tx, mut rx) = mpsc::unbounded::<ExtensionCall>();
692            let extension_task = async move {
693                while let Some(call) = rx.next().await {
694                    (call)(&mut extension, &mut store).await;
695                }
696            };
697
698            anyhow::Ok((
699                extension_task,
700                manifest.clone(),
701                this.work_dir.join(manifest.id.as_ref()).into(),
702                tx,
703                zed_api_version,
704            ))
705        };
706
707        cx.spawn(async move |cx| {
708            let (zed_api_version, component) = compile_task.await?;
709
710            // Run wasi-dependent operations on tokio.
711            // wasmtime_wasi internally uses tokio for I/O operations.
712            let (extension_task, manifest, work_dir, tx, zed_api_version) =
713                gpui_tokio::Tokio::spawn(cx, load_extension(zed_api_version, component)).await??;
714
715            // Run the extension message loop on tokio since extension
716            // calls may invoke wasi functions that require a tokio runtime.
717            let task = Arc::new(gpui_tokio::Tokio::spawn(cx, extension_task));
718
719            Ok(WasmExtension {
720                manifest,
721                work_dir,
722                tx,
723                zed_api_version,
724                _task: task,
725            })
726        })
727    }
728
729    async fn build_wasi_ctx(&self, manifest: &Arc<ExtensionManifest>) -> Result<WasiCtx> {
730        let extension_work_dir = self.work_dir.join(manifest.id.as_ref());
731        self.fs
732            .create_dir(&extension_work_dir)
733            .await
734            .context("failed to create extension work dir")?;
735
736        let file_perms = wasmtime_wasi::FilePerms::all();
737        let dir_perms = wasmtime_wasi::DirPerms::all();
738        let path = SanitizedPath::new(&extension_work_dir).to_string();
739        #[cfg(target_os = "windows")]
740        let path = path.replace('\\', "/");
741
742        let mut ctx = WasiCtxBuilder::new();
743        ctx.inherit_stdio()
744            .env("PWD", &path)
745            .env("RUST_BACKTRACE", "full");
746
747        ctx.preopened_dir(&path, ".", dir_perms, file_perms)?;
748        ctx.preopened_dir(&path, &path, dir_perms, file_perms)?;
749
750        Ok(ctx.build())
751    }
752
753    pub async fn writeable_path_from_extension(
754        &self,
755        id: &Arc<str>,
756        path: &Path,
757    ) -> Result<PathBuf> {
758        let canonical_work_dir = self
759            .fs
760            .canonicalize(&self.work_dir)
761            .await
762            .with_context(|| format!("canonicalizing work dir {:?}", self.work_dir))?;
763        let extension_work_dir = canonical_work_dir.join(id.as_ref());
764
765        let absolute = if path.is_relative() {
766            extension_work_dir.join(path)
767        } else {
768            path.to_path_buf()
769        };
770
771        let normalized = util::paths::normalize_lexically(&absolute)
772            .map_err(|_| anyhow!("path {path:?} escapes its parent"))?;
773
774        // Canonicalize the nearest existing ancestor to resolve any symlinks
775        // in the on-disk portion of the path. Components beyond that ancestor
776        // are re-appended, which lets this work for destinations that don't
777        // exist yet (e.g. nested directories created by tar extraction).
778        let mut existing = normalized.as_path();
779        let mut tail_components = Vec::new();
780        let canonical_prefix = loop {
781            match self.fs.canonicalize(existing).await {
782                Ok(canonical) => break canonical,
783                Err(_) => {
784                    if let Some(file_name) = existing.file_name() {
785                        tail_components.push(file_name.to_owned());
786                    }
787                    existing = existing
788                        .parent()
789                        .context(format!("cannot resolve path {path:?}"))?;
790                }
791            }
792        };
793
794        let mut resolved = canonical_prefix;
795        for component in tail_components.into_iter().rev() {
796            resolved.push(component);
797        }
798
799        anyhow::ensure!(
800            resolved.starts_with(&extension_work_dir),
801            "cannot write to path {resolved:?}",
802        );
803        Ok(resolved)
804    }
805}
806
807pub fn parse_wasm_extension_version(extension_id: &str, wasm_bytes: &[u8]) -> Result<Version> {
808    let mut version = None;
809
810    for part in wasmparser::Parser::new(0).parse_all(wasm_bytes) {
811        if let wasmparser::Payload::CustomSection(s) =
812            part.context("error parsing wasm extension")?
813            && s.name() == "zed:api-version"
814        {
815            version = parse_wasm_extension_version_custom_section(s.data());
816            if version.is_none() {
817                bail!(
818                    "extension {} has invalid zed:api-version section: {:?}",
819                    extension_id,
820                    s.data()
821                );
822            }
823        }
824    }
825
826    // The reason we wait until we're done parsing all of the Wasm bytes to return the version
827    // is to work around a panic that can happen inside of Wasmtime when the bytes are invalid.
828    //
829    // By parsing the entirety of the Wasm bytes before we return, we're able to detect this problem
830    // earlier as an `Err` rather than as a panic.
831    version.with_context(|| format!("extension {extension_id} has no zed:api-version section"))
832}
833
834fn parse_wasm_extension_version_custom_section(data: &[u8]) -> Option<Version> {
835    if data.len() == 6 {
836        Some(Version::new(
837            u16::from_be_bytes([data[0], data[1]]) as _,
838            u16::from_be_bytes([data[2], data[3]]) as _,
839            u16::from_be_bytes([data[4], data[5]]) as _,
840        ))
841    } else {
842        None
843    }
844}
845
846impl WasmExtension {
847    pub async fn load(
848        extension_dir: &Path,
849        manifest: &Arc<ExtensionManifest>,
850        wasm_host: Arc<WasmHost>,
851        cx: &AsyncApp,
852    ) -> Result<Self> {
853        let path = extension_dir.join("extension.wasm");
854
855        let mut wasm_file = wasm_host
856            .fs
857            .open_sync(&path)
858            .await
859            .context(format!("opening wasm file, path: {path:?}"))?;
860
861        let mut wasm_bytes = Vec::new();
862        wasm_file
863            .read_to_end(&mut wasm_bytes)
864            .context(format!("reading wasm file, path: {path:?}"))?;
865
866        wasm_host
867            .load_extension(wasm_bytes, manifest, cx)
868            .await
869            .with_context(|| format!("loading wasm extension: {}", manifest.id))
870    }
871
872    pub async fn call<T, Fn>(&self, f: Fn) -> Result<T>
873    where
874        T: 'static + Send,
875        Fn: 'static
876            + Send
877            + for<'a> FnOnce(&'a mut Extension, &'a mut Store<WasmState>) -> BoxFuture<'a, T>,
878    {
879        let (return_tx, return_rx) = oneshot::channel();
880        self.tx
881            .unbounded_send(Box::new(move |extension, store| {
882                async {
883                    let result = f(extension, store).await;
884                    return_tx.send(result).ok();
885                }
886                .boxed()
887            }))
888            .map_err(|_| {
889                anyhow!(
890                    "wasm extension channel should not be closed yet, extension {} (id {})",
891                    self.manifest.name,
892                    self.manifest.id,
893                )
894            })?;
895        return_rx.await.with_context(|| {
896            format!(
897                "wasm extension channel, extension {} (id {})",
898                self.manifest.name, self.manifest.id,
899            )
900        })
901    }
902}
903
904impl WasmState {
905    fn on_main_thread<T, Fn>(&self, f: Fn) -> impl 'static + Future<Output = T>
906    where
907        T: 'static + Send,
908        Fn: 'static + Send + for<'a> FnOnce(&'a mut AsyncApp) -> LocalBoxFuture<'a, T>,
909    {
910        let (return_tx, return_rx) = oneshot::channel();
911        self.host
912            .main_thread_message_tx
913            .clone()
914            .unbounded_send(Box::new(move |cx| {
915                async {
916                    let result = f(cx).await;
917                    return_tx.send(result).ok();
918                }
919                .boxed_local()
920            }))
921            .unwrap_or_else(|_| {
922                panic!(
923                    "main thread message channel should not be closed yet, extension {} (id {})",
924                    self.manifest.name, self.manifest.id,
925                )
926            });
927        let name = self.manifest.name.clone();
928        let id = self.manifest.id.clone();
929        async move {
930            return_rx.await.unwrap_or_else(|_| {
931                panic!("main thread message channel, extension {name} (id {id})")
932            })
933        }
934    }
935
936    fn work_dir(&self) -> PathBuf {
937        self.host.work_dir.join(self.manifest.id.as_ref())
938    }
939
940    fn extension_error(&self, message: String) -> anyhow::Error {
941        anyhow!(
942            "from extension \"{}\" version {}: {}",
943            self.manifest.name,
944            self.manifest.version,
945            message
946        )
947    }
948}
949
950impl wasmtime::component::HasData for WasmState {
951    type Data<'a> = &'a mut WasmState;
952}
953
954impl WasiView for WasmState {
955    fn ctx(&mut self) -> WasiCtxView<'_> {
956        WasiCtxView {
957            ctx: &mut self.ctx,
958            table: &mut self.table,
959        }
960    }
961}
962
963/// Wrapper around a mini-moka bounded cache for storing incremental compilation artifacts.
964/// Since wasm modules have many similar elements, this can save us a lot of work at the
965/// cost of a small memory footprint. However, we don't want this to be unbounded, so we use
966/// a LFU/LRU cache to evict less used cache entries.
967#[derive(Debug)]
968struct IncrementalCompilationCache {
969    cache: Cache<Vec<u8>, Vec<u8>>,
970}
971
972impl IncrementalCompilationCache {
973    fn new() -> Self {
974        let cache = Cache::builder()
975            // Cap this at 32 MB for now. Our extensions turn into roughly 512kb in the cache,
976            // which means we could store 64 completely novel extensions in the cache, but in
977            // practice we will more than that, which is more than enough for our use case.
978            .max_capacity(32 * 1024 * 1024)
979            .weigher(|k: &Vec<u8>, v: &Vec<u8>| (k.len() + v.len()).try_into().unwrap_or(u32::MAX))
980            .build();
981        Self { cache }
982    }
983}
984
985impl CacheStore for IncrementalCompilationCache {
986    fn get(&self, key: &[u8]) -> Option<Cow<'_, [u8]>> {
987        self.cache.get(key).map(|v| v.into())
988    }
989
990    fn insert(&self, key: &[u8], value: Vec<u8>) -> bool {
991        self.cache.insert(key.to_vec(), value);
992        true
993    }
994}
995
996#[cfg(test)]
997mod tests {
998    use super::*;
999    use extension::ExtensionHostProxy;
1000    use fs::FakeFs;
1001    use gpui::TestAppContext;
1002    use http_client::FakeHttpClient;
1003    use node_runtime::NodeRuntime;
1004    use serde_json::json;
1005    use settings::SettingsStore;
1006
1007    fn init_test(cx: &mut TestAppContext) {
1008        cx.update(|cx| {
1009            let store = SettingsStore::test(cx);
1010            cx.set_global(store);
1011            release_channel::init(semver::Version::new(0, 0, 0), cx);
1012            extension::init(cx);
1013            gpui_tokio::init(cx);
1014        });
1015    }
1016
1017    #[gpui::test]
1018    async fn test_writeable_path_rejects_escape_attempts(cx: &mut TestAppContext) {
1019        init_test(cx);
1020
1021        let fs = FakeFs::new(cx.executor());
1022        fs.insert_tree(
1023            "/work",
1024            json!({
1025                "test-extension": {
1026                    "legit.txt": "legitimate content"
1027                }
1028            }),
1029        )
1030        .await;
1031        fs.insert_tree("/outside", json!({ "secret.txt": "sensitive data" }))
1032            .await;
1033        fs.insert_symlink("/work/test-extension/escape", PathBuf::from("/outside"))
1034            .await;
1035
1036        let host = cx.update(|cx| {
1037            WasmHost::new(
1038                fs.clone(),
1039                FakeHttpClient::with_200_response(),
1040                NodeRuntime::unavailable(),
1041                Arc::new(ExtensionHostProxy::default()),
1042                PathBuf::from("/work"),
1043                cx,
1044            )
1045        });
1046
1047        let extension_id: Arc<str> = "test-extension".into();
1048
1049        // A path traversing through a symlink that points outside the work dir
1050        // must be rejected. Canonicalization resolves the symlink before the
1051        // prefix check, so this is caught.
1052        let result = host
1053            .writeable_path_from_extension(
1054                &extension_id,
1055                Path::new("/work/test-extension/escape/secret.txt"),
1056            )
1057            .await;
1058        assert!(
1059            result.is_err(),
1060            "symlink escape should be rejected, but got: {result:?}",
1061        );
1062
1063        // A path using `..` to escape the extension work dir must be rejected.
1064        let result = host
1065            .writeable_path_from_extension(
1066                &extension_id,
1067                Path::new("/work/test-extension/../../outside/secret.txt"),
1068            )
1069            .await;
1070        assert!(
1071            result.is_err(),
1072            "parent traversal escape should be rejected, but got: {result:?}",
1073        );
1074
1075        // A legitimate path within the extension work dir should succeed.
1076        let result = host
1077            .writeable_path_from_extension(
1078                &extension_id,
1079                Path::new("/work/test-extension/legit.txt"),
1080            )
1081            .await;
1082        assert!(
1083            result.is_ok(),
1084            "legitimate path should be accepted, but got: {result:?}",
1085        );
1086
1087        // A relative path with non-existent intermediate directories should
1088        // succeed, mirroring the integration test pattern where an extension
1089        // downloads a tar to e.g. "gleam-v1.2.3" (creating the directory)
1090        // and then references "gleam-v1.2.3/gleam" inside it.
1091        let result = host
1092            .writeable_path_from_extension(&extension_id, Path::new("new-dir/nested/binary"))
1093            .await;
1094        assert!(
1095            result.is_ok(),
1096            "relative path with non-existent parents should be accepted, but got: {result:?}",
1097        );
1098
1099        // A symlink deeper than the immediate parent must still be caught.
1100        // Here "escape" is a symlink to /outside, so "escape/deep/file.txt"
1101        // has multiple non-existent components beyond the symlink.
1102        let result = host
1103            .writeable_path_from_extension(&extension_id, Path::new("escape/deep/nested/file.txt"))
1104            .await;
1105        assert!(
1106            result.is_err(),
1107            "symlink escape through deep non-existent path should be rejected, but got: {result:?}",
1108        );
1109    }
1110}
1111