Add tests

crates/agent2/src/tests/mod.rs

@@ -673,15 +673,6 @@ async fn test_resume_after_tool_use_limit(cx: &mut TestAppContext) {
             "}
         )
     });
-
-    // Ensure we error if calling resume when tool use limit was *not* reached.
-    let error = thread
-        .update(cx, |thread, cx| thread.resume(cx))
-        .unwrap_err();
-    assert_eq!(
-        error.to_string(),
-        "can only resume after tool use limit is reached"
-    )
 }
 
 #[gpui::test]
@@ -2116,6 +2107,7 @@ async fn test_send_retry_on_error(cx: &mut TestAppContext) {
 
     fake_model.send_last_completion_stream_text_chunk("Hey!");
     fake_model.end_last_completion_stream();
+    cx.run_until_parked();
 
     let mut retry_events = Vec::new();
     while let Some(Ok(event)) = events.next().await {
@@ -2149,6 +2141,104 @@ async fn test_send_retry_on_error(cx: &mut TestAppContext) {
     });
 }
 
+#[gpui::test]
+async fn test_send_retry_cancels_tool_calls_on_error(cx: &mut TestAppContext) {
+    let ThreadTest { thread, model, .. } = setup(cx, TestModel::Fake).await;
+    let fake_model = model.as_fake();
+
+    thread
+        .update(cx, |thread, cx| {
+            thread.set_completion_mode(agent_settings::CompletionMode::Burn, cx);
+            thread.add_tool(EchoTool);
+            thread.send(UserMessageId::new(), ["Call the echo tool!"], cx)
+        })
+        .unwrap();
+    cx.run_until_parked();
+
+    let tool_use_1 = LanguageModelToolUse {
+        id: "tool_1".into(),
+        name: EchoTool::name().into(),
+        raw_input: json!({"text": "test"}).to_string(),
+        input: json!({"text": "test"}),
+        is_input_complete: true,
+    };
+    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
+        tool_use_1.clone(),
+    ));
+    fake_model.send_last_completion_stream_error(LanguageModelCompletionError::ServerOverloaded {
+        provider: LanguageModelProviderName::new("Anthropic"),
+        retry_after: Some(Duration::from_secs(3)),
+    });
+    fake_model.end_last_completion_stream();
+    cx.run_until_parked();
+    thread.read_with(cx, |thread, _cx| {
+        assert_eq!(
+            thread.to_markdown(),
+            indoc! {"
+            ## User
+
+            Call the echo tool!
+        "}
+        )
+    });
+
+    cx.executor().advance_clock(Duration::from_secs(3));
+    cx.run_until_parked();
+    let completion = fake_model.pending_completions().pop().unwrap();
+    assert_eq!(
+        completion.messages[1..],
+        vec![LanguageModelRequestMessage {
+            role: Role::User,
+            content: vec!["Call the echo tool!".into()],
+            cache: true
+        }]
+    );
+
+    let tool_use_2 = LanguageModelToolUse {
+        id: "tool_2".into(),
+        name: EchoTool::name().into(),
+        raw_input: json!({"text": "test"}).to_string(),
+        input: json!({"text": "test"}),
+        is_input_complete: true,
+    };
+    let tool_result_2 = LanguageModelToolResult {
+        tool_use_id: "tool_2".into(),
+        tool_name: EchoTool::name().into(),
+        is_error: false,
+        content: "test".into(),
+        output: Some("test".into()),
+    };
+    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
+        tool_use_2.clone(),
+    ));
+    fake_model
+        .send_last_completion_stream_event(LanguageModelCompletionEvent::Stop(StopReason::EndTurn));
+    fake_model.end_last_completion_stream();
+    cx.run_until_parked();
+
+    let completion = fake_model.pending_completions().pop().unwrap();
+    assert_eq!(
+        completion.messages[1..],
+        vec![
+            LanguageModelRequestMessage {
+                role: Role::User,
+                content: vec!["Call the echo tool!".into()],
+                cache: false
+            },
+            LanguageModelRequestMessage {
+                role: Role::Assistant,
+                content: vec![MessageContent::ToolUse(tool_use_2.clone())],
+                cache: false
+            },
+            LanguageModelRequestMessage {
+                role: Role::User,
+                content: vec![MessageContent::ToolResult(tool_result_2)],
+                cache: true
+            }
+        ]
+    );
+}
+
 #[gpui::test]
 async fn test_send_max_retries_exceeded(cx: &mut TestAppContext) {
     let ThreadTest { thread, model, .. } = setup(cx, TestModel::Fake).await;

crates/agent2/src/thread.rs

@@ -27,8 +27,7 @@ use futures::{
 };
 use git::repository::DiffType;
 use gpui::{
-    App, AppContext, AsyncApp, Context, Entity, EventEmitter, FutureExt as _, SharedString, Task,
+    App, AppContext, AsyncApp, Context, Entity, EventEmitter, SharedString, Task, WeakEntity,
-    WeakEntity,
 };
 use language_model::{
     LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelExt,