search_command.rs

  1use anyhow::Result;
  2use assistant_slash_command::{
  3    ArgumentCompletion, SlashCommand, SlashCommandOutput, SlashCommandOutputSection,
  4    SlashCommandResult,
  5};
  6use feature_flags::FeatureFlag;
  7use gpui::{AppContext, Task, WeakView};
  8use language::{CodeLabel, LspAdapterDelegate};
  9use semantic_index::{LoadedSearchResult, SemanticDb};
 10use std::{
 11    fmt::Write,
 12    sync::{atomic::AtomicBool, Arc},
 13};
 14use ui::{prelude::*, IconName};
 15use workspace::Workspace;
 16
 17use crate::slash_command::create_label_for_command;
 18use crate::slash_command::file_command::{build_entry_output_section, codeblock_fence_for_path};
 19
 20pub(crate) struct SearchSlashCommandFeatureFlag;
 21
 22impl FeatureFlag for SearchSlashCommandFeatureFlag {
 23    const NAME: &'static str = "search-slash-command";
 24
 25    fn enabled_for_staff() -> bool {
 26        false
 27    }
 28}
 29
 30pub(crate) struct SearchSlashCommand;
 31
 32impl SlashCommand for SearchSlashCommand {
 33    fn name(&self) -> String {
 34        "search".into()
 35    }
 36
 37    fn label(&self, cx: &AppContext) -> CodeLabel {
 38        create_label_for_command("search", &["--n"], cx)
 39    }
 40
 41    fn description(&self) -> String {
 42        "Search your project semantically".into()
 43    }
 44
 45    fn icon(&self) -> IconName {
 46        IconName::SearchCode
 47    }
 48
 49    fn menu_text(&self) -> String {
 50        self.description()
 51    }
 52
 53    fn requires_argument(&self) -> bool {
 54        true
 55    }
 56
 57    fn complete_argument(
 58        self: Arc<Self>,
 59        _arguments: &[String],
 60        _cancel: Arc<AtomicBool>,
 61        _workspace: Option<WeakView<Workspace>>,
 62        _cx: &mut WindowContext,
 63    ) -> Task<Result<Vec<ArgumentCompletion>>> {
 64        Task::ready(Ok(Vec::new()))
 65    }
 66
 67    fn run(
 68        self: Arc<Self>,
 69        arguments: &[String],
 70        _context_slash_command_output_sections: &[SlashCommandOutputSection<language::Anchor>],
 71        _context_buffer: language::BufferSnapshot,
 72        workspace: WeakView<Workspace>,
 73        _delegate: Option<Arc<dyn LspAdapterDelegate>>,
 74        cx: &mut WindowContext,
 75    ) -> Task<SlashCommandResult> {
 76        let Some(workspace) = workspace.upgrade() else {
 77            return Task::ready(Err(anyhow::anyhow!("workspace was dropped")));
 78        };
 79        if arguments.is_empty() {
 80            return Task::ready(Err(anyhow::anyhow!("missing search query")));
 81        };
 82
 83        let mut limit = None;
 84        let mut query = String::new();
 85        for part in arguments {
 86            if let Some(parameter) = part.strip_prefix("--") {
 87                if let Ok(count) = parameter.parse::<usize>() {
 88                    limit = Some(count);
 89                    continue;
 90                }
 91            }
 92
 93            query.push_str(part);
 94            query.push(' ');
 95        }
 96        query.pop();
 97
 98        if query.is_empty() {
 99            return Task::ready(Err(anyhow::anyhow!("missing search query")));
100        }
101
102        let project = workspace.read(cx).project().clone();
103        let fs = project.read(cx).fs().clone();
104        let Some(project_index) =
105            cx.update_global(|index: &mut SemanticDb, cx| index.project_index(project, cx))
106        else {
107            return Task::ready(Err(anyhow::anyhow!("no project indexer")));
108        };
109
110        cx.spawn(|cx| async move {
111            let results = project_index
112                .read_with(&cx, |project_index, cx| {
113                    project_index.search(vec![query.clone()], limit.unwrap_or(5), cx)
114                })?
115                .await?;
116
117            let loaded_results = SemanticDb::load_results(results, &fs, &cx).await?;
118
119            let output = cx
120                .background_executor()
121                .spawn(async move {
122                    let mut text = format!("Search results for {query}:\n");
123                    let mut sections = Vec::new();
124                    for loaded_result in &loaded_results {
125                        add_search_result_section(loaded_result, &mut text, &mut sections);
126                    }
127
128                    let query = SharedString::from(query);
129                    sections.push(SlashCommandOutputSection {
130                        range: 0..text.len(),
131                        icon: IconName::MagnifyingGlass,
132                        label: query,
133                        metadata: None,
134                    });
135
136                    SlashCommandOutput {
137                        text,
138                        sections,
139                        run_commands_in_text: false,
140                    }
141                    .to_event_stream()
142                })
143                .await;
144
145            Ok(output)
146        })
147    }
148}
149
150pub fn add_search_result_section(
151    loaded_result: &LoadedSearchResult,
152    text: &mut String,
153    sections: &mut Vec<SlashCommandOutputSection<usize>>,
154) {
155    let LoadedSearchResult {
156        path,
157        full_path,
158        excerpt_content,
159        row_range,
160        ..
161    } = loaded_result;
162    let section_start_ix = text.len();
163    text.push_str(&codeblock_fence_for_path(
164        Some(&path),
165        Some(row_range.clone()),
166    ));
167
168    text.push_str(&excerpt_content);
169    if !text.ends_with('\n') {
170        text.push('\n');
171    }
172    writeln!(text, "```\n").unwrap();
173    let section_end_ix = text.len() - 1;
174    sections.push(build_entry_output_section(
175        section_start_ix..section_end_ix,
176        Some(&full_path),
177        false,
178        Some(row_range.start() + 1..row_range.end() + 1),
179    ));
180}