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        _context_slash_command_output_sections: &[SlashCommandOutputSection<language::Anchor>],
 64        _context_buffer: language::BufferSnapshot,
 65        workspace: WeakView<Workspace>,
 66        _delegate: Option<Arc<dyn LspAdapterDelegate>>,
 67        cx: &mut WindowContext,
 68    ) -> Task<Result<SlashCommandOutput>> {
 69        let Some(workspace) = workspace.upgrade() else {
 70            return Task::ready(Err(anyhow::anyhow!("workspace was dropped")));
 71        };
 72        if arguments.is_empty() {
 73            return Task::ready(Err(anyhow::anyhow!("missing search query")));
 74        };
 75
 76        let mut limit = None;
 77        let mut query = String::new();
 78        for part in arguments {
 79            if let Some(parameter) = part.strip_prefix("--") {
 80                if let Ok(count) = parameter.parse::<usize>() {
 81                    limit = Some(count);
 82                    continue;
 83                }
 84            }
 85
 86            query.push_str(part);
 87            query.push(' ');
 88        }
 89        query.pop();
 90
 91        if query.is_empty() {
 92            return Task::ready(Err(anyhow::anyhow!("missing search query")));
 93        }
 94
 95        let project = workspace.read(cx).project().clone();
 96        let fs = project.read(cx).fs().clone();
 97        let Some(project_index) =
 98            cx.update_global(|index: &mut SemanticDb, cx| index.project_index(project, cx))
 99        else {
100            return Task::ready(Err(anyhow::anyhow!("no project indexer")));
101        };
102
103        cx.spawn(|cx| async move {
104            let results = project_index
105                .read_with(&cx, |project_index, cx| {
106                    project_index.search(query.clone(), limit.unwrap_or(5), cx)
107                })?
108                .await?;
109
110            let mut loaded_results = Vec::new();
111            for result in results {
112                let (full_path, file_content) =
113                    result.worktree.read_with(&cx, |worktree, _cx| {
114                        let entry_abs_path = worktree.abs_path().join(&result.path);
115                        let mut entry_full_path = PathBuf::from(worktree.root_name());
116                        entry_full_path.push(&result.path);
117                        let file_content = async {
118                            let entry_abs_path = entry_abs_path;
119                            fs.load(&entry_abs_path).await
120                        };
121                        (entry_full_path, file_content)
122                    })?;
123                if let Some(file_content) = file_content.await.log_err() {
124                    loaded_results.push((result, full_path, file_content));
125                }
126            }
127
128            let output = cx
129                .background_executor()
130                .spawn(async move {
131                    let mut text = format!("Search results for {query}:\n");
132                    let mut sections = Vec::new();
133                    for (result, full_path, file_content) in loaded_results {
134                        let range_start = result.range.start.min(file_content.len());
135                        let range_end = result.range.end.min(file_content.len());
136
137                        let start_row = file_content[0..range_start].matches('\n').count() as u32;
138                        let end_row = file_content[0..range_end].matches('\n').count() as u32;
139                        let start_line_byte_offset = file_content[0..range_start]
140                            .rfind('\n')
141                            .map(|pos| pos + 1)
142                            .unwrap_or_default();
143                        let end_line_byte_offset = file_content[range_end..]
144                            .find('\n')
145                            .map(|pos| range_end + pos)
146                            .unwrap_or_else(|| file_content.len());
147
148                        let section_start_ix = text.len();
149                        text.push_str(&codeblock_fence_for_path(
150                            Some(&result.path),
151                            Some(start_row..end_row),
152                        ));
153
154                        let mut excerpt =
155                            file_content[start_line_byte_offset..end_line_byte_offset].to_string();
156                        LineEnding::normalize(&mut excerpt);
157                        text.push_str(&excerpt);
158                        writeln!(text, "\n```\n").unwrap();
159                        let section_end_ix = text.len() - 1;
160                        sections.push(build_entry_output_section(
161                            section_start_ix..section_end_ix,
162                            Some(&full_path),
163                            false,
164                            Some(start_row + 1..end_row + 1),
165                        ));
166                    }
167
168                    let query = SharedString::from(query);
169                    sections.push(SlashCommandOutputSection {
170                        range: 0..text.len(),
171                        icon: IconName::MagnifyingGlass,
172                        label: query,
173                        metadata: None,
174                    });
175
176                    SlashCommandOutput {
177                        text,
178                        sections,
179                        run_commands_in_text: false,
180                    }
181                })
182                .await;
183
184            Ok(output)
185        })
186    }
187}