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

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

@ -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);
{
let [_, _, Mention::Text { content, uri, .. }] = contents.as_slice() else {
panic!("Unexpected mentions");
};
pretty_assertions::assert_eq!(content, "1");
pretty_assertions::assert_eq!(
contents[2],
Mention::Text {
content: "1".into(),
uri: format!("{url_one}?symbol=MySymbol#L1:1").parse().unwrap(),
}
uri,
&format!("{url_one}?symbol=MySymbol#L1:1")
.parse::<MentionUri>()
.unwrap()
);
}
cx.run_until_parked();

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

@ -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
});