diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index 54cc6296d5..180cc88390 100644 --- a/crates/agent/src/thread.rs +++ b/crates/agent/src/thread.rs @@ -396,6 +396,7 @@ pub struct Thread { remaining_turns: u32, configured_model: Option, profile: AgentProfile, + last_error_context: Option<(Arc, CompletionIntent)>, } #[derive(Clone, Debug)] @@ -489,10 +490,11 @@ impl Thread { retry_state: None, message_feedback: HashMap::default(), last_auto_capture_at: None, + last_error_context: None, last_received_chunk_at: None, request_callback: None, remaining_turns: u32::MAX, - configured_model, + configured_model: configured_model.clone(), profile: AgentProfile::new(profile_id, tools), } } @@ -613,6 +615,7 @@ impl Thread { feedback: None, message_feedback: HashMap::default(), last_auto_capture_at: None, + last_error_context: None, last_received_chunk_at: None, request_callback: None, remaining_turns: u32::MAX, @@ -1264,9 +1267,58 @@ impl Thread { self.flush_notifications(model.clone(), intent, cx); - let request = self.to_completion_request(model.clone(), intent, cx); + let _checkpoint = self.finalize_pending_checkpoint(cx); + self.stream_completion( + self.to_completion_request(model.clone(), intent, cx), + model, + intent, + window, + cx, + ); + } - self.stream_completion(request, model, intent, window, cx); + pub fn retry_last_completion( + &mut self, + window: Option, + cx: &mut Context, + ) { + // Clear any existing error state + self.retry_state = None; + + // Use the last error context if available, otherwise fall back to configured model + let (model, intent) = if let Some((model, intent)) = self.last_error_context.take() { + (model, intent) + } else if let Some(configured_model) = self.configured_model.as_ref() { + let model = configured_model.model.clone(); + let intent = if self.has_pending_tool_uses() { + CompletionIntent::ToolResults + } else { + CompletionIntent::UserPrompt + }; + (model, intent) + } else if let Some(configured_model) = self.get_or_init_configured_model(cx) { + let model = configured_model.model.clone(); + let intent = if self.has_pending_tool_uses() { + CompletionIntent::ToolResults + } else { + CompletionIntent::UserPrompt + }; + (model, intent) + } else { + return; + }; + + self.send_to_model(model, intent, window, cx); + } + + pub fn enable_burn_mode_and_retry( + &mut self, + window: Option, + cx: &mut Context, + ) { + self.completion_mode = CompletionMode::Burn; + cx.emit(ThreadEvent::ProfileChanged); + self.retry_last_completion(window, cx); } pub fn used_tools_since_last_user_message(&self) -> bool { @@ -2222,6 +2274,23 @@ impl Thread { window: Option, cx: &mut Context, ) -> bool { + // Store context for the Retry button + self.last_error_context = Some((model.clone(), intent)); + + // Only auto-retry if Burn Mode is enabled + if self.completion_mode != CompletionMode::Burn { + // Show error with retry options + cx.emit(ThreadEvent::ShowError(ThreadError::RetryableError { + message: format!( + "{}\n\nTo automatically retry when similar errors happen, enable Burn Mode.", + error + ) + .into(), + can_enable_burn_mode: true, + })); + return false; + } + let Some(strategy) = strategy.or_else(|| Self::get_retry_strategy(error)) else { return false; }; @@ -2302,6 +2371,13 @@ impl Thread { // Stop generating since we're giving up on retrying. self.pending_completions.clear(); + // Show error alongside a Retry button, but no + // Enable Burn Mode button (since it's already enabled) + cx.emit(ThreadEvent::ShowError(ThreadError::RetryableError { + message: format!("Failed after retrying: {}", error).into(), + can_enable_burn_mode: false, + })); + false } } @@ -3212,6 +3288,11 @@ pub enum ThreadError { header: SharedString, message: SharedString, }, + #[error("Retryable error: {message}")] + RetryableError { + message: SharedString, + can_enable_burn_mode: bool, + }, } #[derive(Debug, Clone)] @@ -4167,6 +4248,11 @@ fn main() {{ let project = create_test_project(cx, json!({})).await; let (_, _, thread, _, _base_model) = setup_test_environment(cx, project.clone()).await; + // Enable Burn Mode to allow retries + thread.update(cx, |thread, _| { + thread.set_completion_mode(CompletionMode::Burn); + }); + // Create model that returns overloaded error let model = Arc::new(ErrorInjector::new(TestError::Overloaded)); @@ -4240,6 +4326,11 @@ fn main() {{ let project = create_test_project(cx, json!({})).await; let (_, _, thread, _, _base_model) = setup_test_environment(cx, project.clone()).await; + // Enable Burn Mode to allow retries + thread.update(cx, |thread, _| { + thread.set_completion_mode(CompletionMode::Burn); + }); + // Create model that returns internal server error let model = Arc::new(ErrorInjector::new(TestError::InternalServerError)); @@ -4316,6 +4407,11 @@ fn main() {{ let project = create_test_project(cx, json!({})).await; let (_, _, thread, _, _base_model) = setup_test_environment(cx, project.clone()).await; + // Enable Burn Mode to allow retries + thread.update(cx, |thread, _| { + thread.set_completion_mode(CompletionMode::Burn); + }); + // Create model that returns internal server error let model = Arc::new(ErrorInjector::new(TestError::InternalServerError)); @@ -4423,6 +4519,11 @@ fn main() {{ let project = create_test_project(cx, json!({})).await; let (_, _, thread, _, _base_model) = setup_test_environment(cx, project.clone()).await; + // Enable Burn Mode to allow retries + thread.update(cx, |thread, _| { + thread.set_completion_mode(CompletionMode::Burn); + }); + // Create model that returns overloaded error let model = Arc::new(ErrorInjector::new(TestError::Overloaded)); @@ -4509,6 +4610,11 @@ fn main() {{ let project = create_test_project(cx, json!({})).await; let (_, _, thread, _, _base_model) = setup_test_environment(cx, project.clone()).await; + // Enable Burn Mode to allow retries + thread.update(cx, |thread, _| { + thread.set_completion_mode(CompletionMode::Burn); + }); + // We'll use a wrapper to switch behavior after first failure struct RetryTestModel { inner: Arc, @@ -4677,6 +4783,11 @@ fn main() {{ let project = create_test_project(cx, json!({})).await; let (_, _, thread, _, _base_model) = setup_test_environment(cx, project.clone()).await; + // Enable Burn Mode to allow retries + thread.update(cx, |thread, _| { + thread.set_completion_mode(CompletionMode::Burn); + }); + // Create a model that fails once then succeeds struct FailOnceModel { inner: Arc, @@ -4838,6 +4949,11 @@ fn main() {{ let project = create_test_project(cx, json!({})).await; let (_, _, thread, _, _base_model) = setup_test_environment(cx, project.clone()).await; + // Enable Burn Mode to allow retries + thread.update(cx, |thread, _| { + thread.set_completion_mode(CompletionMode::Burn); + }); + // Create a model that returns rate limit error with retry_after struct RateLimitModel { inner: Arc, @@ -5111,6 +5227,79 @@ fn main() {{ ); } + #[gpui::test] + async fn test_no_retry_without_burn_mode(cx: &mut TestAppContext) { + init_test_settings(cx); + + let project = create_test_project(cx, json!({})).await; + let (_, _, thread, _, _base_model) = setup_test_environment(cx, project.clone()).await; + + // Ensure we're in Normal mode (not Burn mode) + thread.update(cx, |thread, _| { + thread.set_completion_mode(CompletionMode::Normal); + }); + + // Track error events + let error_events = Arc::new(Mutex::new(Vec::new())); + let error_events_clone = error_events.clone(); + + let _subscription = thread.update(cx, |_, cx| { + cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| { + if let ThreadEvent::ShowError(error) = event { + error_events_clone.lock().push(error.clone()); + } + }) + }); + + // Create model that returns overloaded error + let model = Arc::new(ErrorInjector::new(TestError::Overloaded)); + + // Insert a user message + thread.update(cx, |thread, cx| { + thread.insert_user_message("Hello!", ContextLoadResult::default(), None, vec![], cx); + }); + + // Start completion + thread.update(cx, |thread, cx| { + thread.send_to_model(model.clone(), CompletionIntent::UserPrompt, None, cx); + }); + + cx.run_until_parked(); + + // Verify no retry state was created + thread.read_with(cx, |thread, _| { + assert!( + thread.retry_state.is_none(), + "Should not have retry state in Normal mode" + ); + }); + + // Check that a retryable error was reported + let errors = error_events.lock(); + assert!(!errors.is_empty(), "Should have received an error event"); + + if let ThreadError::RetryableError { + message: _, + can_enable_burn_mode, + } = &errors[0] + { + assert!( + *can_enable_burn_mode, + "Error should indicate burn mode can be enabled" + ); + } else { + panic!("Expected RetryableError, got {:?}", errors[0]); + } + + // Verify the thread is no longer generating + thread.read_with(cx, |thread, _| { + assert!( + !thread.is_generating(), + "Should not be generating after error without retry" + ); + }); + } + #[gpui::test] async fn test_retry_cancelled_on_stop(cx: &mut TestAppContext) { init_test_settings(cx); @@ -5118,6 +5307,11 @@ fn main() {{ let project = create_test_project(cx, json!({})).await; let (_, _, thread, _, _base_model) = setup_test_environment(cx, project.clone()).await; + // Enable Burn Mode to allow retries + thread.update(cx, |thread, _| { + thread.set_completion_mode(CompletionMode::Burn); + }); + // Create model that returns overloaded error let model = Arc::new(ErrorInjector::new(TestError::Overloaded)); diff --git a/crates/agent_ui/src/agent_panel.rs b/crates/agent_ui/src/agent_panel.rs index 895a499502..139e32f835 100644 --- a/crates/agent_ui/src/agent_panel.rs +++ b/crates/agent_ui/src/agent_panel.rs @@ -65,8 +65,9 @@ use theme::ThemeSettings; use time::UtcOffset; use ui::utils::WithRemSize; use ui::{ - Banner, Callout, CheckboxWithLabel, ContextMenu, ElevationIndex, KeyBinding, PopoverMenu, - PopoverMenuHandle, ProgressBar, Tab, Tooltip, Vector, VectorName, prelude::*, + Banner, Button, Callout, CheckboxWithLabel, ContextMenu, ElevationIndex, IconPosition, + KeyBinding, PopoverMenu, PopoverMenuHandle, ProgressBar, Tab, Tooltip, Vector, VectorName, + prelude::*, }; use util::ResultExt as _; use workspace::{ @@ -2977,6 +2978,21 @@ impl AgentPanel { .size(IconSize::Small) .color(Color::Error); + let retry_button = Button::new("retry", "Retry") + .icon(IconName::RotateCw) + .icon_position(IconPosition::Start) + .on_click({ + let thread = thread.clone(); + move |_, window, cx| { + thread.update(cx, |thread, cx| { + thread.clear_last_error(); + thread.thread().update(cx, |thread, cx| { + thread.retry_last_completion(Some(window.window_handle()), cx); + }); + }); + } + }); + div() .border_t_1() .border_color(cx.theme().colors().border) @@ -2985,13 +3001,72 @@ impl AgentPanel { .icon(icon) .title(header) .description(message.clone()) - .primary_action(self.dismiss_error_button(thread, cx)) - .secondary_action(self.create_copy_button(message_with_header)) + .primary_action(retry_button) + .secondary_action(self.dismiss_error_button(thread, cx)) + .tertiary_action(self.create_copy_button(message_with_header)) .bg_color(self.error_callout_bg(cx)), ) .into_any_element() } + fn render_retryable_error( + &self, + message: SharedString, + can_enable_burn_mode: bool, + thread: &Entity, + cx: &mut Context, + ) -> AnyElement { + let icon = Icon::new(IconName::XCircle) + .size(IconSize::Small) + .color(Color::Error); + + let retry_button = Button::new("retry", "Retry") + .icon(IconName::RotateCw) + .icon_position(IconPosition::Start) + .on_click({ + let thread = thread.clone(); + move |_, window, cx| { + thread.update(cx, |thread, cx| { + thread.clear_last_error(); + thread.thread().update(cx, |thread, cx| { + thread.retry_last_completion(Some(window.window_handle()), cx); + }); + }); + } + }); + + let mut callout = Callout::new() + .icon(icon) + .title("Error") + .description(message.clone()) + .bg_color(self.error_callout_bg(cx)) + .primary_action(retry_button); + + if can_enable_burn_mode { + let burn_mode_button = Button::new("enable_burn_retry", "Enable Burn Mode and Retry") + .icon(IconName::ZedBurnMode) + .icon_position(IconPosition::Start) + .on_click({ + let thread = thread.clone(); + move |_, window, cx| { + thread.update(cx, |thread, cx| { + thread.clear_last_error(); + thread.thread().update(cx, |thread, cx| { + thread.enable_burn_mode_and_retry(Some(window.window_handle()), cx); + }); + }); + } + }); + callout = callout.secondary_action(burn_mode_button); + } + + div() + .border_t_1() + .border_color(cx.theme().colors().border) + .child(callout) + .into_any_element() + } + fn render_prompt_editor( &self, context_editor: &Entity, @@ -3233,6 +3308,15 @@ impl Render for AgentPanel { ThreadError::Message { header, message } => { self.render_error_message(header, message, thread, cx) } + ThreadError::RetryableError { + message, + can_enable_burn_mode, + } => self.render_retryable_error( + message, + can_enable_burn_mode, + thread, + cx, + ), }) .into_any(), )