Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add tracked buffers for agent2 mentions (#36608)

Browse source 
Release Notes:

- N/A
This commit is contained in:
Cole Miller 2025-08-20 14:01:18 -04:00 • committed by GitHub
parent 8334cdb358
commit 41e28a7185
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 107 additions and 57 deletions
Download patch file Download diff file Expand all files Collapse all files

151
crates/agent_ui/src/acp/message_editor.rs
View file

@ -34,7 +34,7 @@ use settings::Settings;
use std::{ use std::{
cell::Cell, cell::Cell,
ffi::OsStr, ffi::OsStr,
fmt::{Display, Write}, fmt::Write,
ops::Range, ops::Range,
path::{Path, PathBuf}, path::{Path, PathBuf},
rc::Rc, rc::Rc,
@ -391,30 +391,33 @@ impl MessageEditor {
let rope = buffer let rope = buffer
.read_with(cx, |buffer, _cx| buffer.as_rope().clone()) .read_with(cx, |buffer, _cx| buffer.as_rope().clone())
.log_err()?; .log_err()?;
Some(rope) Some((rope, buffer))
}); });
cx.background_spawn(async move { cx.background_spawn(async move {
let rope = rope_task.await?; let (rope, buffer) = rope_task.await?;
Some((rel_path, full_path, rope.to_string())) Some((rel_path, full_path, rope.to_string(), buffer))
}) })
})) }))
})?; })?;
let contents = cx let contents = cx
.background_spawn(async move { .background_spawn(async move {
let contents = descendants_future.await.into_iter().flatten(); let (contents, tracked_buffers) = descendants_future
contents.collect() .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; .await;
anyhow::Ok(contents) anyhow::Ok(contents)
}); });
let task = cx let task = cx
.spawn(async move |_, _| { .spawn(async move |_, _| task.await.map_err(|e| e.to_string()))
task.await
.map(|contents| DirectoryContents(contents).to_string())
.map_err(|e| e.to_string())
})
.shared(); .shared();
self.mention_set self.mention_set
@ -663,7 +666,7 @@ impl MessageEditor {
&self, &self,
window: &mut Window, window: &mut Window,
cx: &mut Context<Self>, cx: &mut Context<Self>,
) -> Task<Result<Vec<acp::ContentBlock>>> { ) -> Task<Result<(Vec<acp::ContentBlock>, Vec<Entity<Buffer>>)>> {
let contents = let contents =
self.mention_set self.mention_set
.contents(&self.project, self.prompt_store.as_ref(), window, cx); .contents(&self.project, self.prompt_store.as_ref(), window, cx);
@ -672,6 +675,7 @@ impl MessageEditor {
cx.spawn(async move |_, cx| { cx.spawn(async move |_, cx| {
let contents = contents.await?; let contents = contents.await?;
let mut all_tracked_buffers = Vec::new();
editor.update(cx, |editor, cx| { editor.update(cx, |editor, cx| {
let mut ix = 0; let mut ix = 0;
@ -702,7 +706,12 @@ impl MessageEditor {
chunks.push(chunk); chunks.push(chunk);
} }
let chunk = match mention { 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 { acp::ContentBlock::Resource(acp::EmbeddedResource {
annotations: None, annotations: None,
resource: acp::EmbeddedResourceResource::TextResourceContents( 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()); .add_fetch_result(url, Task::ready(Ok(text)).shared());
} }
MentionUri::Directory { abs_path } => { 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); self.mention_set.directories.insert(abs_path, task);
} }
MentionUri::File { .. } MentionUri::File { .. }
@ -1153,16 +1162,13 @@ impl MessageEditor {
} }
} }
struct DirectoryContents(Arc<[(Arc<Path>, PathBuf, String)]>); fn render_directory_contents(entries: Vec<(Arc<Path>, PathBuf, String)>) -> String {
let mut output = String::new();
impl Display for DirectoryContents { for (_relative_path, full_path, content) in entries {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { let fence = codeblock_fence_for_path(Some(&full_path), None);
for (_relative_path, full_path, content) in self.0.iter() { write!(output, "\n{fence}\n{content}\n```").unwrap();
let fence = codeblock_fence_for_path(Some(full_path), None);
write!(f, "\n{fence}\n{content}\n```")?;
}
Ok(())
} }
output
} }
impl Focusable for MessageEditor { impl Focusable for MessageEditor {
@ -1328,7 +1334,11 @@ impl Render for ImageHover {
#[derive(Debug, Eq, PartialEq)] #[derive(Debug, Eq, PartialEq)]
pub enum Mention { pub enum Mention {
Text { uri: MentionUri, content: String }, Text {
uri: MentionUri,
content: String,
tracked_buffers: Vec<Entity<Buffer>>,
},
Image(MentionImage), Image(MentionImage),
} }
@ -1346,7 +1356,7 @@ pub struct MentionSet {
images: HashMap<CreaseId, Shared<Task<Result<MentionImage, String>>>>, images: HashMap<CreaseId, Shared<Task<Result<MentionImage, String>>>>,
thread_summaries: HashMap<acp::SessionId, Shared<Task<Result<SharedString, String>>>>, thread_summaries: HashMap<acp::SessionId, Shared<Task<Result<SharedString, String>>>>,
text_thread_summaries: HashMap<PathBuf, Shared<Task<Result<String, 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 { impl MentionSet {
@ -1382,6 +1392,7 @@ impl MentionSet {
self.fetch_results.clear(); self.fetch_results.clear();
self.thread_summaries.clear(); self.thread_summaries.clear();
self.text_thread_summaries.clear(); self.text_thread_summaries.clear();
self.directories.clear();
self.uri_by_crease_id self.uri_by_crease_id
.drain() .drain()
.map(|(id, _)| id) .map(|(id, _)| id)
@ -1424,7 +1435,14 @@ impl MentionSet {
let buffer = buffer_task?.await?; let buffer = buffer_task?.await?;
let content = buffer.read_with(cx, |buffer, _cx| buffer.text())?; 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 } => { MentionUri::Directory { abs_path } => {
@ -1433,11 +1451,14 @@ impl MentionSet {
}; };
let uri = uri.clone(); let uri = uri.clone();
cx.spawn(async move |_| { cx.spawn(async move |_| {
let (content, tracked_buffers) =
content.await.map_err(|e| anyhow::anyhow!("{e}"))?;
Ok(( Ok((
crease_id, crease_id,
Mention::Text { Mention::Text {
uri, uri,
content: content.await.map_err(|e| anyhow::anyhow!("{e}"))?, content,
tracked_buffers,
}, },
)) ))
}) })
@ -1473,7 +1494,14 @@ impl MentionSet {
.collect() .collect()
})?; })?;
anyhow::Ok((crease_id, Mention::Text { uri, content })) anyhow::Ok((
crease_id,
Mention::Text {
uri,
content,
tracked_buffers: vec![buffer],
},
))
}) })
} }
MentionUri::Thread { id, .. } => { MentionUri::Thread { id, .. } => {
@ -1490,6 +1518,7 @@ impl MentionSet {
.await .await
.map_err(|e| anyhow::anyhow!("{e}"))? .map_err(|e| anyhow::anyhow!("{e}"))?
.to_string(), .to_string(),
tracked_buffers: Vec::new(),
}, },
)) ))
}) })
@ -1505,6 +1534,7 @@ impl MentionSet {
Mention::Text { Mention::Text {
uri, uri,
content: content.await.map_err(|e| anyhow::anyhow!("{e}"))?, content: content.await.map_err(|e| anyhow::anyhow!("{e}"))?,
tracked_buffers: Vec::new(),
}, },
)) ))
}) })
@ -1518,7 +1548,14 @@ impl MentionSet {
cx.spawn(async move |_| { cx.spawn(async move |_| {
// TODO: report load errors instead of just logging // TODO: report load errors instead of just logging
let text = text_task.await?; 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 } => { MentionUri::Fetch { url } => {
@ -1532,6 +1569,7 @@ impl MentionSet {
Mention::Text { Mention::Text {
uri, uri,
content: content.await.map_err(|e| anyhow::anyhow!("{e}"))?, content: content.await.map_err(|e| anyhow::anyhow!("{e}"))?,
tracked_buffers: Vec::new(),
}, },
)) ))
}) })
@ -1703,6 +1741,7 @@ impl Addon for MessageEditorAddon {
mod tests { mod tests {
use std::{ops::Range, path::Path, sync::Arc}; use std::{ops::Range, path::Path, sync::Arc};
use acp_thread::MentionUri;
use agent_client_protocol as acp; use agent_client_protocol as acp;
use agent2::HistoryStore; use agent2::HistoryStore;
use assistant_context::ContextStore; use assistant_context::ContextStore;
@ -1815,7 +1854,7 @@ mod tests {
editor.backspace(&Default::default(), window, cx); editor.backspace(&Default::default(), window, cx);
}); });
let content = message_editor let (content, _) = message_editor
.update_in(cx, |message_editor, window, cx| { .update_in(cx, |message_editor, window, cx| {
message_editor.contents(window, cx) message_editor.contents(window, cx)
}) })
@ -2046,13 +2085,13 @@ mod tests {
.into_values() .into_values()
.collect::<Vec<_>>(); .collect::<Vec<_>>();
pretty_assertions::assert_eq!( {
contents, let [Mention::Text { content, uri, .. }] = contents.as_slice() else {
[Mention::Text { panic!("Unexpected mentions");
content: "1".into(), };
uri: url_one.parse().unwrap() pretty_assertions::assert_eq!(content, "1");
}] pretty_assertions::assert_eq!(uri, &url_one.parse::<MentionUri>().unwrap());
); }
cx.simulate_input(" "); cx.simulate_input(" ");
@ -2098,15 +2137,15 @@ mod tests {
.into_values() .into_values()
.collect::<Vec<_>>(); .collect::<Vec<_>>();
assert_eq!(contents.len(), 2);
let url_eight = uri!("file:///dir/b/eight.txt"); let url_eight = uri!("file:///dir/b/eight.txt");
pretty_assertions::assert_eq!(
contents[1], {
Mention::Text { let [_, Mention::Text { content, uri, .. }] = contents.as_slice() else {
content: "8".to_string(), panic!("Unexpected mentions");
uri: url_eight.parse().unwrap(), };
} pretty_assertions::assert_eq!(content, "8");
); pretty_assertions::assert_eq!(uri, &url_eight.parse::<MentionUri>().unwrap());
}
editor.update(&mut cx, |editor, cx| { editor.update(&mut cx, |editor, cx| {
assert_eq!( assert_eq!(
@ -2208,14 +2247,18 @@ mod tests {
.into_values() .into_values()
.collect::<Vec<_>>(); .collect::<Vec<_>>();
assert_eq!(contents.len(), 3); {
pretty_assertions::assert_eq!( let [_, _, Mention::Text { content, uri, .. }] = contents.as_slice() else {
contents[2], panic!("Unexpected mentions");
Mention::Text { };
content: "1".into(), pretty_assertions::assert_eq!(content, "1");
uri: format!("{url_one}?symbol=MySymbol#L1:1").parse().unwrap(), pretty_assertions::assert_eq!(
} uri,
); &format!("{url_one}?symbol=MySymbol#L1:1")
.parse::<MentionUri>()
.unwrap()
);
}
cx.run_until_parked(); cx.run_until_parked();

13
crates/agent_ui/src/acp/thread_view.rs
View file

@ -739,7 +739,7 @@ impl AcpThreadView {
fn send_impl( fn send_impl(
&mut self, &mut self,
contents: Task<anyhow::Result<Vec<acp::ContentBlock>>>, contents: Task<anyhow::Result<(Vec<acp::ContentBlock>, Vec<Entity<Buffer>>)>>,
window: &mut Window, window: &mut Window,
cx: &mut Context<Self>, cx: &mut Context<Self>,
) { ) {
@ -751,7 +751,7 @@ impl AcpThreadView {
return; return;
}; };
let task = cx.spawn_in(window, async move |this, cx| { let task = cx.spawn_in(window, async move |this, cx| {
let contents = contents.await?; let (contents, tracked_buffers) = contents.await?;
if contents.is_empty() { if contents.is_empty() {
return Ok(()); return Ok(());
@ -764,7 +764,14 @@ impl AcpThreadView {
message_editor.clear(window, cx); 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 send.await
}); });