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