Add label_for_completion to extension API (#10175)

This PR adds the ability for extensions to implement
`label_for_completion` to customize completions coming back from the
language server.

We've used the Gleam extension as a motivating example, adding
`label_for_completion` support to it.

Release Notes:

- N/A

---------

Co-authored-by: Max <max@zed.dev>
Co-authored-by: Max Brunsfeld <maxbrunsfeld@gmail.com>

crates/extension/src/wasm_host/wit.rs

@@ -1,20 +1,28 @@
 mod since_v0_0_1;
 mod since_v0_0_4;
+mod since_v0_0_6;
 
-use super::{wasm_engine, WasmState};
-use anyhow::{Context, Result};
-use language::LspAdapterDelegate;
-use semantic_version::SemanticVersion;
 use std::ops::RangeInclusive;
 use std::sync::Arc;
+
+use anyhow::bail;
+use anyhow::{Context, Result};
+use language::{LanguageServerName, LspAdapterDelegate};
+use semantic_version::SemanticVersion;
 use wasmtime::{
     component::{Component, Instance, Linker, Resource},
     Store,
 };
 
-use since_v0_0_4 as latest;
+use super::{wasm_engine, WasmState};
 
-pub use latest::{Command, LanguageServerConfig};
+use since_v0_0_6 as latest;
+
+pub use latest::{
+    zed::extension::lsp::{Completion, CompletionKind, InsertTextFormat},
+    CodeLabel, CodeLabelSpan, Command, Range,
+};
+pub use since_v0_0_4::LanguageServerConfig;
 
 pub fn new_linker(
     f: impl Fn(&mut Linker<WasmState>, fn(&mut WasmState) -> &mut WasmState) -> Result<()>,
@@ -41,6 +49,7 @@ pub fn wasm_api_version_range() -> RangeInclusive<SemanticVersion> {
 }
 
 pub enum Extension {
+    V006(since_v0_0_6::Extension),
     V004(since_v0_0_4::Extension),
     V001(since_v0_0_1::Extension),
 }
@@ -51,16 +60,13 @@ impl Extension {
         version: SemanticVersion,
         component: &Component,
     ) -> Result<(Self, Instance)> {
-        if version < latest::MIN_VERSION {
-            let (extension, instance) = since_v0_0_1::Extension::instantiate_async(
-                store,
-                &component,
-                since_v0_0_1::linker(),
-            )
-            .await
-            .context("failed to instantiate wasm extension")?;
-            Ok((Self::V001(extension), instance))
-        } else {
+        if version >= latest::MIN_VERSION {
+            let (extension, instance) =
+                latest::Extension::instantiate_async(store, &component, latest::linker())
+                    .await
+                    .context("failed to instantiate wasm extension")?;
+            Ok((Self::V006(extension), instance))
+        } else if version >= since_v0_0_4::MIN_VERSION {
             let (extension, instance) = since_v0_0_4::Extension::instantiate_async(
                 store,
                 &component,
@@ -69,11 +75,21 @@ impl Extension {
             .await
             .context("failed to instantiate wasm extension")?;
             Ok((Self::V004(extension), instance))
+        } else {
+            let (extension, instance) = since_v0_0_1::Extension::instantiate_async(
+                store,
+                &component,
+                since_v0_0_1::linker(),
+            )
+            .await
+            .context("failed to instantiate wasm extension")?;
+            Ok((Self::V001(extension), instance))
         }
     }
 
     pub async fn call_init_extension(&self, store: &mut Store<WasmState>) -> Result<()> {
         match self {
+            Extension::V006(ext) => ext.call_init_extension(store).await,
             Extension::V004(ext) => ext.call_init_extension(store).await,
             Extension::V001(ext) => ext.call_init_extension(store).await,
         }
@@ -82,14 +98,19 @@ impl Extension {
     pub async fn call_language_server_command(
         &self,
         store: &mut Store<WasmState>,
+        language_server_id: &LanguageServerName,
         config: &LanguageServerConfig,
         resource: Resource<Arc<dyn LspAdapterDelegate>>,
     ) -> Result<Result<Command, String>> {
         match self {
-            Extension::V004(ext) => {
-                ext.call_language_server_command(store, config, resource)
+            Extension::V006(ext) => {
+                ext.call_language_server_command(store, &language_server_id.0, resource)
                     .await
             }
+            Extension::V004(ext) => Ok(ext
+                .call_language_server_command(store, config, resource)
+                .await?
+                .map(|command| command.into())),
             Extension::V001(ext) => Ok(ext
                 .call_language_server_command(store, &config.clone().into(), resource)
                 .await?
@@ -100,10 +121,19 @@ impl Extension {
     pub async fn call_language_server_initialization_options(
         &self,
         store: &mut Store<WasmState>,
+        language_server_id: &LanguageServerName,
         config: &LanguageServerConfig,
         resource: Resource<Arc<dyn LspAdapterDelegate>>,
     ) -> Result<Result<Option<String>, String>> {
         match self {
+            Extension::V006(ext) => {
+                ext.call_language_server_initialization_options(
+                    store,
+                    &language_server_id.0,
+                    resource,
+                )
+                .await
+            }
             Extension::V004(ext) => {
                 ext.call_language_server_initialization_options(store, config, resource)
                     .await
@@ -118,6 +148,23 @@ impl Extension {
             }
         }
     }
+
+    pub async fn call_labels_for_completions(
+        &self,
+        store: &mut Store<WasmState>,
+        language_server_id: &LanguageServerName,
+        completions: Vec<latest::Completion>,
+    ) -> Result<Result<Vec<Option<CodeLabel>>, String>> {
+        match self {
+            Extension::V001(_) | Extension::V004(_) => {
+                bail!("unsupported function: 'labels_for_completions'")
+            }
+            Extension::V006(ext) => {
+                ext.call_labels_for_completions(store, &language_server_id.0, &completions)
+                    .await
+            }
+        }
+    }
 }
 
 trait ToWasmtimeResult<T> {