acp: Eagerly load all kinds of mentions (#36741)

This PR makes it so that all kinds of @-mentions start loading their
context as soon as they are confirmed. Previously, we were waiting to
load the context for file, symbol, selection, and rule mentions until
the user's message was sent. By kicking off loading immediately for
these kinds of context, we can support adding selections from unsaved
buffers, and we make the semantics of @-mentions more consistent.

Loading all kinds of context eagerly also makes it possible to simplify
the structure of the MentionSet and the code around it. Now MentionSet
is just a single hash map, all the management of creases happens in a
uniform way in `MessageEditor::confirm_completion`, and the helper
methods for loading different kinds of context are much more focused and
orthogonal.

Release Notes:

- N/A

---------

Co-authored-by: Conrad <conrad@zed.dev>

crates/acp_thread/src/mention.rs

@@ -5,7 +5,7 @@ use prompt_store::{PromptId, UserPromptId};
 use serde::{Deserialize, Serialize};
 use std::{
     fmt,
-    ops::Range,
+    ops::RangeInclusive,
     path::{Path, PathBuf},
     str::FromStr,
 };
@@ -17,13 +17,14 @@ pub enum MentionUri {
     File {
         abs_path: PathBuf,
     },
+    PastedImage,
     Directory {
         abs_path: PathBuf,
     },
     Symbol {
-        path: PathBuf,
+        abs_path: PathBuf,
         name: String,
-        line_range: Range<u32>,
+        line_range: RangeInclusive<u32>,
     },
     Thread {
         id: acp::SessionId,
@@ -38,8 +39,9 @@ pub enum MentionUri {
         name: String,
     },
     Selection {
-        path: PathBuf,
+        #[serde(default, skip_serializing_if = "Option::is_none")]
-        line_range: Range<u32>,
+        abs_path: Option<PathBuf>,
+        line_range: RangeInclusive<u32>,
     },
     Fetch {
         url: Url,
@@ -48,36 +50,44 @@ pub enum MentionUri {
 
 impl MentionUri {
     pub fn parse(input: &str) -> Result<Self> {
+        fn parse_line_range(fragment: &str) -> Result<RangeInclusive<u32>> {
+            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 range = start
+                .parse::<u32>()
+                .context("Parsing line range start")?
+                .checked_sub(1)
+                .context("Line numbers should be 1-based")?
+                ..=end
+                    .parse::<u32>()
+                    .context("Parsing line range end")?
+                    .checked_sub(1)
+                    .context("Line numbers should be 1-based")?;
+            Ok(range)
+        }
+
         let url = url::Url::parse(input)?;
         let path = url.path();
         match url.scheme() {
             "file" => {
                 let path = url.to_file_path().ok().context("Extracting file path")?;
                 if let Some(fragment) = url.fragment() {
-                    let range = fragment
+                    let line_range = parse_line_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::<u32>()
-                        .context("Parsing line range start")?
-                        .checked_sub(1)
-                        .context("Line numbers should be 1-based")?
-                        ..end
-                            .parse::<u32>()
-                            .context("Parsing line range end")?
-                            .checked_sub(1)
-                            .context("Line numbers should be 1-based")?;
                     if let Some(name) = single_query_param(&url, "symbol")? {
                         Ok(Self::Symbol {
                             name,
-                            path,
+                            abs_path: path,
                             line_range,
                         })
                     } else {
-                        Ok(Self::Selection { path, line_range })
+                        Ok(Self::Selection {
+                            abs_path: Some(path),
+                            line_range,
+                        })
                     }
                 } else if input.ends_with("/") {
                     Ok(Self::Directory { abs_path: path })
@@ -105,6 +115,17 @@ impl MentionUri {
                         id: rule_id.into(),
                         name,
                     })
+                } else if path.starts_with("/agent/pasted-image") {
+                    Ok(Self::PastedImage)
+                } else if path.starts_with("/agent/untitled-buffer") {
+                    let fragment = url
+                        .fragment()
+                        .context("Missing fragment for untitled buffer selection")?;
+                    let line_range = parse_line_range(fragment)?;
+                    Ok(Self::Selection {
+                        abs_path: None,
+                        line_range,
+                    })
                 } else {
                     bail!("invalid zed url: {:?}", input);
                 }
@@ -121,13 +142,16 @@ impl MentionUri {
                 .unwrap_or_default()
                 .to_string_lossy()
                 .into_owned(),
+            MentionUri::PastedImage => "Image".to_string(),
             MentionUri::Symbol { name, .. } => name.clone(),
             MentionUri::Thread { name, .. } => name.clone(),
             MentionUri::TextThread { name, .. } => name.clone(),
             MentionUri::Rule { name, .. } => name.clone(),
             MentionUri::Selection {
-                path, line_range, ..
+                abs_path: path,
-            } => selection_name(path, line_range),
+                line_range,
+                ..
+            } => selection_name(path.as_deref(), line_range),
             MentionUri::Fetch { url } => url.to_string(),
         }
     }
@@ -137,6 +161,7 @@ impl MentionUri {
             MentionUri::File { abs_path } => {
                 FileIcons::get_icon(abs_path, cx).unwrap_or_else(|| IconName::File.path().into())
             }
+            MentionUri::PastedImage => IconName::Image.path().into(),
             MentionUri::Directory { .. } => FileIcons::get_folder_icon(false, cx)
                 .unwrap_or_else(|| IconName::Folder.path().into()),
             MentionUri::Symbol { .. } => IconName::Code.path().into(),
@@ -157,29 +182,40 @@ impl MentionUri {
             MentionUri::File { abs_path } => {
                 Url::from_file_path(abs_path).expect("mention path should be absolute")
             }
+            MentionUri::PastedImage => Url::parse("zed:///agent/pasted-image").unwrap(),
             MentionUri::Directory { abs_path } => {
                 Url::from_directory_path(abs_path).expect("mention path should be absolute")
             }
             MentionUri::Symbol {
-                path,
+                abs_path,
                 name,
                 line_range,
             } => {
-                let mut url = Url::from_file_path(path).expect("mention path should be absolute");
+                let mut url =
+                    Url::from_file_path(abs_path).expect("mention path should be absolute");
                 url.query_pairs_mut().append_pair("symbol", name);
                 url.set_fragment(Some(&format!(
                     "L{}:{}",
-                    line_range.start + 1,
+                    line_range.start() + 1,
-                    line_range.end + 1
+                    line_range.end() + 1
                 )));
                 url
             }
-            MentionUri::Selection { path, line_range } => {
+            MentionUri::Selection {
-                let mut url = Url::from_file_path(path).expect("mention path should be absolute");
+                abs_path: path,
+                line_range,
+            } => {
+                let mut url = if let Some(path) = path {
+                    Url::from_file_path(path).expect("mention path should be absolute")
+                } else {
+                    let mut url = Url::parse("zed:///").unwrap();
+                    url.set_path("/agent/untitled-buffer");
+                    url
+                };
                 url.set_fragment(Some(&format!(
                     "L{}:{}",
-                    line_range.start + 1,
+                    line_range.start() + 1,
-                    line_range.end + 1
+                    line_range.end() + 1
                 )));
                 url
             }
@@ -191,7 +227,10 @@ impl MentionUri {
             }
             MentionUri::TextThread { path, name } => {
                 let mut url = Url::parse("zed:///").unwrap();
-                url.set_path(&format!("/agent/text-thread/{}", path.to_string_lossy()));
+                url.set_path(&format!(
+                    "/agent/text-thread/{}",
+                    path.to_string_lossy().trim_start_matches('/')
+                ));
                 url.query_pairs_mut().append_pair("name", name);
                 url
             }
@@ -237,12 +276,14 @@ fn single_query_param(url: &Url, name: &'static str) -> Result<Option<String>> {
     }
 }
 
-pub fn selection_name(path: &Path, line_range: &Range<u32>) -> String {
+pub fn selection_name(path: Option<&Path>, line_range: &RangeInclusive<u32>) -> String {
     format!(
         "{} ({}:{})",
-        path.file_name().unwrap_or_default().display(),
+        path.and_then(|path| path.file_name())
-        line_range.start + 1,
+            .unwrap_or("Untitled".as_ref())
-        line_range.end + 1
+            .display(),
+        *line_range.start() + 1,
+        *line_range.end() + 1
     )
 }
 
@@ -302,14 +343,14 @@ mod tests {
         let parsed = MentionUri::parse(symbol_uri).unwrap();
         match &parsed {
             MentionUri::Symbol {
-                path,
+                abs_path: path,
                 name,
                 line_range,
             } => {
                 assert_eq!(path.to_str().unwrap(), path!("/path/to/file.rs"));
                 assert_eq!(name, "MySymbol");
-                assert_eq!(line_range.start, 9);
+                assert_eq!(line_range.start(), &9);
-                assert_eq!(line_range.end, 19);
+                assert_eq!(line_range.end(), &19);
             }
             _ => panic!("Expected Symbol variant"),
         }
@@ -321,16 +362,39 @@ mod tests {
         let selection_uri = uri!("file:///path/to/file.rs#L5:15");
         let parsed = MentionUri::parse(selection_uri).unwrap();
         match &parsed {
-            MentionUri::Selection { path, line_range } => {
+            MentionUri::Selection {
-                assert_eq!(path.to_str().unwrap(), path!("/path/to/file.rs"));
+                abs_path: path,
-                assert_eq!(line_range.start, 4);
+                line_range,
-                assert_eq!(line_range.end, 14);
+            } => {
+                assert_eq!(
+                    path.as_ref().unwrap().to_str().unwrap(),
+                    path!("/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().to_string(), selection_uri);
     }
 
+    #[test]
+    fn test_parse_untitled_selection_uri() {
+        let selection_uri = uri!("zed:///agent/untitled-buffer#L1:10");
+        let parsed = MentionUri::parse(selection_uri).unwrap();
+        match &parsed {
+            MentionUri::Selection {
+                abs_path: None,
+                line_range,
+            } => {
+                assert_eq!(line_range.start(), &0);
+                assert_eq!(line_range.end(), &9);
+            }
+            _ => panic!("Expected Selection variant without path"),
+        }
+        assert_eq!(parsed.to_uri().to_string(), selection_uri);
+    }
+
     #[test]
     fn test_parse_thread_uri() {
         let thread_uri = "zed:///agent/thread/session123?name=Thread+name";

crates/agent/src/thread_store.rs

@@ -893,8 +893,19 @@ impl ThreadsDatabase {
 
         let needs_migration_from_heed = mdb_path.exists();
 
-        let connection = if *ZED_STATELESS || cfg!(any(feature = "test-support", test)) {
+        let connection = if *ZED_STATELESS {
             Connection::open_memory(Some("THREAD_FALLBACK_DB"))
+        } else if cfg!(any(feature = "test-support", test)) {
+            // rust stores the name of the test on the current thread.
+            // We use this to automatically create a database that will
+            // be shared within the test (for the test_retrieve_old_thread)
+            // but not with concurrent tests.
+            let thread = std::thread::current();
+            let test_name = thread.name();
+            Connection::open_memory(Some(&format!(
+                "THREAD_FALLBACK_{}",
+                test_name.unwrap_or_default()
+            )))
         } else {
             Connection::open_file(&sqlite_path.to_string_lossy())
         };

crates/agent2/src/db.rs

@@ -266,8 +266,19 @@ impl ThreadsDatabase {
     }
 
     pub fn new(executor: BackgroundExecutor) -> Result<Self> {
-        let connection = if *ZED_STATELESS || cfg!(any(feature = "test-support", test)) {
+        let connection = if *ZED_STATELESS {
             Connection::open_memory(Some("THREAD_FALLBACK_DB"))
+        } else if cfg!(any(feature = "test-support", test)) {
+            // rust stores the name of the test on the current thread.
+            // We use this to automatically create a database that will
+            // be shared within the test (for the test_retrieve_old_thread)
+            // but not with concurrent tests.
+            let thread = std::thread::current();
+            let test_name = thread.name();
+            Connection::open_memory(Some(&format!(
+                "THREAD_FALLBACK_{}",
+                test_name.unwrap_or_default()
+            )))
         } else {
             let threads_dir = paths::data_dir().join("threads");
             std::fs::create_dir_all(&threads_dir)?;

crates/agent2/src/thread.rs

@@ -45,14 +45,15 @@ use schemars::{JsonSchema, Schema};
 use serde::{Deserialize, Serialize};
 use settings::{Settings, update_settings_file};
 use smol::stream::StreamExt;
+use std::fmt::Write;
 use std::{
     collections::BTreeMap,
+    ops::RangeInclusive,
     path::Path,
     sync::Arc,
     time::{Duration, Instant},
 };
-use std::{fmt::Write, ops::Range};
+use util::{ResultExt, debug_panic, markdown::MarkdownCodeBlock};
-use util::{ResultExt, markdown::MarkdownCodeBlock};
 use uuid::Uuid;
 
 const TOOL_CANCELED_MESSAGE: &str = "Tool canceled by user";
@@ -187,6 +188,7 @@ impl UserMessage {
         const OPEN_FILES_TAG: &str = "<files>";
         const OPEN_DIRECTORIES_TAG: &str = "<directories>";
         const OPEN_SYMBOLS_TAG: &str = "<symbols>";
+        const OPEN_SELECTIONS_TAG: &str = "<selections>";
         const OPEN_THREADS_TAG: &str = "<threads>";
         const OPEN_FETCH_TAG: &str = "<fetched_urls>";
         const OPEN_RULES_TAG: &str =
@@ -195,6 +197,7 @@ impl UserMessage {
         let mut file_context = OPEN_FILES_TAG.to_string();
         let mut directory_context = OPEN_DIRECTORIES_TAG.to_string();
         let mut symbol_context = OPEN_SYMBOLS_TAG.to_string();
+        let mut selection_context = OPEN_SELECTIONS_TAG.to_string();
         let mut thread_context = OPEN_THREADS_TAG.to_string();
         let mut fetch_context = OPEN_FETCH_TAG.to_string();
         let mut rules_context = OPEN_RULES_TAG.to_string();
@@ -211,7 +214,7 @@ impl UserMessage {
                     match uri {
                         MentionUri::File { abs_path } => {
                             write!(
-                                &mut symbol_context,
+                                &mut file_context,
                                 "\n{}",
                                 MarkdownCodeBlock {
                                     tag: &codeblock_tag(abs_path, None),
@@ -220,17 +223,19 @@ impl UserMessage {
                             )
                             .ok();
                         }
+                        MentionUri::PastedImage => {
+                            debug_panic!("pasted image URI should not be used in mention content")
+                        }
                         MentionUri::Directory { .. } => {
                             write!(&mut directory_context, "\n{}\n", content).ok();
                         }
                         MentionUri::Symbol {
-                            path, line_range, ..
+                            abs_path: path,
-                        }
+                            line_range,
-                        | MentionUri::Selection {
+                            ..
-                            path, line_range, ..
                         } => {
                             write!(
-                                &mut rules_context,
+                                &mut symbol_context,
                                 "\n{}",
                                 MarkdownCodeBlock {
                                     tag: &codeblock_tag(path, Some(line_range)),
@@ -239,6 +244,24 @@ impl UserMessage {
                             )
                             .ok();
                         }
+                        MentionUri::Selection {
+                            abs_path: path,
+                            line_range,
+                            ..
+                        } => {
+                            write!(
+                                &mut selection_context,
+                                "\n{}",
+                                MarkdownCodeBlock {
+                                    tag: &codeblock_tag(
+                                        path.as_deref().unwrap_or("Untitled".as_ref()),
+                                        Some(line_range)
+                                    ),
+                                    text: content
+                                }
+                            )
+                            .ok();
+                        }
                         MentionUri::Thread { .. } => {
                             write!(&mut thread_context, "\n{}\n", content).ok();
                         }
@@ -291,6 +314,13 @@ impl UserMessage {
                 .push(language_model::MessageContent::Text(symbol_context));
         }
 
+        if selection_context.len() > OPEN_SELECTIONS_TAG.len() {
+            selection_context.push_str("</selections>\n");
+            message
+                .content
+                .push(language_model::MessageContent::Text(selection_context));
+        }
+
         if thread_context.len() > OPEN_THREADS_TAG.len() {
             thread_context.push_str("</threads>\n");
             message
@@ -326,7 +356,7 @@ impl UserMessage {
     }
 }
 
-fn codeblock_tag(full_path: &Path, line_range: Option<&Range<u32>>) -> String {
+fn codeblock_tag(full_path: &Path, line_range: Option<&RangeInclusive<u32>>) -> String {
     let mut result = String::new();
 
     if let Some(extension) = full_path.extension().and_then(|ext| ext.to_str()) {
@@ -336,10 +366,10 @@ fn codeblock_tag(full_path: &Path, line_range: Option<&Range<u32>>) -> String {
     let _ = write!(result, "{}", full_path.display());
 
     if let Some(range) = line_range {
-        if range.start == range.end {
+        if range.start() == range.end() {
-            let _ = write!(result, ":{}", range.start + 1);
+            let _ = write!(result, ":{}", range.start() + 1);
         } else {
-            let _ = write!(result, ":{}-{}", range.start + 1, range.end + 1);
+            let _ = write!(result, ":{}-{}", range.start() + 1, range.end() + 1);
         }
     }
 

crates/agent_ui/src/acp/completion_provider.rs

@@ -247,9 +247,9 @@ impl ContextPickerCompletionProvider {
 
         let abs_path = project.read(cx).absolute_path(&symbol.path, cx)?;
         let uri = MentionUri::Symbol {
-            path: abs_path,
+            abs_path,
             name: symbol.name.clone(),
-            line_range: symbol.range.start.0.row..symbol.range.end.0.row,
+            line_range: symbol.range.start.0.row..=symbol.range.end.0.row,
         };
         let new_text = format!("{} ", uri.as_link());
         let new_text_len = new_text.len();

crates/agent_ui/src/acp/thread_view.rs

@@ -274,6 +274,7 @@ pub struct AcpThreadView {
     edits_expanded: bool,
     plan_expanded: bool,
     editor_expanded: bool,
+    terminal_expanded: bool,
     editing_message: Option<usize>,
     prompt_capabilities: Rc<Cell<PromptCapabilities>>,
     _cancel_task: Option<Task<()>>,
@@ -384,6 +385,7 @@ impl AcpThreadView {
             edits_expanded: false,
             plan_expanded: false,
             editor_expanded: false,
+            terminal_expanded: true,
             history_store,
             hovered_recent_history_item: None,
             prompt_capabilities,
@@ -835,7 +837,7 @@ impl AcpThreadView {
 
         let contents = self
             .message_editor
-            .update(cx, |message_editor, cx| message_editor.contents(window, cx));
+            .update(cx, |message_editor, cx| message_editor.contents(cx));
         self.send_impl(contents, window, cx)
     }
 
@@ -848,7 +850,7 @@ impl AcpThreadView {
 
         let contents = self
             .message_editor
-            .update(cx, |message_editor, cx| message_editor.contents(window, cx));
+            .update(cx, |message_editor, cx| message_editor.contents(cx));
 
         cx.spawn_in(window, async move |this, cx| {
             cancelled.await;
@@ -956,8 +958,7 @@ impl AcpThreadView {
             return;
         };
 
-        let contents =
+        let contents = message_editor.update(cx, |message_editor, cx| message_editor.contents(cx));
-            message_editor.update(cx, |message_editor, cx| message_editor.contents(window, cx));
 
         let task = cx.foreground_executor().spawn(async move {
             rewind.await?;
@@ -1690,9 +1691,10 @@ impl AcpThreadView {
             matches!(tool_call.kind, acp::ToolKind::Edit) || tool_call.diffs().next().is_some();
         let use_card_layout = needs_confirmation || is_edit;
 
-        let is_collapsible = !tool_call.content.is_empty() && !needs_confirmation;
+        let is_collapsible = !tool_call.content.is_empty() && !use_card_layout;
 
-        let is_open = needs_confirmation || self.expanded_tool_calls.contains(&tool_call.id);
+        let is_open =
+            needs_confirmation || is_edit || self.expanded_tool_calls.contains(&tool_call.id);
 
         let gradient_overlay = |color: Hsla| {
             div()
@@ -2162,8 +2164,6 @@ impl AcpThreadView {
             .map(|path| format!("{}", path.display()))
             .unwrap_or_else(|| "current directory".to_string());
 
-        let is_expanded = self.expanded_tool_calls.contains(&tool_call.id);
-
         let header = h_flex()
             .id(SharedString::from(format!(
                 "terminal-tool-header-{}",
@@ -2297,19 +2297,12 @@ impl AcpThreadView {
                         "terminal-tool-disclosure-{}",
                         terminal.entity_id()
                     )),
-                    is_expanded,
+                    self.terminal_expanded,
                 )
                 .opened_icon(IconName::ChevronUp)
                 .closed_icon(IconName::ChevronDown)
-                .on_click(cx.listener({
+                .on_click(cx.listener(move |this, _event, _window, _cx| {
-                    let id = tool_call.id.clone();
+                    this.terminal_expanded = !this.terminal_expanded;
-                    move |this, _event, _window, _cx| {
-                        if is_expanded {
-                            this.expanded_tool_calls.remove(&id);
-                        } else {
-                            this.expanded_tool_calls.insert(id.clone());
-                        }
-                    }
                 })),
             );
 
@@ -2318,7 +2311,7 @@ impl AcpThreadView {
             .read(cx)
             .entry(entry_ix)
             .and_then(|entry| entry.terminal(terminal));
-        let show_output = is_expanded && terminal_view.is_some();
+        let show_output = self.terminal_expanded && terminal_view.is_some();
 
         v_flex()
             .mb_2()
@@ -3655,6 +3648,7 @@ impl AcpThreadView {
                         .open_path(path, None, true, window, cx)
                         .detach_and_log_err(cx);
                 }
+                MentionUri::PastedImage => {}
                 MentionUri::Directory { abs_path } => {
                     let project = workspace.project();
                     let Some(entry) = project.update(cx, |project, cx| {
@@ -3669,9 +3663,14 @@ impl AcpThreadView {
                     });
                 }
                 MentionUri::Symbol {
-                    path, line_range, ..
+                    abs_path: path,
+                    line_range,
+                    ..
                 }
-                | MentionUri::Selection { path, line_range } => {
+                | MentionUri::Selection {
+                    abs_path: Some(path),
+                    line_range,
+                } => {
                     let project = workspace.project();
                     let Some((path, _)) = project.update(cx, |project, cx| {
                         let path = project.find_project_path(path, cx)?;
@@ -3687,8 +3686,8 @@ impl AcpThreadView {
                             let Some(editor) = item.await?.downcast::<Editor>() else {
                                 return Ok(());
                             };
-                            let range =
+                            let range = Point::new(*line_range.start(), 0)
-                                Point::new(line_range.start, 0)..Point::new(line_range.start, 0);
+                                ..Point::new(*line_range.start(), 0);
                             editor
                                 .update_in(cx, |editor, window, cx| {
                                     editor.change_selections(
@@ -3703,6 +3702,7 @@ impl AcpThreadView {
                         })
                         .detach_and_log_err(cx);
                 }
+                MentionUri::Selection { abs_path: None, .. } => {}
                 MentionUri::Thread { id, name } => {
                     if let Some(panel) = workspace.panel::<AgentPanel>(cx) {
                         panel.update(cx, |panel, cx| {