From 7bf6cc058ca778e528c1c1986b4c78728d615bfb Mon Sep 17 00:00:00 2001 From: Cole Miller Date: Sat, 23 Aug 2025 01:21:20 -0400 Subject: [PATCH] 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 --- crates/acp_thread/src/mention.rs | 154 +- crates/agent/src/thread_store.rs | 13 +- crates/agent2/src/db.rs | 13 +- crates/agent2/src/thread.rs | 54 +- .../agent_ui/src/acp/completion_provider.rs | 4 +- crates/agent_ui/src/acp/message_editor.rs | 1252 +++++++---------- crates/agent_ui/src/acp/thread_view.rs | 46 +- 7 files changed, 699 insertions(+), 837 deletions(-) diff --git a/crates/acp_thread/src/mention.rs b/crates/acp_thread/src/mention.rs index a1e713cffa..6fa0887e22 100644 --- a/crates/acp_thread/src/mention.rs +++ b/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, + line_range: RangeInclusive, }, Thread { id: acp::SessionId, @@ -38,8 +39,9 @@ pub enum MentionUri { name: String, }, Selection { - path: PathBuf, - line_range: Range, + #[serde(default, skip_serializing_if = "Option::is_none")] + abs_path: Option, + line_range: RangeInclusive, }, Fetch { url: Url, @@ -48,36 +50,44 @@ pub enum MentionUri { impl MentionUri { pub fn parse(input: &str) -> Result { + fn parse_line_range(fragment: &str) -> Result> { + 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::() + .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")?; + 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 - .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 line_range = parse_line_range(fragment)?; 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, .. - } => selection_name(path, line_range), + abs_path: path, + 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.end + 1 + line_range.start() + 1, + line_range.end() + 1 ))); url } - MentionUri::Selection { path, line_range } => { - let mut url = Url::from_file_path(path).expect("mention path should be absolute"); + MentionUri::Selection { + 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.end + 1 + line_range.start() + 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> { } } -pub fn selection_name(path: &Path, line_range: &Range) -> String { +pub fn selection_name(path: Option<&Path>, line_range: &RangeInclusive) -> String { format!( "{} ({}:{})", - path.file_name().unwrap_or_default().display(), - line_range.start + 1, - line_range.end + 1 + path.and_then(|path| path.file_name()) + .unwrap_or("Untitled".as_ref()) + .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.end, 19); + assert_eq!(line_range.start(), &9); + 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 } => { - assert_eq!(path.to_str().unwrap(), path!("/path/to/file.rs")); - assert_eq!(line_range.start, 4); - assert_eq!(line_range.end, 14); + MentionUri::Selection { + abs_path: path, + line_range, + } => { + 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"; diff --git a/crates/agent/src/thread_store.rs b/crates/agent/src/thread_store.rs index 45e551dbdf..cba2457566 100644 --- a/crates/agent/src/thread_store.rs +++ b/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()) }; diff --git a/crates/agent2/src/db.rs b/crates/agent2/src/db.rs index 1b88955a24..e7d31c0c7a 100644 --- a/crates/agent2/src/db.rs +++ b/crates/agent2/src/db.rs @@ -266,8 +266,19 @@ impl ThreadsDatabase { } pub fn new(executor: BackgroundExecutor) -> Result { - 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)?; diff --git a/crates/agent2/src/thread.rs b/crates/agent2/src/thread.rs index c89e5875f9..6d616f73fc 100644 --- a/crates/agent2/src/thread.rs +++ b/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, markdown::MarkdownCodeBlock}; +use util::{ResultExt, debug_panic, 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 = ""; const OPEN_DIRECTORIES_TAG: &str = ""; const OPEN_SYMBOLS_TAG: &str = ""; + const OPEN_SELECTIONS_TAG: &str = ""; const OPEN_THREADS_TAG: &str = ""; const OPEN_FETCH_TAG: &str = ""; 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, .. - } - | MentionUri::Selection { - path, line_range, .. + abs_path: 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("\n"); + message + .content + .push(language_model::MessageContent::Text(selection_context)); + } + if thread_context.len() > OPEN_THREADS_TAG.len() { thread_context.push_str("\n"); message @@ -326,7 +356,7 @@ impl UserMessage { } } -fn codeblock_tag(full_path: &Path, line_range: Option<&Range>) -> String { +fn codeblock_tag(full_path: &Path, line_range: Option<&RangeInclusive>) -> 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>) -> String { let _ = write!(result, "{}", full_path.display()); if let Some(range) = line_range { - if range.start == range.end { - let _ = write!(result, ":{}", range.start + 1); + if range.start() == range.end() { + 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); } } diff --git a/crates/agent_ui/src/acp/completion_provider.rs b/crates/agent_ui/src/acp/completion_provider.rs index 22a9ea6773..5b40967069 100644 --- a/crates/agent_ui/src/acp/completion_provider.rs +++ b/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(); diff --git a/crates/agent_ui/src/acp/message_editor.rs b/crates/agent_ui/src/acp/message_editor.rs index b82f1968a7..5a8abeb07a 100644 --- a/crates/agent_ui/src/acp/message_editor.rs +++ b/crates/agent_ui/src/acp/message_editor.rs @@ -6,7 +6,7 @@ use acp_thread::{MentionUri, selection_name}; use agent_client_protocol as acp; use agent_servers::AgentServer; use agent2::HistoryStore; -use anyhow::{Context as _, Result, anyhow}; +use anyhow::{Result, anyhow}; use assistant_slash_commands::codeblock_fence_for_path; use collections::{HashMap, HashSet}; use editor::{ @@ -17,8 +17,8 @@ use editor::{ display_map::{Crease, CreaseId, FoldId}, }; use futures::{ - FutureExt as _, TryFutureExt as _, - future::{Shared, join_all, try_join_all}, + FutureExt as _, + future::{Shared, join_all}, }; use gpui::{ AppContext, ClipboardEntry, Context, Entity, EventEmitter, FocusHandle, Focusable, @@ -28,14 +28,14 @@ use gpui::{ use language::{Buffer, Language}; use language_model::LanguageModelImage; use project::{CompletionIntent, Project, ProjectItem, ProjectPath, Worktree}; -use prompt_store::PromptStore; +use prompt_store::{PromptId, PromptStore}; use rope::Point; use settings::Settings; use std::{ cell::Cell, ffi::OsStr, fmt::Write, - ops::Range, + ops::{Range, RangeInclusive}, path::{Path, PathBuf}, rc::Rc, sync::Arc, @@ -49,12 +49,8 @@ use ui::{ Render, SelectableButton, SharedString, Styled, TextSize, TintColor, Toggleable, Window, div, h_flex, px, }; -use url::Url; -use util::ResultExt; -use workspace::{ - Toast, Workspace, - notifications::{NotificationId, NotifyResultExt as _}, -}; +use util::{ResultExt, debug_panic}; +use workspace::{Workspace, notifications::NotifyResultExt as _}; use zed_actions::agent::Chat; const PARSE_SLASH_COMMAND_DEBOUNCE: Duration = Duration::from_millis(50); @@ -219,9 +215,9 @@ impl MessageEditor { pub fn mentions(&self) -> HashSet { self.mention_set - .uri_by_crease_id + .mentions .values() - .cloned() + .map(|(uri, _)| uri.clone()) .collect() } @@ -246,132 +242,168 @@ impl MessageEditor { else { return Task::ready(()); }; + let end_anchor = snapshot + .buffer_snapshot + .anchor_before(start_anchor.to_offset(&snapshot.buffer_snapshot) + content_len + 1); - if let MentionUri::File { abs_path, .. } = &mention_uri { - let extension = abs_path - .extension() - .and_then(OsStr::to_str) - .unwrap_or_default(); - - if Img::extensions().contains(&extension) && !extension.contains("svg") { - if !self.prompt_capabilities.get().image { - struct ImagesNotAllowed; - - let end_anchor = snapshot.buffer_snapshot.anchor_before( - start_anchor.to_offset(&snapshot.buffer_snapshot) + content_len + 1, - ); - - self.editor.update(cx, |editor, cx| { - // Remove mention - editor.edit([((start_anchor..end_anchor), "")], cx); - }); - - self.workspace - .update(cx, |workspace, cx| { - workspace.show_toast( - Toast::new( - NotificationId::unique::(), - "This agent does not support images yet", - ) - .autohide(), - cx, - ); - }) - .ok(); - return Task::ready(()); - } - - let project = self.project.clone(); - let Some(project_path) = project - .read(cx) - .project_path_for_absolute_path(abs_path, cx) - else { - return Task::ready(()); - }; - let image = cx - .spawn(async move |_, cx| { - let image = project - .update(cx, |project, cx| project.open_image(project_path, cx)) - .map_err(|e| e.to_string())? - .await - .map_err(|e| e.to_string())?; - image - .read_with(cx, |image, _cx| image.image.clone()) - .map_err(|e| e.to_string()) - }) - .shared(); - let Some(crease_id) = insert_crease_for_image( - *excerpt_id, - start, - content_len, - Some(abs_path.as_path().into()), - image.clone(), - self.editor.clone(), - window, - cx, - ) else { - return Task::ready(()); - }; - return self.confirm_mention_for_image( - crease_id, - start_anchor, - Some(abs_path.clone()), - image, - window, - cx, - ); - } - } - - let Some(crease_id) = crate::context_picker::insert_crease_for_mention( - *excerpt_id, - start, - content_len, - crease_text, - mention_uri.icon_path(cx), - self.editor.clone(), - window, - cx, - ) else { + let crease_id = if let MentionUri::File { abs_path } = &mention_uri + && let Some(extension) = abs_path.extension() + && let Some(extension) = extension.to_str() + && Img::extensions().contains(&extension) + && !extension.contains("svg") + { + let Some(project_path) = self + .project + .read(cx) + .project_path_for_absolute_path(&abs_path, cx) + else { + log::error!("project path not found"); + return Task::ready(()); + }; + let image = self + .project + .update(cx, |project, cx| project.open_image(project_path, cx)); + let image = cx + .spawn(async move |_, cx| { + let image = image.await.map_err(|e| e.to_string())?; + let image = image + .update(cx, |image, _| image.image.clone()) + .map_err(|e| e.to_string())?; + Ok(image) + }) + .shared(); + insert_crease_for_image( + *excerpt_id, + start, + content_len, + Some(abs_path.as_path().into()), + image, + self.editor.clone(), + window, + cx, + ) + } else { + crate::context_picker::insert_crease_for_mention( + *excerpt_id, + start, + content_len, + crease_text, + mention_uri.icon_path(cx), + self.editor.clone(), + window, + cx, + ) + }; + let Some(crease_id) = crease_id else { return Task::ready(()); }; - match mention_uri { - MentionUri::Fetch { url } => { - self.confirm_mention_for_fetch(crease_id, start_anchor, url, window, cx) + let task = match mention_uri.clone() { + MentionUri::Fetch { url } => self.confirm_mention_for_fetch(url, cx), + MentionUri::Directory { abs_path } => self.confirm_mention_for_directory(abs_path, cx), + MentionUri::Thread { id, .. } => self.confirm_mention_for_thread(id, cx), + MentionUri::TextThread { path, .. } => self.confirm_mention_for_text_thread(path, cx), + MentionUri::File { abs_path } => self.confirm_mention_for_file(abs_path, cx), + MentionUri::Symbol { + abs_path, + line_range, + .. + } => self.confirm_mention_for_symbol(abs_path, line_range, cx), + MentionUri::Rule { id, .. } => self.confirm_mention_for_rule(id, cx), + MentionUri::PastedImage => { + debug_panic!("pasted image URI should not be included in completions"); + Task::ready(Err(anyhow!( + "pasted imaged URI should not be included in completions" + ))) } - MentionUri::Directory { abs_path } => { - self.confirm_mention_for_directory(crease_id, start_anchor, abs_path, window, cx) + MentionUri::Selection { .. } => { + // Handled elsewhere + debug_panic!("unexpected selection URI"); + Task::ready(Err(anyhow!("unexpected selection URI"))) } - MentionUri::Thread { id, name } => { - self.confirm_mention_for_thread(crease_id, start_anchor, id, name, window, cx) + }; + let task = cx + .spawn(async move |_, _| task.await.map_err(|e| e.to_string())) + .shared(); + self.mention_set + .mentions + .insert(crease_id, (mention_uri, task.clone())); + + // Notify the user if we failed to load the mentioned context + cx.spawn_in(window, async move |this, cx| { + if task.await.notify_async_err(cx).is_none() { + this.update(cx, |this, cx| { + this.editor.update(cx, |editor, cx| { + // Remove mention + editor.edit([(start_anchor..end_anchor, "")], cx); + }); + this.mention_set.mentions.remove(&crease_id); + }) + .ok(); } - MentionUri::TextThread { path, name } => self.confirm_mention_for_text_thread( - crease_id, - start_anchor, - path, - name, - window, - cx, - ), - MentionUri::File { .. } - | MentionUri::Symbol { .. } - | MentionUri::Rule { .. } - | MentionUri::Selection { .. } => { - self.mention_set.insert_uri(crease_id, mention_uri.clone()); - Task::ready(()) + }) + } + + fn confirm_mention_for_file( + &mut self, + abs_path: PathBuf, + cx: &mut Context, + ) -> Task> { + let Some(project_path) = self + .project + .read(cx) + .project_path_for_absolute_path(&abs_path, cx) + else { + return Task::ready(Err(anyhow!("project path not found"))); + }; + let extension = abs_path + .extension() + .and_then(OsStr::to_str) + .unwrap_or_default(); + + if Img::extensions().contains(&extension) && !extension.contains("svg") { + if !self.prompt_capabilities.get().image { + return Task::ready(Err(anyhow!("This agent does not support images yet"))); } + let task = self + .project + .update(cx, |project, cx| project.open_image(project_path, cx)); + return cx.spawn(async move |_, cx| { + let image = task.await?; + let image = image.update(cx, |image, _| image.image.clone())?; + let format = image.format; + let image = cx + .update(|cx| LanguageModelImage::from_image(image, cx))? + .await; + if let Some(image) = image { + Ok(Mention::Image(MentionImage { + data: image.source, + format, + })) + } else { + Err(anyhow!("Failed to convert image")) + } + }); } + + let buffer = self + .project + .update(cx, |project, cx| project.open_buffer(project_path, cx)); + cx.spawn(async move |_, cx| { + let buffer = buffer.await?; + let mention = buffer.update(cx, |buffer, cx| Mention::Text { + content: buffer.text(), + tracked_buffers: vec![cx.entity()], + })?; + anyhow::Ok(mention) + }) } fn confirm_mention_for_directory( &mut self, - crease_id: CreaseId, - anchor: Anchor, abs_path: PathBuf, - window: &mut Window, cx: &mut Context, - ) -> Task<()> { + ) -> Task> { fn collect_files_in_path(worktree: &Worktree, path: &Path) -> Vec<(Arc, PathBuf)> { let mut files = Vec::new(); @@ -386,24 +418,21 @@ impl MessageEditor { files } - let uri = MentionUri::Directory { - abs_path: abs_path.clone(), - }; let Some(project_path) = self .project .read(cx) .project_path_for_absolute_path(&abs_path, cx) else { - return Task::ready(()); + return Task::ready(Err(anyhow!("project path not found"))); }; let Some(entry) = self.project.read(cx).entry_for_path(&project_path, cx) else { - return Task::ready(()); + return Task::ready(Err(anyhow!("project entry not found"))); }; let Some(worktree) = self.project.read(cx).worktree_for_entry(entry.id, cx) else { - return Task::ready(()); + return Task::ready(Err(anyhow!("worktree not found"))); }; let project = self.project.clone(); - let task = cx.spawn(async move |_, cx| { + cx.spawn(async move |_, cx| { let directory_path = entry.path.clone(); let worktree_id = worktree.read_with(cx, |worktree, _| worktree.id())?; @@ -453,89 +482,83 @@ impl MessageEditor { ((rel_path, full_path, rope), buffer) }) .unzip(); - (render_directory_contents(contents), tracked_buffers) + Mention::Text { + content: render_directory_contents(contents), + tracked_buffers, + } }) .await; anyhow::Ok(contents) - }); - let task = cx - .spawn(async move |_, _| task.await.map_err(|e| e.to_string())) - .shared(); - - self.mention_set - .directories - .insert(abs_path.clone(), task.clone()); - - let editor = self.editor.clone(); - cx.spawn_in(window, async move |this, cx| { - if task.await.notify_async_err(cx).is_some() { - this.update(cx, |this, _| { - this.mention_set.insert_uri(crease_id, uri); - }) - .ok(); - } else { - editor - .update(cx, |editor, cx| { - editor.display_map.update(cx, |display_map, cx| { - display_map.unfold_intersecting(vec![anchor..anchor], true, cx); - }); - editor.remove_creases([crease_id], cx); - }) - .ok(); - this.update(cx, |this, _cx| { - this.mention_set.directories.remove(&abs_path); - }) - .ok(); - } }) } fn confirm_mention_for_fetch( &mut self, - crease_id: CreaseId, - anchor: Anchor, url: url::Url, - window: &mut Window, cx: &mut Context, - ) -> Task<()> { - let Some(http_client) = self + ) -> Task> { + let http_client = match self .workspace - .update(cx, |workspace, _cx| workspace.client().http_client()) - .ok() - else { - return Task::ready(()); + .update(cx, |workspace, _| workspace.client().http_client()) + { + Ok(http_client) => http_client, + Err(e) => return Task::ready(Err(e)), }; - - let url_string = url.to_string(); - let fetch = cx - .background_executor() - .spawn(async move { - fetch_url_content(http_client, url_string) - .map_err(|e| e.to_string()) - .await + cx.background_executor().spawn(async move { + let content = fetch_url_content(http_client, url.to_string()).await?; + Ok(Mention::Text { + content, + tracked_buffers: Vec::new(), }) - .shared(); - self.mention_set - .add_fetch_result(url.clone(), fetch.clone()); + }) + } - cx.spawn_in(window, async move |this, cx| { - let fetch = fetch.await.notify_async_err(cx); - this.update(cx, |this, cx| { - if fetch.is_some() { - this.mention_set - .insert_uri(crease_id, MentionUri::Fetch { url }); - } else { - // Remove crease if we failed to fetch - this.editor.update(cx, |editor, cx| { - editor.display_map.update(cx, |display_map, cx| { - display_map.unfold_intersecting(vec![anchor..anchor], true, cx); - }); - editor.remove_creases([crease_id], cx); - }); - this.mention_set.fetch_results.remove(&url); + fn confirm_mention_for_symbol( + &mut self, + abs_path: PathBuf, + line_range: RangeInclusive, + cx: &mut Context, + ) -> Task> { + let Some(project_path) = self + .project + .read(cx) + .project_path_for_absolute_path(&abs_path, cx) + else { + return Task::ready(Err(anyhow!("project path not found"))); + }; + let buffer = self + .project + .update(cx, |project, cx| project.open_buffer(project_path, cx)); + cx.spawn(async move |_, cx| { + let buffer = buffer.await?; + let mention = buffer.update(cx, |buffer, cx| { + let start = Point::new(*line_range.start(), 0).min(buffer.max_point()); + let end = Point::new(*line_range.end() + 1, 0).min(buffer.max_point()); + let content = buffer.text_for_range(start..end).collect(); + Mention::Text { + content, + tracked_buffers: vec![cx.entity()], } + })?; + anyhow::Ok(mention) + }) + } + + fn confirm_mention_for_rule( + &mut self, + id: PromptId, + cx: &mut Context, + ) -> Task> { + let Some(prompt_store) = self.prompt_store.clone() else { + return Task::ready(Err(anyhow!("missing prompt store"))); + }; + let prompt = prompt_store.read(cx).load(id, cx); + cx.spawn(async move |_, _| { + let prompt = prompt.await?; + Ok(Mention::Text { + content: prompt, + tracked_buffers: Vec::new(), }) - .ok(); }) } @@ -560,24 +583,24 @@ impl MessageEditor { let range = snapshot.anchor_after(offset + range_to_fold.start) ..snapshot.anchor_after(offset + range_to_fold.end); - // TODO support selections from buffers with no path - let Some(project_path) = buffer.read(cx).project_path(cx) else { - continue; - }; - let Some(abs_path) = self.project.read(cx).absolute_path(&project_path, cx) else { - continue; - }; + let abs_path = buffer + .read(cx) + .project_path(cx) + .and_then(|project_path| self.project.read(cx).absolute_path(&project_path, cx)); let snapshot = buffer.read(cx).snapshot(); + let text = snapshot + .text_for_range(selection_range.clone()) + .collect::(); let point_range = selection_range.to_point(&snapshot); - let line_range = point_range.start.row..point_range.end.row; + let line_range = point_range.start.row..=point_range.end.row; let uri = MentionUri::Selection { - path: abs_path.clone(), + abs_path: abs_path.clone(), line_range: line_range.clone(), }; let crease = crate::context_picker::crease_for_mention( - selection_name(&abs_path, &line_range).into(), + selection_name(abs_path.as_deref(), &line_range).into(), uri.icon_path(cx), range, self.editor.downgrade(), @@ -589,132 +612,69 @@ impl MessageEditor { crease_ids.first().copied().unwrap() }); - self.mention_set.insert_uri(crease_id, uri); + self.mention_set.mentions.insert( + crease_id, + ( + uri, + Task::ready(Ok(Mention::Text { + content: text, + tracked_buffers: vec![buffer], + })) + .shared(), + ), + ); } } fn confirm_mention_for_thread( &mut self, - crease_id: CreaseId, - anchor: Anchor, id: acp::SessionId, - name: String, - window: &mut Window, cx: &mut Context, - ) -> Task<()> { - let uri = MentionUri::Thread { - id: id.clone(), - name, - }; + ) -> Task> { let server = Rc::new(agent2::NativeAgentServer::new( self.project.read(cx).fs().clone(), self.history_store.clone(), )); let connection = server.connect(Path::new(""), &self.project, cx); - let load_summary = cx.spawn({ - let id = id.clone(); - async move |_, cx| { - let agent = connection.await?; - let agent = agent.downcast::().unwrap(); - let summary = agent - .0 - .update(cx, |agent, cx| agent.thread_summary(id, cx))? - .await?; - anyhow::Ok(summary) - } - }); - let task = cx - .spawn(async move |_, _| load_summary.await.map_err(|e| format!("{e}"))) - .shared(); - - self.mention_set.insert_thread(id.clone(), task.clone()); - self.mention_set.insert_uri(crease_id, uri); - - let editor = self.editor.clone(); - cx.spawn_in(window, async move |this, cx| { - if task.await.notify_async_err(cx).is_none() { - editor - .update(cx, |editor, cx| { - editor.display_map.update(cx, |display_map, cx| { - display_map.unfold_intersecting(vec![anchor..anchor], true, cx); - }); - editor.remove_creases([crease_id], cx); - }) - .ok(); - this.update(cx, |this, _| { - this.mention_set.thread_summaries.remove(&id); - this.mention_set.uri_by_crease_id.remove(&crease_id); - }) - .ok(); - } + cx.spawn(async move |_, cx| { + let agent = connection.await?; + let agent = agent.downcast::().unwrap(); + let summary = agent + .0 + .update(cx, |agent, cx| agent.thread_summary(id, cx))? + .await?; + anyhow::Ok(Mention::Text { + content: summary.to_string(), + tracked_buffers: Vec::new(), + }) }) } fn confirm_mention_for_text_thread( &mut self, - crease_id: CreaseId, - anchor: Anchor, path: PathBuf, - name: String, - window: &mut Window, cx: &mut Context, - ) -> Task<()> { - let uri = MentionUri::TextThread { - path: path.clone(), - name, - }; + ) -> Task> { let context = self.history_store.update(cx, |text_thread_store, cx| { text_thread_store.load_text_thread(path.as_path().into(), cx) }); - let task = cx - .spawn(async move |_, cx| { - let context = context.await.map_err(|e| e.to_string())?; - let xml = context - .update(cx, |context, cx| context.to_xml(cx)) - .map_err(|e| e.to_string())?; - Ok(xml) + cx.spawn(async move |_, cx| { + let context = context.await?; + let xml = context.update(cx, |context, cx| context.to_xml(cx))?; + Ok(Mention::Text { + content: xml, + tracked_buffers: Vec::new(), }) - .shared(); - - self.mention_set - .insert_text_thread(path.clone(), task.clone()); - - let editor = self.editor.clone(); - cx.spawn_in(window, async move |this, cx| { - if task.await.notify_async_err(cx).is_some() { - this.update(cx, |this, _| { - this.mention_set.insert_uri(crease_id, uri); - }) - .ok(); - } else { - editor - .update(cx, |editor, cx| { - editor.display_map.update(cx, |display_map, cx| { - display_map.unfold_intersecting(vec![anchor..anchor], true, cx); - }); - editor.remove_creases([crease_id], cx); - }) - .ok(); - this.update(cx, |this, _| { - this.mention_set.text_thread_summaries.remove(&path); - }) - .ok(); - } }) } pub fn contents( &self, - window: &mut Window, cx: &mut Context, ) -> Task, Vec>)>> { - let contents = self.mention_set.contents( - &self.project, - self.prompt_store.as_ref(), - &self.prompt_capabilities.get(), - window, - cx, - ); + let contents = self + .mention_set + .contents(&self.prompt_capabilities.get(), cx); let editor = self.editor.clone(); let prevent_slash_commands = self.prevent_slash_commands; @@ -729,7 +689,7 @@ impl MessageEditor { editor.display_map.update(cx, |map, cx| { let snapshot = map.snapshot(cx); for (crease_id, crease) in snapshot.crease_snapshot.creases() { - let Some(mention) = contents.get(&crease_id) else { + let Some((uri, mention)) = contents.get(&crease_id) else { continue; }; @@ -747,7 +707,6 @@ impl MessageEditor { } let chunk = match mention { Mention::Text { - uri, content, tracked_buffers, } => { @@ -764,17 +723,25 @@ impl MessageEditor { }) } Mention::Image(mention_image) => { + let uri = match uri { + MentionUri::File { .. } => Some(uri.to_uri().to_string()), + MentionUri::PastedImage => None, + other => { + debug_panic!( + "unexpected mention uri for image: {:?}", + other + ); + None + } + }; acp::ContentBlock::Image(acp::ImageContent { annotations: None, data: mention_image.data.to_string(), mime_type: mention_image.format.mime_type().into(), - uri: mention_image - .abs_path - .as_ref() - .map(|path| format!("file://{}", path.display())), + uri, }) } - Mention::UriOnly(uri) => { + Mention::UriOnly => { acp::ContentBlock::ResourceLink(acp::ResourceLink { name: uri.name(), uri: uri.to_uri().to_string(), @@ -813,7 +780,13 @@ impl MessageEditor { pub fn clear(&mut self, window: &mut Window, cx: &mut Context) { self.editor.update(cx, |editor, cx| { editor.clear(window, cx); - editor.remove_creases(self.mention_set.drain(), cx) + editor.remove_creases( + self.mention_set + .mentions + .drain() + .map(|(crease_id, _)| crease_id), + cx, + ) }); } @@ -853,7 +826,7 @@ impl MessageEditor { } cx.stop_propagation(); - let replacement_text = "image"; + let replacement_text = MentionUri::PastedImage.as_link().to_string(); for image in images { let (excerpt_id, text_anchor, multibuffer_anchor) = self.editor.update(cx, |message_editor, cx| { @@ -876,24 +849,62 @@ impl MessageEditor { }); let content_len = replacement_text.len(); - let Some(anchor) = multibuffer_anchor else { - return; + let Some(start_anchor) = multibuffer_anchor else { + continue; }; - let task = Task::ready(Ok(Arc::new(image))).shared(); + let end_anchor = self.editor.update(cx, |editor, cx| { + let snapshot = editor.buffer().read(cx).snapshot(cx); + snapshot.anchor_before(start_anchor.to_offset(&snapshot) + content_len) + }); + let image = Arc::new(image); let Some(crease_id) = insert_crease_for_image( excerpt_id, text_anchor, content_len, None.clone(), - task.clone(), + Task::ready(Ok(image.clone())).shared(), self.editor.clone(), window, cx, ) else { - return; + continue; }; - self.confirm_mention_for_image(crease_id, anchor, None, task, window, cx) - .detach(); + let task = cx + .spawn_in(window, { + async move |_, cx| { + let format = image.format; + let image = cx + .update(|_, cx| LanguageModelImage::from_image(image, cx)) + .map_err(|e| e.to_string())? + .await; + if let Some(image) = image { + Ok(Mention::Image(MentionImage { + data: image.source, + format, + })) + } else { + Err("Failed to convert image".into()) + } + } + }) + .shared(); + + self.mention_set + .mentions + .insert(crease_id, (MentionUri::PastedImage, task.clone())); + + cx.spawn_in(window, async move |this, cx| { + if task.await.notify_async_err(cx).is_none() { + this.update(cx, |this, cx| { + this.editor.update(cx, |editor, cx| { + editor.edit([(start_anchor..end_anchor, "")], cx); + }); + this.mention_set.mentions.remove(&crease_id); + }) + .ok(); + } + }) + .detach(); } } @@ -995,67 +1006,6 @@ impl MessageEditor { }) } - fn confirm_mention_for_image( - &mut self, - crease_id: CreaseId, - anchor: Anchor, - abs_path: Option, - image: Shared, String>>>, - window: &mut Window, - cx: &mut Context, - ) -> Task<()> { - let editor = self.editor.clone(); - let task = cx - .spawn_in(window, { - let abs_path = abs_path.clone(); - async move |_, cx| { - let image = image.await?; - let format = image.format; - let image = cx - .update(|_, cx| LanguageModelImage::from_image(image, cx)) - .map_err(|e| e.to_string())? - .await; - if let Some(image) = image { - Ok(MentionImage { - abs_path, - data: image.source, - format, - }) - } else { - Err("Failed to convert image".into()) - } - } - }) - .shared(); - - self.mention_set.insert_image(crease_id, task.clone()); - - cx.spawn_in(window, async move |this, cx| { - if task.await.notify_async_err(cx).is_some() { - if let Some(abs_path) = abs_path.clone() { - this.update(cx, |this, _cx| { - this.mention_set - .insert_uri(crease_id, MentionUri::File { abs_path }); - }) - .ok(); - } - } else { - editor - .update(cx, |editor, cx| { - editor.display_map.update(cx, |display_map, cx| { - display_map.unfold_intersecting(vec![anchor..anchor], true, cx); - }); - editor.remove_creases([crease_id], cx); - }) - .ok(); - this.update(cx, |this, _cx| { - this.mention_set.images.remove(&crease_id); - }) - .ok(); - } - }) - } - pub fn set_mode(&mut self, mode: EditorMode, cx: &mut Context) { self.editor.update(cx, |editor, cx| { editor.set_mode(mode); @@ -1073,7 +1023,6 @@ impl MessageEditor { let mut text = String::new(); let mut mentions = Vec::new(); - let mut images = Vec::new(); for chunk in message { match chunk { @@ -1084,26 +1033,58 @@ impl MessageEditor { resource: acp::EmbeddedResourceResource::TextResourceContents(resource), .. }) => { - if let Some(mention_uri) = MentionUri::parse(&resource.uri).log_err() { - let start = text.len(); - write!(&mut text, "{}", mention_uri.as_link()).ok(); - let end = text.len(); - mentions.push((start..end, mention_uri, resource.text)); - } + let Some(mention_uri) = MentionUri::parse(&resource.uri).log_err() else { + continue; + }; + let start = text.len(); + write!(&mut text, "{}", mention_uri.as_link()).ok(); + let end = text.len(); + mentions.push(( + start..end, + mention_uri, + Mention::Text { + content: resource.text, + tracked_buffers: Vec::new(), + }, + )); } acp::ContentBlock::ResourceLink(resource) => { if let Some(mention_uri) = MentionUri::parse(&resource.uri).log_err() { let start = text.len(); write!(&mut text, "{}", mention_uri.as_link()).ok(); let end = text.len(); - mentions.push((start..end, mention_uri, resource.uri)); + mentions.push((start..end, mention_uri, Mention::UriOnly)); } } - acp::ContentBlock::Image(content) => { + acp::ContentBlock::Image(acp::ImageContent { + uri, + data, + mime_type, + annotations: _, + }) => { + let mention_uri = if let Some(uri) = uri { + MentionUri::parse(&uri) + } else { + Ok(MentionUri::PastedImage) + }; + let Some(mention_uri) = mention_uri.log_err() else { + continue; + }; + let Some(format) = ImageFormat::from_mime_type(&mime_type) else { + log::error!("failed to parse MIME type for image: {mime_type:?}"); + continue; + }; let start = text.len(); - text.push_str("image"); + write!(&mut text, "{}", mention_uri.as_link()).ok(); let end = text.len(); - images.push((start..end, content)); + mentions.push(( + start..end, + mention_uri, + Mention::Image(MentionImage { + data: data.into(), + format, + }), + )); } acp::ContentBlock::Audio(_) | acp::ContentBlock::Resource(_) => {} } @@ -1114,9 +1095,9 @@ impl MessageEditor { editor.buffer().read(cx).snapshot(cx) }); - for (range, mention_uri, text) in mentions { + for (range, mention_uri, mention) in mentions { let anchor = snapshot.anchor_before(range.start); - let crease_id = crate::context_picker::insert_crease_for_mention( + let Some(crease_id) = crate::context_picker::insert_crease_for_mention( anchor.excerpt_id, anchor.text_anchor, range.end - range.start, @@ -1125,77 +1106,14 @@ impl MessageEditor { self.editor.clone(), window, cx, - ); - - if let Some(crease_id) = crease_id { - self.mention_set.insert_uri(crease_id, mention_uri.clone()); - } - - match mention_uri { - MentionUri::Thread { id, .. } => { - self.mention_set - .insert_thread(id, Task::ready(Ok(text.into())).shared()); - } - MentionUri::TextThread { path, .. } => { - self.mention_set - .insert_text_thread(path, Task::ready(Ok(text)).shared()); - } - MentionUri::Fetch { url } => { - self.mention_set - .add_fetch_result(url, Task::ready(Ok(text)).shared()); - } - MentionUri::Directory { abs_path } => { - let task = Task::ready(Ok((text, Vec::new()))).shared(); - self.mention_set.directories.insert(abs_path, task); - } - MentionUri::File { .. } - | MentionUri::Symbol { .. } - | MentionUri::Rule { .. } - | MentionUri::Selection { .. } => {} - } - } - for (range, content) in images { - let Some(format) = ImageFormat::from_mime_type(&content.mime_type) else { + ) else { continue; }; - let anchor = snapshot.anchor_before(range.start); - let abs_path = content - .uri - .as_ref() - .and_then(|uri| uri.strip_prefix("file://").map(|s| Path::new(s).into())); - let name = content - .uri - .as_ref() - .and_then(|uri| { - uri.strip_prefix("file://") - .and_then(|path| Path::new(path).file_name()) - }) - .map(|name| name.to_string_lossy().to_string()) - .unwrap_or("Image".to_owned()); - let crease_id = crate::context_picker::insert_crease_for_mention( - anchor.excerpt_id, - anchor.text_anchor, - range.end - range.start, - name.into(), - IconName::Image.path().into(), - self.editor.clone(), - window, - cx, + self.mention_set.mentions.insert( + crease_id, + (mention_uri.clone(), Task::ready(Ok(mention)).shared()), ); - let data: SharedString = content.data.to_string().into(); - - if let Some(crease_id) = crease_id { - self.mention_set.insert_image( - crease_id, - Task::ready(Ok(MentionImage { - abs_path, - data, - format, - })) - .shared(), - ); - } } cx.notify(); } @@ -1425,289 +1343,60 @@ impl Render for ImageHover { } } -#[derive(Debug, Eq, PartialEq)] +#[derive(Debug, Clone, Eq, PartialEq)] pub enum Mention { Text { - uri: MentionUri, content: String, tracked_buffers: Vec>, }, Image(MentionImage), - UriOnly(MentionUri), + UriOnly, } #[derive(Clone, Debug, Eq, PartialEq)] pub struct MentionImage { - pub abs_path: Option, pub data: SharedString, pub format: ImageFormat, } #[derive(Default)] pub struct MentionSet { - uri_by_crease_id: HashMap, - fetch_results: HashMap>>>, - images: HashMap>>>, - thread_summaries: HashMap>>>, - text_thread_summaries: HashMap>>>, - directories: HashMap>), String>>>>, + mentions: HashMap>>)>, } impl MentionSet { - pub fn insert_uri(&mut self, crease_id: CreaseId, uri: MentionUri) { - self.uri_by_crease_id.insert(crease_id, uri); - } - - pub fn add_fetch_result(&mut self, url: Url, content: Shared>>) { - self.fetch_results.insert(url, content); - } - - pub fn insert_image( - &mut self, - crease_id: CreaseId, - task: Shared>>, - ) { - self.images.insert(crease_id, task); - } - - fn insert_thread( - &mut self, - id: acp::SessionId, - task: Shared>>, - ) { - self.thread_summaries.insert(id, task); - } - - fn insert_text_thread(&mut self, path: PathBuf, task: Shared>>) { - self.text_thread_summaries.insert(path, task); - } - - pub fn contents( + fn contents( &self, - project: &Entity, - prompt_store: Option<&Entity>, prompt_capabilities: &acp::PromptCapabilities, - _window: &mut Window, cx: &mut App, - ) -> Task>> { + ) -> Task>> { if !prompt_capabilities.embedded_context { let mentions = self - .uri_by_crease_id + .mentions .iter() - .map(|(crease_id, uri)| (*crease_id, Mention::UriOnly(uri.clone()))) + .map(|(crease_id, (uri, _))| (*crease_id, (uri.clone(), Mention::UriOnly))) .collect(); return Task::ready(Ok(mentions)); } - let mut processed_image_creases = HashSet::default(); - - let mut contents = self - .uri_by_crease_id - .iter() - .map(|(&crease_id, uri)| { - match uri { - MentionUri::File { abs_path, .. } => { - let uri = uri.clone(); - let abs_path = abs_path.to_path_buf(); - - if let Some(task) = self.images.get(&crease_id).cloned() { - processed_image_creases.insert(crease_id); - return cx.spawn(async move |_| { - let image = task.await.map_err(|e| anyhow!("{e}"))?; - anyhow::Ok((crease_id, Mention::Image(image))) - }); - } - - let buffer_task = project.update(cx, |project, cx| { - let path = project - .find_project_path(abs_path, 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())?; - - anyhow::Ok(( - crease_id, - Mention::Text { - uri, - content, - tracked_buffers: vec![buffer], - }, - )) - }) - } - MentionUri::Directory { abs_path } => { - let Some(content) = self.directories.get(abs_path).cloned() else { - return Task::ready(Err(anyhow!("missing directory load task"))); - }; - let uri = uri.clone(); - cx.spawn(async move |_| { - let (content, tracked_buffers) = - content.await.map_err(|e| anyhow::anyhow!("{e}"))?; - Ok(( - crease_id, - Mention::Text { - uri, - content, - tracked_buffers, - }, - )) - }) - } - MentionUri::Symbol { - path, line_range, .. - } - | MentionUri::Selection { - path, line_range, .. - } => { - 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::Text { - uri, - content, - tracked_buffers: vec![buffer], - }, - )) - }) - } - MentionUri::Thread { id, .. } => { - let Some(content) = self.thread_summaries.get(id).cloned() else { - return Task::ready(Err(anyhow!("missing thread summary"))); - }; - let uri = uri.clone(); - cx.spawn(async move |_| { - Ok(( - crease_id, - Mention::Text { - uri, - content: content - .await - .map_err(|e| anyhow::anyhow!("{e}"))? - .to_string(), - tracked_buffers: Vec::new(), - }, - )) - }) - } - MentionUri::TextThread { path, .. } => { - let Some(content) = self.text_thread_summaries.get(path).cloned() else { - return Task::ready(Err(anyhow!("missing text thread summary"))); - }; - let uri = uri.clone(); - cx.spawn(async move |_| { - Ok(( - crease_id, - Mention::Text { - uri, - content: content.await.map_err(|e| anyhow::anyhow!("{e}"))?, - tracked_buffers: Vec::new(), - }, - )) - }) - } - MentionUri::Rule { id: prompt_id, .. } => { - let Some(prompt_store) = prompt_store else { - return Task::ready(Err(anyhow!("missing prompt store"))); - }; - let text_task = prompt_store.read(cx).load(*prompt_id, cx); - let uri = uri.clone(); - cx.spawn(async move |_| { - // TODO: report load errors instead of just logging - let text = text_task.await?; - anyhow::Ok(( - crease_id, - Mention::Text { - uri, - content: text, - tracked_buffers: Vec::new(), - }, - )) - }) - } - MentionUri::Fetch { url } => { - let Some(content) = self.fetch_results.get(url).cloned() else { - return Task::ready(Err(anyhow!("missing fetch result"))); - }; - let uri = uri.clone(); - cx.spawn(async move |_| { - Ok(( - crease_id, - Mention::Text { - uri, - content: content.await.map_err(|e| anyhow::anyhow!("{e}"))?, - tracked_buffers: Vec::new(), - }, - )) - }) - } - } - }) - .collect::>(); - - // Handle images that didn't have a mention URI (because they were added by the paste handler). - contents.extend(self.images.iter().filter_map(|(crease_id, image)| { - if processed_image_creases.contains(crease_id) { - return None; - } - let crease_id = *crease_id; - let image = image.clone(); - Some(cx.spawn(async move |_| { - Ok(( - crease_id, - Mention::Image(image.await.map_err(|e| anyhow::anyhow!("{e}"))?), - )) - })) - })); - + let mentions = self.mentions.clone(); cx.spawn(async move |_cx| { - let contents = try_join_all(contents).await?.into_iter().collect(); - anyhow::Ok(contents) + let mut contents = HashMap::default(); + for (crease_id, (mention_uri, task)) in mentions { + contents.insert( + crease_id, + (mention_uri, task.await.map_err(|e| anyhow!("{e}"))?), + ); + } + Ok(contents) }) } - pub fn drain(&mut self) -> impl Iterator { - self.fetch_results.clear(); - self.thread_summaries.clear(); - self.text_thread_summaries.clear(); - self.directories.clear(); - self.uri_by_crease_id - .drain() - .map(|(id, _)| id) - .chain(self.images.drain().map(|(id, _)| id)) - } - - pub fn remove_invalid(&mut self, snapshot: EditorSnapshot) { + fn remove_invalid(&mut self, snapshot: EditorSnapshot) { for (crease_id, crease) in snapshot.crease_snapshot.creases() { if !crease.range().start.is_valid(&snapshot.buffer_snapshot) { - self.uri_by_crease_id.remove(&crease_id); + self.mentions.remove(&crease_id); } } } @@ -1969,9 +1658,7 @@ mod tests { }); let (content, _) = message_editor - .update_in(cx, |message_editor, window, cx| { - message_editor.contents(window, cx) - }) + .update(cx, |message_editor, cx| message_editor.contents(cx)) .await .unwrap(); @@ -2038,7 +1725,8 @@ mod tests { "six.txt": "6", "seven.txt": "7", "eight.txt": "8", - } + }, + "x.png": "", }), ) .await; @@ -2222,14 +1910,10 @@ mod tests { }; let contents = message_editor - .update_in(&mut cx, |message_editor, window, cx| { - message_editor.mention_set().contents( - &project, - None, - &all_prompt_capabilities, - window, - cx, - ) + .update(&mut cx, |message_editor, cx| { + message_editor + .mention_set() + .contents(&all_prompt_capabilities, cx) }) .await .unwrap() @@ -2237,7 +1921,7 @@ mod tests { .collect::>(); { - let [Mention::Text { content, uri, .. }] = contents.as_slice() else { + let [(uri, Mention::Text { content, .. })] = contents.as_slice() else { panic!("Unexpected mentions"); }; pretty_assertions::assert_eq!(content, "1"); @@ -2245,14 +1929,10 @@ mod tests { } let contents = message_editor - .update_in(&mut cx, |message_editor, window, cx| { - message_editor.mention_set().contents( - &project, - None, - &acp::PromptCapabilities::default(), - window, - cx, - ) + .update(&mut cx, |message_editor, cx| { + message_editor + .mention_set() + .contents(&acp::PromptCapabilities::default(), cx) }) .await .unwrap() @@ -2260,7 +1940,7 @@ mod tests { .collect::>(); { - let [Mention::UriOnly(uri)] = contents.as_slice() else { + let [(uri, Mention::UriOnly)] = contents.as_slice() else { panic!("Unexpected mentions"); }; pretty_assertions::assert_eq!(uri, &url_one.parse::().unwrap()); @@ -2300,14 +1980,10 @@ mod tests { cx.run_until_parked(); let contents = message_editor - .update_in(&mut cx, |message_editor, window, cx| { - message_editor.mention_set().contents( - &project, - None, - &all_prompt_capabilities, - window, - cx, - ) + .update(&mut cx, |message_editor, cx| { + message_editor + .mention_set() + .contents(&all_prompt_capabilities, cx) }) .await .unwrap() @@ -2317,7 +1993,7 @@ mod tests { let url_eight = uri!("file:///dir/b/eight.txt"); { - let [_, Mention::Text { content, uri, .. }] = contents.as_slice() else { + let [_, (uri, Mention::Text { content, .. })] = contents.as_slice() else { panic!("Unexpected mentions"); }; pretty_assertions::assert_eq!(content, "8"); @@ -2414,14 +2090,10 @@ mod tests { }); let contents = message_editor - .update_in(&mut cx, |message_editor, window, cx| { - message_editor.mention_set().contents( - &project, - None, - &all_prompt_capabilities, - window, - cx, - ) + .update(&mut cx, |message_editor, cx| { + message_editor + .mention_set() + .contents(&all_prompt_capabilities, cx) }) .await .unwrap() @@ -2429,7 +2101,7 @@ mod tests { .collect::>(); { - let [_, _, Mention::Text { content, uri, .. }] = contents.as_slice() else { + let [_, _, (uri, Mention::Text { content, .. })] = contents.as_slice() else { panic!("Unexpected mentions"); }; pretty_assertions::assert_eq!(content, "1"); @@ -2444,11 +2116,85 @@ mod tests { cx.run_until_parked(); editor.read_with(&cx, |editor, cx| { - assert_eq!( - editor.text(cx), - format!("Lorem [@one.txt]({url_one}) Ipsum [@eight.txt]({url_eight}) [@MySymbol]({url_one}?symbol=MySymbol#L1:1) ") - ); - }); + assert_eq!( + editor.text(cx), + format!("Lorem [@one.txt]({url_one}) Ipsum [@eight.txt]({url_eight}) [@MySymbol]({url_one}?symbol=MySymbol#L1:1) ") + ); + }); + + // Try to mention an "image" file that will fail to load + cx.simulate_input("@file x.png"); + + editor.update(&mut cx, |editor, cx| { + assert_eq!( + editor.text(cx), + format!("Lorem [@one.txt]({url_one}) Ipsum [@eight.txt]({url_eight}) [@MySymbol]({url_one}?symbol=MySymbol#L1:1) @file x.png") + ); + assert!(editor.has_visible_completions_menu()); + assert_eq!(current_completion_labels(editor), &["x.png dir/"]); + }); + + editor.update_in(&mut cx, |editor, window, cx| { + editor.confirm_completion(&editor::actions::ConfirmCompletion::default(), window, cx); + }); + + // Getting the message contents fails + message_editor + .update(&mut cx, |message_editor, cx| { + message_editor + .mention_set() + .contents(&all_prompt_capabilities, cx) + }) + .await + .expect_err("Should fail to load x.png"); + + cx.run_until_parked(); + + // Mention was removed + editor.read_with(&cx, |editor, cx| { + assert_eq!( + editor.text(cx), + format!("Lorem [@one.txt]({url_one}) Ipsum [@eight.txt]({url_eight}) [@MySymbol]({url_one}?symbol=MySymbol#L1:1) ") + ); + }); + + // Once more + cx.simulate_input("@file x.png"); + + editor.update(&mut cx, |editor, cx| { + assert_eq!( + editor.text(cx), + format!("Lorem [@one.txt]({url_one}) Ipsum [@eight.txt]({url_eight}) [@MySymbol]({url_one}?symbol=MySymbol#L1:1) @file x.png") + ); + assert!(editor.has_visible_completions_menu()); + assert_eq!(current_completion_labels(editor), &["x.png dir/"]); + }); + + editor.update_in(&mut cx, |editor, window, cx| { + editor.confirm_completion(&editor::actions::ConfirmCompletion::default(), window, cx); + }); + + // This time don't immediately get the contents, just let the confirmed completion settle + cx.run_until_parked(); + + // Mention was removed + editor.read_with(&cx, |editor, cx| { + assert_eq!( + editor.text(cx), + format!("Lorem [@one.txt]({url_one}) Ipsum [@eight.txt]({url_eight}) [@MySymbol]({url_one}?symbol=MySymbol#L1:1) ") + ); + }); + + // Now getting the contents succeeds, because the invalid mention was removed + let contents = message_editor + .update(&mut cx, |message_editor, cx| { + message_editor + .mention_set() + .contents(&all_prompt_capabilities, cx) + }) + .await + .unwrap(); + assert_eq!(contents.len(), 3); } fn fold_ranges(editor: &Editor, cx: &mut App) -> Vec> { diff --git a/crates/agent_ui/src/acp/thread_view.rs b/crates/agent_ui/src/acp/thread_view.rs index 0e1d4123b9..3ad1234e22 100644 --- a/crates/agent_ui/src/acp/thread_view.rs +++ b/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, prompt_capabilities: Rc>, _cancel_task: Option>, @@ -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 = - message_editor.update(cx, |message_editor, cx| message_editor.contents(window, cx)); + let contents = message_editor.update(cx, |message_editor, cx| message_editor.contents(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({ - let id = tool_call.id.clone(); - move |this, _event, _window, _cx| { - if is_expanded { - this.expanded_tool_calls.remove(&id); - } else { - this.expanded_tool_calls.insert(id.clone()); - } - } + .on_click(cx.listener(move |this, _event, _window, _cx| { + this.terminal_expanded = !this.terminal_expanded; })), ); @@ -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::() else { return Ok(()); }; - let range = - Point::new(line_range.start, 0)..Point::new(line_range.start, 0); + let range = 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::(cx) { panel.update(cx, |panel, cx| {