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    ) -> Result<DebugAdapterBinary> {
383        self.call(|extension, store| {
384            async move {
385                let dap_binary = extension
386                    .call_get_dap_binary(store, dap_name, config, user_installed_path)
387                    .await?
388                    .map_err(|err| anyhow!("{err:?}"))?;
389                let dap_binary = dap_binary.try_into()?;
390                Ok(dap_binary)
391            }
392            .boxed()
393        })
394        .await
395    }
396}
397
398pub struct WasmState {
399    manifest: Arc<ExtensionManifest>,
400    pub table: ResourceTable,
401    ctx: wasi::WasiCtx,
402    pub host: Arc<WasmHost>,
403}
404
405type MainThreadCall = Box<dyn Send + for<'a> FnOnce(&'a mut AsyncApp) -> LocalBoxFuture<'a, ()>>;
406
407type ExtensionCall = Box<
408    dyn Send + for<'a> FnOnce(&'a mut Extension, &'a mut Store<WasmState>) -> BoxFuture<'a, ()>,
409>;
410
411fn wasm_engine() -> wasmtime::Engine {
412    static WASM_ENGINE: OnceLock<wasmtime::Engine> = OnceLock::new();
413
414    WASM_ENGINE
415        .get_or_init(|| {
416            let mut config = wasmtime::Config::new();
417            config.wasm_component_model(true);
418            config.async_support(true);
419            wasmtime::Engine::new(&config).unwrap()
420        })
421        .clone()
422}
423
424impl WasmHost {
425    pub fn new(
426        fs: Arc<dyn Fs>,
427        http_client: Arc<dyn HttpClient>,
428        node_runtime: NodeRuntime,
429        proxy: Arc<ExtensionHostProxy>,
430        work_dir: PathBuf,
431        cx: &mut App,
432    ) -> Arc<Self> {
433        let (tx, mut rx) = mpsc::unbounded::<MainThreadCall>();
434        let task = cx.spawn(async move |cx| {
435            while let Some(message) = rx.next().await {
436                message(cx).await;
437            }
438        });
439        Arc::new(Self {
440            engine: wasm_engine(),
441            fs,
442            work_dir,
443            http_client,
444            node_runtime,
445            proxy,
446            release_channel: ReleaseChannel::global(cx),
447            _main_thread_message_task: task,
448            main_thread_message_tx: tx,
449        })
450    }
451
452    pub fn load_extension(
453        self: &Arc<Self>,
454        wasm_bytes: Vec<u8>,
455        manifest: &Arc<ExtensionManifest>,
456        executor: BackgroundExecutor,
457    ) -> Task<Result<WasmExtension>> {
458        let this = self.clone();
459        let manifest = manifest.clone();
460        executor.clone().spawn(async move {
461            let zed_api_version = parse_wasm_extension_version(&manifest.id, &wasm_bytes)?;
462
463            let component = Component::from_binary(&this.engine, &wasm_bytes)
464                .context("failed to compile wasm component")?;
465
466            let mut store = wasmtime::Store::new(
467                &this.engine,
468                WasmState {
469                    ctx: this.build_wasi_ctx(&manifest).await?,
470                    manifest: manifest.clone(),
471                    table: ResourceTable::new(),
472                    host: this.clone(),
473                },
474            );
475
476            let mut extension = Extension::instantiate_async(
477                &mut store,
478                this.release_channel,
479                zed_api_version,
480                &component,
481            )
482            .await?;
483
484            extension
485                .call_init_extension(&mut store)
486                .await
487                .context("failed to initialize wasm extension")?;
488
489            let (tx, mut rx) = mpsc::unbounded::<ExtensionCall>();
490            executor
491                .spawn(async move {
492                    while let Some(call) = rx.next().await {
493                        (call)(&mut extension, &mut store).await;
494                    }
495                })
496                .detach();
497
498            Ok(WasmExtension {
499                manifest: manifest.clone(),
500                work_dir: this.work_dir.join(manifest.id.as_ref()).into(),
501                tx,
502                zed_api_version,
503            })
504        })
505    }
506
507    async fn build_wasi_ctx(&self, manifest: &Arc<ExtensionManifest>) -> Result<wasi::WasiCtx> {
508        let extension_work_dir = self.work_dir.join(manifest.id.as_ref());
509        self.fs
510            .create_dir(&extension_work_dir)
511            .await
512            .context("failed to create extension work dir")?;
513
514        let file_perms = wasi::FilePerms::all();
515        let dir_perms = wasi::DirPerms::all();
516
517        Ok(wasi::WasiCtxBuilder::new()
518            .inherit_stdio()
519            .preopened_dir(&extension_work_dir, ".", dir_perms, file_perms)?
520            .preopened_dir(
521                &extension_work_dir,
522                extension_work_dir.to_string_lossy(),
523                dir_perms,
524                file_perms,
525            )?
526            .env("PWD", extension_work_dir.to_string_lossy())
527            .env("RUST_BACKTRACE", "full")
528            .build())
529    }
530
531    pub fn writeable_path_from_extension(&self, id: &Arc<str>, path: &Path) -> Result<PathBuf> {
532        let extension_work_dir = self.work_dir.join(id.as_ref());
533        let path = normalize_path(&extension_work_dir.join(path));
534        if path.starts_with(&extension_work_dir) {
535            Ok(path)
536        } else {
537            Err(anyhow!("cannot write to path {}", path.display()))
538        }
539    }
540}
541
542pub fn parse_wasm_extension_version(
543    extension_id: &str,
544    wasm_bytes: &[u8],
545) -> Result<SemanticVersion> {
546    let mut version = None;
547
548    for part in wasmparser::Parser::new(0).parse_all(wasm_bytes) {
549        if let wasmparser::Payload::CustomSection(s) =
550            part.context("error parsing wasm extension")?
551        {
552            if s.name() == "zed:api-version" {
553                version = parse_wasm_extension_version_custom_section(s.data());
554                if version.is_none() {
555                    bail!(
556                        "extension {} has invalid zed:api-version section: {:?}",
557                        extension_id,
558                        s.data()
559                    );
560                }
561            }
562        }
563    }
564
565    // The reason we wait until we're done parsing all of the Wasm bytes to return the version
566    // is to work around a panic that can happen inside of Wasmtime when the bytes are invalid.
567    //
568    // By parsing the entirety of the Wasm bytes before we return, we're able to detect this problem
569    // earlier as an `Err` rather than as a panic.
570    version.ok_or_else(|| anyhow!("extension {} has no zed:api-version section", extension_id))
571}
572
573fn parse_wasm_extension_version_custom_section(data: &[u8]) -> Option<SemanticVersion> {
574    if data.len() == 6 {
575        Some(SemanticVersion::new(
576            u16::from_be_bytes([data[0], data[1]]) as _,
577            u16::from_be_bytes([data[2], data[3]]) as _,
578            u16::from_be_bytes([data[4], data[5]]) as _,
579        ))
580    } else {
581        None
582    }
583}
584
585impl WasmExtension {
586    pub async fn load(
587        extension_dir: PathBuf,
588        manifest: &Arc<ExtensionManifest>,
589        wasm_host: Arc<WasmHost>,
590        cx: &AsyncApp,
591    ) -> Result<Self> {
592        let path = extension_dir.join("extension.wasm");
593
594        let mut wasm_file = wasm_host
595            .fs
596            .open_sync(&path)
597            .await
598            .context("failed to open wasm file")?;
599
600        let mut wasm_bytes = Vec::new();
601        wasm_file
602            .read_to_end(&mut wasm_bytes)
603            .context("failed to read wasm")?;
604
605        wasm_host
606            .load_extension(wasm_bytes, manifest, cx.background_executor().clone())
607            .await
608            .with_context(|| format!("failed to load wasm extension {}", manifest.id))
609    }
610
611    pub async fn call<T, Fn>(&self, f: Fn) -> T
612    where
613        T: 'static + Send,
614        Fn: 'static
615            + Send
616            + for<'a> FnOnce(&'a mut Extension, &'a mut Store<WasmState>) -> BoxFuture<'a, T>,
617    {
618        let (return_tx, return_rx) = oneshot::channel();
619        self.tx
620            .clone()
621            .unbounded_send(Box::new(move |extension, store| {
622                async {
623                    let result = f(extension, store).await;
624                    return_tx.send(result).ok();
625                }
626                .boxed()
627            }))
628            .expect("wasm extension channel should not be closed yet");
629        return_rx.await.expect("wasm extension channel")
630    }
631}
632
633impl WasmState {
634    fn on_main_thread<T, Fn>(&self, f: Fn) -> impl 'static + Future<Output = T>
635    where
636        T: 'static + Send,
637        Fn: 'static + Send + for<'a> FnOnce(&'a mut AsyncApp) -> LocalBoxFuture<'a, T>,
638    {
639        let (return_tx, return_rx) = oneshot::channel();
640        self.host
641            .main_thread_message_tx
642            .clone()
643            .unbounded_send(Box::new(move |cx| {
644                async {
645                    let result = f(cx).await;
646                    return_tx.send(result).ok();
647                }
648                .boxed_local()
649            }))
650            .expect("main thread message channel should not be closed yet");
651        async move { return_rx.await.expect("main thread message channel") }
652    }
653
654    fn work_dir(&self) -> PathBuf {
655        self.host.work_dir.join(self.manifest.id.as_ref())
656    }
657}
658
659impl wasi::WasiView for WasmState {
660    fn table(&mut self) -> &mut ResourceTable {
661        &mut self.table
662    }
663
664    fn ctx(&mut self) -> &mut wasi::WasiCtx {
665        &mut self.ctx
666    }
667}