Use Extension trait when registering extension context servers (#21070)

This PR updates the extension context server registration to go through
the `Extension` trait for interacting with extensions rather than going
through the `WasmHost` directly.

Release Notes:

- N/A

crates/extension_host/src/extension_host.rs

@@ -149,8 +149,8 @@ pub trait ExtensionRegistrationHooks: Send + Sync + 'static {
 
     fn register_context_server(
         &self,
+        _extension: Arc<dyn Extension>,
         _id: Arc<str>,
-        _extension: WasmExtension,
         _cx: &mut AppContext,
     ) {
     }
@@ -1284,8 +1284,8 @@ impl ExtensionStore {
 
                     for (id, _context_server_entry) in &manifest.context_servers {
                         this.registration_hooks.register_context_server(
+                            extension.clone(),
                             id.clone(),
-                            wasm_extension.clone(),
                             cx,
                         );
                     }

crates/extension_host/src/wasm_host.rs

@@ -4,7 +4,7 @@ use crate::{ExtensionManifest, ExtensionRegistrationHooks};
 use anyhow::{anyhow, bail, Context as _, Result};
 use async_trait::async_trait;
 use extension::{
-    CodeLabel, Command, Completion, KeyValueStoreDelegate, SlashCommand,
+    CodeLabel, Command, Completion, KeyValueStoreDelegate, ProjectDelegate, SlashCommand,
     SlashCommandArgumentCompletion, SlashCommandOutput, Symbol, WorktreeDelegate,
 };
 use fs::{normalize_path, Fs};
@@ -34,7 +34,6 @@ use wasmtime::{
 };
 use wasmtime_wasi::{self as wasi, WasiView};
 use wit::Extension;
-pub use wit::ExtensionProject;
 
 pub struct WasmHost {
     engine: Engine,
@@ -238,6 +237,25 @@ impl extension::Extension for WasmExtension {
         .await
     }
 
+    async fn context_server_command(
+        &self,
+        context_server_id: Arc<str>,
+        project: Arc<dyn ProjectDelegate>,
+    ) -> Result<Command> {
+        self.call(|extension, store| {
+            async move {
+                let project_resource = store.data_mut().table().push(project)?;
+                let command = extension
+                    .call_context_server_command(store, context_server_id.clone(), project_resource)
+                    .await?
+                    .map_err(|err| anyhow!("{err}"))?;
+                anyhow::Ok(command.into())
+            }
+            .boxed()
+        })
+        .await
+    }
+
     async fn suggest_docs_packages(&self, provider: Arc<str>) -> Result<Vec<String>> {
         self.call(|extension, store| {
             async move {

crates/extension_host/src/wasm_host/wit/since_v0_2_0.rs

@@ -8,7 +8,7 @@ use async_compression::futures::bufread::GzipDecoder;
 use async_tar::Archive;
 use async_trait::async_trait;
 use context_servers::manager::ContextServerSettings;
-use extension::{KeyValueStoreDelegate, WorktreeDelegate};
+use extension::{KeyValueStoreDelegate, ProjectDelegate, WorktreeDelegate};
 use futures::{io::BufReader, FutureExt as _};
 use futures::{lock::Mutex, AsyncReadExt};
 use language::{language_settings::AllLanguageSettings, LanguageName, LanguageServerBinaryStatus};
@@ -44,13 +44,10 @@ mod settings {
 }
 
 pub type ExtensionWorktree = Arc<dyn WorktreeDelegate>;
+pub type ExtensionProject = Arc<dyn ProjectDelegate>;
 pub type ExtensionKeyValueStore = Arc<dyn KeyValueStoreDelegate>;
 pub type ExtensionHttpResponseStream = Arc<Mutex<::http_client::Response<AsyncBody>>>;
 
-pub struct ExtensionProject {
-    pub worktree_ids: Vec<u64>,
-}
-
 pub fn linker() -> &'static Linker<WasmState> {
     static LINKER: OnceLock<Linker<WasmState>> = OnceLock::new();
     LINKER.get_or_init(|| super::new_linker(Extension::add_to_linker))
@@ -273,7 +270,7 @@ impl HostProject for WasmState {
         project: Resource<ExtensionProject>,
     ) -> wasmtime::Result<Vec<u64>> {
         let project = self.table.get(&project)?;
-        Ok(project.worktree_ids.clone())
+        Ok(project.worktree_ids())
     }
 
     fn drop(&mut self, _project: Resource<Project>) -> Result<()> {