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

Cargo.lock

@@ -4123,9 +4123,11 @@ dependencies = [
  "anyhow",
  "async-compression",
  "async-tar",
+ "async-trait",
  "collections",
  "fs",
  "futures 0.3.30",
+ "gpui",
  "http_client",
  "language",
  "log",
@@ -4219,6 +4221,7 @@ dependencies = [
  "db",
  "editor",
  "env_logger 0.11.5",
+ "extension",
  "extension_host",
  "fs",
  "futures 0.3.30",
@@ -6005,7 +6008,7 @@ dependencies = [
  "cargo_metadata",
  "collections",
  "derive_more",
- "extension_host",
+ "extension",
  "fs",
  "futures 0.3.30",
  "fuzzy",

crates/extension/Cargo.toml

@@ -15,9 +15,11 @@ path = "src/extension.rs"
 anyhow.workspace = true
 async-compression.workspace = true
 async-tar.workspace = true
+async-trait.workspace = true
 collections.workspace = true
 fs.workspace = true
 futures.workspace = true
+gpui.workspace = true
 http_client.workspace = true
 language.workspace = true
 log.workspace = true

crates/extension/src/extension.rs

@@ -1,11 +1,38 @@
 pub mod extension_builder;
 mod extension_manifest;
 
+use std::path::Path;
+use std::sync::Arc;
+
 use anyhow::{anyhow, bail, Context as _, Result};
+use async_trait::async_trait;
+use gpui::Task;
 use semantic_version::SemanticVersion;
 
 pub use crate::extension_manifest::*;
 
+pub trait KeyValueStoreDelegate: Send + Sync + 'static {
+    fn insert(&self, key: String, docs: String) -> Task<Result<()>>;
+}
+
+#[async_trait]
+pub trait Extension: Send + Sync + 'static {
+    /// Returns the [`ExtensionManifest`] for this extension.
+    fn manifest(&self) -> Arc<ExtensionManifest>;
+
+    /// Returns the path to this extension's working directory.
+    fn work_dir(&self) -> Arc<Path>;
+
+    async fn suggest_docs_packages(&self, provider: Arc<str>) -> Result<Vec<String>>;
+
+    async fn index_docs(
+        &self,
+        provider: Arc<str>,
+        package_name: Arc<str>,
+        kv_store: Arc<dyn KeyValueStoreDelegate>,
+    ) -> Result<()>;
+}
+
 pub fn parse_wasm_extension_version(
     extension_id: &str,
     wasm_bytes: &[u8],

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 {

crates/extensions_ui/Cargo.toml

@@ -23,6 +23,7 @@ collections.workspace = true
 context_servers.workspace = true
 db.workspace = true
 editor.workspace = true
+extension.workspace = true
 extension_host.workspace = true
 fs.workspace = true
 futures.workspace = true

crates/extensions_ui/src/extension_indexed_docs_provider.rs

@@ -1,79 +0,0 @@
-use std::path::PathBuf;
-use std::sync::Arc;
-
-use anyhow::{anyhow, Result};
-use async_trait::async_trait;
-use futures::FutureExt;
-use indexed_docs::{IndexedDocsDatabase, IndexedDocsProvider, PackageName, ProviderId};
-use wasmtime_wasi::WasiView;
-
-use extension_host::wasm_host::{WasmExtension, WasmHost};
-
-pub struct ExtensionIndexedDocsProvider {
-    pub(crate) extension: WasmExtension,
-    pub(crate) host: Arc<WasmHost>,
-    pub(crate) id: ProviderId,
-}
-
-#[async_trait]
-impl IndexedDocsProvider for ExtensionIndexedDocsProvider {
-    fn id(&self) -> ProviderId {
-        self.id.clone()
-    }
-
-    fn database_path(&self) -> PathBuf {
-        let mut database_path = self.host.work_dir.clone();
-        database_path.push(self.extension.manifest.id.as_ref());
-        database_path.push("docs");
-        database_path.push(format!("{}.0.mdb", self.id));
-
-        database_path
-    }
-
-    async fn suggest_packages(&self) -> Result<Vec<PackageName>> {
-        self.extension
-            .call({
-                let id = self.id.clone();
-                |extension, store| {
-                    async move {
-                        let packages = extension
-                            .call_suggest_docs_packages(store, id.as_ref())
-                            .await?
-                            .map_err(|err| anyhow!("{err:?}"))?;
-
-                        Ok(packages
-                            .into_iter()
-                            .map(|package| PackageName::from(package.as_str()))
-                            .collect())
-                    }
-                    .boxed()
-                }
-            })
-            .await
-    }
-
-    async fn index(&self, package: PackageName, database: Arc<IndexedDocsDatabase>) -> Result<()> {
-        self.extension
-            .call({
-                let id = self.id.clone();
-                |extension, store| {
-                    async move {
-                        let database_resource = store.data_mut().table().push(database as _)?;
-                        extension
-                            .call_index_docs(
-                                store,
-                                id.as_ref(),
-                                package.as_ref(),
-                                database_resource,
-                            )
-                            .await?
-                            .map_err(|err| anyhow!("{err:?}"))?;
-
-                        anyhow::Ok(())
-                    }
-                    .boxed()
-                }
-            })
-            .await
-    }
-}

crates/extensions_ui/src/extension_registration_hooks.rs

@@ -3,17 +3,18 @@ use std::{path::PathBuf, sync::Arc};
 use anyhow::Result;
 use assistant_slash_command::SlashCommandRegistry;
 use context_servers::ContextServerFactoryRegistry;
+use extension::Extension;
 use extension_host::{extension_lsp_adapter::ExtensionLspAdapter, wasm_host};
 use fs::Fs;
 use gpui::{AppContext, BackgroundExecutor, Task};
-use indexed_docs::{IndexedDocsRegistry, ProviderId};
+use indexed_docs::{ExtensionIndexedDocsProvider, IndexedDocsRegistry, ProviderId};
 use language::{LanguageRegistry, LanguageServerBinaryStatus, LoadedLanguage};
 use snippet_provider::SnippetRegistry;
 use theme::{ThemeRegistry, ThemeSettings};
 use ui::SharedString;
 
 use crate::extension_context_server::ExtensionContextServer;
-use crate::{extension_indexed_docs_provider, extension_slash_command::ExtensionSlashCommand};
+use crate::extension_slash_command::ExtensionSlashCommand;
 
 pub struct ConcreteExtensionRegistrationHooks {
     slash_command_registry: Arc<SlashCommandRegistry>,
@@ -99,19 +100,12 @@ impl extension_host::ExtensionRegistrationHooks for ConcreteExtensionRegistratio
             );
     }
 
-    fn register_docs_provider(
-        &self,
-        extension: wasm_host::WasmExtension,
-        host: Arc<wasm_host::WasmHost>,
-        provider_id: Arc<str>,
-    ) {
-        self.indexed_docs_registry.register_provider(Box::new(
-            extension_indexed_docs_provider::ExtensionIndexedDocsProvider {
+    fn register_docs_provider(&self, extension: Arc<dyn Extension>, provider_id: Arc<str>) {
+        self.indexed_docs_registry
+            .register_provider(Box::new(ExtensionIndexedDocsProvider::new(
                 extension,
-                host,
-                id: ProviderId(provider_id),
-            },
-        ));
+                ProviderId(provider_id),
+            )));
     }
 
     fn register_snippets(&self, path: &PathBuf, snippet_contents: &str) -> Result<()> {

crates/extensions_ui/src/extensions_ui.rs

@@ -1,6 +1,5 @@
 mod components;
 mod extension_context_server;
-mod extension_indexed_docs_provider;
 mod extension_registration_hooks;
 mod extension_slash_command;
 mod extension_suggest;

crates/indexed_docs/Cargo.toml

@@ -17,6 +17,7 @@ async-trait.workspace = true
 cargo_metadata.workspace = true
 collections.workspace = true
 derive_more.workspace = true
+extension.workspace = true
 fs.workspace = true
 futures.workspace = true
 fuzzy.workspace = true
@@ -30,7 +31,6 @@ paths.workspace = true
 serde.workspace = true
 strum.workspace = true
 util.workspace = true
-extension_host.workspace = true
 
 [dev-dependencies]
 indoc.workspace = true

crates/indexed_docs/src/extension_indexed_docs_provider.rs

@@ -0,0 +1,53 @@
+use std::path::PathBuf;
+use std::sync::Arc;
+
+use anyhow::Result;
+use async_trait::async_trait;
+use extension::Extension;
+
+use crate::{IndexedDocsDatabase, IndexedDocsProvider, PackageName, ProviderId};
+
+pub struct ExtensionIndexedDocsProvider {
+    extension: Arc<dyn Extension>,
+    id: ProviderId,
+}
+
+impl ExtensionIndexedDocsProvider {
+    pub fn new(extension: Arc<dyn Extension>, id: ProviderId) -> Self {
+        Self { extension, id }
+    }
+}
+
+#[async_trait]
+impl IndexedDocsProvider for ExtensionIndexedDocsProvider {
+    fn id(&self) -> ProviderId {
+        self.id.clone()
+    }
+
+    fn database_path(&self) -> PathBuf {
+        let mut database_path = PathBuf::from(self.extension.work_dir().as_ref());
+        database_path.push(self.extension.manifest().id.as_ref());
+        database_path.push("docs");
+        database_path.push(format!("{}.0.mdb", self.id));
+
+        database_path
+    }
+
+    async fn suggest_packages(&self) -> Result<Vec<PackageName>> {
+        let packages = self
+            .extension
+            .suggest_docs_packages(self.id.0.clone())
+            .await?;
+
+        Ok(packages
+            .into_iter()
+            .map(|package| PackageName::from(package.as_str()))
+            .collect())
+    }
+
+    async fn index(&self, package: PackageName, database: Arc<IndexedDocsDatabase>) -> Result<()> {
+        self.extension
+            .index_docs(self.id.0.clone(), package.as_ref().into(), database)
+            .await
+    }
+}

crates/indexed_docs/src/indexed_docs.rs

@@ -1,7 +1,9 @@
+mod extension_indexed_docs_provider;
 mod providers;
 mod registry;
 mod store;
 
+pub use crate::extension_indexed_docs_provider::ExtensionIndexedDocsProvider;
 pub use crate::providers::rustdoc::*;
 pub use crate::registry::*;
 pub use crate::store::*;

crates/indexed_docs/src/providers/rustdoc.rs

@@ -2,7 +2,6 @@ mod item;
 mod to_markdown;
 
 use cargo_metadata::MetadataCommand;
-use extension_host::DocsDatabase;
 use futures::future::BoxFuture;
 pub use item::*;
 use parking_lot::RwLock;
@@ -209,7 +208,7 @@ impl IndexedDocsProvider for DocsDotRsProvider {
 
 async fn index_rustdoc(
     package: PackageName,
-    database: Arc<dyn DocsDatabase>,
+    database: Arc<IndexedDocsDatabase>,
     fetch_page: impl Fn(&PackageName, Option<&RustdocItem>) -> BoxFuture<'static, Result<Option<String>>>
         + Send
         + Sync,

crates/indexed_docs/src/store.rs

@@ -324,10 +324,8 @@ impl IndexedDocsDatabase {
             Ok(any)
         })
     }
-}
 
-impl extension_host::DocsDatabase for IndexedDocsDatabase {
-    fn insert(&self, key: String, docs: String) -> Task<Result<()>> {
+    pub fn insert(&self, key: String, docs: String) -> Task<Result<()>> {
         let env = self.env.clone();
         let entries = self.entries;
 
@@ -339,3 +337,9 @@ impl extension_host::DocsDatabase for IndexedDocsDatabase {
         })
     }
 }
+
+impl extension::KeyValueStoreDelegate for IndexedDocsDatabase {
+    fn insert(&self, key: String, docs: String) -> Task<Result<()>> {
+        IndexedDocsDatabase::insert(&self, key, docs)
+    }
+}