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

Finalize diff on task drop

Browse source
This commit is contained in:
Antonio Scandurra 2025-08-26 11:15:08 +02:00
parent afd2c3aa76
commit 0fb6c1ee09
3 changed files with 259 additions and 142 deletions
Download patch file Download diff file Expand all files Collapse all files

24
crates/agent2/src/thread.rs
View file

@ -2459,6 +2459,30 @@ impl ToolCallEventStreamReceiver {
} }
} }
pub async fn expect_update_fields(&mut self) -> acp::ToolCallUpdateFields {
let event = self.0.next().await;
if let Some(Ok(ThreadEvent::ToolCallUpdate(acp_thread::ToolCallUpdate::UpdateFields(
update,
)))) = event
{
update.fields
} else {
panic!("Expected update fields but got: {:?}", event);
}
}
pub async fn expect_diff(&mut self) -> Entity<acp_thread::Diff> {
let event = self.0.next().await;
if let Some(Ok(ThreadEvent::ToolCallUpdate(acp_thread::ToolCallUpdate::UpdateDiff(
update,
)))) = event
{
update.diff
} else {
panic!("Expected diff but got: {:?}", event);
}
}
pub async fn expect_terminal(&mut self) -> Entity<acp_thread::Terminal> { pub async fn expect_terminal(&mut self) -> Entity<acp_thread::Terminal> {
let event = self.0.next().await; let event = self.0.next().await;
if let Some(Ok(ThreadEvent::ToolCallUpdate(acp_thread::ToolCallUpdate::UpdateTerminal( if let Some(Ok(ThreadEvent::ToolCallUpdate(acp_thread::ToolCallUpdate::UpdateTerminal(

346
crates/agent2/src/tools/edit_file_tool.rs
View file

@ -272,164 +272,164 @@ impl AgentTool for EditFileTool {
.await?; .await?;
let diff = cx.new(|cx| Diff::new(buffer.clone(), cx))?; let diff = cx.new(|cx| Diff::new(buffer.clone(), cx))?;
let result = async { event_stream.update_diff(diff.clone());
event_stream.update_diff(diff.clone()); let _finalize_diff = util::defer({
let diff = diff.downgrade();
let mut cx = cx.clone();
move || {
diff.update(&mut cx, |diff, cx| diff.finalize(cx)).ok();
}
});
let old_snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot())?; let old_snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot())?;
let old_text = cx let old_text = cx
.background_spawn({ .background_spawn({
let old_snapshot = old_snapshot.clone(); let old_snapshot = old_snapshot.clone();
async move { Arc::new(old_snapshot.text()) } async move { Arc::new(old_snapshot.text()) }
}) })
.await; .await;
let (output, mut events) = if matches!(input.mode, EditFileMode::Edit) { let (output, mut events) = if matches!(input.mode, EditFileMode::Edit) {
edit_agent.edit( edit_agent.edit(
buffer.clone(), buffer.clone(),
input.display_description.clone(), input.display_description.clone(),
&request, &request,
cx, cx,
) )
} else { } else {
edit_agent.overwrite( edit_agent.overwrite(
buffer.clone(), buffer.clone(),
input.display_description.clone(), input.display_description.clone(),
&request, &request,
cx, cx,
) )
}; };
let mut hallucinated_old_text = false; let mut hallucinated_old_text = false;
let mut ambiguous_ranges = Vec::new(); let mut ambiguous_ranges = Vec::new();
let mut emitted_location = false; let mut emitted_location = false;
while let Some(event) = events.next().await { while let Some(event) = events.next().await {
match event { match event {
EditAgentOutputEvent::Edited(range) => { EditAgentOutputEvent::Edited(range) => {
if !emitted_location { if !emitted_location {
let line = buffer.update(cx, |buffer, _cx| { let line = buffer.update(cx, |buffer, _cx| {
range.start.to_point(&buffer.snapshot()).row range.start.to_point(&buffer.snapshot()).row
}).ok(); }).ok();
if let Some(abs_path) = abs_path.clone() { if let Some(abs_path) = abs_path.clone() {
event_stream.update_fields(ToolCallUpdateFields { event_stream.update_fields(ToolCallUpdateFields {
locations: Some(vec![ToolCallLocation { path: abs_path, line }]), locations: Some(vec![ToolCallLocation { path: abs_path, line }]),
..Default::default() ..Default::default()
}); });
}
emitted_location = true;
} }
}, emitted_location = true;
EditAgentOutputEvent::UnresolvedEditRange => hallucinated_old_text = true,
EditAgentOutputEvent::AmbiguousEditRange(ranges) => ambiguous_ranges = ranges,
EditAgentOutputEvent::ResolvingEditRange(range) => {
diff.update(cx, |card, cx| card.reveal_range(range.clone(), cx))?;
// if !emitted_location {
// let line = buffer.update(cx, |buffer, _cx| {
// range.start.to_point(&buffer.snapshot()).row
// }).ok();
// if let Some(abs_path) = abs_path.clone() {
// event_stream.update_fields(ToolCallUpdateFields {
// locations: Some(vec![ToolCallLocation { path: abs_path, line }]),
// ..Default::default()
// });
// }
// }
} }
},
EditAgentOutputEvent::UnresolvedEditRange => hallucinated_old_text = true,
EditAgentOutputEvent::AmbiguousEditRange(ranges) => ambiguous_ranges = ranges,
EditAgentOutputEvent::ResolvingEditRange(range) => {
diff.update(cx, |card, cx| card.reveal_range(range.clone(), cx))?;
// if !emitted_location {
// let line = buffer.update(cx, |buffer, _cx| {
// range.start.to_point(&buffer.snapshot()).row
// }).ok();
// if let Some(abs_path) = abs_path.clone() {
// event_stream.update_fields(ToolCallUpdateFields {
// locations: Some(vec![ToolCallLocation { path: abs_path, line }]),
// ..Default::default()
// });
// }
// }
} }
} }
}
// If format_on_save is enabled, format the buffer // If format_on_save is enabled, format the buffer
let format_on_save_enabled = buffer let format_on_save_enabled = buffer
.read_with(cx, |buffer, cx| { .read_with(cx, |buffer, cx| {
let settings = language_settings::language_settings( let settings = language_settings::language_settings(
buffer.language().map(|l| l.name()), buffer.language().map(|l| l.name()),
buffer.file(), buffer.file(),
cx, cx,
); );
settings.format_on_save != FormatOnSave::Off settings.format_on_save != FormatOnSave::Off
}) })
.unwrap_or(false); .unwrap_or(false);
let edit_agent_output = output.await?; let edit_agent_output = output.await?;
if format_on_save_enabled {
action_log.update(cx, |log, cx| {
log.buffer_edited(buffer.clone(), cx);
})?;
let format_task = project.update(cx, |project, cx| {
project.format(
HashSet::from_iter([buffer.clone()]),
LspFormatTarget::Buffers,
false, // Don't push to history since the tool did it.
FormatTrigger::Save,
cx,
)
})?;
format_task.await.log_err();
}
project
.update(cx, |project, cx| project.save_buffer(buffer.clone(), cx))?
.await?;
if format_on_save_enabled {
action_log.update(cx, |log, cx| { action_log.update(cx, |log, cx| {
log.buffer_edited(buffer.clone(), cx); log.buffer_edited(buffer.clone(), cx);
})?; })?;
let new_snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot())?; let format_task = project.update(cx, |project, cx| {
let (new_text, unified_diff) = cx project.format(
.background_spawn({ HashSet::from_iter([buffer.clone()]),
let new_snapshot = new_snapshot.clone(); LspFormatTarget::Buffers,
let old_text = old_text.clone(); false, // Don't push to history since the tool did it.
async move { FormatTrigger::Save,
let new_text = new_snapshot.text(); cx,
let diff = language::unified_diff(&old_text, &new_text); )
(new_text, diff) })?;
} format_task.await.log_err();
}) }
.await;
let input_path = input.path.display(); project
if unified_diff.is_empty() { .update(cx, |project, cx| project.save_buffer(buffer.clone(), cx))?
anyhow::ensure!( .await?;
!hallucinated_old_text,
formatdoc! {"
Some edits were produced but none of them could be applied.
Read the relevant sections of {input_path} again so that
I can perform the requested edits.
"}
);
anyhow::ensure!(
ambiguous_ranges.is_empty(),
{
let line_numbers = ambiguous_ranges
.iter()
.map(|range| range.start.to_string())
.collect::<Vec<_>>()
.join(", ");
formatdoc! {"
<old_text> matches more than one position in the file (lines: {line_numbers}). Read the
relevant sections of {input_path} again and extend <old_text> so
that I can perform the requested edits.
"}
}
);
}
Ok(EditFileToolOutput { action_log.update(cx, |log, cx| {
input_path: input.path, log.buffer_edited(buffer.clone(), cx);
new_text, })?;
old_text,
diff: unified_diff, let new_snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot())?;
edit_agent_output, let (new_text, unified_diff) = cx
.background_spawn({
let new_snapshot = new_snapshot.clone();
let old_text = old_text.clone();
async move {
let new_text = new_snapshot.text();
let diff = language::unified_diff(&old_text, &new_text);
(new_text, diff)
}
}) })
}.await; .await;
// Always finalize the diff, regardless of whether the operation succeeded or failed let input_path = input.path.display();
diff.update(cx, |diff, cx| diff.finalize(cx)).ok(); if unified_diff.is_empty() {
anyhow::ensure!(
!hallucinated_old_text,
formatdoc! {"
Some edits were produced but none of them could be applied.
Read the relevant sections of {input_path} again so that
I can perform the requested edits.
"}
);
anyhow::ensure!(
ambiguous_ranges.is_empty(),
{
let line_numbers = ambiguous_ranges
.iter()
.map(|range| range.start.to_string())
.collect::<Vec<_>>()
.join(", ");
formatdoc! {"
<old_text> matches more than one position in the file (lines: {line_numbers}). Read the
relevant sections of {input_path} again and extend <old_text> so
that I can perform the requested edits.
"}
}
);
}
result Ok(EditFileToolOutput {
input_path: input.path,
new_text,
old_text,
diff: unified_diff,
edit_agent_output,
})
}) })
} }
@ -1550,6 +1550,76 @@ mod tests {
); );
} }
#[gpui::test]
async fn test_diff_finalization(cx: &mut TestAppContext) {
init_test(cx);
let fs = project::FakeFs::new(cx.executor());
fs.insert_tree("/", json!({"main.rs": ""})).await;
let project = Project::test(fs.clone(), [path!("/").as_ref()], cx).await;
let languages = project.read_with(cx, |project, _cx| project.languages().clone());
let context_server_registry =
cx.new(|cx| ContextServerRegistry::new(project.read(cx).context_server_store(), cx));
let model = Arc::new(FakeLanguageModel::default());
let thread = cx.new(|cx| {
Thread::new(
project.clone(),
cx.new(|_cx| ProjectContext::default()),
context_server_registry.clone(),
Templates::new(),
Some(model.clone()),
cx,
)
});
// Ensure the diff is finalized if an error occurs while editing.
{
model.forbid_requests();
let tool = Arc::new(EditFileTool::new(thread.downgrade(), languages.clone()));
let (stream_tx, mut stream_rx) = ToolCallEventStream::test();
let edit = cx.update(|cx| {
tool.run(
EditFileToolInput {
display_description: "Edit file".into(),
path: "/main.rs".into(),
mode: EditFileMode::Edit,
},
stream_tx,
cx,
)
});
stream_rx.expect_update_fields().await;
let diff = stream_rx.expect_diff().await;
diff.read_with(cx, |diff, _| assert!(matches!(diff, Diff::Pending(_))));
edit.await.unwrap_err();
diff.read_with(cx, |diff, _| assert!(matches!(diff, Diff::Finalized(_))));
model.allow_requests();
}
// Ensure the diff is finalized if the tool call gets dropped.
{
let tool = Arc::new(EditFileTool::new(thread.downgrade(), languages.clone()));
let (stream_tx, mut stream_rx) = ToolCallEventStream::test();
let edit = cx.update(|cx| {
tool.run(
EditFileToolInput {
display_description: "Edit file".into(),
path: "/main.rs".into(),
mode: EditFileMode::Edit,
},
stream_tx,
cx,
)
});
stream_rx.expect_update_fields().await;
let diff = stream_rx.expect_diff().await;
diff.read_with(cx, |diff, _| assert!(matches!(diff, Diff::Pending(_))));
drop(edit);
cx.run_until_parked();
diff.read_with(cx, |diff, _| assert!(matches!(diff, Diff::Finalized(_))));
}
}
fn init_test(cx: &mut TestAppContext) { fn init_test(cx: &mut TestAppContext) {
cx.update(|cx| { cx.update(|cx| {
let settings_store = SettingsStore::test(cx); let settings_store = SettingsStore::test(cx);

31
crates/language_model/src/fake_provider.rs
View file

@ -4,12 +4,16 @@ use crate::{
LanguageModelProviderId, LanguageModelProviderName, LanguageModelProviderState, LanguageModelProviderId, LanguageModelProviderName, LanguageModelProviderState,
LanguageModelRequest, LanguageModelToolChoice, LanguageModelRequest, LanguageModelToolChoice,
}; };
use anyhow::anyhow;
use futures::{FutureExt, channel::mpsc, future::BoxFuture, stream::BoxStream}; use futures::{FutureExt, channel::mpsc, future::BoxFuture, stream::BoxStream};
use gpui::{AnyView, App, AsyncApp, Entity, Task, Window}; use gpui::{AnyView, App, AsyncApp, Entity, Task, Window};
use http_client::Result; use http_client::Result;
use parking_lot::Mutex; use parking_lot::Mutex;
use smol::stream::StreamExt; use smol::stream::StreamExt;
use std::sync::Arc; use std::sync::{
Arc,
atomic::{AtomicBool, Ordering::SeqCst},
};
#[derive(Clone)] #[derive(Clone)]
pub struct FakeLanguageModelProvider { pub struct FakeLanguageModelProvider {
@ -106,6 +110,7 @@ pub struct FakeLanguageModel {
>, >,
)>, )>,
>, >,
forbid_requests: AtomicBool,
} }
impl Default for FakeLanguageModel { impl Default for FakeLanguageModel {
@ -114,11 +119,20 @@ impl Default for FakeLanguageModel {
provider_id: LanguageModelProviderId::from("fake".to_string()), provider_id: LanguageModelProviderId::from("fake".to_string()),
provider_name: LanguageModelProviderName::from("Fake".to_string()), provider_name: LanguageModelProviderName::from("Fake".to_string()),
current_completion_txs: Mutex::new(Vec::new()), current_completion_txs: Mutex::new(Vec::new()),
forbid_requests: AtomicBool::new(false),
} }
} }
} }
impl FakeLanguageModel { impl FakeLanguageModel {
pub fn allow_requests(&self) {
self.forbid_requests.store(false, SeqCst);
}
pub fn forbid_requests(&self) {
self.forbid_requests.store(true, SeqCst);
}
pub fn pending_completions(&self) -> Vec<LanguageModelRequest> { pub fn pending_completions(&self) -> Vec<LanguageModelRequest> {
self.current_completion_txs self.current_completion_txs
.lock() .lock()
@ -251,9 +265,18 @@ impl LanguageModel for FakeLanguageModel {
LanguageModelCompletionError, LanguageModelCompletionError,
>, >,
> { > {
let (tx, rx) = mpsc::unbounded(); if self.forbid_requests.load(SeqCst) {
self.current_completion_txs.lock().push((request, tx)); async move {
async move { Ok(rx.boxed()) }.boxed() Err(LanguageModelCompletionError::Other(anyhow!(
"requests are forbidden"
)))
}
.boxed()
} else {
let (tx, rx) = mpsc::unbounded();
self.current_completion_txs.lock().push((request, tx));
async move { Ok(rx.boxed()) }.boxed()
}
} }
fn as_fake(&self) -> &Self { fn as_fake(&self) -> &Self {