docs_command.rs

  1use std::path::Path;
  2use std::sync::atomic::AtomicBool;
  3use std::sync::Arc;
  4use std::time::Duration;
  5
  6use anyhow::{anyhow, bail, Result};
  7use assistant_slash_command::{
  8    ArgumentCompletion, SlashCommand, SlashCommandOutput, SlashCommandOutputSection,
  9};
 10use feature_flags::FeatureFlag;
 11use gpui::{AppContext, BackgroundExecutor, Model, Task, WeakView};
 12use indexed_docs::{
 13    DocsDotRsProvider, IndexedDocsRegistry, IndexedDocsStore, LocalRustdocProvider, PackageName,
 14    ProviderId,
 15};
 16use language::LspAdapterDelegate;
 17use project::{Project, ProjectPath};
 18use ui::prelude::*;
 19use util::{maybe, ResultExt};
 20use workspace::Workspace;
 21
 22pub(crate) struct DocsSlashCommandFeatureFlag;
 23
 24impl FeatureFlag for DocsSlashCommandFeatureFlag {
 25    const NAME: &'static str = "docs-slash-command";
 26}
 27
 28pub(crate) struct DocsSlashCommand;
 29
 30impl DocsSlashCommand {
 31    pub const NAME: &'static str = "docs";
 32
 33    fn path_to_cargo_toml(project: Model<Project>, cx: &mut AppContext) -> Option<Arc<Path>> {
 34        let worktree = project.read(cx).worktrees(cx).next()?;
 35        let worktree = worktree.read(cx);
 36        let entry = worktree.entry_for_path("Cargo.toml")?;
 37        let path = ProjectPath {
 38            worktree_id: worktree.id(),
 39            path: entry.path.clone(),
 40        };
 41        Some(Arc::from(
 42            project.read(cx).absolute_path(&path, cx)?.as_path(),
 43        ))
 44    }
 45
 46    /// Ensures that the indexed doc providers for Rust are registered.
 47    ///
 48    /// Ideally we would do this sooner, but we need to wait until we're able to
 49    /// access the workspace so we can read the project.
 50    fn ensure_rust_doc_providers_are_registered(
 51        &self,
 52        workspace: Option<WeakView<Workspace>>,
 53        cx: &mut AppContext,
 54    ) {
 55        let indexed_docs_registry = IndexedDocsRegistry::global(cx);
 56        if indexed_docs_registry
 57            .get_provider_store(LocalRustdocProvider::id())
 58            .is_none()
 59        {
 60            let index_provider_deps = maybe!({
 61                let workspace = workspace.clone().ok_or_else(|| anyhow!("no workspace"))?;
 62                let workspace = workspace
 63                    .upgrade()
 64                    .ok_or_else(|| anyhow!("workspace was dropped"))?;
 65                let project = workspace.read(cx).project().clone();
 66                let fs = project.read(cx).fs().clone();
 67                let cargo_workspace_root = Self::path_to_cargo_toml(project, cx)
 68                    .and_then(|path| path.parent().map(|path| path.to_path_buf()))
 69                    .ok_or_else(|| anyhow!("no Cargo workspace root found"))?;
 70
 71                anyhow::Ok((fs, cargo_workspace_root))
 72            });
 73
 74            if let Some((fs, cargo_workspace_root)) = index_provider_deps.log_err() {
 75                indexed_docs_registry.register_provider(Box::new(LocalRustdocProvider::new(
 76                    fs,
 77                    cargo_workspace_root,
 78                )));
 79            }
 80        }
 81
 82        if indexed_docs_registry
 83            .get_provider_store(DocsDotRsProvider::id())
 84            .is_none()
 85        {
 86            let http_client = maybe!({
 87                let workspace = workspace.ok_or_else(|| anyhow!("no workspace"))?;
 88                let workspace = workspace
 89                    .upgrade()
 90                    .ok_or_else(|| anyhow!("workspace was dropped"))?;
 91                let project = workspace.read(cx).project().clone();
 92                anyhow::Ok(project.read(cx).client().http_client().clone())
 93            });
 94
 95            if let Some(http_client) = http_client.log_err() {
 96                indexed_docs_registry
 97                    .register_provider(Box::new(DocsDotRsProvider::new(http_client)));
 98            }
 99        }
100    }
101
102    /// Runs just-in-time indexing for a given package, in case the slash command
103    /// is run without any entries existing in the index.
104    fn run_just_in_time_indexing(
105        store: Arc<IndexedDocsStore>,
106        key: String,
107        package: PackageName,
108        executor: BackgroundExecutor,
109    ) -> Task<()> {
110        executor.clone().spawn(async move {
111            let (prefix, needs_full_index) = if let Some((prefix, _)) = key.split_once('*') {
112                // If we have a wildcard in the search, we want to wait until
113                // we've completely finished indexing so we get a full set of
114                // results for the wildcard.
115                (prefix.to_string(), true)
116            } else {
117                (key, false)
118            };
119
120            // If we already have some entries, we assume that we've indexed the package before
121            // and don't need to do it again.
122            let has_any_entries = store
123                .any_with_prefix(prefix.clone())
124                .await
125                .unwrap_or_default();
126            if has_any_entries {
127                return ();
128            };
129
130            let index_task = store.clone().index(package.clone());
131
132            if needs_full_index {
133                _ = index_task.await;
134            } else {
135                loop {
136                    executor.timer(Duration::from_millis(200)).await;
137
138                    if store
139                        .any_with_prefix(prefix.clone())
140                        .await
141                        .unwrap_or_default()
142                        || !store.is_indexing(&package)
143                    {
144                        break;
145                    }
146                }
147            }
148        })
149    }
150}
151
152impl SlashCommand for DocsSlashCommand {
153    fn name(&self) -> String {
154        Self::NAME.into()
155    }
156
157    fn description(&self) -> String {
158        "insert docs".into()
159    }
160
161    fn menu_text(&self) -> String {
162        "Insert Documentation".into()
163    }
164
165    fn requires_argument(&self) -> bool {
166        true
167    }
168
169    fn complete_argument(
170        self: Arc<Self>,
171        query: String,
172        _cancel: Arc<AtomicBool>,
173        workspace: Option<WeakView<Workspace>>,
174        cx: &mut AppContext,
175    ) -> Task<Result<Vec<ArgumentCompletion>>> {
176        self.ensure_rust_doc_providers_are_registered(workspace, cx);
177
178        let indexed_docs_registry = IndexedDocsRegistry::global(cx);
179        let args = DocsSlashCommandArgs::parse(&query);
180        let store = args
181            .provider()
182            .ok_or_else(|| anyhow!("no docs provider specified"))
183            .and_then(|provider| IndexedDocsStore::try_global(provider, cx));
184        cx.background_executor().spawn(async move {
185            fn build_completions(
186                provider: ProviderId,
187                items: Vec<String>,
188            ) -> Vec<ArgumentCompletion> {
189                items
190                    .into_iter()
191                    .map(|item| ArgumentCompletion {
192                        label: item.clone(),
193                        new_text: format!("{provider} {item}"),
194                        run_command: true,
195                    })
196                    .collect()
197            }
198
199            match args {
200                DocsSlashCommandArgs::NoProvider => {
201                    let providers = indexed_docs_registry.list_providers();
202                    if providers.is_empty() {
203                        return Ok(vec![ArgumentCompletion {
204                            label: "No available docs providers.".to_string(),
205                            new_text: String::new(),
206                            run_command: false,
207                        }]);
208                    }
209
210                    Ok(providers
211                        .into_iter()
212                        .map(|provider| ArgumentCompletion {
213                            label: provider.to_string(),
214                            new_text: provider.to_string(),
215                            run_command: false,
216                        })
217                        .collect())
218                }
219                DocsSlashCommandArgs::SearchPackageDocs {
220                    provider,
221                    package,
222                    index,
223                } => {
224                    let store = store?;
225
226                    if index {
227                        // We don't need to hold onto this task, as the `IndexedDocsStore` will hold it
228                        // until it completes.
229                        drop(store.clone().index(package.as_str().into()));
230                    }
231
232                    let items = store.search(package).await;
233                    Ok(build_completions(provider, items))
234                }
235                DocsSlashCommandArgs::SearchItemDocs {
236                    provider,
237                    item_path,
238                    ..
239                } => {
240                    let store = store?;
241                    let items = store.search(item_path).await;
242                    Ok(build_completions(provider, items))
243                }
244            }
245        })
246    }
247
248    fn run(
249        self: Arc<Self>,
250        argument: Option<&str>,
251        _workspace: WeakView<Workspace>,
252        _delegate: Option<Arc<dyn LspAdapterDelegate>>,
253        cx: &mut WindowContext,
254    ) -> Task<Result<SlashCommandOutput>> {
255        let Some(argument) = argument else {
256            return Task::ready(Err(anyhow!("missing argument")));
257        };
258
259        let args = DocsSlashCommandArgs::parse(argument);
260        let executor = cx.background_executor().clone();
261        let task = cx.background_executor().spawn({
262            let store = args
263                .provider()
264                .ok_or_else(|| anyhow!("no docs provider specified"))
265                .and_then(|provider| IndexedDocsStore::try_global(provider, cx));
266            async move {
267                let (provider, key) = match args.clone() {
268                    DocsSlashCommandArgs::NoProvider => bail!("no docs provider specified"),
269                    DocsSlashCommandArgs::SearchPackageDocs {
270                        provider, package, ..
271                    } => (provider, package),
272                    DocsSlashCommandArgs::SearchItemDocs {
273                        provider,
274                        item_path,
275                        ..
276                    } => (provider, item_path),
277                };
278
279                if key.trim().is_empty() {
280                    let package_term = match provider.as_ref() {
281                        "docs-rs" | "rustdoc" => "crate",
282                        _ => "package",
283                    };
284
285                    bail!("no {package_term} name provided");
286                }
287
288                let store = store?;
289
290                if let Some(package) = args.package() {
291                    Self::run_just_in_time_indexing(store.clone(), key.clone(), package, executor)
292                        .await;
293                }
294
295                let (text, ranges) = if let Some((prefix, _)) = key.split_once('*') {
296                    let docs = store.load_many_by_prefix(prefix.to_string()).await?;
297
298                    let mut text = String::new();
299                    let mut ranges = Vec::new();
300
301                    for (key, docs) in docs {
302                        let prev_len = text.len();
303
304                        text.push_str(&docs.0);
305                        text.push_str("\n");
306                        ranges.push((key, prev_len..text.len()));
307                        text.push_str("\n");
308                    }
309
310                    (text, ranges)
311                } else {
312                    let item_docs = store.load(key.clone()).await?;
313                    let text = item_docs.to_string();
314                    let range = 0..text.len();
315
316                    (text, vec![(key, range)])
317                };
318
319                anyhow::Ok((provider, text, ranges))
320            }
321        });
322
323        cx.foreground_executor().spawn(async move {
324            let (provider, text, ranges) = task.await?;
325            Ok(SlashCommandOutput {
326                text,
327                sections: ranges
328                    .into_iter()
329                    .map(|(key, range)| SlashCommandOutputSection {
330                        range,
331                        icon: IconName::FileDoc,
332                        label: format!("docs ({provider}): {key}",).into(),
333                    })
334                    .collect(),
335                run_commands_in_text: false,
336            })
337        })
338    }
339}
340
341fn is_item_path_delimiter(char: char) -> bool {
342    !char.is_alphanumeric() && char != '-' && char != '_'
343}
344
345#[derive(Debug, PartialEq, Clone)]
346pub(crate) enum DocsSlashCommandArgs {
347    NoProvider,
348    SearchPackageDocs {
349        provider: ProviderId,
350        package: String,
351        index: bool,
352    },
353    SearchItemDocs {
354        provider: ProviderId,
355        package: String,
356        item_path: String,
357    },
358}
359
360impl DocsSlashCommandArgs {
361    pub fn parse(argument: &str) -> Self {
362        let Some((provider, argument)) = argument.split_once(' ') else {
363            return Self::NoProvider;
364        };
365
366        let provider = ProviderId(provider.into());
367
368        if let Some((package, rest)) = argument.split_once(is_item_path_delimiter) {
369            if rest.trim().is_empty() {
370                Self::SearchPackageDocs {
371                    provider,
372                    package: package.to_owned(),
373                    index: true,
374                }
375            } else {
376                Self::SearchItemDocs {
377                    provider,
378                    package: package.to_owned(),
379                    item_path: argument.to_owned(),
380                }
381            }
382        } else {
383            Self::SearchPackageDocs {
384                provider,
385                package: argument.to_owned(),
386                index: false,
387            }
388        }
389    }
390
391    pub fn provider(&self) -> Option<ProviderId> {
392        match self {
393            Self::NoProvider => None,
394            Self::SearchPackageDocs { provider, .. } | Self::SearchItemDocs { provider, .. } => {
395                Some(provider.clone())
396            }
397        }
398    }
399
400    pub fn package(&self) -> Option<PackageName> {
401        match self {
402            Self::NoProvider => None,
403            Self::SearchPackageDocs { package, .. } | Self::SearchItemDocs { package, .. } => {
404                Some(package.as_str().into())
405            }
406        }
407    }
408}
409
410#[cfg(test)]
411mod tests {
412    use super::*;
413
414    #[test]
415    fn test_parse_docs_slash_command_args() {
416        assert_eq!(
417            DocsSlashCommandArgs::parse(""),
418            DocsSlashCommandArgs::NoProvider
419        );
420        assert_eq!(
421            DocsSlashCommandArgs::parse("rustdoc"),
422            DocsSlashCommandArgs::NoProvider
423        );
424
425        assert_eq!(
426            DocsSlashCommandArgs::parse("rustdoc "),
427            DocsSlashCommandArgs::SearchPackageDocs {
428                provider: ProviderId("rustdoc".into()),
429                package: "".into(),
430                index: false
431            }
432        );
433        assert_eq!(
434            DocsSlashCommandArgs::parse("gleam "),
435            DocsSlashCommandArgs::SearchPackageDocs {
436                provider: ProviderId("gleam".into()),
437                package: "".into(),
438                index: false
439            }
440        );
441
442        assert_eq!(
443            DocsSlashCommandArgs::parse("rustdoc gpui"),
444            DocsSlashCommandArgs::SearchPackageDocs {
445                provider: ProviderId("rustdoc".into()),
446                package: "gpui".into(),
447                index: false,
448            }
449        );
450        assert_eq!(
451            DocsSlashCommandArgs::parse("gleam gleam_stdlib"),
452            DocsSlashCommandArgs::SearchPackageDocs {
453                provider: ProviderId("gleam".into()),
454                package: "gleam_stdlib".into(),
455                index: false
456            }
457        );
458
459        // Adding an item path delimiter indicates we can start indexing.
460        assert_eq!(
461            DocsSlashCommandArgs::parse("rustdoc gpui:"),
462            DocsSlashCommandArgs::SearchPackageDocs {
463                provider: ProviderId("rustdoc".into()),
464                package: "gpui".into(),
465                index: true,
466            }
467        );
468        assert_eq!(
469            DocsSlashCommandArgs::parse("gleam gleam_stdlib/"),
470            DocsSlashCommandArgs::SearchPackageDocs {
471                provider: ProviderId("gleam".into()),
472                package: "gleam_stdlib".into(),
473                index: true
474            }
475        );
476
477        assert_eq!(
478            DocsSlashCommandArgs::parse("rustdoc gpui::foo::bar::Baz"),
479            DocsSlashCommandArgs::SearchItemDocs {
480                provider: ProviderId("rustdoc".into()),
481                package: "gpui".into(),
482                item_path: "gpui::foo::bar::Baz".into()
483            }
484        );
485        assert_eq!(
486            DocsSlashCommandArgs::parse("gleam gleam_stdlib/gleam/int"),
487            DocsSlashCommandArgs::SearchItemDocs {
488                provider: ProviderId("gleam".into()),
489                package: "gleam_stdlib".into(),
490                item_path: "gleam_stdlib/gleam/int".into()
491            }
492        );
493    }
494}