From 41e28a71855c9e5595d3764423e56517e5315931 Mon Sep 17 00:00:00 2001 From: Cole Miller Date: Wed, 20 Aug 2025 14:01:18 -0400 Subject: [PATCH] Add tracked buffers for agent2 mentions (#36608) Release Notes: - N/A --- crates/agent_ui/src/acp/message_editor.rs | 151 ++++++++++++++-------- crates/agent_ui/src/acp/thread_view.rs | 13 +- 2 files changed, 107 insertions(+), 57 deletions(-) diff --git a/crates/agent_ui/src/acp/message_editor.rs b/crates/agent_ui/src/acp/message_editor.rs index a50e33dc31..ccd33c9247 100644 --- a/crates/agent_ui/src/acp/message_editor.rs +++ b/crates/agent_ui/src/acp/message_editor.rs @@ -34,7 +34,7 @@ use settings::Settings; use std::{ cell::Cell, ffi::OsStr, - fmt::{Display, Write}, + fmt::Write, ops::Range, path::{Path, PathBuf}, rc::Rc, @@ -391,30 +391,33 @@ impl MessageEditor { let rope = buffer .read_with(cx, |buffer, _cx| buffer.as_rope().clone()) .log_err()?; - Some(rope) + Some((rope, buffer)) }); cx.background_spawn(async move { - let rope = rope_task.await?; - Some((rel_path, full_path, rope.to_string())) + let (rope, buffer) = rope_task.await?; + Some((rel_path, full_path, rope.to_string(), buffer)) }) })) })?; let contents = cx .background_spawn(async move { - let contents = descendants_future.await.into_iter().flatten(); - contents.collect() + let (contents, tracked_buffers) = descendants_future + .await + .into_iter() + .flatten() + .map(|(rel_path, full_path, rope, buffer)| { + ((rel_path, full_path, rope), buffer) + }) + .unzip(); + (render_directory_contents(contents), tracked_buffers) }) .await; anyhow::Ok(contents) }); let task = cx - .spawn(async move |_, _| { - task.await - .map(|contents| DirectoryContents(contents).to_string()) - .map_err(|e| e.to_string()) - }) + .spawn(async move |_, _| task.await.map_err(|e| e.to_string())) .shared(); self.mention_set @@ -663,7 +666,7 @@ impl MessageEditor { &self, window: &mut Window, cx: &mut Context, - ) -> Task>> { + ) -> Task, Vec>)>> { let contents = self.mention_set .contents(&self.project, self.prompt_store.as_ref(), window, cx); @@ -672,6 +675,7 @@ impl MessageEditor { cx.spawn(async move |_, cx| { let contents = contents.await?; + let mut all_tracked_buffers = Vec::new(); editor.update(cx, |editor, cx| { let mut ix = 0; @@ -702,7 +706,12 @@ impl MessageEditor { chunks.push(chunk); } let chunk = match mention { - Mention::Text { uri, content } => { + Mention::Text { + uri, + content, + tracked_buffers, + } => { + all_tracked_buffers.extend(tracked_buffers.iter().cloned()); acp::ContentBlock::Resource(acp::EmbeddedResource { annotations: None, resource: acp::EmbeddedResourceResource::TextResourceContents( @@ -745,7 +754,7 @@ impl MessageEditor { } }); - chunks + (chunks, all_tracked_buffers) }) }) } @@ -1043,7 +1052,7 @@ impl MessageEditor { .add_fetch_result(url, Task::ready(Ok(text)).shared()); } MentionUri::Directory { abs_path } => { - let task = Task::ready(Ok(text)).shared(); + let task = Task::ready(Ok((text, Vec::new()))).shared(); self.mention_set.directories.insert(abs_path, task); } MentionUri::File { .. } @@ -1153,16 +1162,13 @@ impl MessageEditor { } } -struct DirectoryContents(Arc<[(Arc, PathBuf, String)]>); - -impl Display for DirectoryContents { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - for (_relative_path, full_path, content) in self.0.iter() { - let fence = codeblock_fence_for_path(Some(full_path), None); - write!(f, "\n{fence}\n{content}\n```")?; - } - Ok(()) +fn render_directory_contents(entries: Vec<(Arc, PathBuf, String)>) -> String { + let mut output = String::new(); + for (_relative_path, full_path, content) in entries { + let fence = codeblock_fence_for_path(Some(&full_path), None); + write!(output, "\n{fence}\n{content}\n```").unwrap(); } + output } impl Focusable for MessageEditor { @@ -1328,7 +1334,11 @@ impl Render for ImageHover { #[derive(Debug, Eq, PartialEq)] pub enum Mention { - Text { uri: MentionUri, content: String }, + Text { + uri: MentionUri, + content: String, + tracked_buffers: Vec>, + }, Image(MentionImage), } @@ -1346,7 +1356,7 @@ pub struct MentionSet { images: HashMap>>>, thread_summaries: HashMap>>>, text_thread_summaries: HashMap>>>, - directories: HashMap>>>, + directories: HashMap>), String>>>>, } impl MentionSet { @@ -1382,6 +1392,7 @@ impl MentionSet { 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) @@ -1424,7 +1435,14 @@ impl MentionSet { let buffer = buffer_task?.await?; let content = buffer.read_with(cx, |buffer, _cx| buffer.text())?; - anyhow::Ok((crease_id, Mention::Text { uri, content })) + anyhow::Ok(( + crease_id, + Mention::Text { + uri, + content, + tracked_buffers: vec![buffer], + }, + )) }) } MentionUri::Directory { abs_path } => { @@ -1433,11 +1451,14 @@ impl MentionSet { }; 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: content.await.map_err(|e| anyhow::anyhow!("{e}"))?, + content, + tracked_buffers, }, )) }) @@ -1473,7 +1494,14 @@ impl MentionSet { .collect() })?; - anyhow::Ok((crease_id, Mention::Text { uri, content })) + anyhow::Ok(( + crease_id, + Mention::Text { + uri, + content, + tracked_buffers: vec![buffer], + }, + )) }) } MentionUri::Thread { id, .. } => { @@ -1490,6 +1518,7 @@ impl MentionSet { .await .map_err(|e| anyhow::anyhow!("{e}"))? .to_string(), + tracked_buffers: Vec::new(), }, )) }) @@ -1505,6 +1534,7 @@ impl MentionSet { Mention::Text { uri, content: content.await.map_err(|e| anyhow::anyhow!("{e}"))?, + tracked_buffers: Vec::new(), }, )) }) @@ -1518,7 +1548,14 @@ impl MentionSet { 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 })) + anyhow::Ok(( + crease_id, + Mention::Text { + uri, + content: text, + tracked_buffers: Vec::new(), + }, + )) }) } MentionUri::Fetch { url } => { @@ -1532,6 +1569,7 @@ impl MentionSet { Mention::Text { uri, content: content.await.map_err(|e| anyhow::anyhow!("{e}"))?, + tracked_buffers: Vec::new(), }, )) }) @@ -1703,6 +1741,7 @@ impl Addon for MessageEditorAddon { mod tests { use std::{ops::Range, path::Path, sync::Arc}; + use acp_thread::MentionUri; use agent_client_protocol as acp; use agent2::HistoryStore; use assistant_context::ContextStore; @@ -1815,7 +1854,7 @@ mod tests { editor.backspace(&Default::default(), window, cx); }); - let content = message_editor + let (content, _) = message_editor .update_in(cx, |message_editor, window, cx| { message_editor.contents(window, cx) }) @@ -2046,13 +2085,13 @@ mod tests { .into_values() .collect::>(); - pretty_assertions::assert_eq!( - contents, - [Mention::Text { - content: "1".into(), - uri: url_one.parse().unwrap() - }] - ); + { + let [Mention::Text { content, uri, .. }] = contents.as_slice() else { + panic!("Unexpected mentions"); + }; + pretty_assertions::assert_eq!(content, "1"); + pretty_assertions::assert_eq!(uri, &url_one.parse::().unwrap()); + } cx.simulate_input(" "); @@ -2098,15 +2137,15 @@ mod tests { .into_values() .collect::>(); - assert_eq!(contents.len(), 2); let url_eight = uri!("file:///dir/b/eight.txt"); - pretty_assertions::assert_eq!( - contents[1], - Mention::Text { - content: "8".to_string(), - uri: url_eight.parse().unwrap(), - } - ); + + { + let [_, Mention::Text { content, uri, .. }] = contents.as_slice() else { + panic!("Unexpected mentions"); + }; + pretty_assertions::assert_eq!(content, "8"); + pretty_assertions::assert_eq!(uri, &url_eight.parse::().unwrap()); + } editor.update(&mut cx, |editor, cx| { assert_eq!( @@ -2208,14 +2247,18 @@ mod tests { .into_values() .collect::>(); - assert_eq!(contents.len(), 3); - pretty_assertions::assert_eq!( - contents[2], - Mention::Text { - content: "1".into(), - uri: format!("{url_one}?symbol=MySymbol#L1:1").parse().unwrap(), - } - ); + { + let [_, _, Mention::Text { content, uri, .. }] = contents.as_slice() else { + panic!("Unexpected mentions"); + }; + pretty_assertions::assert_eq!(content, "1"); + pretty_assertions::assert_eq!( + uri, + &format!("{url_one}?symbol=MySymbol#L1:1") + .parse::() + .unwrap() + ); + } cx.run_until_parked(); diff --git a/crates/agent_ui/src/acp/thread_view.rs b/crates/agent_ui/src/acp/thread_view.rs index 4ce55cce56..14f9cacd15 100644 --- a/crates/agent_ui/src/acp/thread_view.rs +++ b/crates/agent_ui/src/acp/thread_view.rs @@ -739,7 +739,7 @@ impl AcpThreadView { fn send_impl( &mut self, - contents: Task>>, + contents: Task, Vec>)>>, window: &mut Window, cx: &mut Context, ) { @@ -751,7 +751,7 @@ impl AcpThreadView { return; }; let task = cx.spawn_in(window, async move |this, cx| { - let contents = contents.await?; + let (contents, tracked_buffers) = contents.await?; if contents.is_empty() { return Ok(()); @@ -764,7 +764,14 @@ impl AcpThreadView { message_editor.clear(window, cx); }); })?; - let send = thread.update(cx, |thread, cx| thread.send(contents, cx))?; + let send = thread.update(cx, |thread, cx| { + thread.action_log().update(cx, |action_log, cx| { + for buffer in tracked_buffers { + action_log.buffer_read(buffer, cx) + } + }); + thread.send(contents, cx) + })?; send.await });