wasm_host.rs

  1pub mod wit;
  2
  3use crate::ExtensionManifest;
  4use anyhow::{Context as _, Result, anyhow, bail};
  5use async_trait::async_trait;
  6use extension::{
  7    CodeLabel, Command, Completion, ContextServerConfiguration, DebugAdapterBinary,
  8    DebugTaskDefinition, ExtensionHostProxy, KeyValueStoreDelegate, ProjectDelegate, SlashCommand,
  9    SlashCommandArgumentCompletion, SlashCommandOutput, Symbol, WorktreeDelegate,
 10};
 11use fs::{Fs, normalize_path};
 12use futures::future::LocalBoxFuture;
 13use futures::{
 14    Future, FutureExt, StreamExt as _,
 15    channel::{
 16        mpsc::{self, UnboundedSender},
 17        oneshot,
 18    },
 19    future::BoxFuture,
 20};
 21use gpui::{App, AsyncApp, BackgroundExecutor, Task};
 22use http_client::HttpClient;
 23use language::LanguageName;
 24use lsp::LanguageServerName;
 25use node_runtime::NodeRuntime;
 26use release_channel::ReleaseChannel;
 27use semantic_version::SemanticVersion;
 28use std::{
 29    path::{Path, PathBuf},
 30    sync::{Arc, OnceLock},
 31};
 32use wasmtime::{
 33    Engine, Store,
 34    component::{Component, ResourceTable},
 35};
 36use wasmtime_wasi::{self as wasi, WasiView};
 37use wit::Extension;
 38
 39pub struct WasmHost {
 40    engine: Engine,
 41    release_channel: ReleaseChannel,
 42    http_client: Arc<dyn HttpClient>,
 43    node_runtime: NodeRuntime,
 44    pub(crate) proxy: Arc<ExtensionHostProxy>,
 45    fs: Arc<dyn Fs>,
 46    pub work_dir: PathBuf,
 47    _main_thread_message_task: Task<()>,
 48    main_thread_message_tx: mpsc::UnboundedSender<MainThreadCall>,
 49}
 50
 51#[derive(Clone)]
 52pub struct WasmExtension {
 53    tx: UnboundedSender<ExtensionCall>,
 54    pub manifest: Arc<ExtensionManifest>,
 55    pub work_dir: Arc<Path>,
 56    #[allow(unused)]
 57    pub zed_api_version: SemanticVersion,
 58}
 59
 60#[async_trait]
 61impl extension::Extension for WasmExtension {
 62    fn manifest(&self) -> Arc<ExtensionManifest> {
 63        self.manifest.clone()
 64    }
 65
 66    fn work_dir(&self) -> Arc<Path> {
 67        self.work_dir.clone()
 68    }
 69
 70    async fn language_server_command(
 71        &self,
 72        language_server_id: LanguageServerName,
 73        language_name: LanguageName,
 74        worktree: Arc<dyn WorktreeDelegate>,
 75    ) -> Result<Command> {
 76        self.call(|extension, store| {
 77            async move {
 78                let resource = store.data_mut().table().push(worktree)?;
 79                let command = extension
 80                    .call_language_server_command(
 81                        store,
 82                        &language_server_id,
 83                        &language_name,
 84                        resource,
 85                    )
 86                    .await?
 87                    .map_err(|err| anyhow!("{err}"))?;
 88
 89                Ok(command.into())
 90            }
 91            .boxed()
 92        })
 93        .await
 94    }
 95
 96    async fn language_server_initialization_options(
 97        &self,
 98        language_server_id: LanguageServerName,
 99        language_name: LanguageName,
100        worktree: Arc<dyn WorktreeDelegate>,
101    ) -> Result<Option<String>> {
102        self.call(|extension, store| {
103            async move {
104                let resource = store.data_mut().table().push(worktree)?;
105                let options = extension
106                    .call_language_server_initialization_options(
107                        store,
108                        &language_server_id,
109                        &language_name,
110                        resource,
111                    )
112                    .await?
113                    .map_err(|err| anyhow!("{err}"))?;
114                anyhow::Ok(options)
115            }
116            .boxed()
117        })
118        .await
119    }
120
121    async fn language_server_workspace_configuration(
122        &self,
123        language_server_id: LanguageServerName,
124        worktree: Arc<dyn WorktreeDelegate>,
125    ) -> Result<Option<String>> {
126        self.call(|extension, store| {
127            async move {
128                let resource = store.data_mut().table().push(worktree)?;
129                let options = extension
130                    .call_language_server_workspace_configuration(
131                        store,
132                        &language_server_id,
133                        resource,
134                    )
135                    .await?
136                    .map_err(|err| anyhow!("{err}"))?;
137                anyhow::Ok(options)
138            }
139            .boxed()
140        })
141        .await
142    }
143
144    async fn language_server_additional_initialization_options(
145        &self,
146        language_server_id: LanguageServerName,
147        target_language_server_id: LanguageServerName,
148        worktree: Arc<dyn WorktreeDelegate>,
149    ) -> Result<Option<String>> {
150        self.call(|extension, store| {
151            async move {
152                let resource = store.data_mut().table().push(worktree)?;
153                let options = extension
154                    .call_language_server_additional_initialization_options(
155                        store,
156                        &language_server_id,
157                        &target_language_server_id,
158                        resource,
159                    )
160                    .await?
161                    .map_err(|err| anyhow!("{err}"))?;
162                anyhow::Ok(options)
163            }
164            .boxed()
165        })
166        .await
167    }
168
169    async fn language_server_additional_workspace_configuration(
170        &self,
171        language_server_id: LanguageServerName,
172        target_language_server_id: LanguageServerName,
173        worktree: Arc<dyn WorktreeDelegate>,
174    ) -> Result<Option<String>> {
175        self.call(|extension, store| {
176            async move {
177                let resource = store.data_mut().table().push(worktree)?;
178                let options = extension
179                    .call_language_server_additional_workspace_configuration(
180                        store,
181                        &language_server_id,
182                        &target_language_server_id,
183                        resource,
184                    )
185                    .await?
186                    .map_err(|err| anyhow!("{err}"))?;
187                anyhow::Ok(options)
188            }
189            .boxed()
190        })
191        .await
192    }
193
194    async fn labels_for_completions(
195        &self,
196        language_server_id: LanguageServerName,
197        completions: Vec<Completion>,
198    ) -> Result<Vec<Option<CodeLabel>>> {
199        self.call(|extension, store| {
200            async move {
201                let labels = extension
202                    .call_labels_for_completions(
203                        store,
204                        &language_server_id,
205                        completions.into_iter().map(Into::into).collect(),
206                    )
207                    .await?
208                    .map_err(|err| anyhow!("{err}"))?;
209
210                Ok(labels
211                    .into_iter()
212                    .map(|label| label.map(Into::into))
213                    .collect())
214            }
215            .boxed()
216        })
217        .await
218    }
219
220    async fn labels_for_symbols(
221        &self,
222        language_server_id: LanguageServerName,
223        symbols: Vec<Symbol>,
224    ) -> Result<Vec<Option<CodeLabel>>> {
225        self.call(|extension, store| {
226            async move {
227                let labels = extension
228                    .call_labels_for_symbols(
229                        store,
230                        &language_server_id,
231                        symbols.into_iter().map(Into::into).collect(),
232                    )
233                    .await?
234                    .map_err(|err| anyhow!("{err}"))?;
235
236                Ok(labels
237                    .into_iter()
238                    .map(|label| label.map(Into::into))
239                    .collect())
240            }
241            .boxed()
242        })
243        .await
244    }
245
246    async fn complete_slash_command_argument(
247        &self,
248        command: SlashCommand,
249        arguments: Vec<String>,
250    ) -> Result<Vec<SlashCommandArgumentCompletion>> {
251        self.call(|extension, store| {
252            async move {
253                let completions = extension
254                    .call_complete_slash_command_argument(store, &command.into(), &arguments)
255                    .await?
256                    .map_err(|err| anyhow!("{err}"))?;
257
258                Ok(completions.into_iter().map(Into::into).collect())
259            }
260            .boxed()
261        })
262        .await
263    }
264
265    async fn run_slash_command(
266        &self,
267        command: SlashCommand,
268        arguments: Vec<String>,
269        delegate: Option<Arc<dyn WorktreeDelegate>>,
270    ) -> Result<SlashCommandOutput> {
271        self.call(|extension, store| {
272            async move {
273                let resource = if let Some(delegate) = delegate {
274                    Some(store.data_mut().table().push(delegate)?)
275                } else {
276                    None
277                };
278
279                let output = extension
280                    .call_run_slash_command(store, &command.into(), &arguments, resource)
281                    .await?
282                    .map_err(|err| anyhow!("{err}"))?;
283
284                Ok(output.into())
285            }
286            .boxed()
287        })
288        .await
289    }
290
291    async fn context_server_command(
292        &self,
293        context_server_id: Arc<str>,
294        project: Arc<dyn ProjectDelegate>,
295    ) -> Result<Command> {
296        self.call(|extension, store| {
297            async move {
298                let project_resource = store.data_mut().table().push(project)?;
299                let command = extension
300                    .call_context_server_command(store, context_server_id.clone(), project_resource)
301                    .await?
302                    .map_err(|err| anyhow!("{err}"))?;
303                anyhow::Ok(command.into())
304            }
305            .boxed()
306        })
307        .await
308    }
309
310    async fn context_server_configuration(
311        &self,
312        context_server_id: Arc<str>,
313        project: Arc<dyn ProjectDelegate>,
314    ) -> Result<Option<ContextServerConfiguration>> {
315        self.call(|extension, store| {
316            async move {
317                let project_resource = store.data_mut().table().push(project)?;
318                let Some(configuration) = extension
319                    .call_context_server_configuration(
320                        store,
321                        context_server_id.clone(),
322                        project_resource,
323                    )
324                    .await?
325                    .map_err(|err| anyhow!("{err}"))?
326                else {
327                    return Ok(None);
328                };
329
330                Ok(Some(configuration.try_into()?))
331            }
332            .boxed()
333        })
334        .await
335    }
336
337    async fn suggest_docs_packages(&self, provider: Arc<str>) -> Result<Vec<String>> {
338        self.call(|extension, store| {
339            async move {
340                let packages = extension
341                    .call_suggest_docs_packages(store, provider.as_ref())
342                    .await?
343                    .map_err(|err| anyhow!("{err:?}"))?;
344
345                Ok(packages)
346            }
347            .boxed()
348        })
349        .await
350    }
351
352    async fn index_docs(
353        &self,
354        provider: Arc<str>,
355        package_name: Arc<str>,
356        kv_store: Arc<dyn KeyValueStoreDelegate>,
357    ) -> Result<()> {
358        self.call(|extension, store| {
359            async move {
360                let kv_store_resource = store.data_mut().table().push(kv_store)?;
361                extension
362                    .call_index_docs(
363                        store,
364                        provider.as_ref(),
365                        package_name.as_ref(),
366                        kv_store_resource,
367                    )
368                    .await?
369                    .map_err(|err| anyhow!("{err:?}"))?;
370
371                anyhow::Ok(())
372            }
373            .boxed()
374        })
375        .await
376    }
377    async fn get_dap_binary(
378        &self,
379        dap_name: Arc<str>,
380        config: DebugTaskDefinition,
381        user_installed_path: Option<PathBuf>,
382        worktree: Arc<dyn WorktreeDelegate>,
383    ) -> Result<DebugAdapterBinary> {
384        self.call(|extension, store| {
385            async move {
386                let resource = store.data_mut().table().push(worktree)?;
387                let dap_binary = extension
388                    .call_get_dap_binary(store, dap_name, config, user_installed_path, resource)
389                    .await?
390                    .map_err(|err| anyhow!("{err:?}"))?;
391                let dap_binary = dap_binary.try_into()?;
392                Ok(dap_binary)
393            }
394            .boxed()
395        })
396        .await
397    }
398}
399
400pub struct WasmState {
401    manifest: Arc<ExtensionManifest>,
402    pub table: ResourceTable,
403    ctx: wasi::WasiCtx,
404    pub host: Arc<WasmHost>,
405}
406
407type MainThreadCall = Box<dyn Send + for<'a> FnOnce(&'a mut AsyncApp) -> LocalBoxFuture<'a, ()>>;
408
409type ExtensionCall = Box<
410    dyn Send + for<'a> FnOnce(&'a mut Extension, &'a mut Store<WasmState>) -> BoxFuture<'a, ()>,
411>;
412
413fn wasm_engine() -> wasmtime::Engine {
414    static WASM_ENGINE: OnceLock<wasmtime::Engine> = OnceLock::new();
415
416    WASM_ENGINE
417        .get_or_init(|| {
418            let mut config = wasmtime::Config::new();
419            config.wasm_component_model(true);
420            config.async_support(true);
421            wasmtime::Engine::new(&config).unwrap()
422        })
423        .clone()
424}
425
426impl WasmHost {
427    pub fn new(
428        fs: Arc<dyn Fs>,
429        http_client: Arc<dyn HttpClient>,
430        node_runtime: NodeRuntime,
431        proxy: Arc<ExtensionHostProxy>,
432        work_dir: PathBuf,
433        cx: &mut App,
434    ) -> Arc<Self> {
435        let (tx, mut rx) = mpsc::unbounded::<MainThreadCall>();
436        let task = cx.spawn(async move |cx| {
437            while let Some(message) = rx.next().await {
438                message(cx).await;
439            }
440        });
441        Arc::new(Self {
442            engine: wasm_engine(),
443            fs,
444            work_dir,
445            http_client,
446            node_runtime,
447            proxy,
448            release_channel: ReleaseChannel::global(cx),
449            _main_thread_message_task: task,
450            main_thread_message_tx: tx,
451        })
452    }
453
454    pub fn load_extension(
455        self: &Arc<Self>,
456        wasm_bytes: Vec<u8>,
457        manifest: &Arc<ExtensionManifest>,
458        executor: BackgroundExecutor,
459    ) -> Task<Result<WasmExtension>> {
460        let this = self.clone();
461        let manifest = manifest.clone();
462        executor.clone().spawn(async move {
463            let zed_api_version = parse_wasm_extension_version(&manifest.id, &wasm_bytes)?;
464
465            let component = Component::from_binary(&this.engine, &wasm_bytes)
466                .context("failed to compile wasm component")?;
467
468            let mut store = wasmtime::Store::new(
469                &this.engine,
470                WasmState {
471                    ctx: this.build_wasi_ctx(&manifest).await?,
472                    manifest: manifest.clone(),
473                    table: ResourceTable::new(),
474                    host: this.clone(),
475                },
476            );
477
478            let mut extension = Extension::instantiate_async(
479                &mut store,
480                this.release_channel,
481                zed_api_version,
482                &component,
483            )
484            .await?;
485
486            extension
487                .call_init_extension(&mut store)
488                .await
489                .context("failed to initialize wasm extension")?;
490
491            let (tx, mut rx) = mpsc::unbounded::<ExtensionCall>();
492            executor
493                .spawn(async move {
494                    while let Some(call) = rx.next().await {
495                        (call)(&mut extension, &mut store).await;
496                    }
497                })
498                .detach();
499
500            Ok(WasmExtension {
501                manifest: manifest.clone(),
502                work_dir: this.work_dir.join(manifest.id.as_ref()).into(),
503                tx,
504                zed_api_version,
505            })
506        })
507    }
508
509    async fn build_wasi_ctx(&self, manifest: &Arc<ExtensionManifest>) -> Result<wasi::WasiCtx> {
510        let extension_work_dir = self.work_dir.join(manifest.id.as_ref());
511        self.fs
512            .create_dir(&extension_work_dir)
513            .await
514            .context("failed to create extension work dir")?;
515
516        let file_perms = wasi::FilePerms::all();
517        let dir_perms = wasi::DirPerms::all();
518
519        Ok(wasi::WasiCtxBuilder::new()
520            .inherit_stdio()
521            .preopened_dir(&extension_work_dir, ".", dir_perms, file_perms)?
522            .preopened_dir(
523                &extension_work_dir,
524                extension_work_dir.to_string_lossy(),
525                dir_perms,
526                file_perms,
527            )?
528            .env("PWD", extension_work_dir.to_string_lossy())
529            .env("RUST_BACKTRACE", "full")
530            .build())
531    }
532
533    pub fn writeable_path_from_extension(&self, id: &Arc<str>, path: &Path) -> Result<PathBuf> {
534        let extension_work_dir = self.work_dir.join(id.as_ref());
535        let path = normalize_path(&extension_work_dir.join(path));
536        if path.starts_with(&extension_work_dir) {
537            Ok(path)
538        } else {
539            Err(anyhow!("cannot write to path {}", path.display()))
540        }
541    }
542}
543
544pub fn parse_wasm_extension_version(
545    extension_id: &str,
546    wasm_bytes: &[u8],
547) -> Result<SemanticVersion> {
548    let mut version = None;
549
550    for part in wasmparser::Parser::new(0).parse_all(wasm_bytes) {
551        if let wasmparser::Payload::CustomSection(s) =
552            part.context("error parsing wasm extension")?
553        {
554            if s.name() == "zed:api-version" {
555                version = parse_wasm_extension_version_custom_section(s.data());
556                if version.is_none() {
557                    bail!(
558                        "extension {} has invalid zed:api-version section: {:?}",
559                        extension_id,
560                        s.data()
561                    );
562                }
563            }
564        }
565    }
566
567    // The reason we wait until we're done parsing all of the Wasm bytes to return the version
568    // is to work around a panic that can happen inside of Wasmtime when the bytes are invalid.
569    //
570    // By parsing the entirety of the Wasm bytes before we return, we're able to detect this problem
571    // earlier as an `Err` rather than as a panic.
572    version.ok_or_else(|| anyhow!("extension {} has no zed:api-version section", extension_id))
573}
574
575fn parse_wasm_extension_version_custom_section(data: &[u8]) -> Option<SemanticVersion> {
576    if data.len() == 6 {
577        Some(SemanticVersion::new(
578            u16::from_be_bytes([data[0], data[1]]) as _,
579            u16::from_be_bytes([data[2], data[3]]) as _,
580            u16::from_be_bytes([data[4], data[5]]) as _,
581        ))
582    } else {
583        None
584    }
585}
586
587impl WasmExtension {
588    pub async fn load(
589        extension_dir: PathBuf,
590        manifest: &Arc<ExtensionManifest>,
591        wasm_host: Arc<WasmHost>,
592        cx: &AsyncApp,
593    ) -> Result<Self> {
594        let path = extension_dir.join("extension.wasm");
595
596        let mut wasm_file = wasm_host
597            .fs
598            .open_sync(&path)
599            .await
600            .context("failed to open wasm file")?;
601
602        let mut wasm_bytes = Vec::new();
603        wasm_file
604            .read_to_end(&mut wasm_bytes)
605            .context("failed to read wasm")?;
606
607        wasm_host
608            .load_extension(wasm_bytes, manifest, cx.background_executor().clone())
609            .await
610            .with_context(|| format!("failed to load wasm extension {}", manifest.id))
611    }
612
613    pub async fn call<T, Fn>(&self, f: Fn) -> T
614    where
615        T: 'static + Send,
616        Fn: 'static
617            + Send
618            + for<'a> FnOnce(&'a mut Extension, &'a mut Store<WasmState>) -> BoxFuture<'a, T>,
619    {
620        let (return_tx, return_rx) = oneshot::channel();
621        self.tx
622            .clone()
623            .unbounded_send(Box::new(move |extension, store| {
624                async {
625                    let result = f(extension, store).await;
626                    return_tx.send(result).ok();
627                }
628                .boxed()
629            }))
630            .expect("wasm extension channel should not be closed yet");
631        return_rx.await.expect("wasm extension channel")
632    }
633}
634
635impl WasmState {
636    fn on_main_thread<T, Fn>(&self, f: Fn) -> impl 'static + Future<Output = T>
637    where
638        T: 'static + Send,
639        Fn: 'static + Send + for<'a> FnOnce(&'a mut AsyncApp) -> LocalBoxFuture<'a, T>,
640    {
641        let (return_tx, return_rx) = oneshot::channel();
642        self.host
643            .main_thread_message_tx
644            .clone()
645            .unbounded_send(Box::new(move |cx| {
646                async {
647                    let result = f(cx).await;
648                    return_tx.send(result).ok();
649                }
650                .boxed_local()
651            }))
652            .expect("main thread message channel should not be closed yet");
653        async move { return_rx.await.expect("main thread message channel") }
654    }
655
656    fn work_dir(&self) -> PathBuf {
657        self.host.work_dir.join(self.manifest.id.as_ref())
658    }
659}
660
661impl wasi::WasiView for WasmState {
662    fn table(&mut self) -> &mut ResourceTable {
663        &mut self.table
664    }
665
666    fn ctx(&mut self) -> &mut wasi::WasiCtx {
667        &mut self.ctx
668    }
669}