Extract ExtensionIndexedDocsProvider to indexed_docs crate (#20607)

This PR extracts the `ExtensionIndexedDocsProvider` implementation to
the `indexed_docs` crate.

To achieve this, we introduce a new `Extension` trait that provides an
abstracted interface for calling an extension. This trait resides in the
`extension` crate, which has minimal dependencies and can be depended on
by other crates, like `indexed_docs`.

We're then able to implement the `ExtensionIndexedDocsProvider` without
having any knowledge of the Wasm-specific internals of the extension
system.

Release Notes:

- N/A

crates/extension_host/src/extension_host.rs

@@ -9,6 +9,7 @@ use async_tar::Archive;
 use client::{telemetry::Telemetry, Client, ExtensionMetadata, GetExtensionsResponse};
 use collections::{btree_map, BTreeMap, HashSet};
 use extension::extension_builder::{CompileExtensionOptions, ExtensionBuilder};
+use extension::Extension;
 pub use extension::ExtensionManifest;
 use fs::{Fs, RemoveOptions};
 use futures::{
@@ -90,10 +91,6 @@ pub fn is_version_compatible(
     true
 }
 
-pub trait DocsDatabase: Send + Sync + 'static {
-    fn insert(&self, key: String, docs: String) -> Task<Result<()>>;
-}
-
 pub trait ExtensionRegistrationHooks: Send + Sync + 'static {
     fn remove_user_themes(&self, _themes: Vec<SharedString>) {}
 
@@ -149,13 +146,7 @@ pub trait ExtensionRegistrationHooks: Send + Sync + 'static {
     ) {
     }
 
-    fn register_docs_provider(
-        &self,
-        _extension: WasmExtension,
-        _host: Arc<WasmHost>,
-        _provider_id: Arc<str>,
-    ) {
-    }
+    fn register_docs_provider(&self, _extension: Arc<dyn Extension>, _provider_id: Arc<str>) {}
 
     fn register_snippets(&self, _path: &PathBuf, _snippet_contents: &str) -> Result<()> {
         Ok(())
@@ -1238,6 +1229,8 @@ impl ExtensionStore {
                 this.reload_complete_senders.clear();
 
                 for (manifest, wasm_extension) in &wasm_extensions {
+                    let extension = Arc::new(wasm_extension.clone());
+
                     for (language_server_id, language_server_config) in &manifest.language_servers {
                         for language in language_server_config.languages() {
                             this.registration_hooks.register_lsp_adapter(
@@ -1280,11 +1273,8 @@ impl ExtensionStore {
                     }
 
                     for (provider_id, _provider) in &manifest.indexed_docs_providers {
-                        this.registration_hooks.register_docs_provider(
-                            wasm_extension.clone(),
-                            this.wasm_host.clone(),
-                            provider_id.clone(),
-                        );
+                        this.registration_hooks
+                            .register_docs_provider(extension.clone(), provider_id.clone());
                     }
                 }
 

crates/extension_host/src/wasm_host.rs

@@ -2,6 +2,8 @@ pub mod wit;
 
 use crate::{ExtensionManifest, ExtensionRegistrationHooks};
 use anyhow::{anyhow, bail, Context as _, Result};
+use async_trait::async_trait;
+use extension::KeyValueStoreDelegate;
 use fs::{normalize_path, Fs};
 use futures::future::LocalBoxFuture;
 use futures::{
@@ -25,7 +27,7 @@ use wasmtime::{
     component::{Component, ResourceTable},
     Engine, Store,
 };
-use wasmtime_wasi as wasi;
+use wasmtime_wasi::{self as wasi, WasiView};
 use wit::Extension;
 pub use wit::{ExtensionProject, SlashCommand};
 
@@ -45,10 +47,63 @@ pub struct WasmHost {
 pub struct WasmExtension {
     tx: UnboundedSender<ExtensionCall>,
     pub manifest: Arc<ExtensionManifest>,
+    pub work_dir: Arc<Path>,
     #[allow(unused)]
     pub zed_api_version: SemanticVersion,
 }
 
+#[async_trait]
+impl extension::Extension for WasmExtension {
+    fn manifest(&self) -> Arc<ExtensionManifest> {
+        self.manifest.clone()
+    }
+
+    fn work_dir(&self) -> Arc<Path> {
+        self.work_dir.clone()
+    }
+
+    async fn suggest_docs_packages(&self, provider: Arc<str>) -> Result<Vec<String>> {
+        self.call(|extension, store| {
+            async move {
+                let packages = extension
+                    .call_suggest_docs_packages(store, provider.as_ref())
+                    .await?
+                    .map_err(|err| anyhow!("{err:?}"))?;
+
+                Ok(packages)
+            }
+            .boxed()
+        })
+        .await
+    }
+
+    async fn index_docs(
+        &self,
+        provider: Arc<str>,
+        package_name: Arc<str>,
+        kv_store: Arc<dyn KeyValueStoreDelegate>,
+    ) -> Result<()> {
+        self.call(|extension, store| {
+            async move {
+                let kv_store_resource = store.data_mut().table().push(kv_store)?;
+                extension
+                    .call_index_docs(
+                        store,
+                        provider.as_ref(),
+                        package_name.as_ref(),
+                        kv_store_resource,
+                    )
+                    .await?
+                    .map_err(|err| anyhow!("{err:?}"))?;
+
+                anyhow::Ok(())
+            }
+            .boxed()
+        })
+        .await
+    }
+}
+
 pub struct WasmState {
     manifest: Arc<ExtensionManifest>,
     pub table: ResourceTable,
@@ -152,6 +207,7 @@ impl WasmHost {
 
             Ok(WasmExtension {
                 manifest: manifest.clone(),
+                work_dir: this.work_dir.clone().into(),
                 tx,
                 zed_api_version,
             })

crates/extension_host/src/wasm_host/wit.rs

@@ -3,12 +3,11 @@ mod since_v0_0_4;
 mod since_v0_0_6;
 mod since_v0_1_0;
 mod since_v0_2_0;
+use extension::KeyValueStoreDelegate;
 use lsp::LanguageServerName;
 use release_channel::ReleaseChannel;
 use since_v0_2_0 as latest;
 
-use crate::DocsDatabase;
-
 use super::{wasm_engine, WasmState};
 use anyhow::{anyhow, Context, Result};
 use language::LspAdapterDelegate;
@@ -422,15 +421,15 @@ impl Extension {
         store: &mut Store<WasmState>,
         provider: &str,
         package_name: &str,
-        database: Resource<Arc<dyn DocsDatabase>>,
+        kv_store: Resource<Arc<dyn KeyValueStoreDelegate>>,
     ) -> Result<Result<(), String>> {
         match self {
             Extension::V020(ext) => {
-                ext.call_index_docs(store, provider, package_name, database)
+                ext.call_index_docs(store, provider, package_name, kv_store)
                     .await
             }
             Extension::V010(ext) => {
-                ext.call_index_docs(store, provider, package_name, database)
+                ext.call_index_docs(store, provider, package_name, kv_store)
                     .await
             }
             Extension::V001(_) | Extension::V004(_) | Extension::V006(_) => {

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

@@ -1,11 +1,11 @@
 use crate::wasm_host::{wit::ToWasmtimeResult, WasmState};
-use crate::DocsDatabase;
 use ::http_client::{AsyncBody, HttpRequestExt};
 use ::settings::{Settings, WorktreeId};
 use anyhow::{anyhow, bail, Context, Result};
 use async_compression::futures::bufread::GzipDecoder;
 use async_tar::Archive;
 use async_trait::async_trait;
+use extension::KeyValueStoreDelegate;
 use futures::{io::BufReader, FutureExt as _};
 use futures::{lock::Mutex, AsyncReadExt};
 use language::LanguageName;
@@ -48,7 +48,7 @@ mod settings {
 }
 
 pub type ExtensionWorktree = Arc<dyn LspAdapterDelegate>;
-pub type ExtensionKeyValueStore = Arc<dyn DocsDatabase>;
+pub type ExtensionKeyValueStore = Arc<dyn KeyValueStoreDelegate>;
 pub type ExtensionHttpResponseStream = Arc<Mutex<::http_client::Response<AsyncBody>>>;
 
 pub fn linker() -> &'static Linker<WasmState> {

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

@@ -1,5 +1,4 @@
 use crate::wasm_host::{wit::ToWasmtimeResult, WasmState};
-use crate::DocsDatabase;
 use ::http_client::{AsyncBody, HttpRequestExt};
 use ::settings::{Settings, WorktreeId};
 use anyhow::{anyhow, bail, Context, Result};
@@ -7,6 +6,7 @@ use async_compression::futures::bufread::GzipDecoder;
 use async_tar::Archive;
 use async_trait::async_trait;
 use context_servers::manager::ContextServerSettings;
+use extension::KeyValueStoreDelegate;
 use futures::{io::BufReader, FutureExt as _};
 use futures::{lock::Mutex, AsyncReadExt};
 use language::{
@@ -45,7 +45,7 @@ mod settings {
 }
 
 pub type ExtensionWorktree = Arc<dyn LspAdapterDelegate>;
-pub type ExtensionKeyValueStore = Arc<dyn DocsDatabase>;
+pub type ExtensionKeyValueStore = Arc<dyn KeyValueStoreDelegate>;
 pub type ExtensionHttpResponseStream = Arc<Mutex<::http_client::Response<AsyncBody>>>;
 
 pub struct ExtensionProject {