Add tracked buffers for agent2 mentions (#36608)

Cole Miller created

Release Notes:

- N/A

Change summary

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(-)

Detailed changes

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<Self>,
-    ) -> Task<Result<Vec<acp::ContentBlock>>> {
+    ) -> Task<Result<(Vec<acp::ContentBlock>, Vec<Entity<Buffer>>)>> {
         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<Path>, 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<Path>, 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<Entity<Buffer>>,
+    },
     Image(MentionImage),
 }
 
@@ -1346,7 +1356,7 @@ pub struct MentionSet {
     images: HashMap<CreaseId, Shared<Task<Result<MentionImage, String>>>>,
     thread_summaries: HashMap<acp::SessionId, Shared<Task<Result<SharedString, String>>>>,
     text_thread_summaries: HashMap<PathBuf, Shared<Task<Result<String, String>>>>,
-    directories: HashMap<PathBuf, Shared<Task<Result<String, String>>>>,
+    directories: HashMap<PathBuf, Shared<Task<Result<(String, Vec<Entity<Buffer>>), 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::<Vec<_>>();
 
-        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::<MentionUri>().unwrap());
+        }
 
         cx.simulate_input(" ");
 
@@ -2098,15 +2137,15 @@ mod tests {
             .into_values()
             .collect::<Vec<_>>();
 
-        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::<MentionUri>().unwrap());
+        }
 
         editor.update(&mut cx, |editor, cx| {
             assert_eq!(
@@ -2208,14 +2247,18 @@ mod tests {
             .into_values()
             .collect::<Vec<_>>();
 
-        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::<MentionUri>()
+                    .unwrap()
+            );
+        }
 
         cx.run_until_parked();
 

crates/agent_ui/src/acp/thread_view.rs 🔗

@@ -739,7 +739,7 @@ impl AcpThreadView {
 
     fn send_impl(
         &mut self,
-        contents: Task<anyhow::Result<Vec<acp::ContentBlock>>>,
+        contents: Task<anyhow::Result<(Vec<acp::ContentBlock>, Vec<Entity<Buffer>>)>>,
         window: &mut Window,
         cx: &mut Context<Self>,
     ) {
@@ -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
         });