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::SemanticDb;
 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 Some(project_index) =
 96            cx.update_global(|index: &mut SemanticDb, cx| index.project_index(project, cx))
 97        else {
 98            return Task::ready(Err(anyhow::anyhow!("no project indexer")));
 99        };
100
101        cx.spawn(|cx| async move {
102            let results = project_index
103                .read_with(&cx, |project_index, cx| {
104                    project_index.search(query.clone(), limit.unwrap_or(5), cx)
105                })?
106                .await?;
107
108            let mut loaded_results = Vec::new();
109            for result in results {
110                let (full_path, file_content) =
111                    result.worktree.read_with(&cx, |worktree, _cx| {
112                        let entry_abs_path = worktree.abs_path().join(&result.path);
113                        let mut entry_full_path = PathBuf::from(worktree.root_name());
114                        entry_full_path.push(&result.path);
115                        let file_content = async {
116                            let entry_abs_path = entry_abs_path;
117                            fs.load(&entry_abs_path).await
118                        };
119                        (entry_full_path, file_content)
120                    })?;
121                if let Some(file_content) = file_content.await.log_err() {
122                    loaded_results.push((result, full_path, file_content));
123                }
124            }
125
126            let output = cx
127                .background_executor()
128                .spawn(async move {
129                    let mut text = format!("Search results for {query}:\n");
130                    let mut sections = Vec::new();
131                    for (result, full_path, file_content) in loaded_results {
132                        let range_start = result.range.start.min(file_content.len());
133                        let range_end = result.range.end.min(file_content.len());
134
135                        let start_row = file_content[0..range_start].matches('\n').count() as u32;
136                        let end_row = file_content[0..range_end].matches('\n').count() as u32;
137                        let start_line_byte_offset = file_content[0..range_start]
138                            .rfind('\n')
139                            .map(|pos| pos + 1)
140                            .unwrap_or_default();
141                        let end_line_byte_offset = file_content[range_end..]
142                            .find('\n')
143                            .map(|pos| range_end + pos)
144                            .unwrap_or_else(|| file_content.len());
145
146                        let section_start_ix = text.len();
147                        text.push_str(&codeblock_fence_for_path(
148                            Some(&result.path),
149                            Some(start_row..end_row),
150                        ));
151
152                        let mut excerpt =
153                            file_content[start_line_byte_offset..end_line_byte_offset].to_string();
154                        LineEnding::normalize(&mut excerpt);
155                        text.push_str(&excerpt);
156                        writeln!(text, "\n```\n").unwrap();
157                        let section_end_ix = text.len() - 1;
158                        sections.push(build_entry_output_section(
159                            section_start_ix..section_end_ix,
160                            Some(&full_path),
161                            false,
162                            Some(start_row + 1..end_row + 1),
163                        ));
164                    }
165
166                    let query = SharedString::from(query);
167                    sections.push(SlashCommandOutputSection {
168                        range: 0..text.len(),
169                        icon: IconName::MagnifyingGlass,
170                        label: query,
171                    });
172
173                    SlashCommandOutput {
174                        text,
175                        sections,
176                        run_commands_in_text: false,
177                    }
178                })
179                .await;
180
181            Ok(output)
182        })
183    }
184}