assistant: Populate `/docs rustdoc` with workspace crates (#16172)

Marshall Bowers created 1 year ago

This PR makes the `/docs` slash command populate the list with all of
the workspace crates when using the local rustdoc provider.

The workspace crates are shown in the search results when a crate is not
already indexed:

<img width="577" alt="Screenshot 2024-08-13 at 2 18 39 PM"
src="https://github.com/user-attachments/assets/39bee576-8e1a-4b21-a9f8-7951ebae4cc3">

These crates are shown with an `(unindexed)` suffix to convey this:

<img width="570" alt="Screenshot 2024-08-13 at 2 18 45 PM"
src="https://github.com/user-attachments/assets/4eeb07f7-378f-44d4-ae11-4ffe45a23964">

Release Notes:

- N/A

Change summary

Cargo.lock                                         |  1 
crates/assistant/src/slash_command/docs_command.rs | 26 ++++++++++++
crates/indexed_docs/Cargo.toml                     |  1 
crates/indexed_docs/src/providers/rustdoc.rs       | 34 +++++++++++++++
4 files changed, 61 insertions(+), 1 deletion(-)

Detailed changes

Cargo.lock

@@ -5469,6 +5469,7 @@ version = "0.1.0"
 dependencies = [
  "anyhow",
  "async-trait",
+ "cargo_metadata",
  "collections",
  "derive_more",
  "fs",

crates/assistant/src/slash_command/docs_command.rs

@@ -72,6 +72,9 @@ impl DocsSlashCommand {
             });
 
             if let Some((fs, cargo_workspace_root)) = index_provider_deps.log_err() {
+                // List the workspace crates once to prime the cache.
+                LocalRustdocProvider::list_workspace_crates().ok();
+
                 indexed_docs_registry.register_provider(Box::new(LocalRustdocProvider::new(
                     fs,
                     cargo_workspace_root,
@@ -230,6 +233,29 @@ impl SlashCommand for DocsSlashCommand {
                     }
 
                     let items = store.search(package).await;
+
+                    if provider == LocalRustdocProvider::id() {
+                        let items = build_completions(provider.clone(), items);
+                        let workspace_crates = LocalRustdocProvider::list_workspace_crates()?;
+
+                        let mut all_items = items;
+                        let workspace_crate_completions = workspace_crates
+                            .into_iter()
+                            .filter(|crate_name| {
+                                !all_items
+                                    .iter()
+                                    .any(|item| item.label.as_str() == crate_name.as_ref())
+                            })
+                            .map(|crate_name| ArgumentCompletion {
+                                label: format!("{crate_name} (unindexed)"),
+                                new_text: format!("{provider} {crate_name}"),
+                                run_command: true,
+                            })
+                            .collect::<Vec<_>>();
+                        all_items.extend(workspace_crate_completions);
+                        return Ok(all_items);
+                    }
+
                     if items.is_empty() {
                         if provider == DocsDotRsProvider::id() {
                             return Ok(std::iter::once(ArgumentCompletion {

crates/indexed_docs/Cargo.toml

@@ -14,6 +14,7 @@ path = "src/indexed_docs.rs"
 [dependencies]
 anyhow.workspace = true
 async-trait.workspace = true
+cargo_metadata.workspace = true
 collections.workspace = true
 derive_more.workspace = true
 fs.workspace = true

crates/indexed_docs/src/providers/rustdoc.rs

@@ -1,12 +1,16 @@
 mod item;
 mod to_markdown;
 
+use cargo_metadata::MetadataCommand;
 use futures::future::BoxFuture;
 pub use item::*;
+use parking_lot::RwLock;
 pub use to_markdown::convert_rustdoc_to_markdown;
 
+use std::collections::BTreeSet;
 use std::path::PathBuf;
-use std::sync::Arc;
+use std::sync::{Arc, LazyLock};
+use std::time::{Duration, Instant};
 
 use anyhow::{bail, Context, Result};
 use async_trait::async_trait;
@@ -40,6 +44,34 @@ impl LocalRustdocProvider {
             cargo_workspace_root,
         }
     }
+
+    /// Returns the list of all crates in the Cargo workspace.
+    ///
+    /// Includes the list of workspace crates as well as all dependency crates.
+    pub fn list_workspace_crates() -> Result<Vec<Arc<str>>> {
+        static WORKSPACE_CRATES: LazyLock<RwLock<Option<(BTreeSet<Arc<str>>, Instant)>>> =
+            LazyLock::new(|| RwLock::new(None));
+
+        if let Some((crates, fetched_at)) = &*WORKSPACE_CRATES.read() {
+            if fetched_at.elapsed() < Duration::from_secs(300) {
+                return Ok(crates.iter().cloned().collect());
+            }
+        }
+
+        let workspace = MetadataCommand::new()
+            .exec()
+            .context("failed to load cargo metadata")?;
+
+        let workspace_crates = workspace
+            .packages
+            .into_iter()
+            .map(|package| package.name.into())
+            .collect::<BTreeSet<_>>();
+
+        *WORKSPACE_CRATES.write() = Some((workspace_crates.clone(), Instant::now()));
+
+        Ok(workspace_crates.iter().cloned().collect())
+    }
 }
 
 #[async_trait]