assistant: Allow `/docs` to perform JIT indexing when run (#14768)

Marshall Bowers created 2 years ago

This PR updates the `/docs` slash command with the ability to
just-in-time index a package when there are not yet any results in the
index.

When running a `/docs` slash command, we fist check to see if there are
any results in the index that would match the search.

If there are, we go ahead and return them, as we do today.

However, if there are not yet any results we kick off an indexing task
as part of the command execution to fetch the results.

Release Notes:

- N/A

Change summary

crates/assistant/src/slash_command/docs_command.rs | 63 +++++++++++++++
crates/indexed_docs/src/store.rs                   | 24 ++++++
2 files changed, 84 insertions(+), 3 deletions(-)

Detailed changes

crates/assistant/src/slash_command/docs_command.rs

@@ -1,12 +1,13 @@
 use std::path::Path;
 use std::sync::atomic::AtomicBool;
 use std::sync::Arc;
+use std::time::Duration;
 
 use anyhow::{anyhow, bail, Result};
 use assistant_slash_command::{
     ArgumentCompletion, SlashCommand, SlashCommandOutput, SlashCommandOutputSection,
 };
-use gpui::{AppContext, Model, Task, WeakView};
+use gpui::{AppContext, BackgroundExecutor, Model, Task, WeakView};
 use indexed_docs::{
     DocsDotRsProvider, IndexedDocsRegistry, IndexedDocsStore, LocalRustdocProvider, PackageName,
     ProviderId,
@@ -90,6 +91,55 @@ impl DocsSlashCommand {
             }
         }
     }
+
+    /// Runs just-in-time indexing for a given package, in case the slash command
+    /// is run without any entries existing in the index.
+    fn run_just_in_time_indexing(
+        store: Arc<IndexedDocsStore>,
+        key: String,
+        package: PackageName,
+        executor: BackgroundExecutor,
+    ) -> Task<()> {
+        executor.clone().spawn(async move {
+            let (prefix, needs_full_index) = if let Some((prefix, _)) = key.split_once('*') {
+                // If we have a wildcard in the search, we want to wait until
+                // we've completely finished indexing so we get a full set of
+                // results for the wildcard.
+                (prefix.to_string(), true)
+            } else {
+                (key, false)
+            };
+
+            // If we already have some entries, we assume that we've indexed the package before
+            // and don't need to do it again.
+            let has_any_entries = store
+                .any_with_prefix(prefix.clone())
+                .await
+                .unwrap_or_default();
+            if has_any_entries {
+                return ();
+            };
+
+            let index_task = store.clone().index(package.clone());
+
+            if needs_full_index {
+                _ = index_task.await;
+            } else {
+                loop {
+                    executor.timer(Duration::from_millis(200)).await;
+
+                    if store
+                        .any_with_prefix(prefix.clone())
+                        .await
+                        .unwrap_or_default()
+                        || !store.is_indexing(&package)
+                    {
+                        break;
+                    }
+                }
+            }
+        })
+    }
 }
 
 impl SlashCommand for DocsSlashCommand {
@@ -200,13 +250,14 @@ impl SlashCommand for DocsSlashCommand {
         };
 
         let args = DocsSlashCommandArgs::parse(argument);
+        let executor = cx.background_executor().clone();
         let task = cx.background_executor().spawn({
             let store = args
                 .provider()
                 .ok_or_else(|| anyhow!("no docs provider specified"))
                 .and_then(|provider| IndexedDocsStore::try_global(provider, cx));
             async move {
-                let (provider, key) = match args {
+                let (provider, key) = match args.clone() {
                     DocsSlashCommandArgs::NoProvider => bail!("no docs provider specified"),
                     DocsSlashCommandArgs::SearchPackageDocs {
                         provider, package, ..
@@ -219,6 +270,12 @@ impl SlashCommand for DocsSlashCommand {
                 };
 
                 let store = store?;
+
+                if let Some(package) = args.package() {
+                    Self::run_just_in_time_indexing(store.clone(), key.clone(), package, executor)
+                        .await;
+                }
+
                 let (text, ranges) = if let Some((prefix, _)) = key.split_once('*') {
                     let docs = store.load_many_by_prefix(prefix.to_string()).await?;
 
@@ -269,7 +326,7 @@ fn is_item_path_delimiter(char: char) -> bool {
     !char.is_alphanumeric() && char != '-' && char != '_'
 }
 
-#[derive(Debug, PartialEq)]
+#[derive(Debug, PartialEq, Clone)]
 pub(crate) enum DocsSlashCommandArgs {
     NoProvider,
     SearchPackageDocs {

crates/indexed_docs/src/store.rs

@@ -112,6 +112,16 @@ impl IndexedDocsStore {
             .await
     }
 
+    /// Returns whether any entries exist with the given prefix.
+    pub async fn any_with_prefix(&self, prefix: String) -> Result<bool> {
+        self.database_future
+            .clone()
+            .await
+            .map_err(|err| anyhow!(err))?
+            .any_with_prefix(prefix)
+            .await
+    }
+
     pub fn index(
         self: Arc<Self>,
         package: PackageName,
@@ -288,6 +298,20 @@ impl IndexedDocsDatabase {
         })
     }
 
+    /// Returns whether any entries exist with the given prefix.
+    pub fn any_with_prefix(&self, prefix: String) -> Task<Result<bool>> {
+        let env = self.env.clone();
+        let entries = self.entries;
+
+        self.executor.spawn(async move {
+            let txn = env.read_txn()?;
+            let any = entries
+                .iter(&txn)?
+                .any(|entry| entry.map_or(false, |(key, _value)| key.starts_with(&prefix)));
+            Ok(any)
+        })
+    }
+
     pub fn insert(&self, key: String, docs: String) -> Task<Result<()>> {
         let env = self.env.clone();
         let entries = self.entries;