Provide wasm extensions with APIs needed for using pre-installed LSP binaries (#9085)

In this PR, we've added two new methods that LSP extensions can call:
* `shell_env()`, for retrieving the environment variables set in the
user's default shell in the worktree
* `which(command)`, for looking up paths to an executable (accounting
for the user's shell env in the worktree)

To test this out, we moved the `uiua` language support into an
extension. We went ahead and removed the built-in support, since this
language is extremely obscure. Sorry @mikayla-maki. To continue coding
in Uiua in Zed, for now you can `Add Dev Extension` from the extensions
pane, and select the `extensions/uiua` directory in the Zed repo. Very
soon, we'll support publishing these extensions so that you'll be able
to just install it normally.

Release Notes:

- N/A

---------

Co-authored-by: Marshall <marshall@zed.dev>

crates/extension/Cargo.toml

@@ -37,7 +37,8 @@ settings.workspace = true
 theme.workspace = true
 toml.workspace = true
 util.workspace = true
-wasmtime = { workspace = true, features = ["async"] }
+wasm-encoder.workspace = true
+wasmtime.workspace = true
 wasmtime-wasi.workspace = true
 wasmparser.workspace = true
 wit-component.workspace = true

crates/extension/src/build_extension.rs

@@ -6,13 +6,16 @@ use async_tar::Archive;
 use futures::io::BufReader;
 use futures::AsyncReadExt;
 use serde::Deserialize;
+use std::mem;
 use std::{
     env, fs,
     path::{Path, PathBuf},
     process::{Command, Stdio},
     sync::Arc,
 };
-use util::http::{AsyncBody, HttpClient};
+use util::http::{self, AsyncBody, HttpClient};
+use wasm_encoder::{ComponentSectionId, Encode as _, RawSection, Section as _};
+use wasmparser::Parser;
 use wit_component::ComponentEncoder;
 
 /// Currently, we compile with Rust's `wasm32-wasi` target, which works with WASI `preview1`.
@@ -59,8 +62,11 @@ struct CargoTomlPackage {
 }
 
 impl ExtensionBuilder {
-    pub fn new(cache_dir: PathBuf, http: Arc<dyn HttpClient>) -> Self {
-        Self { cache_dir, http }
+    pub fn new(cache_dir: PathBuf) -> Self {
+        Self {
+            cache_dir,
+            http: http::client(),
+        }
     }
 
     pub async fn compile_extension(
@@ -138,6 +144,10 @@ impl ExtensionBuilder {
             .encode()
             .context("failed to encode wasm component")?;
 
+        let component_bytes = self
+            .strip_custom_sections(&component_bytes)
+            .context("failed to strip debug sections from wasm component")?;
+
         fs::write(extension_dir.join("extension.wasm"), &component_bytes)
             .context("failed to write extension.wasm")?;
 
@@ -310,7 +320,7 @@ impl ExtensionBuilder {
     async fn install_wasi_preview1_adapter_if_needed(&self) -> Result<Vec<u8>> {
         let cache_path = self.cache_dir.join("wasi_snapshot_preview1.reactor.wasm");
         if let Ok(content) = fs::read(&cache_path) {
-            if wasmparser::Parser::is_core_wasm(&content) {
+            if Parser::is_core_wasm(&content) {
                 return Ok(content);
             }
         }
@@ -333,7 +343,7 @@ impl ExtensionBuilder {
         fs::write(&cache_path, &content)
             .with_context(|| format!("failed to save file {}", cache_path.display()))?;
 
-        if !wasmparser::Parser::is_core_wasm(&content) {
+        if !Parser::is_core_wasm(&content) {
             bail!("downloaded wasi adapter is invalid");
         }
         Ok(content)
@@ -379,4 +389,68 @@ impl ExtensionBuilder {
 
         Ok(clang_path)
     }
+
+    // This was adapted from:
+    // https://github.com/bytecodealliance/wasm-tools/1791a8f139722e9f8679a2bd3d8e423e55132b22/src/bin/wasm-tools/strip.rs
+    fn strip_custom_sections(&self, input: &Vec<u8>) -> Result<Vec<u8>> {
+        use wasmparser::Payload::*;
+
+        let strip_custom_section = |name: &str| name.starts_with(".debug");
+
+        let mut output = Vec::new();
+        let mut stack = Vec::new();
+
+        for payload in Parser::new(0).parse_all(input) {
+            let payload = payload?;
+
+            // Track nesting depth, so that we don't mess with inner producer sections:
+            match payload {
+                Version { encoding, .. } => {
+                    output.extend_from_slice(match encoding {
+                        wasmparser::Encoding::Component => &wasm_encoder::Component::HEADER,
+                        wasmparser::Encoding::Module => &wasm_encoder::Module::HEADER,
+                    });
+                }
+                ModuleSection { .. } | ComponentSection { .. } => {
+                    stack.push(mem::take(&mut output));
+                    continue;
+                }
+                End { .. } => {
+                    let mut parent = match stack.pop() {
+                        Some(c) => c,
+                        None => break,
+                    };
+                    if output.starts_with(&wasm_encoder::Component::HEADER) {
+                        parent.push(ComponentSectionId::Component as u8);
+                        output.encode(&mut parent);
+                    } else {
+                        parent.push(ComponentSectionId::CoreModule as u8);
+                        output.encode(&mut parent);
+                    }
+                    output = parent;
+                }
+                _ => {}
+            }
+
+            match &payload {
+                CustomSection(c) => {
+                    if strip_custom_section(c.name()) {
+                        continue;
+                    }
+                }
+
+                _ => {}
+            }
+
+            if let Some((id, range)) = payload.as_section() {
+                RawSection {
+                    id,
+                    data: &input[range],
+                }
+                .append_to(&mut output);
+            }
+        }
+
+        Ok(output)
+    }
 }

crates/extension/src/extension_store.rs

@@ -193,7 +193,7 @@ impl ExtensionStore {
             extension_index: Default::default(),
             installed_dir,
             index_path,
-            builder: Arc::new(ExtensionBuilder::new(build_dir, http_client.clone())),
+            builder: Arc::new(ExtensionBuilder::new(build_dir)),
             outstanding_operations: Default::default(),
             modified_extensions: Default::default(),
             reload_complete_senders: Vec::new(),
@@ -545,7 +545,7 @@ impl ExtensionStore {
                         builder
                             .compile_extension(
                                 &extension_source_path,
-                                CompileExtensionOptions { release: true },
+                                CompileExtensionOptions { release: false },
                             )
                             .await
                     }

crates/extension/src/extension_store_test.rs

@@ -7,10 +7,7 @@ use collections::BTreeMap;
 use fs::{FakeFs, Fs, RealFs};
 use futures::{io::BufReader, AsyncReadExt, StreamExt};
 use gpui::{Context, TestAppContext};
-use language::{
-    Language, LanguageConfig, LanguageMatcher, LanguageRegistry, LanguageServerBinaryStatus,
-    LanguageServerName,
-};
+use language::{LanguageMatcher, LanguageRegistry, LanguageServerBinaryStatus, LanguageServerName};
 use node_runtime::FakeNodeRuntime;
 use parking_lot::Mutex;
 use project::Project;
@@ -573,19 +570,6 @@ async fn test_extension_store_with_gleam_extension(cx: &mut TestAppContext) {
         })
         .await
         .unwrap();
-    project.update(cx, |project, cx| {
-        project.set_language_for_buffer(
-            &buffer,
-            Arc::new(Language::new(
-                LanguageConfig {
-                    name: "Gleam".into(),
-                    ..Default::default()
-                },
-                None,
-            )),
-            cx,
-        )
-    });
 
     let fake_server = fake_servers.next().await.unwrap();
     let expected_server_path = extensions_dir.join("work/gleam/gleam-v1.2.3/gleam");

crates/extension/src/wasm_host.rs

@@ -3,7 +3,7 @@ use anyhow::{anyhow, bail, Context as _, Result};
 use async_compression::futures::bufread::GzipDecoder;
 use async_tar::Archive;
 use async_trait::async_trait;
-use fs::Fs;
+use fs::{normalize_path, Fs};
 use futures::{
     channel::{
         mpsc::{self, UnboundedSender},
@@ -72,8 +72,6 @@ type ExtensionCall = Box<
 
 static WASM_ENGINE: OnceLock<wasmtime::Engine> = OnceLock::new();
 
-const EXTENSION_WORK_DIR_PATH: &str = "/zed/work";
-
 impl WasmHost {
     pub fn new(
         fs: Arc<dyn Fs>,
@@ -181,11 +179,12 @@ impl WasmHost {
             .await
             .context("failed to create extension work dir")?;
 
-        let work_dir_preopen = Dir::open_ambient_dir(extension_work_dir, ambient_authority())
+        let work_dir_preopen = Dir::open_ambient_dir(&extension_work_dir, ambient_authority())
             .context("failed to preopen extension work directory")?;
         let current_dir_preopen = work_dir_preopen
             .try_clone()
             .context("failed to preopen extension current directory")?;
+        let extension_work_dir = extension_work_dir.to_string_lossy();
 
         let perms = wasi::FilePerms::all();
         let dir_perms = wasi::DirPerms::all();
@@ -193,26 +192,24 @@ impl WasmHost {
         Ok(wasi::WasiCtxBuilder::new()
             .inherit_stdio()
             .preopened_dir(current_dir_preopen, dir_perms, perms, ".")
-            .preopened_dir(work_dir_preopen, dir_perms, perms, EXTENSION_WORK_DIR_PATH)
-            .env("PWD", EXTENSION_WORK_DIR_PATH)
-            .env("RUST_BACKTRACE", "1")
+            .preopened_dir(work_dir_preopen, dir_perms, perms, &extension_work_dir)
+            .env("PWD", &extension_work_dir)
+            .env("RUST_BACKTRACE", "full")
             .build())
     }
 
     pub fn path_from_extension(&self, id: &Arc<str>, path: &Path) -> PathBuf {
-        self.writeable_path_from_extension(id, path)
-            .unwrap_or_else(|| path.to_path_buf())
+        let extension_work_dir = self.work_dir.join(id.as_ref());
+        normalize_path(&extension_work_dir.join(path))
     }
 
-    pub fn writeable_path_from_extension(&self, id: &Arc<str>, path: &Path) -> Option<PathBuf> {
-        let path = path.strip_prefix(EXTENSION_WORK_DIR_PATH).unwrap_or(path);
-        if path.is_relative() {
-            let mut result = self.work_dir.clone();
-            result.push(id.as_ref());
-            result.extend(path);
-            Some(result)
+    pub fn writeable_path_from_extension(&self, id: &Arc<str>, path: &Path) -> Result<PathBuf> {
+        let extension_work_dir = self.work_dir.join(id.as_ref());
+        let path = normalize_path(&extension_work_dir.join(path));
+        if path.starts_with(&extension_work_dir) {
+            Ok(path)
         } else {
-            None
+            Err(anyhow!("cannot write to path {}", path.display()))
         }
     }
 }
@@ -252,13 +249,6 @@ impl WasmExtension {
     }
 }
 
-impl WasmState {
-    pub fn writeable_path_from_extension(&self, path: &Path) -> Option<PathBuf> {
-        self.host
-            .writeable_path_from_extension(&self.manifest.id, path)
-    }
-}
-
 #[async_trait]
 impl wit::HostWorktree for WasmState {
     async fn read_text_file(
@@ -273,6 +263,26 @@ impl wit::HostWorktree for WasmState {
             .map_err(|error| error.to_string()))
     }
 
+    async fn shell_env(
+        &mut self,
+        delegate: Resource<Arc<dyn LspAdapterDelegate>>,
+    ) -> wasmtime::Result<wit::EnvVars> {
+        let delegate = self.table.get(&delegate)?;
+        Ok(delegate.shell_env().await.into_iter().collect())
+    }
+
+    async fn which(
+        &mut self,
+        delegate: Resource<Arc<dyn LspAdapterDelegate>>,
+        binary_name: String,
+    ) -> wasmtime::Result<Option<String>> {
+        let delegate = self.table.get(&delegate)?;
+        Ok(delegate
+            .which(binary_name.as_ref())
+            .await
+            .map(|path| path.to_string_lossy().to_string()))
+    }
+
     fn drop(&mut self, _worktree: Resource<wit::Worktree>) -> Result<()> {
         // we only ever hand out borrows of worktrees
         Ok(())
@@ -395,8 +405,8 @@ impl wit::ExtensionImports for WasmState {
             this.host.fs.create_dir(&extension_work_dir).await?;
 
             let destination_path = this
-                .writeable_path_from_extension(&path)
-                .ok_or_else(|| anyhow!("cannot write to path {:?}", path))?;
+                .host
+                .writeable_path_from_extension(&this.manifest.id, &path)?;
 
             let mut response = this
                 .host