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 util::paths::SanitizedPath;
 41use wasmtime::{
 42    CacheStore, Engine, Store,
 43    component::{Component, ResourceTable},
 44};
 45use wasmtime_wasi::{self as wasi, 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: SemanticVersion,
 69}
 70
 71impl Drop for WasmExtension {
 72    fn drop(&mut self) {
 73        self.tx.close_channel();
 74    }
 75}
 76
 77#[async_trait]
 78impl extension::Extension for WasmExtension {
 79    fn manifest(&self) -> Arc<ExtensionManifest> {
 80        self.manifest.clone()
 81    }
 82
 83    fn work_dir(&self) -> Arc<Path> {
 84        self.work_dir.clone()
 85    }
 86
 87    async fn language_server_command(
 88        &self,
 89        language_server_id: LanguageServerName,
 90        language_name: LanguageName,
 91        worktree: Arc<dyn WorktreeDelegate>,
 92    ) -> Result<Command> {
 93        self.call(|extension, store| {
 94            async move {
 95                let resource = store.data_mut().table().push(worktree)?;
 96                let command = extension
 97                    .call_language_server_command(
 98                        store,
 99                        &language_server_id,
100                        &language_name,
101                        resource,
102                    )
103                    .await?
104                    .map_err(|err| store.data().extension_error(err))?;
105
106                Ok(command.into())
107            }
108            .boxed()
109        })
110        .await?
111    }
112
113    async fn language_server_initialization_options(
114        &self,
115        language_server_id: LanguageServerName,
116        language_name: LanguageName,
117        worktree: Arc<dyn WorktreeDelegate>,
118    ) -> Result<Option<String>> {
119        self.call(|extension, store| {
120            async move {
121                let resource = store.data_mut().table().push(worktree)?;
122                let options = extension
123                    .call_language_server_initialization_options(
124                        store,
125                        &language_server_id,
126                        &language_name,
127                        resource,
128                    )
129                    .await?
130                    .map_err(|err| store.data().extension_error(err))?;
131                anyhow::Ok(options)
132            }
133            .boxed()
134        })
135        .await?
136    }
137
138    async fn language_server_workspace_configuration(
139        &self,
140        language_server_id: LanguageServerName,
141        worktree: Arc<dyn WorktreeDelegate>,
142    ) -> Result<Option<String>> {
143        self.call(|extension, store| {
144            async move {
145                let resource = store.data_mut().table().push(worktree)?;
146                let options = extension
147                    .call_language_server_workspace_configuration(
148                        store,
149                        &language_server_id,
150                        resource,
151                    )
152                    .await?
153                    .map_err(|err| store.data().extension_error(err))?;
154                anyhow::Ok(options)
155            }
156            .boxed()
157        })
158        .await?
159    }
160
161    async fn language_server_additional_initialization_options(
162        &self,
163        language_server_id: LanguageServerName,
164        target_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                let options = extension
171                    .call_language_server_additional_initialization_options(
172                        store,
173                        &language_server_id,
174                        &target_language_server_id,
175                        resource,
176                    )
177                    .await?
178                    .map_err(|err| store.data().extension_error(err))?;
179                anyhow::Ok(options)
180            }
181            .boxed()
182        })
183        .await?
184    }
185
186    async fn language_server_additional_workspace_configuration(
187        &self,
188        language_server_id: LanguageServerName,
189        target_language_server_id: LanguageServerName,
190        worktree: Arc<dyn WorktreeDelegate>,
191    ) -> Result<Option<String>> {
192        self.call(|extension, store| {
193            async move {
194                let resource = store.data_mut().table().push(worktree)?;
195                let options = extension
196                    .call_language_server_additional_workspace_configuration(
197                        store,
198                        &language_server_id,
199                        &target_language_server_id,
200                        resource,
201                    )
202                    .await?
203                    .map_err(|err| store.data().extension_error(err))?;
204                anyhow::Ok(options)
205            }
206            .boxed()
207        })
208        .await?
209    }
210
211    async fn labels_for_completions(
212        &self,
213        language_server_id: LanguageServerName,
214        completions: Vec<Completion>,
215    ) -> Result<Vec<Option<CodeLabel>>> {
216        self.call(|extension, store| {
217            async move {
218                let labels = extension
219                    .call_labels_for_completions(
220                        store,
221                        &language_server_id,
222                        completions.into_iter().map(Into::into).collect(),
223                    )
224                    .await?
225                    .map_err(|err| store.data().extension_error(err))?;
226
227                Ok(labels
228                    .into_iter()
229                    .map(|label| label.map(Into::into))
230                    .collect())
231            }
232            .boxed()
233        })
234        .await?
235    }
236
237    async fn labels_for_symbols(
238        &self,
239        language_server_id: LanguageServerName,
240        symbols: Vec<Symbol>,
241    ) -> Result<Vec<Option<CodeLabel>>> {
242        self.call(|extension, store| {
243            async move {
244                let labels = extension
245                    .call_labels_for_symbols(
246                        store,
247                        &language_server_id,
248                        symbols.into_iter().map(Into::into).collect(),
249                    )
250                    .await?
251                    .map_err(|err| store.data().extension_error(err))?;
252
253                Ok(labels
254                    .into_iter()
255                    .map(|label| label.map(Into::into))
256                    .collect())
257            }
258            .boxed()
259        })
260        .await?
261    }
262
263    async fn complete_slash_command_argument(
264        &self,
265        command: SlashCommand,
266        arguments: Vec<String>,
267    ) -> Result<Vec<SlashCommandArgumentCompletion>> {
268        self.call(|extension, store| {
269            async move {
270                let completions = extension
271                    .call_complete_slash_command_argument(store, &command.into(), &arguments)
272                    .await?
273                    .map_err(|err| store.data().extension_error(err))?;
274
275                Ok(completions.into_iter().map(Into::into).collect())
276            }
277            .boxed()
278        })
279        .await?
280    }
281
282    async fn run_slash_command(
283        &self,
284        command: SlashCommand,
285        arguments: Vec<String>,
286        delegate: Option<Arc<dyn WorktreeDelegate>>,
287    ) -> Result<SlashCommandOutput> {
288        self.call(|extension, store| {
289            async move {
290                let resource = if let Some(delegate) = delegate {
291                    Some(store.data_mut().table().push(delegate)?)
292                } else {
293                    None
294                };
295
296                let output = extension
297                    .call_run_slash_command(store, &command.into(), &arguments, resource)
298                    .await?
299                    .map_err(|err| store.data().extension_error(err))?;
300
301                Ok(output.into())
302            }
303            .boxed()
304        })
305        .await?
306    }
307
308    async fn context_server_command(
309        &self,
310        context_server_id: Arc<str>,
311        project: Arc<dyn ProjectDelegate>,
312    ) -> Result<Command> {
313        self.call(|extension, store| {
314            async move {
315                let project_resource = store.data_mut().table().push(project)?;
316                let command = extension
317                    .call_context_server_command(store, context_server_id.clone(), project_resource)
318                    .await?
319                    .map_err(|err| store.data().extension_error(err))?;
320                anyhow::Ok(command.into())
321            }
322            .boxed()
323        })
324        .await?
325    }
326
327    async fn context_server_configuration(
328        &self,
329        context_server_id: Arc<str>,
330        project: Arc<dyn ProjectDelegate>,
331    ) -> Result<Option<ContextServerConfiguration>> {
332        self.call(|extension, store| {
333            async move {
334                let project_resource = store.data_mut().table().push(project)?;
335                let Some(configuration) = extension
336                    .call_context_server_configuration(
337                        store,
338                        context_server_id.clone(),
339                        project_resource,
340                    )
341                    .await?
342                    .map_err(|err| store.data().extension_error(err))?
343                else {
344                    return Ok(None);
345                };
346
347                Ok(Some(configuration.try_into()?))
348            }
349            .boxed()
350        })
351        .await?
352    }
353
354    async fn suggest_docs_packages(&self, provider: Arc<str>) -> Result<Vec<String>> {
355        self.call(|extension, store| {
356            async move {
357                let packages = extension
358                    .call_suggest_docs_packages(store, provider.as_ref())
359                    .await?
360                    .map_err(|err| store.data().extension_error(err))?;
361
362                Ok(packages)
363            }
364            .boxed()
365        })
366        .await?
367    }
368
369    async fn index_docs(
370        &self,
371        provider: Arc<str>,
372        package_name: Arc<str>,
373        kv_store: Arc<dyn KeyValueStoreDelegate>,
374    ) -> Result<()> {
375        self.call(|extension, store| {
376            async move {
377                let kv_store_resource = store.data_mut().table().push(kv_store)?;
378                extension
379                    .call_index_docs(
380                        store,
381                        provider.as_ref(),
382                        package_name.as_ref(),
383                        kv_store_resource,
384                    )
385                    .await?
386                    .map_err(|err| store.data().extension_error(err))?;
387
388                anyhow::Ok(())
389            }
390            .boxed()
391        })
392        .await?
393    }
394
395    async fn get_dap_binary(
396        &self,
397        dap_name: Arc<str>,
398        config: DebugTaskDefinition,
399        user_installed_path: Option<PathBuf>,
400        worktree: Arc<dyn WorktreeDelegate>,
401    ) -> Result<DebugAdapterBinary> {
402        self.call(|extension, store| {
403            async move {
404                let resource = store.data_mut().table().push(worktree)?;
405                let dap_binary = extension
406                    .call_get_dap_binary(store, dap_name, config, user_installed_path, resource)
407                    .await?
408                    .map_err(|err| store.data().extension_error(err))?;
409                let dap_binary = dap_binary.try_into()?;
410                Ok(dap_binary)
411            }
412            .boxed()
413        })
414        .await?
415    }
416    async fn dap_request_kind(
417        &self,
418        dap_name: Arc<str>,
419        config: serde_json::Value,
420    ) -> Result<StartDebuggingRequestArgumentsRequest> {
421        self.call(|extension, store| {
422            async move {
423                let kind = extension
424                    .call_dap_request_kind(store, dap_name, config)
425                    .await?
426                    .map_err(|err| store.data().extension_error(err))?;
427                Ok(kind.into())
428            }
429            .boxed()
430        })
431        .await?
432    }
433
434    async fn dap_config_to_scenario(&self, config: ZedDebugConfig) -> Result<DebugScenario> {
435        self.call(|extension, store| {
436            async move {
437                let kind = extension
438                    .call_dap_config_to_scenario(store, config)
439                    .await?
440                    .map_err(|err| store.data().extension_error(err))?;
441                Ok(kind)
442            }
443            .boxed()
444        })
445        .await?
446    }
447
448    async fn dap_locator_create_scenario(
449        &self,
450        locator_name: String,
451        build_config_template: TaskTemplate,
452        resolved_label: String,
453        debug_adapter_name: String,
454    ) -> Result<Option<DebugScenario>> {
455        self.call(|extension, store| {
456            async move {
457                extension
458                    .call_dap_locator_create_scenario(
459                        store,
460                        locator_name,
461                        build_config_template,
462                        resolved_label,
463                        debug_adapter_name,
464                    )
465                    .await
466            }
467            .boxed()
468        })
469        .await?
470    }
471    async fn run_dap_locator(
472        &self,
473        locator_name: String,
474        config: SpawnInTerminal,
475    ) -> Result<DebugRequest> {
476        self.call(|extension, store| {
477            async move {
478                extension
479                    .call_run_dap_locator(store, locator_name, config)
480                    .await?
481                    .map_err(|err| store.data().extension_error(err))
482            }
483            .boxed()
484        })
485        .await?
486    }
487}
488
489pub struct WasmState {
490    manifest: Arc<ExtensionManifest>,
491    pub table: ResourceTable,
492    ctx: wasi::WasiCtx,
493    pub host: Arc<WasmHost>,
494    pub(crate) capability_granter: CapabilityGranter,
495}
496
497type MainThreadCall = Box<dyn Send + for<'a> FnOnce(&'a mut AsyncApp) -> LocalBoxFuture<'a, ()>>;
498
499type ExtensionCall = Box<
500    dyn Send + for<'a> FnOnce(&'a mut Extension, &'a mut Store<WasmState>) -> BoxFuture<'a, ()>,
501>;
502
503fn wasm_engine(executor: &BackgroundExecutor) -> wasmtime::Engine {
504    static WASM_ENGINE: OnceLock<wasmtime::Engine> = OnceLock::new();
505    WASM_ENGINE
506        .get_or_init(|| {
507            let mut config = wasmtime::Config::new();
508            config.wasm_component_model(true);
509            config.async_support(true);
510            config
511                .enable_incremental_compilation(cache_store())
512                .unwrap();
513            // Async support introduces the issue that extension execution happens during `Future::poll`,
514            // which could block an async thread.
515            // https://docs.rs/wasmtime/latest/wasmtime/struct.Config.html#execution-in-poll
516            //
517            // Epoch interruption is a lightweight mechanism to allow the extensions to yield control
518            // back to the executor at regular intervals.
519            config.epoch_interruption(true);
520
521            let engine = wasmtime::Engine::new(&config).unwrap();
522
523            // It might be safer to do this on a non-async thread to make sure it makes progress
524            // regardless of if extensions are blocking.
525            // However, due to our current setup, this isn't a likely occurrence and we'd rather
526            // not have a dedicated thread just for this. If it becomes an issue, we can consider
527            // creating a separate thread for epoch interruption.
528            let engine_ref = engine.weak();
529            executor
530                .spawn(async move {
531                    // Somewhat arbitrary interval, as it isn't a guaranteed interval.
532                    // But this is a rough upper bound for how long the extension execution can block on
533                    // `Future::poll`.
534                    const EPOCH_INTERVAL: Duration = Duration::from_millis(100);
535                    let mut timer = Timer::interval(EPOCH_INTERVAL);
536                    while (timer.next().await).is_some() {
537                        // Exit the loop and thread once the engine is dropped.
538                        let Some(engine) = engine_ref.upgrade() else {
539                            break;
540                        };
541                        engine.increment_epoch();
542                    }
543                })
544                .detach();
545
546            engine
547        })
548        .clone()
549}
550
551fn cache_store() -> Arc<IncrementalCompilationCache> {
552    static CACHE_STORE: LazyLock<Arc<IncrementalCompilationCache>> =
553        LazyLock::new(|| Arc::new(IncrementalCompilationCache::new()));
554    CACHE_STORE.clone()
555}
556
557impl WasmHost {
558    pub fn new(
559        fs: Arc<dyn Fs>,
560        http_client: Arc<dyn HttpClient>,
561        node_runtime: NodeRuntime,
562        proxy: Arc<ExtensionHostProxy>,
563        work_dir: PathBuf,
564        cx: &mut App,
565    ) -> Arc<Self> {
566        let (tx, mut rx) = mpsc::unbounded::<MainThreadCall>();
567        let task = cx.spawn(async move |cx| {
568            while let Some(message) = rx.next().await {
569                message(cx).await;
570            }
571        });
572        Arc::new(Self {
573            engine: wasm_engine(cx.background_executor()),
574            fs,
575            work_dir,
576            http_client,
577            node_runtime,
578            proxy,
579            release_channel: ReleaseChannel::global(cx),
580            granted_capabilities: vec![
581                ExtensionCapability::ProcessExec(ProcessExecCapability {
582                    command: "*".to_string(),
583                    args: vec!["**".to_string()],
584                }),
585                ExtensionCapability::DownloadFile(DownloadFileCapability {
586                    host: "*".to_string(),
587                    path: vec!["**".to_string()],
588                }),
589                ExtensionCapability::NpmInstallPackage(NpmInstallPackageCapability {
590                    package: "*".to_string(),
591                }),
592            ],
593            _main_thread_message_task: task,
594            main_thread_message_tx: tx,
595        })
596    }
597
598    pub fn load_extension(
599        self: &Arc<Self>,
600        wasm_bytes: Vec<u8>,
601        manifest: &Arc<ExtensionManifest>,
602        executor: BackgroundExecutor,
603    ) -> Task<Result<WasmExtension>> {
604        let this = self.clone();
605        let manifest = manifest.clone();
606        executor.clone().spawn(async move {
607            let zed_api_version = parse_wasm_extension_version(&manifest.id, &wasm_bytes)?;
608
609            let component = Component::from_binary(&this.engine, &wasm_bytes)
610                .context("failed to compile wasm component")?;
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        let path = SanitizedPath::new(&extension_work_dir).to_string();
670        #[cfg(target_os = "windows")]
671        let path = path.replace('\\', "/");
672
673        let mut ctx = wasi::WasiCtxBuilder::new();
674        ctx.inherit_stdio()
675            .env("PWD", &path)
676            .env("RUST_BACKTRACE", "full");
677
678        ctx.preopened_dir(&path, ".", dir_perms, file_perms)?;
679        ctx.preopened_dir(&path, &path, dir_perms, file_perms)?;
680
681        Ok(ctx.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}