agent2: Show retry button for errors

Co-authored-by: Antonio Scandurra <me@as-cii.com>

crates/acp_thread/src/acp_thread.rs

@@ -1373,6 +1373,10 @@ impl AcpThread {
         })
     }
 
+    pub fn can_resume(&self, cx: &App) -> bool {
+        self.connection.resume(&self.session_id, cx).is_some()
+    }
+
     pub fn resume(&mut self, cx: &mut Context<Self>) -> BoxFuture<'static, Result<()>> {
         self.run_turn(cx, async move |this, cx| {
             this.update(cx, |this, cx| {
@@ -2659,7 +2663,7 @@ mod tests {
         fn truncate(
             &self,
             session_id: &acp::SessionId,
-            _cx: &mut App,
+            _cx: &App,
         ) -> Option<Rc<dyn AgentSessionTruncate>> {
             Some(Rc::new(FakeAgentSessionEditor {
                 _session_id: session_id.clone(),

crates/acp_thread/src/connection.rs

@@ -43,7 +43,7 @@ pub trait AgentConnection {
     fn resume(
         &self,
         _session_id: &acp::SessionId,
-        _cx: &mut App,
+        _cx: &App,
     ) -> Option<Rc<dyn AgentSessionResume>> {
         None
     }
@@ -53,7 +53,7 @@ pub trait AgentConnection {
     fn truncate(
         &self,
         _session_id: &acp::SessionId,
-        _cx: &mut App,
+        _cx: &App,
     ) -> Option<Rc<dyn AgentSessionTruncate>> {
         None
     }
@@ -61,7 +61,7 @@ pub trait AgentConnection {
     fn set_title(
         &self,
         _session_id: &acp::SessionId,
-        _cx: &mut App,
+        _cx: &App,
     ) -> Option<Rc<dyn AgentSessionSetTitle>> {
         None
     }
@@ -439,7 +439,7 @@ mod test_support {
         fn truncate(
             &self,
             _session_id: &agent_client_protocol::SessionId,
-            _cx: &mut App,
+            _cx: &App,
         ) -> Option<Rc<dyn AgentSessionTruncate>> {
             Some(Rc::new(StubAgentSessionEditor))
         }

crates/agent2/src/agent.rs

@@ -936,7 +936,7 @@ impl acp_thread::AgentConnection for NativeAgentConnection {
     fn resume(
         &self,
         session_id: &acp::SessionId,
-        _cx: &mut App,
+        _cx: &App,
     ) -> Option<Rc<dyn acp_thread::AgentSessionResume>> {
         Some(Rc::new(NativeAgentSessionResume {
             connection: self.clone(),
@@ -956,9 +956,9 @@ impl acp_thread::AgentConnection for NativeAgentConnection {
     fn truncate(
         &self,
         session_id: &agent_client_protocol::SessionId,
-        cx: &mut App,
+        cx: &App,
     ) -> Option<Rc<dyn acp_thread::AgentSessionTruncate>> {
-        self.0.update(cx, |agent, _cx| {
+        self.0.read_with(cx, |agent, _cx| {
             agent.sessions.get(session_id).map(|session| {
                 Rc::new(NativeAgentSessionEditor {
                     thread: session.thread.clone(),
@@ -971,7 +971,7 @@ impl acp_thread::AgentConnection for NativeAgentConnection {
     fn set_title(
         &self,
         session_id: &acp::SessionId,
-        _cx: &mut App,
+        _cx: &App,
     ) -> Option<Rc<dyn acp_thread::AgentSessionSetTitle>> {
         Some(Rc::new(NativeAgentSessionSetTitle {
             connection: self.clone(),

crates/agent2/src/thread.rs

@@ -27,7 +27,8 @@ use futures::{
 };
 use git::repository::DiffType;
 use gpui::{
-    App, AppContext, AsyncApp, Context, Entity, EventEmitter, SharedString, Task, WeakEntity,
+    App, AppContext, AsyncApp, Context, Entity, EventEmitter, FutureExt as _, SharedString, Task,
+    WeakEntity,
 };
 use language_model::{
     LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelExt,
@@ -1085,11 +1086,6 @@ impl Thread {
         &mut self,
         cx: &mut Context<Self>,
     ) -> Result<mpsc::UnboundedReceiver<Result<ThreadEvent>>> {
-        anyhow::ensure!(
-            self.tool_use_limit_reached,
-            "can only resume after tool use limit is reached"
-        );
-
         self.messages.push(Message::Resume);
         cx.notify();
 
@@ -1216,12 +1212,13 @@ impl Thread {
         cx: &mut AsyncApp,
     ) -> Result<()> {
         log::debug!("Stream completion started successfully");
-        let request = this.update(cx, |this, cx| {
-            this.build_completion_request(completion_intent, cx)
-        })??;
 
         let mut attempt = None;
         'retry: loop {
+            let request = this.update(cx, |this, cx| {
+                this.build_completion_request(completion_intent, cx)
+            })??;
+
             telemetry::event!(
                 "Agent Thread Completion",
                 thread_id = this.read_with(cx, |this, _| this.id.to_string())?,
@@ -1236,7 +1233,7 @@ impl Thread {
                 attempt.unwrap_or(0)
             );
             let mut events = model
-                .stream_completion(request.clone(), cx)
+                .stream_completion(request, cx)
                 .await
                 .map_err(|error| anyhow!(error))?;
             let mut tool_results = FuturesUnordered::new();
@@ -1291,6 +1288,7 @@ impl Thread {
                             started_at: Instant::now(),
                             duration: delay,
                         });
+                        this.update(cx, |this, cx| this.flush_pending_message(cx))?;
 
                         cx.background_executor().timer(delay).await;
                         continue 'retry;
@@ -1737,6 +1735,10 @@ impl Thread {
             return;
         };
 
+        if message.content.is_empty() {
+            return;
+        }
+
         for content in &message.content {
             let AgentMessageContent::ToolUse(tool_use) = content else {
                 continue;

crates/agent_ui/src/acp/thread_view.rs

@@ -820,6 +820,9 @@ impl AcpThreadView {
         let Some(thread) = self.thread() else {
             return;
         };
+        if !thread.read(cx).can_resume(cx) {
+            return;
+        }
 
         let task = thread.update(cx, |thread, cx| thread.resume(cx));
         cx.spawn(async move |this, cx| {
@@ -4469,12 +4472,53 @@ impl AcpThreadView {
     }
 
     fn render_any_thread_error(&self, error: SharedString, cx: &mut Context<'_, Self>) -> Callout {
+        let can_resume = self
+            .thread()
+            .map_or(false, |thread| thread.read(cx).can_resume(cx));
+
+        let can_enable_burn_mode = self.as_native_thread(cx).map_or(false, |thread| {
+            let thread = thread.read(cx);
+            let supports_burn_mode = thread
+                .model()
+                .map_or(false, |model| model.supports_burn_mode());
+            supports_burn_mode && thread.completion_mode() == CompletionMode::Normal
+        });
+
         Callout::new()
             .severity(Severity::Error)
             .title("Error")
             .icon(IconName::XCircle)
             .description(error.clone())
-            .actions_slot(self.create_copy_button(error.to_string()))
+            .actions_slot(
+                h_flex()
+                    .gap_0p5()
+                    .when(can_resume && can_enable_burn_mode, |this| {
+                        this.child(
+                            Button::new("enable-burn-mode-and-retry", "Enable Burn Mode and Retry")
+                                .icon(IconName::ZedBurnMode)
+                                .icon_position(IconPosition::Start)
+                                .icon_size(IconSize::Small)
+                                .label_size(LabelSize::Small)
+                                .on_click(cx.listener(|this, _, window, cx| {
+                                    this.toggle_burn_mode(&ToggleBurnMode, window, cx);
+                                    this.resume_chat(cx);
+                                })),
+                        )
+                    })
+                    .when(can_resume, |this| {
+                        this.child(
+                            Button::new("retry", "Retry")
+                                .icon(IconName::RotateCw)
+                                .icon_position(IconPosition::Start)
+                                .icon_size(IconSize::Small)
+                                .label_size(LabelSize::Small)
+                                .on_click(cx.listener(|this, _, _window, cx| {
+                                    this.resume_chat(cx);
+                                })),
+                        )
+                    })
+                    .child(self.create_copy_button(error.to_string())),
+            )
             .dismiss_action(self.dismiss_error_button(cx))
     }