From 5da39b39fbe9ca87971456206681f740d686f272 Mon Sep 17 00:00:00 2001 From: Marshall Bowers Date: Fri, 22 Nov 2024 13:21:30 -0500 Subject: [PATCH] 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 --- Cargo.lock | 1 - crates/extension/src/extension.rs | 10 ++++ crates/extension/src/types.rs | 1 + crates/extension_host/src/extension_host.rs | 4 +- crates/extension_host/src/wasm_host.rs | 22 ++++++- .../src/wasm_host/wit/since_v0_2_0.rs | 9 +-- crates/extensions_ui/Cargo.toml | 1 - .../src/extension_registration_hooks.rs | 57 ++++++++----------- 8 files changed, 59 insertions(+), 46 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 75d69fdcf9a4f..514338c59053c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4237,7 +4237,6 @@ dependencies = [ "ui", "util", "vim_mode_setting", - "wasmtime-wasi", "workspace", "zed_actions", ] diff --git a/crates/extension/src/extension.rs b/crates/extension/src/extension.rs index a3c275c5377c7..fe9b49909b9dd 100644 --- a/crates/extension/src/extension.rs +++ b/crates/extension/src/extension.rs @@ -25,6 +25,10 @@ pub trait WorktreeDelegate: Send + Sync + 'static { async fn shell_env(&self) -> Vec<(String, String)>; } +pub trait ProjectDelegate: Send + Sync + 'static { + fn worktree_ids(&self) -> Vec; +} + pub trait KeyValueStoreDelegate: Send + Sync + 'static { fn insert(&self, key: String, docs: String) -> Task>; } @@ -87,6 +91,12 @@ pub trait Extension: Send + Sync + 'static { worktree: Option>, ) -> Result; + async fn context_server_command( + &self, + context_server_id: Arc, + project: Arc, + ) -> Result; + async fn suggest_docs_packages(&self, provider: Arc) -> Result>; async fn index_docs( diff --git a/crates/extension/src/types.rs b/crates/extension/src/types.rs index f4c37b5daf26a..f04d31300fa44 100644 --- a/crates/extension/src/types.rs +++ b/crates/extension/src/types.rs @@ -10,6 +10,7 @@ pub use slash_command::*; pub type EnvVars = Vec<(String, String)>; /// A command. +#[derive(Debug)] pub struct Command { /// The command to execute. pub command: String, diff --git a/crates/extension_host/src/extension_host.rs b/crates/extension_host/src/extension_host.rs index 236c8091b4860..85da8127959e7 100644 --- a/crates/extension_host/src/extension_host.rs +++ b/crates/extension_host/src/extension_host.rs @@ -149,8 +149,8 @@ pub trait ExtensionRegistrationHooks: Send + Sync + 'static { fn register_context_server( &self, + _extension: Arc, _id: Arc, - _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, ); } diff --git a/crates/extension_host/src/wasm_host.rs b/crates/extension_host/src/wasm_host.rs index 54699ac0a1cfd..01c57599a894f 100644 --- a/crates/extension_host/src/wasm_host.rs +++ b/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, + project: Arc, + ) -> Result { + 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) -> Result> { self.call(|extension, store| { async move { diff --git a/crates/extension_host/src/wasm_host/wit/since_v0_2_0.rs b/crates/extension_host/src/wasm_host/wit/since_v0_2_0.rs index 02577abd0e7ae..234eec26ec5ee 100644 --- a/crates/extension_host/src/wasm_host/wit/since_v0_2_0.rs +++ b/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; +pub type ExtensionProject = Arc; pub type ExtensionKeyValueStore = Arc; pub type ExtensionHttpResponseStream = Arc>>; -pub struct ExtensionProject { - pub worktree_ids: Vec, -} - pub fn linker() -> &'static Linker { static LINKER: OnceLock> = OnceLock::new(); LINKER.get_or_init(|| super::new_linker(Extension::add_to_linker)) @@ -273,7 +270,7 @@ impl HostProject for WasmState { project: Resource, ) -> wasmtime::Result> { let project = self.table.get(&project)?; - Ok(project.worktree_ids.clone()) + Ok(project.worktree_ids()) } fn drop(&mut self, _project: Resource) -> Result<()> { diff --git a/crates/extensions_ui/Cargo.toml b/crates/extensions_ui/Cargo.toml index e8de7c3f12920..a219fe4bd4f60 100644 --- a/crates/extensions_ui/Cargo.toml +++ b/crates/extensions_ui/Cargo.toml @@ -41,7 +41,6 @@ theme.workspace = true ui.workspace = true util.workspace = true vim_mode_setting.workspace = true -wasmtime-wasi.workspace = true workspace.workspace = true zed_actions.workspace = true diff --git a/crates/extensions_ui/src/extension_registration_hooks.rs b/crates/extensions_ui/src/extension_registration_hooks.rs index 07a4c1455c855..1b427cd187f6e 100644 --- a/crates/extensions_ui/src/extension_registration_hooks.rs +++ b/crates/extensions_ui/src/extension_registration_hooks.rs @@ -1,13 +1,11 @@ use std::{path::PathBuf, sync::Arc}; -use anyhow::{anyhow, Result}; +use anyhow::Result; use assistant_slash_command::{ExtensionSlashCommand, SlashCommandRegistry}; use context_servers::manager::ServerCommand; use context_servers::ContextServerFactoryRegistry; -use db::smol::future::FutureExt as _; -use extension::Extension; -use extension_host::wasm_host::ExtensionProject; -use extension_host::{extension_lsp_adapter::ExtensionLspAdapter, wasm_host}; +use extension::{Extension, ProjectDelegate}; +use extension_host::extension_lsp_adapter::ExtensionLspAdapter; use fs::Fs; use gpui::{AppContext, BackgroundExecutor, Model, Task}; use indexed_docs::{ExtensionIndexedDocsProvider, IndexedDocsRegistry, ProviderId}; @@ -16,7 +14,16 @@ use lsp::LanguageServerName; use snippet_provider::SnippetRegistry; use theme::{ThemeRegistry, ThemeSettings}; use ui::SharedString; -use wasmtime_wasi::WasiView as _; + +struct ExtensionProject { + worktree_ids: Vec, +} + +impl ProjectDelegate for ExtensionProject { + fn worktree_ids(&self) -> Vec { + self.worktree_ids.clone() + } +} pub struct ConcreteExtensionRegistrationHooks { slash_command_registry: Arc, @@ -72,8 +79,8 @@ impl extension_host::ExtensionRegistrationHooks for ConcreteExtensionRegistratio fn register_context_server( &self, + extension: Arc, id: Arc, - extension: wasm_host::WasmExtension, cx: &mut AppContext, ) { self.context_server_factory_registry @@ -84,42 +91,24 @@ impl extension_host::ExtensionRegistrationHooks for ConcreteExtensionRegistratio move |project, cx| { log::info!( "loading command for context server {id} from extension {}", - extension.manifest.id + extension.manifest().id ); let id = id.clone(); let extension = extension.clone(); cx.spawn(|mut cx| async move { let extension_project = - project.update(&mut cx, |project, cx| ExtensionProject { - worktree_ids: project - .visible_worktrees(cx) - .map(|worktree| worktree.read(cx).id().to_proto()) - .collect(), + project.update(&mut cx, |project, cx| { + Arc::new(ExtensionProject { + worktree_ids: project + .visible_worktrees(cx) + .map(|worktree| worktree.read(cx).id().to_proto()) + .collect(), + }) })?; let command = extension - .call({ - let id = id.clone(); - |extension, store| { - async move { - let project = store - .data_mut() - .table() - .push(extension_project)?; - let command = extension - .call_context_server_command( - store, - id.clone(), - project, - ) - .await? - .map_err(|e| anyhow!("{}", e))?; - anyhow::Ok(command) - } - .boxed() - } - }) + .context_server_command(id.clone(), extension_project) .await?; log::info!("loaded command for context server {id}: {command:?}");