search_command.rs

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