From b4d97c437d57f42f0b046cee4497a1a7a96af5c2 Mon Sep 17 00:00:00 2001 From: Agus Zubiaga Date: Tue, 12 Aug 2025 17:38:28 -0300 Subject: [PATCH] Get contents of symbols and selections Co-authored-by: Cole Miller --- crates/acp_thread/src/mention.rs | 179 +++++++++-- crates/agent2/src/thread.rs | 32 +- .../agent_ui/src/acp/completion_provider.rs | 298 ++++++++++-------- crates/agent_ui/src/acp/thread_view.rs | 4 +- crates/agent_ui/src/agent_panel.rs | 4 +- crates/agent_ui/src/context_picker.rs | 2 +- 6 files changed, 355 insertions(+), 164 deletions(-) diff --git a/crates/acp_thread/src/mention.rs b/crates/acp_thread/src/mention.rs index 61e613a3c8..ca0af26330 100644 --- a/crates/acp_thread/src/mention.rs +++ b/crates/acp_thread/src/mention.rs @@ -1,13 +1,24 @@ -use anyhow::{Result, bail}; -use std::path::PathBuf; +use anyhow::{Context as _, Result, bail}; +use std::{ + ops::Range, + path::{Path, PathBuf}, +}; #[derive(Clone, Debug, PartialEq, Eq)] pub enum MentionUri { File(PathBuf), - Symbol(PathBuf, String), + Symbol { + path: PathBuf, + name: String, + line_range: Range, + }, Thread(String), TextThread(PathBuf), Rule(String), + Selection { + path: PathBuf, + line_range: Range, + }, } impl MentionUri { @@ -17,7 +28,40 @@ impl MentionUri { match url.scheme() { "file" => { if let Some(fragment) = url.fragment() { - Ok(Self::Symbol(path.into(), fragment.into())) + let range = fragment + .strip_prefix("L") + .context("Line range must start with \"L\"")?; + let (start, end) = range + .split_once(":") + .context("Line range must use colon as separator")?; + let line_range = start + .parse::() + .context("Parsing line range start")? + .checked_sub(1) + .context("Line numbers should be 1-based")? + ..end + .parse::() + .context("Parsing line range end")? + .checked_sub(1) + .context("Line numbers should be 1-based")?; + let pairs = url.query_pairs().collect::>(); + match pairs.as_slice() { + [] => Ok(Self::Selection { + path: path.into(), + line_range, + }), + [(k, v)] => { + if k != "symbol" { + bail!("invalid query parameter") + } + Ok(Self::Symbol { + name: v.to_string(), + path: path.into(), + line_range, + }) + } + _ => bail!("too many query pairs"), + } } else { Ok(Self::File(path.into())) } @@ -37,11 +81,18 @@ impl MentionUri { pub fn name(&self) -> String { match self { - MentionUri::File(path) => path.file_name().unwrap().to_string_lossy().into_owned(), - MentionUri::Symbol(_path, name) => name.clone(), + MentionUri::File(path) => path + .file_name() + .unwrap_or_default() + .to_string_lossy() + .into_owned(), + MentionUri::Symbol { name, .. } => name.clone(), MentionUri::Thread(thread) => thread.to_string(), MentionUri::TextThread(thread) => thread.display().to_string(), MentionUri::Rule(rule) => rule.clone(), + MentionUri::Selection { + path, line_range, .. + } => selection_name(path, line_range), } } @@ -57,8 +108,26 @@ impl MentionUri { MentionUri::File(path) => { format!("file://{}", path.display()) } - MentionUri::Symbol(path, name) => { - format!("file://{}#{}", path.display(), name) + MentionUri::Symbol { + path, + name, + line_range, + } => { + format!( + "file://{}?symbol={}#L{}:{}", + path.display(), + name, + line_range.start + 1, + line_range.end + 1, + ) + } + MentionUri::Selection { path, line_range } => { + format!( + "file://{}#L{}:{}", + path.display(), + line_range.start + 1, + line_range.end + 1, + ) } MentionUri::Thread(thread) => { format!("zed:///agent/thread/{}", thread) @@ -73,13 +142,21 @@ impl MentionUri { } } +pub fn selection_name(path: &Path, line_range: &Range) -> String { + format!( + "{} ({}:{})", + path.file_name().unwrap_or_default().display(), + line_range.start + 1, + line_range.end + 1 + ) +} + #[cfg(test)] mod tests { use super::*; #[test] - fn test_mention_uri_parse_and_display() { - // Test file URI + fn test_parse_file_uri() { let file_uri = "file:///path/to/file.rs"; let parsed = MentionUri::parse(file_uri).unwrap(); match &parsed { @@ -87,20 +164,45 @@ mod tests { _ => panic!("Expected File variant"), } assert_eq!(parsed.to_uri(), file_uri); + } - // Test symbol URI - let symbol_uri = "file:///path/to/file.rs#MySymbol"; + #[test] + fn test_parse_symbol_uri() { + let symbol_uri = "file:///path/to/file.rs?symbol=MySymbol#L10:20"; let parsed = MentionUri::parse(symbol_uri).unwrap(); match &parsed { - MentionUri::Symbol(path, symbol) => { + MentionUri::Symbol { + path, + name, + line_range, + } => { assert_eq!(path.to_str().unwrap(), "/path/to/file.rs"); - assert_eq!(symbol, "MySymbol"); + assert_eq!(name, "MySymbol"); + assert_eq!(line_range.start, 9); + assert_eq!(line_range.end, 19); } _ => panic!("Expected Symbol variant"), } assert_eq!(parsed.to_uri(), symbol_uri); + } - // Test thread URI + #[test] + fn test_parse_selection_uri() { + let selection_uri = "file:///path/to/file.rs#L5:15"; + let parsed = MentionUri::parse(selection_uri).unwrap(); + match &parsed { + MentionUri::Selection { path, line_range } => { + assert_eq!(path.to_str().unwrap(), "/path/to/file.rs"); + assert_eq!(line_range.start, 4); + assert_eq!(line_range.end, 14); + } + _ => panic!("Expected Selection variant"), + } + assert_eq!(parsed.to_uri(), selection_uri); + } + + #[test] + fn test_parse_thread_uri() { let thread_uri = "zed:///agent/thread/session123"; let parsed = MentionUri::parse(thread_uri).unwrap(); match &parsed { @@ -108,8 +210,10 @@ mod tests { _ => panic!("Expected Thread variant"), } assert_eq!(parsed.to_uri(), thread_uri); + } - // Test rule URI + #[test] + fn test_parse_rule_uri() { let rule_uri = "zed:///agent/rule/my_rule"; let parsed = MentionUri::parse(rule_uri).unwrap(); match &parsed { @@ -117,11 +221,50 @@ mod tests { _ => panic!("Expected Rule variant"), } assert_eq!(parsed.to_uri(), rule_uri); + } - // Test invalid scheme + #[test] + fn test_invalid_scheme() { assert!(MentionUri::parse("http://example.com").is_err()); + assert!(MentionUri::parse("https://example.com").is_err()); + assert!(MentionUri::parse("ftp://example.com").is_err()); + } - // Test invalid zed path + #[test] + fn test_invalid_zed_path() { assert!(MentionUri::parse("zed:///invalid/path").is_err()); + assert!(MentionUri::parse("zed:///agent/unknown/test").is_err()); + } + + #[test] + fn test_invalid_line_range_format() { + // Missing L prefix + assert!(MentionUri::parse("file:///path/to/file.rs#10:20").is_err()); + + // Missing colon separator + assert!(MentionUri::parse("file:///path/to/file.rs#L1020").is_err()); + + // Invalid numbers + assert!(MentionUri::parse("file:///path/to/file.rs#L10:abc").is_err()); + assert!(MentionUri::parse("file:///path/to/file.rs#Labc:20").is_err()); + } + + #[test] + fn test_invalid_query_parameters() { + // Invalid query parameter name + assert!(MentionUri::parse("file:///path/to/file.rs#L10:20?invalid=test").is_err()); + + // Too many query parameters + assert!( + MentionUri::parse("file:///path/to/file.rs#L10:20?symbol=test&another=param").is_err() + ); + } + + #[test] + fn test_zero_based_line_numbers() { + // Test that 0-based line numbers are rejected (should be 1-based) + assert!(MentionUri::parse("file:///path/to/file.rs#L0:10").is_err()); + assert!(MentionUri::parse("file:///path/to/file.rs#L1:0").is_err()); + assert!(MentionUri::parse("file:///path/to/file.rs#L0:0").is_err()); } } diff --git a/crates/agent2/src/thread.rs b/crates/agent2/src/thread.rs index eea129cc0c..9da047cf08 100644 --- a/crates/agent2/src/thread.rs +++ b/crates/agent2/src/thread.rs @@ -26,8 +26,8 @@ use schemars::{JsonSchema, Schema}; use serde::{Deserialize, Serialize}; use settings::{Settings, update_settings_file}; use smol::stream::StreamExt; -use std::fmt::Write; use std::{cell::RefCell, collections::BTreeMap, path::Path, rc::Rc, sync::Arc}; +use std::{fmt::Write, ops::Range}; use util::{ResultExt, markdown::MarkdownCodeBlock}; #[derive(Debug, Clone)] @@ -1182,17 +1182,33 @@ impl AgentMessage { } MessageContent::Mention { uri, content } => { match uri { - MentionUri::File(path) | MentionUri::Symbol(path, _) => { + MentionUri::File(path) => { write!( &mut symbol_context, "\n{}", MarkdownCodeBlock { - tag: &codeblock_tag(&path), + tag: &codeblock_tag(&path, None), text: &content.to_string(), } ) .ok(); } + MentionUri::Symbol { + path, line_range, .. + } + | MentionUri::Selection { + path, line_range, .. + } => { + write!( + &mut rules_context, + "\n{}", + MarkdownCodeBlock { + tag: &codeblock_tag(&path, Some(line_range)), + text: &content + } + ) + .ok(); + } MentionUri::Thread(_session_id) => { write!(&mut thread_context, "\n{}\n", content).ok(); } @@ -1263,7 +1279,7 @@ impl AgentMessage { } } -fn codeblock_tag(full_path: &Path) -> String { +fn codeblock_tag(full_path: &Path, line_range: Option<&Range>) -> String { let mut result = String::new(); if let Some(extension) = full_path.extension().and_then(|ext| ext.to_str()) { @@ -1272,6 +1288,14 @@ fn codeblock_tag(full_path: &Path) -> String { let _ = write!(result, "{}", full_path.display()); + if let Some(range) = line_range { + if range.start == range.end { + let _ = write!(result, ":{}", range.start + 1); + } else { + let _ = write!(result, ":{}-{}", range.start + 1, range.end + 1); + } + } + result } diff --git a/crates/agent_ui/src/acp/completion_provider.rs b/crates/agent_ui/src/acp/completion_provider.rs index 090df753c1..c45a7a5843 100644 --- a/crates/agent_ui/src/acp/completion_provider.rs +++ b/crates/agent_ui/src/acp/completion_provider.rs @@ -3,16 +3,16 @@ use std::path::{Path, PathBuf}; use std::sync::Arc; use std::sync::atomic::AtomicBool; -use acp_thread::MentionUri; -use agent::context_store::ContextStore; +use acp_thread::{MentionUri, selection_name}; use anyhow::{Context as _, Result}; use collections::{HashMap, HashSet}; use editor::display_map::CreaseId; -use editor::{CompletionProvider, Editor, ExcerptId}; +use editor::{AnchorRangeExt, CompletionProvider, Editor, ExcerptId, ToOffset as _, ToPoint}; use file_icons::FileIcons; use futures::future::try_join_all; use fuzzy::{StringMatch, StringMatchCandidate}; use gpui::{App, Entity, Task, WeakEntity}; +use itertools::Itertools as _; use language::{Buffer, CodeLabel, HighlightId}; use lsp::CompletionContext; use parking_lot::Mutex; @@ -21,14 +21,12 @@ use project::{ }; use prompt_store::PromptStore; use rope::Point; -use text::{Anchor, ToPoint}; +use text::{Anchor, ToPoint as _}; use ui::prelude::*; -use util::ResultExt as _; use workspace::Workspace; use agent::{ - Thread, - context::{AgentContextHandle, AgentContextKey, RULES_ICON}, + context::RULES_ICON, thread_store::{TextThreadStore, ThreadStore}, }; @@ -40,10 +38,9 @@ use crate::context_picker::thread_context_picker::{ ThreadContextEntry, ThreadMatch, search_threads, }; use crate::context_picker::{ - ContextPickerEntry, ContextPickerMode, RecentEntry, available_context_picker_entries, - recent_context_picker_entries, + ContextPickerAction, ContextPickerEntry, ContextPickerMode, RecentEntry, + available_context_picker_entries, recent_context_picker_entries, selection_ranges, }; -use crate::message_editor::ContextCreasesAddon; #[derive(Default)] pub struct MentionSet { @@ -86,10 +83,44 @@ impl MentionSet { anyhow::Ok((crease_id, Mention { uri, content })) }) } - _ => { - // TODO - unimplemented!() + MentionUri::Symbol { + path, line_range, .. } + | MentionUri::Selection { + path, line_range, .. + } => { + let crease_id = *crease_id; + let uri = uri.clone(); + let path_buf = path.clone(); + let line_range = line_range.clone(); + + let buffer_task = project.update(cx, |project, cx| { + let path = project + .find_project_path(&path_buf, cx) + .context("Failed to find project path")?; + anyhow::Ok(project.open_buffer(path, cx)) + }); + + cx.spawn(async move |cx| { + let buffer = buffer_task?.await?; + let content = buffer.read_with(cx, |buffer, _cx| { + buffer + .text_for_range( + Point::new(line_range.start, 0) + ..Point::new( + line_range.end, + buffer.line_len(line_range.end), + ), + ) + .collect() + })?; + + anyhow::Ok((crease_id, Mention { uri, content })) + }) + } + MentionUri::Thread(_) => todo!(), + MentionUri::TextThread(path_buf) => todo!(), + MentionUri::Rule(_) => todo!(), }) .collect::>(); @@ -136,8 +167,8 @@ fn search( cancellation_flag: Arc, recent_entries: Vec, prompt_store: Option>, - thread_store: Option>, - text_thread_context_store: Option>, + thread_store: WeakEntity, + text_thread_context_store: WeakEntity, workspace: Entity, cx: &mut App, ) -> Task> { @@ -168,9 +199,8 @@ fn search( Some(ContextPickerMode::Thread) => { if let Some((thread_store, context_store)) = thread_store - .as_ref() - .and_then(|t| t.upgrade()) - .zip(text_thread_context_store.as_ref().and_then(|t| t.upgrade())) + .upgrade() + .zip(text_thread_context_store.upgrade()) { let search_threads_task = search_threads( query.clone(), @@ -240,14 +270,19 @@ fn search( .collect::>(); matches.extend( - available_context_picker_entries(&prompt_store, &thread_store, &workspace, cx) - .into_iter() - .map(|mode| { - Match::Entry(EntryMatch { - entry: mode, - mat: None, - }) - }), + available_context_picker_entries( + &prompt_store, + &Some(thread_store.clone()), + &workspace, + cx, + ) + .into_iter() + .map(|mode| { + Match::Entry(EntryMatch { + entry: mode, + mat: None, + }) + }), ); Task::ready(matches) @@ -257,8 +292,12 @@ fn search( let search_files_task = search_files(query.clone(), cancellation_flag.clone(), &workspace, cx); - let entries = - available_context_picker_entries(&prompt_store, &thread_store, &workspace, cx); + let entries = available_context_picker_entries( + &prompt_store, + &Some(thread_store.clone()), + &workspace, + cx, + ); let entry_candidates = entries .iter() .enumerate() @@ -352,117 +391,101 @@ impl ContextPickerCompletionProvider { confirm: Some(Arc::new(|_, _, _| true)), }), ContextPickerEntry::Action(action) => { - todo!() - // let (new_text, on_action) = match action { - // ContextPickerAction::AddSelections => { - // let selections = selection_ranges(workspace, cx); + let (new_text, on_action) = match action { + ContextPickerAction::AddSelections => { + let selections = selection_ranges(workspace, cx); - // let selection_infos = selections - // .iter() - // .map(|(buffer, range)| { - // let full_path = buffer - // .read(cx) - // .file() - // .map(|file| file.full_path(cx)) - // .unwrap_or_else(|| PathBuf::from("untitled")); - // let file_name = full_path - // .file_name() - // .unwrap_or_default() - // .to_string_lossy() - // .to_string(); - // let line_range = range.to_point(&buffer.read(cx).snapshot()); + const PLACEHOLDER: &str = "selection "; - // let link = MentionLink::for_selection( - // &file_name, - // &full_path.to_string_lossy(), - // line_range.start.row as usize..line_range.end.row as usize, - // ); - // (file_name, link, line_range) - // }) - // .collect::>(); + let new_text = std::iter::repeat(PLACEHOLDER) + .take(selections.len()) + .chain(std::iter::once("")) + .join(" "); - // let new_text = format!( - // "{} ", - // selection_infos.iter().map(|(_, link, _)| link).join(" ") - // ); + let callback = Arc::new({ + let mention_set = mention_set.clone(); + let selections = selections.clone(); + move |_, window: &mut Window, cx: &mut App| { + let editor = editor.clone(); + let mention_set = mention_set.clone(); + let selections = selections.clone(); + window.defer(cx, move |window, cx| { + let mut current_offset = 0; - // let callback = Arc::new({ - // let context_store = context_store.clone(); - // let selections = selections.clone(); - // let selection_infos = selection_infos.clone(); - // move |_, window: &mut Window, cx: &mut App| { - // context_store.update(cx, |context_store, cx| { - // for (buffer, range) in &selections { - // context_store.add_selection( - // buffer.clone(), - // range.clone(), - // cx, - // ); - // } - // }); + for (buffer, selection_range) in selections { + let snapshot = + editor.read(cx).buffer().read(cx).snapshot(cx); + let Some(start) = snapshot + .anchor_in_excerpt(excerpt_id, source_range.start) + else { + return; + }; - // let editor = editor.clone(); - // let selection_infos = selection_infos.clone(); - // window.defer(cx, move |window, cx| { - // let mut current_offset = 0; - // for (file_name, link, line_range) in selection_infos.iter() { - // let snapshot = - // editor.read(cx).buffer().read(cx).snapshot(cx); - // let Some(start) = snapshot - // .anchor_in_excerpt(excerpt_id, source_range.start) - // else { - // return; - // }; + let offset = start.to_offset(&snapshot) + current_offset; + let text_len = PLACEHOLDER.len() - 1; - // let offset = start.to_offset(&snapshot) + current_offset; - // let text_len = link.len(); + let range = snapshot.anchor_after(offset) + ..snapshot.anchor_after(offset + text_len); - // let range = snapshot.anchor_after(offset) - // ..snapshot.anchor_after(offset + text_len); + let path = buffer + .read(cx) + .file() + .map_or(PathBuf::from("untitled"), |file| { + file.path().to_path_buf() + }); - // let crease = super::crease_for_mention( - // format!( - // "{} ({}-{})", - // file_name, - // line_range.start.row + 1, - // line_range.end.row + 1 - // ) - // .into(), - // IconName::Reader.path().into(), - // range, - // editor.downgrade(), - // ); + let point_range = range.to_point(&snapshot); + let line_range = point_range.start.row..point_range.end.row; + let crease = crate::context_picker::crease_for_mention( + selection_name(&path, &line_range).into(), + IconName::Reader.path().into(), + range, + editor.downgrade(), + ); - // editor.update(cx, |editor, cx| { - // editor.insert_creases(vec![crease.clone()], cx); - // editor.fold_creases(vec![crease], false, window, cx); - // }); + let [crease_id]: [_; 1] = + editor.update(cx, |editor, cx| { + let crease_ids = + editor.insert_creases(vec![crease.clone()], cx); + editor.fold_creases( + vec![crease], + false, + window, + cx, + ); + crease_ids.try_into().unwrap() + }); - // current_offset += text_len + 1; - // } - // }); + mention_set.lock().insert( + crease_id, + MentionUri::Selection { path, line_range }, + ); - // false - // } - // }); + current_offset += text_len + 1; + } + }); - // (new_text, callback) - // } - // }; + false + } + }); - // Some(Completion { - // replace_range: source_range.clone(), - // new_text, - // label: CodeLabel::plain(action.label().to_string(), None), - // icon_path: Some(action.icon().path().into()), - // documentation: None, - // source: project::CompletionSource::Custom, - // insert_text_mode: None, - // // This ensures that when a user accepts this completion, the - // // completion menu will still be shown after "@category " is - // // inserted - // confirm: Some(on_action), - // }) + (new_text, callback) + } + }; + + Some(Completion { + replace_range: source_range.clone(), + new_text, + label: CodeLabel::plain(action.label().to_string(), None), + icon_path: Some(action.icon().path().into()), + documentation: None, + source: project::CompletionSource::Custom, + insert_text_mode: None, + // This ensures that when a user accepts this completion, the + // completion menu will still be shown after "@category " is + // inserted + confirm: Some(on_action), + }) } } } @@ -636,11 +659,12 @@ impl ContextPickerCompletionProvider { let comment_id = cx.theme().syntax().highlight_id("comment").map(HighlightId); let mut label = CodeLabel::plain(symbol.name.clone(), None); - label.push_str(" ", None); - label.push_str(&file_name, comment_id); - label.push_str(&format!(" L{}", symbol.range.start.0.row + 1), comment_id); - let uri = MentionUri::Symbol(full_path.into(), symbol.name.clone()); + let uri = MentionUri::Symbol { + path: full_path.into(), + name: symbol.name.clone(), + line_range: symbol.range.start.0.row..symbol.range.end.0.row, + }; let new_text = format!("{} ", uri.to_link()); let new_text_len = new_text.len(); Some(Completion { @@ -741,20 +765,18 @@ impl CompletionProvider for ContextPickerCompletionProvider { }; let recent_entries = recent_context_picker_entries( - thread_store.clone(), - text_thread_store.clone(), + Some(thread_store.clone()), + Some(text_thread_store.clone()), workspace.clone(), &exclude_paths, &exclude_threads, cx, ); - let prompt_store = thread_store.as_ref().and_then(|thread_store| { - thread_store - .read_with(cx, |thread_store, _cx| thread_store.prompt_store().clone()) - .ok() - .flatten() - }); + let prompt_store = thread_store + .read_with(cx, |thread_store, _cx| thread_store.prompt_store().clone()) + .ok() + .flatten(); let search_task = search( mode, diff --git a/crates/agent_ui/src/acp/thread_view.rs b/crates/agent_ui/src/acp/thread_view.rs index 937c5d1f5e..02d3e1c0d1 100644 --- a/crates/agent_ui/src/acp/thread_view.rs +++ b/crates/agent_ui/src/acp/thread_view.rs @@ -106,7 +106,7 @@ impl AcpThreadView { workspace: WeakEntity, project: Entity, thread_store: WeakEntity, - text_thread_store: WeaksEntity, + text_thread_store: WeakEntity, message_history: Rc>>>, min_lines: usize, max_lines: Option, @@ -3478,6 +3478,8 @@ mod tests { Rc::new(agent), workspace.downgrade(), project, + WeakEntity::new_invalid(), + WeakEntity::new_invalid(), Rc::new(RefCell::new(MessageHistory::default())), 1, None, diff --git a/crates/agent_ui/src/agent_panel.rs b/crates/agent_ui/src/agent_panel.rs index 01e88c3bfa..fbb75f28c0 100644 --- a/crates/agent_ui/src/agent_panel.rs +++ b/crates/agent_ui/src/agent_panel.rs @@ -965,8 +965,8 @@ impl AgentPanel { server, workspace.clone(), project, - thread_store.clone(), - text_thread_store.clone(), + thread_store.downgrade(), + text_thread_store.downgrade(), message_history, MIN_EDITOR_LINES, Some(MAX_EDITOR_LINES), diff --git a/crates/agent_ui/src/context_picker.rs b/crates/agent_ui/src/context_picker.rs index e53b43b306..7dc00bfae2 100644 --- a/crates/agent_ui/src/context_picker.rs +++ b/crates/agent_ui/src/context_picker.rs @@ -742,7 +742,7 @@ fn add_selections_as_context( }) } -fn selection_ranges( +pub(crate) fn selection_ranges( workspace: &Entity, cx: &mut App, ) -> Vec<(Entity, Range)> {