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

agent2: Emit cancellation stop reason on cancel (#36381)

Browse source 
Release Notes:

- N/A

---------

Co-authored-by: Antonio Scandurra <me@as-cii.com>
This commit is contained in:
Ben Brandt 2025-08-18 09:58:30 +02:00 • committed by GitHub
parent b3969ed427
commit ea828c0c59
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 191 additions and 93 deletions
Download patch file Download diff file Expand all files Collapse all files

66
crates/agent2/src/tests/mod.rs
View file

@ -941,7 +941,15 @@ async fn test_cancellation(cx: &mut TestAppContext) {
// Cancel the current send and ensure that the event stream is closed, even
// if one of the tools is still running.
thread.update(cx, |thread, _cx| thread.cancel());
events.collect::<Vec<_>>().await;
let events = events.collect::<Vec<_>>().await;
let last_event = events.last();
assert!(
matches!(
last_event,
Some(Ok(AgentResponseEvent::Stop(acp::StopReason::Canceled)))
),
"unexpected event {last_event:?}"
);
// Ensure we can still send a new message after cancellation.
let events = thread
@ -965,6 +973,62 @@ async fn test_cancellation(cx: &mut TestAppContext) {
assert_eq!(stop_events(events), vec![acp::StopReason::EndTurn]);
}
#[gpui::test]
async fn test_in_progress_send_canceled_by_next_send(cx: &mut TestAppContext) {
let ThreadTest { model, thread, .. } = setup(cx, TestModel::Fake).await;
let fake_model = model.as_fake();
let events_1 = thread.update(cx, |thread, cx| {
thread.send(UserMessageId::new(), ["Hello 1"], cx)
});
cx.run_until_parked();
fake_model.send_last_completion_stream_text_chunk("Hey 1!");
cx.run_until_parked();
let events_2 = thread.update(cx, |thread, cx| {
thread.send(UserMessageId::new(), ["Hello 2"], cx)
});
cx.run_until_parked();
fake_model.send_last_completion_stream_text_chunk("Hey 2!");
fake_model
.send_last_completion_stream_event(LanguageModelCompletionEvent::Stop(StopReason::EndTurn));
fake_model.end_last_completion_stream();
let events_1 = events_1.collect::<Vec<_>>().await;
assert_eq!(stop_events(events_1), vec![acp::StopReason::Canceled]);
let events_2 = events_2.collect::<Vec<_>>().await;
assert_eq!(stop_events(events_2), vec![acp::StopReason::EndTurn]);
}
#[gpui::test]
async fn test_subsequent_successful_sends_dont_cancel(cx: &mut TestAppContext) {
let ThreadTest { model, thread, .. } = setup(cx, TestModel::Fake).await;
let fake_model = model.as_fake();
let events_1 = thread.update(cx, |thread, cx| {
thread.send(UserMessageId::new(), ["Hello 1"], cx)
});
cx.run_until_parked();
fake_model.send_last_completion_stream_text_chunk("Hey 1!");
fake_model
.send_last_completion_stream_event(LanguageModelCompletionEvent::Stop(StopReason::EndTurn));
fake_model.end_last_completion_stream();
let events_1 = events_1.collect::<Vec<_>>().await;
let events_2 = thread.update(cx, |thread, cx| {
thread.send(UserMessageId::new(), ["Hello 2"], cx)
});
cx.run_until_parked();
fake_model.send_last_completion_stream_text_chunk("Hey 2!");
fake_model
.send_last_completion_stream_event(LanguageModelCompletionEvent::Stop(StopReason::EndTurn));
fake_model.end_last_completion_stream();
let events_2 = events_2.collect::<Vec<_>>().await;
assert_eq!(stop_events(events_1), vec![acp::StopReason::EndTurn]);
assert_eq!(stop_events(events_2), vec![acp::StopReason::EndTurn]);
}
#[gpui::test]
async fn test_refusal(cx: &mut TestAppContext) {
let ThreadTest { model, thread, .. } = setup(cx, TestModel::Fake).await;