wasm_host.rs

  1pub mod wit;
  2
  3use crate::ExtensionManifest;
  4use crate::capability_granter::CapabilityGranter;
  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, DownloadFileCapability, ExtensionCapability, ExtensionHostProxy,
 11    KeyValueStoreDelegate, NpmInstallPackageCapability, ProcessExecCapability, ProjectDelegate,
 12    SlashCommand, SlashCommandArgumentCompletion, SlashCommandOutput, Symbol, WorktreeDelegate,
 13};
 14use fs::{Fs, normalize_path};
 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, Timer};
 25use http_client::HttpClient;
 26use language::LanguageName;
 27use lsp::LanguageServerName;
 28use moka::sync::Cache;
 29use node_runtime::NodeRuntime;
 30use release_channel::ReleaseChannel;
 31use semantic_version::SemanticVersion;
 32use std::borrow::Cow;
 33use std::sync::{LazyLock, OnceLock};
 34use std::time::Duration;
 35use std::{
 36    path::{Path, PathBuf},
 37    sync::Arc,
 38};
 39use task::{DebugScenario, SpawnInTerminal, TaskTemplate, ZedDebugConfig};
 40use wasmtime::{
 41    CacheStore, Engine, Store,
 42    component::{Component, ResourceTable},
 43};
 44use wasmtime_wasi::{self as wasi, WasiView};
 45use wit::Extension;
 46
 47pub struct WasmHost {
 48    engine: Engine,
 49    release_channel: ReleaseChannel,
 50    http_client: Arc<dyn HttpClient>,
 51    node_runtime: NodeRuntime,
 52    pub(crate) proxy: Arc<ExtensionHostProxy>,
 53    fs: Arc<dyn Fs>,
 54    pub work_dir: PathBuf,
 55    /// The capabilities granted to extensions running on the host.
 56    pub(crate) granted_capabilities: Vec<ExtensionCapability>,
 57    _main_thread_message_task: Task<()>,
 58    main_thread_message_tx: mpsc::UnboundedSender<MainThreadCall>,
 59}
 60
 61#[derive(Clone, Debug)]
 62pub struct WasmExtension {
 63    tx: UnboundedSender<ExtensionCall>,
 64    pub manifest: Arc<ExtensionManifest>,
 65    pub work_dir: Arc<Path>,
 66    #[allow(unused)]
 67    pub zed_api_version: SemanticVersion,
 68}
 69
 70impl Drop for WasmExtension {
 71    fn drop(&mut self) {
 72        self.tx.close_channel();
 73    }
 74}
 75
 76#[async_trait]
 77impl extension::Extension for WasmExtension {
 78    fn manifest(&self) -> Arc<ExtensionManifest> {
 79        self.manifest.clone()
 80    }
 81
 82    fn work_dir(&self) -> Arc<Path> {
 83        self.work_dir.clone()
 84    }
 85
 86    async fn language_server_command(
 87        &self,
 88        language_server_id: LanguageServerName,
 89        language_name: LanguageName,
 90        worktree: Arc<dyn WorktreeDelegate>,
 91    ) -> Result<Command> {
 92        self.call(|extension, store| {
 93            async move {
 94                let resource = store.data_mut().table().push(worktree)?;
 95                let command = extension
 96                    .call_language_server_command(
 97                        store,
 98                        &language_server_id,
 99                        &language_name,
100                        resource,
101                    )
102                    .await?
103                    .map_err(|err| store.data().extension_error(err))?;
104
105                Ok(command.into())
106            }
107            .boxed()
108        })
109        .await?
110    }
111
112    async fn language_server_initialization_options(
113        &self,
114        language_server_id: LanguageServerName,
115        language_name: LanguageName,
116        worktree: Arc<dyn WorktreeDelegate>,
117    ) -> Result<Option<String>> {
118        self.call(|extension, store| {
119            async move {
120                let resource = store.data_mut().table().push(worktree)?;
121                let options = extension
122                    .call_language_server_initialization_options(
123                        store,
124                        &language_server_id,
125                        &language_name,
126                        resource,
127                    )
128                    .await?
129                    .map_err(|err| store.data().extension_error(err))?;
130                anyhow::Ok(options)
131            }
132            .boxed()
133        })
134        .await?
135    }
136
137    async fn language_server_workspace_configuration(
138        &self,
139        language_server_id: LanguageServerName,
140        worktree: Arc<dyn WorktreeDelegate>,
141    ) -> Result<Option<String>> {
142        self.call(|extension, store| {
143            async move {
144                let resource = store.data_mut().table().push(worktree)?;
145                let options = extension
146                    .call_language_server_workspace_configuration(
147                        store,
148                        &language_server_id,
149                        resource,
150                    )
151                    .await?
152                    .map_err(|err| store.data().extension_error(err))?;
153                anyhow::Ok(options)
154            }
155            .boxed()
156        })
157        .await?
158    }
159
160    async fn language_server_additional_initialization_options(
161        &self,
162        language_server_id: LanguageServerName,
163        target_language_server_id: LanguageServerName,
164        worktree: Arc<dyn WorktreeDelegate>,
165    ) -> Result<Option<String>> {
166        self.call(|extension, store| {
167            async move {
168                let resource = store.data_mut().table().push(worktree)?;
169                let options = extension
170                    .call_language_server_additional_initialization_options(
171                        store,
172                        &language_server_id,
173                        &target_language_server_id,
174                        resource,
175                    )
176                    .await?
177                    .map_err(|err| store.data().extension_error(err))?;
178                anyhow::Ok(options)
179            }
180            .boxed()
181        })
182        .await?
183    }
184
185    async fn language_server_additional_workspace_configuration(
186        &self,
187        language_server_id: LanguageServerName,
188        target_language_server_id: LanguageServerName,
189        worktree: Arc<dyn WorktreeDelegate>,
190    ) -> Result<Option<String>> {
191        self.call(|extension, store| {
192            async move {
193                let resource = store.data_mut().table().push(worktree)?;
194                let options = extension
195                    .call_language_server_additional_workspace_configuration(
196                        store,
197                        &language_server_id,
198                        &target_language_server_id,
199                        resource,
200                    )
201                    .await?
202                    .map_err(|err| store.data().extension_error(err))?;
203                anyhow::Ok(options)
204            }
205            .boxed()
206        })
207        .await?
208    }
209
210    async fn labels_for_completions(
211        &self,
212        language_server_id: LanguageServerName,
213        completions: Vec<Completion>,
214    ) -> Result<Vec<Option<CodeLabel>>> {
215        self.call(|extension, store| {
216            async move {
217                let labels = extension
218                    .call_labels_for_completions(
219                        store,
220                        &language_server_id,
221                        completions.into_iter().map(Into::into).collect(),
222                    )
223                    .await?
224                    .map_err(|err| store.data().extension_error(err))?;
225
226                Ok(labels
227                    .into_iter()
228                    .map(|label| label.map(Into::into))
229                    .collect())
230            }
231            .boxed()
232        })
233        .await?
234    }
235
236    async fn labels_for_symbols(
237        &self,
238        language_server_id: LanguageServerName,
239        symbols: Vec<Symbol>,
240    ) -> Result<Vec<Option<CodeLabel>>> {
241        self.call(|extension, store| {
242            async move {
243                let labels = extension
244                    .call_labels_for_symbols(
245                        store,
246                        &language_server_id,
247                        symbols.into_iter().map(Into::into).collect(),
248                    )
249                    .await?
250                    .map_err(|err| store.data().extension_error(err))?;
251
252                Ok(labels
253                    .into_iter()
254                    .map(|label| label.map(Into::into))
255                    .collect())
256            }
257            .boxed()
258        })
259        .await?
260    }
261
262    async fn complete_slash_command_argument(
263        &self,
264        command: SlashCommand,
265        arguments: Vec<String>,
266    ) -> Result<Vec<SlashCommandArgumentCompletion>> {
267        self.call(|extension, store| {
268            async move {
269                let completions = extension
270                    .call_complete_slash_command_argument(store, &command.into(), &arguments)
271                    .await?
272                    .map_err(|err| store.data().extension_error(err))?;
273
274                Ok(completions.into_iter().map(Into::into).collect())
275            }
276            .boxed()
277        })
278        .await?
279    }
280
281    async fn run_slash_command(
282        &self,
283        command: SlashCommand,
284        arguments: Vec<String>,
285        delegate: Option<Arc<dyn WorktreeDelegate>>,
286    ) -> Result<SlashCommandOutput> {
287        self.call(|extension, store| {
288            async move {
289                let resource = if let Some(delegate) = delegate {
290                    Some(store.data_mut().table().push(delegate)?)
291                } else {
292                    None
293                };
294
295                let output = extension
296                    .call_run_slash_command(store, &command.into(), &arguments, resource)
297                    .await?
298                    .map_err(|err| store.data().extension_error(err))?;
299
300                Ok(output.into())
301            }
302            .boxed()
303        })
304        .await?
305    }
306
307    async fn context_server_command(
308        &self,
309        context_server_id: Arc<str>,
310        project: Arc<dyn ProjectDelegate>,
311    ) -> Result<Command> {
312        self.call(|extension, store| {
313            async move {
314                let project_resource = store.data_mut().table().push(project)?;
315                let command = extension
316                    .call_context_server_command(store, context_server_id.clone(), project_resource)
317                    .await?
318                    .map_err(|err| store.data().extension_error(err))?;
319                anyhow::Ok(command.into())
320            }
321            .boxed()
322        })
323        .await?
324    }
325
326    async fn context_server_configuration(
327        &self,
328        context_server_id: Arc<str>,
329        project: Arc<dyn ProjectDelegate>,
330    ) -> Result<Option<ContextServerConfiguration>> {
331        self.call(|extension, store| {
332            async move {
333                let project_resource = store.data_mut().table().push(project)?;
334                let Some(configuration) = extension
335                    .call_context_server_configuration(
336                        store,
337                        context_server_id.clone(),
338                        project_resource,
339                    )
340                    .await?
341                    .map_err(|err| store.data().extension_error(err))?
342                else {
343                    return Ok(None);
344                };
345
346                Ok(Some(configuration.try_into()?))
347            }
348            .boxed()
349        })
350        .await?
351    }
352
353    async fn suggest_docs_packages(&self, provider: Arc<str>) -> Result<Vec<String>> {
354        self.call(|extension, store| {
355            async move {
356                let packages = extension
357                    .call_suggest_docs_packages(store, provider.as_ref())
358                    .await?
359                    .map_err(|err| store.data().extension_error(err))?;
360
361                Ok(packages)
362            }
363            .boxed()
364        })
365        .await?
366    }
367
368    async fn index_docs(
369        &self,
370        provider: Arc<str>,
371        package_name: Arc<str>,
372        kv_store: Arc<dyn KeyValueStoreDelegate>,
373    ) -> Result<()> {
374        self.call(|extension, store| {
375            async move {
376                let kv_store_resource = store.data_mut().table().push(kv_store)?;
377                extension
378                    .call_index_docs(
379                        store,
380                        provider.as_ref(),
381                        package_name.as_ref(),
382                        kv_store_resource,
383                    )
384                    .await?
385                    .map_err(|err| store.data().extension_error(err))?;
386
387                anyhow::Ok(())
388            }
389            .boxed()
390        })
391        .await?
392    }
393
394    async fn get_dap_binary(
395        &self,
396        dap_name: Arc<str>,
397        config: DebugTaskDefinition,
398        user_installed_path: Option<PathBuf>,
399        worktree: Arc<dyn WorktreeDelegate>,
400    ) -> Result<DebugAdapterBinary> {
401        self.call(|extension, store| {
402            async move {
403                let resource = store.data_mut().table().push(worktree)?;
404                let dap_binary = extension
405                    .call_get_dap_binary(store, dap_name, config, user_installed_path, resource)
406                    .await?
407                    .map_err(|err| store.data().extension_error(err))?;
408                let dap_binary = dap_binary.try_into()?;
409                Ok(dap_binary)
410            }
411            .boxed()
412        })
413        .await?
414    }
415    async fn dap_request_kind(
416        &self,
417        dap_name: Arc<str>,
418        config: serde_json::Value,
419    ) -> Result<StartDebuggingRequestArgumentsRequest> {
420        self.call(|extension, store| {
421            async move {
422                let kind = extension
423                    .call_dap_request_kind(store, dap_name, config)
424                    .await?
425                    .map_err(|err| store.data().extension_error(err))?;
426                Ok(kind.into())
427            }
428            .boxed()
429        })
430        .await?
431    }
432
433    async fn dap_config_to_scenario(&self, config: ZedDebugConfig) -> Result<DebugScenario> {
434        self.call(|extension, store| {
435            async move {
436                let kind = extension
437                    .call_dap_config_to_scenario(store, config)
438                    .await?
439                    .map_err(|err| store.data().extension_error(err))?;
440                Ok(kind)
441            }
442            .boxed()
443        })
444        .await?
445    }
446
447    async fn dap_locator_create_scenario(
448        &self,
449        locator_name: String,
450        build_config_template: TaskTemplate,
451        resolved_label: String,
452        debug_adapter_name: String,
453    ) -> Result<Option<DebugScenario>> {
454        self.call(|extension, store| {
455            async move {
456                extension
457                    .call_dap_locator_create_scenario(
458                        store,
459                        locator_name,
460                        build_config_template,
461                        resolved_label,
462                        debug_adapter_name,
463                    )
464                    .await
465            }
466            .boxed()
467        })
468        .await?
469    }
470    async fn run_dap_locator(
471        &self,
472        locator_name: String,
473        config: SpawnInTerminal,
474    ) -> Result<DebugRequest> {
475        self.call(|extension, store| {
476            async move {
477                extension
478                    .call_run_dap_locator(store, locator_name, config)
479                    .await?
480                    .map_err(|err| store.data().extension_error(err))
481            }
482            .boxed()
483        })
484        .await?
485    }
486}
487
488pub struct WasmState {
489    manifest: Arc<ExtensionManifest>,
490    pub table: ResourceTable,
491    ctx: wasi::WasiCtx,
492    pub host: Arc<WasmHost>,
493    pub(crate) capability_granter: CapabilityGranter,
494}
495
496type MainThreadCall = Box<dyn Send + for<'a> FnOnce(&'a mut AsyncApp) -> LocalBoxFuture<'a, ()>>;
497
498type ExtensionCall = Box<
499    dyn Send + for<'a> FnOnce(&'a mut Extension, &'a mut Store<WasmState>) -> BoxFuture<'a, ()>,
500>;
501
502fn wasm_engine(executor: &BackgroundExecutor) -> wasmtime::Engine {
503    static WASM_ENGINE: OnceLock<wasmtime::Engine> = OnceLock::new();
504    WASM_ENGINE
505        .get_or_init(|| {
506            let mut config = wasmtime::Config::new();
507            config.wasm_component_model(true);
508            config.async_support(true);
509            config
510                .enable_incremental_compilation(cache_store())
511                .unwrap();
512            // Async support introduces the issue that extension execution happens during `Future::poll`,
513            // which could block an async thread.
514            // https://docs.rs/wasmtime/latest/wasmtime/struct.Config.html#execution-in-poll
515            //
516            // Epoch interruption is a lightweight mechanism to allow the extensions to yield control
517            // back to the executor at regular intervals.
518            config.epoch_interruption(true);
519
520            let engine = wasmtime::Engine::new(&config).unwrap();
521
522            // It might be safer to do this on a non-async thread to make sure it makes progress
523            // regardless of if extensions are blocking.
524            // However, due to our current setup, this isn't a likely occurrence and we'd rather
525            // not have a dedicated thread just for this. If it becomes an issue, we can consider
526            // creating a separate thread for epoch interruption.
527            let engine_ref = engine.weak();
528            executor
529                .spawn(async move {
530                    // Somewhat arbitrary interval, as it isn't a guaranteed interval.
531                    // But this is a rough upper bound for how long the extension execution can block on
532                    // `Future::poll`.
533                    const EPOCH_INTERVAL: Duration = Duration::from_millis(100);
534                    let mut timer = Timer::interval(EPOCH_INTERVAL);
535                    while (timer.next().await).is_some() {
536                        // Exit the loop and thread once the engine is dropped.
537                        let Some(engine) = engine_ref.upgrade() else {
538                            break;
539                        };
540                        engine.increment_epoch();
541                    }
542                })
543                .detach();
544
545            engine
546        })
547        .clone()
548}
549
550fn cache_store() -> Arc<IncrementalCompilationCache> {
551    static CACHE_STORE: LazyLock<Arc<IncrementalCompilationCache>> =
552        LazyLock::new(|| Arc::new(IncrementalCompilationCache::new()));
553    CACHE_STORE.clone()
554}
555
556impl WasmHost {
557    pub fn new(
558        fs: Arc<dyn Fs>,
559        http_client: Arc<dyn HttpClient>,
560        node_runtime: NodeRuntime,
561        proxy: Arc<ExtensionHostProxy>,
562        work_dir: PathBuf,
563        cx: &mut App,
564    ) -> Arc<Self> {
565        let (tx, mut rx) = mpsc::unbounded::<MainThreadCall>();
566        let task = cx.spawn(async move |cx| {
567            while let Some(message) = rx.next().await {
568                message(cx).await;
569            }
570        });
571        Arc::new(Self {
572            engine: wasm_engine(cx.background_executor()),
573            fs,
574            work_dir,
575            http_client,
576            node_runtime,
577            proxy,
578            release_channel: ReleaseChannel::global(cx),
579            granted_capabilities: vec![
580                ExtensionCapability::ProcessExec(ProcessExecCapability {
581                    command: "*".to_string(),
582                    args: vec!["**".to_string()],
583                }),
584                ExtensionCapability::DownloadFile(DownloadFileCapability {
585                    host: "*".to_string(),
586                    path: vec!["**".to_string()],
587                }),
588                ExtensionCapability::NpmInstallPackage(NpmInstallPackageCapability {
589                    package: "*".to_string(),
590                }),
591            ],
592            _main_thread_message_task: task,
593            main_thread_message_tx: tx,
594        })
595    }
596
597    pub fn load_extension(
598        self: &Arc<Self>,
599        wasm_bytes: Vec<u8>,
600        manifest: &Arc<ExtensionManifest>,
601        executor: BackgroundExecutor,
602    ) -> Task<Result<WasmExtension>> {
603        let this = self.clone();
604        let manifest = manifest.clone();
605        executor.clone().spawn(async move {
606            let zed_api_version = parse_wasm_extension_version(&manifest.id, &wasm_bytes)?;
607
608            let component = Component::from_binary(&this.engine, &wasm_bytes)
609                .context("failed to compile wasm component")?;
610
611            let mut store = wasmtime::Store::new(
612                &this.engine,
613                WasmState {
614                    ctx: this.build_wasi_ctx(&manifest).await?,
615                    manifest: manifest.clone(),
616                    table: ResourceTable::new(),
617                    host: this.clone(),
618                    capability_granter: CapabilityGranter::new(
619                        this.granted_capabilities.clone(),
620                        manifest.clone(),
621                    ),
622                },
623            );
624            // Store will yield after 1 tick, and get a new deadline of 1 tick after each yield.
625            store.set_epoch_deadline(1);
626            store.epoch_deadline_async_yield_and_update(1);
627
628            let mut extension = Extension::instantiate_async(
629                &executor,
630                &mut store,
631                this.release_channel,
632                zed_api_version,
633                &component,
634            )
635            .await?;
636
637            extension
638                .call_init_extension(&mut store)
639                .await
640                .context("failed to initialize wasm extension")?;
641
642            let (tx, mut rx) = mpsc::unbounded::<ExtensionCall>();
643            executor
644                .spawn(async move {
645                    while let Some(call) = rx.next().await {
646                        (call)(&mut extension, &mut store).await;
647                    }
648                })
649                .detach();
650
651            Ok(WasmExtension {
652                manifest: manifest.clone(),
653                work_dir: this.work_dir.join(manifest.id.as_ref()).into(),
654                tx,
655                zed_api_version,
656            })
657        })
658    }
659
660    async fn build_wasi_ctx(&self, manifest: &Arc<ExtensionManifest>) -> Result<wasi::WasiCtx> {
661        let extension_work_dir = self.work_dir.join(manifest.id.as_ref());
662        self.fs
663            .create_dir(&extension_work_dir)
664            .await
665            .context("failed to create extension work dir")?;
666
667        let file_perms = wasi::FilePerms::all();
668        let dir_perms = wasi::DirPerms::all();
669
670        Ok(wasi::WasiCtxBuilder::new()
671            .inherit_stdio()
672            .preopened_dir(&extension_work_dir, ".", dir_perms, file_perms)?
673            .preopened_dir(
674                &extension_work_dir,
675                extension_work_dir.to_string_lossy(),
676                dir_perms,
677                file_perms,
678            )?
679            .env("PWD", extension_work_dir.to_string_lossy())
680            .env("RUST_BACKTRACE", "full")
681            .build())
682    }
683
684    pub fn writeable_path_from_extension(&self, id: &Arc<str>, path: &Path) -> Result<PathBuf> {
685        let extension_work_dir = self.work_dir.join(id.as_ref());
686        let path = normalize_path(&extension_work_dir.join(path));
687        anyhow::ensure!(
688            path.starts_with(&extension_work_dir),
689            "cannot write to path {path:?}",
690        );
691        Ok(path)
692    }
693}
694
695pub fn parse_wasm_extension_version(
696    extension_id: &str,
697    wasm_bytes: &[u8],
698) -> Result<SemanticVersion> {
699    let mut version = None;
700
701    for part in wasmparser::Parser::new(0).parse_all(wasm_bytes) {
702        if let wasmparser::Payload::CustomSection(s) =
703            part.context("error parsing wasm extension")?
704            && s.name() == "zed:api-version"
705        {
706            version = parse_wasm_extension_version_custom_section(s.data());
707            if version.is_none() {
708                bail!(
709                    "extension {} has invalid zed:api-version section: {:?}",
710                    extension_id,
711                    s.data()
712                );
713            }
714        }
715    }
716
717    // The reason we wait until we're done parsing all of the Wasm bytes to return the version
718    // is to work around a panic that can happen inside of Wasmtime when the bytes are invalid.
719    //
720    // By parsing the entirety of the Wasm bytes before we return, we're able to detect this problem
721    // earlier as an `Err` rather than as a panic.
722    version.with_context(|| format!("extension {extension_id} has no zed:api-version section"))
723}
724
725fn parse_wasm_extension_version_custom_section(data: &[u8]) -> Option<SemanticVersion> {
726    if data.len() == 6 {
727        Some(SemanticVersion::new(
728            u16::from_be_bytes([data[0], data[1]]) as _,
729            u16::from_be_bytes([data[2], data[3]]) as _,
730            u16::from_be_bytes([data[4], data[5]]) as _,
731        ))
732    } else {
733        None
734    }
735}
736
737impl WasmExtension {
738    pub async fn load(
739        extension_dir: &Path,
740        manifest: &Arc<ExtensionManifest>,
741        wasm_host: Arc<WasmHost>,
742        cx: &AsyncApp,
743    ) -> Result<Self> {
744        let path = extension_dir.join("extension.wasm");
745
746        let mut wasm_file = wasm_host
747            .fs
748            .open_sync(&path)
749            .await
750            .context("failed to open wasm file")?;
751
752        let mut wasm_bytes = Vec::new();
753        wasm_file
754            .read_to_end(&mut wasm_bytes)
755            .context("failed to read wasm")?;
756
757        wasm_host
758            .load_extension(wasm_bytes, manifest, cx.background_executor().clone())
759            .await
760            .with_context(|| format!("failed to load wasm extension {}", manifest.id))
761    }
762
763    pub async fn call<T, Fn>(&self, f: Fn) -> Result<T>
764    where
765        T: 'static + Send,
766        Fn: 'static
767            + Send
768            + for<'a> FnOnce(&'a mut Extension, &'a mut Store<WasmState>) -> BoxFuture<'a, T>,
769    {
770        let (return_tx, return_rx) = oneshot::channel();
771        self.tx
772            .unbounded_send(Box::new(move |extension, store| {
773                async {
774                    let result = f(extension, store).await;
775                    return_tx.send(result).ok();
776                }
777                .boxed()
778            }))
779            .map_err(|_| {
780                anyhow!(
781                    "wasm extension channel should not be closed yet, extension {} (id {})",
782                    self.manifest.name,
783                    self.manifest.id,
784                )
785            })?;
786        return_rx.await.with_context(|| {
787            format!(
788                "wasm extension channel, extension {} (id {})",
789                self.manifest.name, self.manifest.id,
790            )
791        })
792    }
793}
794
795impl WasmState {
796    fn on_main_thread<T, Fn>(&self, f: Fn) -> impl 'static + Future<Output = T>
797    where
798        T: 'static + Send,
799        Fn: 'static + Send + for<'a> FnOnce(&'a mut AsyncApp) -> LocalBoxFuture<'a, T>,
800    {
801        let (return_tx, return_rx) = oneshot::channel();
802        self.host
803            .main_thread_message_tx
804            .clone()
805            .unbounded_send(Box::new(move |cx| {
806                async {
807                    let result = f(cx).await;
808                    return_tx.send(result).ok();
809                }
810                .boxed_local()
811            }))
812            .unwrap_or_else(|_| {
813                panic!(
814                    "main thread message channel should not be closed yet, extension {} (id {})",
815                    self.manifest.name, self.manifest.id,
816                )
817            });
818        let name = self.manifest.name.clone();
819        let id = self.manifest.id.clone();
820        async move {
821            return_rx.await.unwrap_or_else(|_| {
822                panic!("main thread message channel, extension {name} (id {id})")
823            })
824        }
825    }
826
827    fn work_dir(&self) -> PathBuf {
828        self.host.work_dir.join(self.manifest.id.as_ref())
829    }
830
831    fn extension_error(&self, message: String) -> anyhow::Error {
832        anyhow!(
833            "from extension \"{}\" version {}: {}",
834            self.manifest.name,
835            self.manifest.version,
836            message
837        )
838    }
839}
840
841impl wasi::WasiView for WasmState {
842    fn table(&mut self) -> &mut ResourceTable {
843        &mut self.table
844    }
845
846    fn ctx(&mut self) -> &mut wasi::WasiCtx {
847        &mut self.ctx
848    }
849}
850
851/// Wrapper around a mini-moka bounded cache for storing incremental compilation artifacts.
852/// Since wasm modules have many similar elements, this can save us a lot of work at the
853/// cost of a small memory footprint. However, we don't want this to be unbounded, so we use
854/// a LFU/LRU cache to evict less used cache entries.
855#[derive(Debug)]
856struct IncrementalCompilationCache {
857    cache: Cache<Vec<u8>, Vec<u8>>,
858}
859
860impl IncrementalCompilationCache {
861    fn new() -> Self {
862        let cache = Cache::builder()
863            // Cap this at 32 MB for now. Our extensions turn into roughly 512kb in the cache,
864            // which means we could store 64 completely novel extensions in the cache, but in
865            // practice we will more than that, which is more than enough for our use case.
866            .max_capacity(32 * 1024 * 1024)
867            .weigher(|k: &Vec<u8>, v: &Vec<u8>| (k.len() + v.len()).try_into().unwrap_or(u32::MAX))
868            .build();
869        Self { cache }
870    }
871}
872
873impl CacheStore for IncrementalCompilationCache {
874    fn get(&self, key: &[u8]) -> Option<Cow<'_, [u8]>> {
875        self.cache.get(key).map(|v| v.into())
876    }
877
878    fn insert(&self, key: &[u8], value: Vec<u8>) -> bool {
879        self.cache.insert(key.to_vec(), value);
880        true
881    }
882}