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

Retry on burn mode (#34669)

Browse source 
Now we only auto-retry if burn mode is enabled. We also show a "Retry"
button (so you don't have to type "continue") if you think that's the
right remedy, and additionally we show a "Retry and Enable Burn Mode"
button if you don't have it enabled.

<img width="484" height="260" alt="Screenshot 2025-07-17 at 6 25 27 PM"
src="https://github.com/user-attachments/assets/dc5bf1f6-8b11-4041-87aa-4f37c95ea9f0"
/>

<img width="478" height="307" alt="Screenshot 2025-07-17 at 6 22 36 PM"
src="https://github.com/user-attachments/assets/1ed6578a-1696-449d-96d1-e447d11959fa"
/>


Release Notes:

- Only auto-retry Agent requests when Burn Mode is enabled
This commit is contained in:
Richard Feldman 2025-07-17 19:04:03 -04:00
parent f9c498318d
commit eeb9e242b4
No known key found for this signature in database
2 changed files with 285 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

200
crates/agent/src/thread.rs
View file

@ -396,6 +396,7 @@ pub struct Thread {
remaining_turns: u32,
configured_model: Option<ConfiguredModel>,
profile: AgentProfile,
last_error_context: Option<(Arc<dyn LanguageModel>, 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<AnyWindowHandle>,
cx: &mut Context<Self>,
) {
// 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<AnyWindowHandle>,
cx: &mut Context<Self>,
) {
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<AnyWindowHandle>,
cx: &mut Context<Self>,
) -> 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<FakeLanguageModel>,
@ -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<FakeLanguageModel>,
@ -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<FakeLanguageModel>,
@ -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));

92
crates/agent_ui/src/agent_panel.rs
View file

@ -64,8 +64,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::{
@ -2913,6 +2914,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)
@ -2921,13 +2937,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<ActiveThread>,
cx: &mut Context<Self>,
) -> 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<TextThreadEditor>,
@ -3169,6 +3244,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(),
)