From a63e50f6cbe36f32c87ced49ad102f7a2d6657f3 Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Tue, 24 Jun 2025 14:32:03 -0400 Subject: [PATCH 01/32] Reduce segment cloning when rendering messages --- crates/agent/src/thread.rs | 7 ++ crates/agent_ui/src/active_thread.rs | 177 ++++++++++++++------------- 2 files changed, 98 insertions(+), 86 deletions(-) diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index a46aa9381ea450..e5c7f0efff75ae 100644 --- a/crates/agent/src/thread.rs +++ b/crates/agent/src/thread.rs @@ -198,6 +198,13 @@ impl MessageSegment { Self::RedactedThinking(_) => false, } } + + pub fn text(&self) -> Option<&str> { + match self { + MessageSegment::Text(text) => Some(text), + _ => None, + } + } } #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] diff --git a/crates/agent_ui/src/active_thread.rs b/crates/agent_ui/src/active_thread.rs index 8df1c88e8a5761..0e7ca9aa897d19 100644 --- a/crates/agent_ui/src/active_thread.rs +++ b/crates/agent_ui/src/active_thread.rs @@ -809,7 +809,12 @@ impl ActiveThread { }; for message in thread.read(cx).messages().cloned().collect::>() { - this.push_message(&message.id, &message.segments, window, cx); + let rendered_message = RenderedMessage::from_segments( + &message.segments, + this.language_registry.clone(), + cx, + ); + this.push_rendered_message(message.id, rendered_message); for tool_use in thread.read(cx).tool_uses_for_message(message.id, cx) { this.render_tool_use_markdown( @@ -875,36 +880,11 @@ impl ActiveThread { &self.text_thread_store } - fn push_message( - &mut self, - id: &MessageId, - segments: &[MessageSegment], - _window: &mut Window, - cx: &mut Context, - ) { + fn push_rendered_message(&mut self, id: MessageId, rendered_message: RenderedMessage) { let old_len = self.messages.len(); - self.messages.push(*id); + self.messages.push(id); self.list_state.splice(old_len..old_len, 1); - - let rendered_message = - RenderedMessage::from_segments(segments, self.language_registry.clone(), cx); - self.rendered_messages_by_id.insert(*id, rendered_message); - } - - fn edited_message( - &mut self, - id: &MessageId, - segments: &[MessageSegment], - _window: &mut Window, - cx: &mut Context, - ) { - let Some(index) = self.messages.iter().position(|message_id| message_id == id) else { - return; - }; - self.list_state.splice(index..index + 1, 1); - let rendered_message = - RenderedMessage::from_segments(segments, self.language_registry.clone(), cx); - self.rendered_messages_by_id.insert(*id, rendered_message); + self.rendered_messages_by_id.insert(id, rendered_message); } fn deleted_message(&mut self, id: &MessageId) { @@ -1037,31 +1017,43 @@ impl ActiveThread { } } ThreadEvent::MessageAdded(message_id) => { - if let Some(message_segments) = self - .thread - .read(cx) - .message(*message_id) - .map(|message| message.segments.clone()) - { - self.push_message(message_id, &message_segments, window, cx); + if let Some(rendered_message) = self.thread.update(cx, |thread, cx| { + thread.message(*message_id).map(|message| { + RenderedMessage::from_segments( + &message.segments, + self.language_registry.clone(), + cx, + ) + }) + }) { + self.push_rendered_message(*message_id, rendered_message); } self.save_thread(cx); cx.notify(); } ThreadEvent::MessageEdited(message_id) => { - if let Some(message_segments) = self - .thread - .read(cx) - .message(*message_id) - .map(|message| message.segments.clone()) - { - self.edited_message(message_id, &message_segments, window, cx); + if let Some(index) = self.messages.iter().position(|id| id == message_id) { + if let Some(rendered_message) = self.thread.update(cx, |thread, cx| { + thread.message(*message_id).map(|message| { + let mut rendered_message = RenderedMessage { + language_registry: self.language_registry.clone(), + segments: Vec::with_capacity(message.segments.len()), + }; + for segment in &message.segments { + rendered_message.push_segment(segment, cx); + } + rendered_message + }) + }) { + self.list_state.splice(index..index + 1, 1); + self.rendered_messages_by_id + .insert(*message_id, rendered_message); + self.scroll_to_bottom(cx); + self.save_thread(cx); + cx.notify(); + } } - - self.scroll_to_bottom(cx); - self.save_thread(cx); - cx.notify(); } ThreadEvent::MessageDeleted(message_id) => { self.deleted_message(message_id); @@ -1311,17 +1303,11 @@ impl ActiveThread { fn start_editing_message( &mut self, message_id: MessageId, - message_segments: &[MessageSegment], + message_text: impl Into>, message_creases: &[MessageCrease], window: &mut Window, cx: &mut Context, ) { - // User message should always consist of a single text segment, - // therefore we can skip returning early if it's not a text segment. - let Some(MessageSegment::Text(message_text)) = message_segments.first() else { - return; - }; - let editor = crate::message_editor::create_editor( self.workspace.clone(), self.context_store.downgrade(), @@ -1333,7 +1319,7 @@ impl ActiveThread { cx, ); editor.update(cx, |editor, cx| { - editor.set_text(message_text.clone(), window, cx); + editor.set_text(message_text, window, cx); insert_message_creases(editor, message_creases, &self.context_store, window, cx); editor.focus_handle(cx).focus(window); editor.move_to_end(&editor::actions::MoveToEnd, window, cx); @@ -1828,8 +1814,6 @@ impl ActiveThread { return div().children(loading_dots).into_any(); } - let message_creases = message.creases.clone(); - let Some(rendered_message) = self.rendered_messages_by_id.get(&message_id) else { return Empty.into_any(); }; @@ -2144,15 +2128,30 @@ impl ActiveThread { }), ) .on_click(cx.listener({ - let message_segments = message.segments.clone(); + let message_creases = message.creases.clone(); move |this, _, window, cx| { - this.start_editing_message( - message_id, - &message_segments, - &message_creases, - window, - cx, - ); + if let Some(message_text) = + this.thread.read(cx).message(message_id).and_then(|message| { + message.segments.first().and_then(|segment| { + match segment { + MessageSegment::Text(message_text) => { + Some(Into::>::into(message_text.as_str())) + } + _ => { + None + } + } + }) + }) + { + this.start_editing_message( + message_id, + message_text, + &message_creases, + window, + cx, + ); + } } })), ), @@ -3826,13 +3825,15 @@ mod tests { }); active_thread.update_in(cx, |active_thread, window, cx| { - active_thread.start_editing_message( - message.id, - message.segments.as_slice(), - message.creases.as_slice(), - window, - cx, - ); + if let Some(message_text) = message.segments.first().and_then(MessageSegment::text) { + active_thread.start_editing_message( + message.id, + message_text, + message.creases.as_slice(), + window, + cx, + ); + } let editor = active_thread .editing_message .as_ref() @@ -3847,13 +3848,15 @@ mod tests { let message = thread.update(cx, |thread, _| thread.message(message.id).cloned().unwrap()); active_thread.update_in(cx, |active_thread, window, cx| { - active_thread.start_editing_message( - message.id, - message.segments.as_slice(), - message.creases.as_slice(), - window, - cx, - ); + if let Some(message_text) = message.segments.first().and_then(MessageSegment::text) { + active_thread.start_editing_message( + message.id, + message_text, + message.creases.as_slice(), + window, + cx, + ); + } let editor = active_thread .editing_message .as_ref() @@ -3935,13 +3938,15 @@ mod tests { // Edit the message while the completion is still running active_thread.update_in(cx, |active_thread, window, cx| { - active_thread.start_editing_message( - message.id, - message.segments.as_slice(), - message.creases.as_slice(), - window, - cx, - ); + if let Some(message_text) = message.segments.first().and_then(MessageSegment::text) { + active_thread.start_editing_message( + message.id, + message_text, + message.creases.as_slice(), + window, + cx, + ); + } let editor = active_thread .editing_message .as_ref() From 0060b30a666ac3a2ff61ac3218a9ef806b7a2599 Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Tue, 24 Jun 2025 12:48:50 -0400 Subject: [PATCH 02/32] Automatically retry on Overloaded --- crates/agent/src/thread.rs | 114 ++++++++++++++++++++++++++- crates/agent_ui/src/active_thread.rs | 28 +++++++ crates/agent_ui/src/agent_diff.rs | 3 +- 3 files changed, 141 insertions(+), 4 deletions(-) diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index e5c7f0efff75ae..6044357ab94353 100644 --- a/crates/agent/src/thread.rs +++ b/crates/agent/src/thread.rs @@ -39,7 +39,12 @@ use proto::Plan; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; use settings::Settings; -use std::{io::Write, ops::Range, sync::Arc, time::Instant}; +use std::{ + io::Write, + ops::Range, + sync::Arc, + time::{Duration, Instant}, +}; use thiserror::Error; use util::{ResultExt as _, post_inc}; use uuid::Uuid; @@ -364,6 +369,8 @@ pub struct Thread { exceeded_window_error: Option, tool_use_limit_reached: bool, feedback: Option, + retry_state: Option, + current_completion_intent: Option, message_feedback: HashMap, last_auto_capture_at: Option, last_received_chunk_at: Option, @@ -375,6 +382,12 @@ pub struct Thread { profile: AgentProfile, } +#[derive(Clone, Debug)] +struct RetryState { + attempt: u32, + max_attempts: u32, +} + #[derive(Clone, Debug, PartialEq, Eq)] pub enum ThreadSummary { Pending, @@ -456,6 +469,8 @@ impl Thread { exceeded_window_error: None, tool_use_limit_reached: false, feedback: None, + retry_state: None, + current_completion_intent: None, message_feedback: HashMap::default(), last_auto_capture_at: None, last_received_chunk_at: None, @@ -522,6 +537,8 @@ impl Thread { detailed_summary_tx, detailed_summary_rx, completion_mode, + retry_state: None, + current_completion_intent: None, messages: serialized .messages .into_iter() @@ -1226,6 +1243,9 @@ impl Thread { self.remaining_turns -= 1; + // Store the intent for potential retries + self.current_completion_intent = Some(intent.clone()); + let request = self.to_completion_request(model.clone(), intent, cx); self.stream_completion(request, model, window, cx); @@ -1460,6 +1480,8 @@ impl Thread { window: Option, cx: &mut Context, ) { + // Clear any previous retry state when starting a new completion + self.retry_state = None; self.tool_use_limit_reached = false; let pending_completion_id = post_inc(&mut self.completion_count); @@ -1864,8 +1886,81 @@ impl Thread { emit_generic_error(error, cx); } LanguageModelKnownError::Overloaded => { - // In the future we will wait and then retry, up to N times. - emit_generic_error(error, cx); + const MAX_RETRY_ATTEMPTS: u32 = 3; + const BASE_RETRY_DELAY_SECS: u64 = 5; + + // Check if we should create or update retry state + let should_retry = if let Some(retry_state) = &mut thread.retry_state { + retry_state.attempt += 1; + retry_state.attempt <= retry_state.max_attempts + } else { + thread.retry_state = Some(RetryState { + attempt: 1, + max_attempts: MAX_RETRY_ATTEMPTS, + }); + true + }; + + if should_retry { + let retry_state = thread.retry_state.as_ref().unwrap(); + let attempt = retry_state.attempt; + let max_attempts = retry_state.max_attempts; + + // Calculate exponential backoff + let delay_secs = BASE_RETRY_DELAY_SECS * 2u64.pow(attempt - 1); + let delay = Duration::from_secs(delay_secs); + + // Add a transient message to inform the user + let retry_message = format!( + "The API is currently overloaded. Automatically retrying in {} seconds... (attempt {}/{})", + delay_secs, + attempt, + max_attempts + ); + + let message_id = thread.insert_message( + Role::System, + vec![MessageSegment::Text(retry_message)], + LoadedContext::default(), + Vec::new(), + true, // is_hidden - won't be persisted + cx, + ); + + cx.emit(ThreadEvent::RetryScheduled { + message_id, + delay, + attempt, + max_attempts, + }); + + // Schedule the retry + let thread_handle = cx.entity().downgrade(); + let model = model.clone(); + let intent = thread.current_completion_intent.clone().unwrap_or(CompletionIntent::UserPrompt); + + cx.spawn(async move |_thread, cx| { + cx.background_executor().timer(delay).await; + + thread_handle.update(cx, |thread, cx| { + // Remove the retry notification message + if let Some(index) = thread.messages.iter().position(|m| m.id == message_id) { + thread.messages.remove(index); + cx.emit(ThreadEvent::MessageDeleted(message_id)); + } + + // Retry the completion + thread.send_to_model(model, intent, window, cx); + }).log_err(); + }).detach(); + + // Don't emit a generic error or cancel the completion yet + return; + } else { + // Max retries exceeded, clear retry state and show error + thread.retry_state = None; + emit_generic_error(error, cx); + } } LanguageModelKnownError::ApiInternalServerError => { // In the future we will retry the request, but only once. @@ -1882,10 +1977,17 @@ impl Thread { emit_generic_error(error, cx); } + // Clear retry state on any other error + thread.retry_state = None; thread.cancel_last_completion(window, cx); } } + // Clear retry state on successful completion + if result.is_ok() && thread.retry_state.is_some() { + thread.retry_state = None; + } + cx.emit(ThreadEvent::Stopped(result.map_err(Arc::new))); if let Some((request_callback, (request, response_events))) = thread @@ -2933,6 +3035,12 @@ pub enum ThreadEvent { CancelEditing, CompletionCanceled, ProfileChanged, + RetryScheduled { + message_id: MessageId, + delay: Duration, + attempt: u32, + max_attempts: u32, + }, } impl EventEmitter for Thread {} diff --git a/crates/agent_ui/src/active_thread.rs b/crates/agent_ui/src/active_thread.rs index 0e7ca9aa897d19..43039b9fa1eb97 100644 --- a/crates/agent_ui/src/active_thread.rs +++ b/crates/agent_ui/src/active_thread.rs @@ -1140,6 +1140,34 @@ impl ActiveThread { self.save_thread(cx); cx.notify(); } + ThreadEvent::RetryScheduled { + message_id, + delay, + attempt, + max_attempts, + } => { + // The retry message is already added to the thread as a system message + // We just need to ensure the UI updates + let message_segments = self + .thread + .read(cx) + .message(*message_id) + .map(|message| message.segments.clone()); + + if let Some(message_segments) = message_segments { + self.push_message(message_id, &message_segments, window, cx); + } + + // Show a notification to the user + let notification_text = format!( + "API overloaded. Retrying in {}s (attempt {}/{})", + delay.as_secs(), + attempt, + max_attempts + ); + self.show_notification(¬ification_text, IconName::ArrowCircle, window, cx); + cx.notify(); + } } } diff --git a/crates/agent_ui/src/agent_diff.rs b/crates/agent_ui/src/agent_diff.rs index aa5e49551b4e1c..15e0b7726cd900 100644 --- a/crates/agent_ui/src/agent_diff.rs +++ b/crates/agent_ui/src/agent_diff.rs @@ -1380,7 +1380,8 @@ impl AgentDiff { | ThreadEvent::ToolConfirmationNeeded | ThreadEvent::ToolUseLimitReached | ThreadEvent::CancelEditing - | ThreadEvent::ProfileChanged => {} + | ThreadEvent::ProfileChanged + | ThreadEvent::RetryScheduled { .. } => {} } } From 0e4a578270eb14bd968623ed2c98a078bfe84db5 Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Mon, 23 Jun 2025 13:47:58 -0400 Subject: [PATCH 03/32] Report provider name in API error --- crates/agent/src/thread.rs | 233 +++++++++++++++++---------- crates/agent_ui/src/active_thread.rs | 4 +- 2 files changed, 154 insertions(+), 83 deletions(-) diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index 6044357ab94353..58e881346451d0 100644 --- a/crates/agent/src/thread.rs +++ b/crates/agent/src/thread.rs @@ -23,11 +23,11 @@ use gpui::{ }; use language_model::{ ConfiguredModel, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent, - LanguageModelId, LanguageModelKnownError, LanguageModelRegistry, LanguageModelRequest, - LanguageModelRequestMessage, LanguageModelRequestTool, LanguageModelToolResult, - LanguageModelToolResultContent, LanguageModelToolUseId, MessageContent, - ModelRequestLimitReachedError, PaymentRequiredError, Role, SelectedModel, StopReason, - TokenUsage, + LanguageModelId, LanguageModelKnownError, LanguageModelProviderName, LanguageModelRegistry, + LanguageModelRequest, LanguageModelRequestMessage, LanguageModelRequestTool, + LanguageModelToolResult, LanguageModelToolResultContent, LanguageModelToolUseId, + MessageContent, ModelRequestLimitReachedError, PaymentRequiredError, Role, SelectedModel, + StopReason, TokenUsage, }; use postage::stream::Stream as _; use project::{ @@ -1881,90 +1881,54 @@ impl Thread { }); cx.notify(); } - LanguageModelKnownError::RateLimitExceeded { .. } => { - // In the future we will report the error to the user, wait retry_after, and then retry. - emit_generic_error(error, cx); + LanguageModelKnownError::RateLimitExceeded { retry_after } => { + let provider_name = model.provider_name(); + let error_message = format!( + "{}'s API rate limit exceeded", + provider_name.0.as_ref() + ); + + if !thread.handle_retryable_error_with_delay( + &error_message, + Some(*retry_after), + model.clone(), + window, + cx, + ) { + emit_generic_error(error, cx); + } } LanguageModelKnownError::Overloaded => { - const MAX_RETRY_ATTEMPTS: u32 = 3; - const BASE_RETRY_DELAY_SECS: u64 = 5; + let provider_name = model.provider_name(); + let error_message = format!( + "{}'s API servers are overloaded right now", + provider_name.0.as_ref() + ); - // Check if we should create or update retry state - let should_retry = if let Some(retry_state) = &mut thread.retry_state { - retry_state.attempt += 1; - retry_state.attempt <= retry_state.max_attempts - } else { - thread.retry_state = Some(RetryState { - attempt: 1, - max_attempts: MAX_RETRY_ATTEMPTS, - }); - true - }; - - if should_retry { - let retry_state = thread.retry_state.as_ref().unwrap(); - let attempt = retry_state.attempt; - let max_attempts = retry_state.max_attempts; - - // Calculate exponential backoff - let delay_secs = BASE_RETRY_DELAY_SECS * 2u64.pow(attempt - 1); - let delay = Duration::from_secs(delay_secs); - - // Add a transient message to inform the user - let retry_message = format!( - "The API is currently overloaded. Automatically retrying in {} seconds... (attempt {}/{})", - delay_secs, - attempt, - max_attempts - ); - - let message_id = thread.insert_message( - Role::System, - vec![MessageSegment::Text(retry_message)], - LoadedContext::default(), - Vec::new(), - true, // is_hidden - won't be persisted - cx, - ); - - cx.emit(ThreadEvent::RetryScheduled { - message_id, - delay, - attempt, - max_attempts, - }); - - // Schedule the retry - let thread_handle = cx.entity().downgrade(); - let model = model.clone(); - let intent = thread.current_completion_intent.clone().unwrap_or(CompletionIntent::UserPrompt); - - cx.spawn(async move |_thread, cx| { - cx.background_executor().timer(delay).await; - - thread_handle.update(cx, |thread, cx| { - // Remove the retry notification message - if let Some(index) = thread.messages.iter().position(|m| m.id == message_id) { - thread.messages.remove(index); - cx.emit(ThreadEvent::MessageDeleted(message_id)); - } - - // Retry the completion - thread.send_to_model(model, intent, window, cx); - }).log_err(); - }).detach(); - - // Don't emit a generic error or cancel the completion yet - return; - } else { - // Max retries exceeded, clear retry state and show error - thread.retry_state = None; + if !thread.handle_retryable_error( + &error_message, + model.clone(), + window, + cx, + ) { emit_generic_error(error, cx); } } LanguageModelKnownError::ApiInternalServerError => { - // In the future we will retry the request, but only once. - emit_generic_error(error, cx); + let provider_name = model.provider_name(); + let error_message = format!( + "{}'s API server reported an internal server error", + provider_name.0.as_ref() + ); + + if !thread.handle_retryable_error( + &error_message, + model.clone(), + window, + cx, + ) { + emit_generic_error(error, cx); + } } LanguageModelKnownError::ReadResponseError(_) | LanguageModelKnownError::DeserializeResponse(_) | @@ -2104,6 +2068,110 @@ impl Thread { }); } + fn handle_retryable_error( + &mut self, + error_message: &str, + model: Arc, + window: Option, + cx: &mut Context, + ) -> bool { + self.handle_retryable_error_with_delay(error_message, None, model, window, cx) + } + + fn handle_retryable_error_with_delay( + &mut self, + error_message: &str, + custom_delay: Option, + model: Arc, + window: Option, + cx: &mut Context, + ) -> bool { + const MAX_RETRY_ATTEMPTS: u32 = 3; + const BASE_RETRY_DELAY_SECS: u64 = 5; + + // Check if we should create or update retry state + let should_retry = if let Some(retry_state) = &mut self.retry_state { + retry_state.attempt += 1; + retry_state.attempt <= retry_state.max_attempts + } else { + self.retry_state = Some(RetryState { + attempt: 1, + max_attempts: MAX_RETRY_ATTEMPTS, + }); + true + }; + + if should_retry { + let retry_state = self.retry_state.as_ref().unwrap(); + let attempt = retry_state.attempt; + let max_attempts = retry_state.max_attempts; + + // Use custom delay if provided (e.g., from rate limit), otherwise exponential backoff + let delay = if let Some(custom_delay) = custom_delay { + custom_delay + } else { + let delay_secs = BASE_RETRY_DELAY_SECS * 2u64.pow(attempt - 1); + Duration::from_secs(delay_secs) + }; + let delay_secs = delay.as_secs(); + + // Add a transient message to inform the user + let retry_message = format!( + "{}. Retrying (attempt {} of {}) in {} seconds...", + error_message, attempt, max_attempts, delay_secs + ); + + let message_id = self.insert_message( + Role::System, + vec![MessageSegment::Text(retry_message)], + LoadedContext::default(), + Vec::new(), + true, // is_hidden - won't be persisted + cx, + ); + + cx.emit(ThreadEvent::RetryScheduled { + message_id, + delay, + attempt, + max_attempts, + provider_name: model.provider_name(), + }); + + // Schedule the retry + let thread_handle = cx.entity().downgrade(); + let intent = self + .current_completion_intent + .clone() + .unwrap_or(CompletionIntent::UserPrompt); + + cx.spawn(async move |_thread, cx| { + cx.background_executor().timer(delay).await; + + thread_handle + .update(cx, |thread, cx| { + // Remove the retry notification message + if let Some(index) = thread.messages.iter().position(|m| m.id == message_id) + { + thread.messages.remove(index); + cx.emit(ThreadEvent::MessageDeleted(message_id)); + } + + // Retry the completion + thread.send_to_model(model, intent, window, cx); + }) + .log_err(); + }) + .detach(); + + true + } else { + // Max retries exceeded, clear retry state + self.retry_state = None; + false + } + } + pub fn start_generating_detailed_summary_if_needed( &mut self, thread_store: WeakEntity, @@ -3040,6 +3108,7 @@ pub enum ThreadEvent { delay: Duration, attempt: u32, max_attempts: u32, + provider_name: LanguageModelProviderName, }, } diff --git a/crates/agent_ui/src/active_thread.rs b/crates/agent_ui/src/active_thread.rs index 43039b9fa1eb97..f718c6549091af 100644 --- a/crates/agent_ui/src/active_thread.rs +++ b/crates/agent_ui/src/active_thread.rs @@ -1145,6 +1145,7 @@ impl ActiveThread { delay, attempt, max_attempts, + provider_name, } => { // The retry message is already added to the thread as a system message // We just need to ensure the UI updates @@ -1160,7 +1161,8 @@ impl ActiveThread { // Show a notification to the user let notification_text = format!( - "API overloaded. Retrying in {}s (attempt {}/{})", + "{} retry in {}s (attempt {}/{})", + provider_name.0.as_ref(), delay.as_secs(), attempt, max_attempts From 7499caf6f09f878715e1a74cd44d45396894191c Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Tue, 24 Jun 2025 12:50:44 -0400 Subject: [PATCH 04/32] Test retries --- crates/agent/src/thread.rs | 570 +++++++++++++++++++++++++++++++++++++ 1 file changed, 570 insertions(+) diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index 58e881346451d0..82b3bae61b4360 100644 --- a/crates/agent/src/thread.rs +++ b/crates/agent/src/thread.rs @@ -3207,14 +3207,23 @@ mod tests { }; use agent_settings::{AgentProfileId, AgentSettings, LanguageModelParameters}; use assistant_tool::ToolRegistry; + use futures::StreamExt; + use futures::future::BoxFuture; + use futures::stream::BoxStream; use gpui::TestAppContext; use icons::IconName; use language_model::fake_provider::{FakeLanguageModel, FakeLanguageModelProvider}; + use language_model::{ + LanguageModelCompletionError, LanguageModelKnownError, LanguageModelName, + LanguageModelProviderId, LanguageModelToolChoice, + }; + use parking_lot::Mutex; use project::{FakeFs, Project}; use prompt_store::PromptBuilder; use serde_json::json; use settings::{Settings, SettingsStore}; use std::sync::Arc; + use std::time::Duration; use theme::ThemeSettings; use util::path; use workspace::Workspace; @@ -3989,6 +3998,567 @@ fn main() {{ } } + // Helper to create a model that returns errors + enum TestError { + Overloaded, + InternalServerError, + } + + struct ErrorInjector { + inner: Arc, + error_type: TestError, + } + + impl ErrorInjector { + fn new(error_type: TestError) -> Self { + Self { + inner: Arc::new(FakeLanguageModel::default()), + error_type, + } + } + } + + impl LanguageModel for ErrorInjector { + fn id(&self) -> LanguageModelId { + self.inner.id() + } + + fn name(&self) -> LanguageModelName { + self.inner.name() + } + + fn provider_id(&self) -> LanguageModelProviderId { + self.inner.provider_id() + } + + fn provider_name(&self) -> LanguageModelProviderName { + self.inner.provider_name() + } + + fn supports_tools(&self) -> bool { + self.inner.supports_tools() + } + + fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool { + self.inner.supports_tool_choice(choice) + } + + fn supports_images(&self) -> bool { + self.inner.supports_images() + } + + fn telemetry_id(&self) -> String { + self.inner.telemetry_id() + } + + fn max_token_count(&self) -> u64 { + self.inner.max_token_count() + } + + fn count_tokens( + &self, + request: LanguageModelRequest, + cx: &App, + ) -> BoxFuture<'static, Result> { + self.inner.count_tokens(request, cx) + } + + fn stream_completion( + &self, + _request: LanguageModelRequest, + _cx: &AsyncApp, + ) -> BoxFuture< + 'static, + Result< + BoxStream< + 'static, + Result, + >, + LanguageModelCompletionError, + >, + > { + let error = match self.error_type { + TestError::Overloaded => LanguageModelKnownError::Overloaded, + TestError::InternalServerError => LanguageModelKnownError::ApiInternalServerError, + }; + async move { + let stream = futures::stream::once(async move { + Err(LanguageModelCompletionError::Other(anyhow::anyhow!(error))) + }); + Ok(stream.boxed()) + } + .boxed() + } + + fn as_fake(&self) -> &FakeLanguageModel { + &self.inner + } + } + + #[gpui::test] + async fn test_retry_on_overloaded_error(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; + + // 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); + }); + + // Track retry events + let retry_events = Arc::new(Mutex::new(Vec::new())); + let retry_events_clone = retry_events.clone(); + + let _subscription = thread.update(cx, |_, cx| { + cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| { + if let ThreadEvent::RetryScheduled { attempt, delay, .. } = event { + retry_events_clone.lock().push((*attempt, *delay)); + } + }) + }); + + // Start completion + thread.update(cx, |thread, cx| { + thread.send_to_model(model.clone(), CompletionIntent::UserPrompt, None, cx); + }); + + cx.run_until_parked(); + + // Check that a retry was scheduled + let events = retry_events.lock(); + assert_eq!(events.len(), 1, "Should have scheduled one retry"); + assert_eq!(events[0].0, 1, "Should be first retry attempt"); + assert_eq!( + events[0].1, + Duration::from_secs(5), + "Should have 5 second delay for first retry" + ); + + // Check that a retry message was added + thread.read_with(cx, |thread, _| { + let mut messages = thread.messages(); + assert!( + messages.any(|msg| { + msg.role == Role::System + && msg.is_hidden + && msg.segments.iter().any(|seg| { + if let MessageSegment::Text(text) = seg { + text.contains("overloaded") + && text.contains("Retrying (attempt 1 of 3)") + } else { + false + } + }) + }), + "Should have added a system retry message" + ); + }); + } + + #[gpui::test] + async fn test_retry_on_internal_server_error(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; + + // Create model that returns internal server error + let model = Arc::new(ErrorInjector::new(TestError::InternalServerError)); + + // Insert a user message + thread.update(cx, |thread, cx| { + thread.insert_user_message("Hello!", ContextLoadResult::default(), None, vec![], cx); + }); + + // Track retry events + let retry_events = Arc::new(Mutex::new(Vec::new())); + let retry_events_clone = retry_events.clone(); + + let _subscription = thread.update(cx, |_, cx| { + cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| { + if let ThreadEvent::RetryScheduled { attempt, delay, .. } = event { + retry_events_clone.lock().push((*attempt, *delay)); + } + }) + }); + + // Start completion + thread.update(cx, |thread, cx| { + thread.send_to_model(model.clone(), CompletionIntent::UserPrompt, None, cx); + }); + + cx.run_until_parked(); + + // Check that a retry was scheduled + let events = retry_events.lock(); + assert_eq!(events.len(), 1, "Should have scheduled one retry"); + assert_eq!(events[0].0, 1, "Should be first retry attempt"); + assert_eq!( + events[0].1, + Duration::from_secs(5), + "Should have 5 second delay for first retry" + ); + + // Check that a retry message was added with provider name + thread.read_with(cx, |thread, _| { + let mut messages = thread.messages(); + assert!( + messages.any(|msg| { + msg.role == Role::System + && msg.is_hidden + && msg.segments.iter().any(|seg| { + if let MessageSegment::Text(text) = seg { + text.contains("internal server error") && + text.contains("Fake") && // Provider name from fake_provider.rs + text.contains("Retrying (attempt 1 of 3)") + } else { + false + } + }) + }), + "Should have added a system retry message with provider name" + ); + }); + } + + #[gpui::test] + async fn test_exponential_backoff_on_retries(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; + + // 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); + }); + + // Track retry events and completion count + let retry_events = Arc::new(Mutex::new(Vec::new())); + let retry_events_clone = retry_events.clone(); + let completion_count = Arc::new(Mutex::new(0)); + let completion_count_clone = completion_count.clone(); + + let _subscription = thread.update(cx, |_, cx| { + cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| match event { + ThreadEvent::RetryScheduled { attempt, delay, .. } => { + retry_events_clone.lock().push((*attempt, *delay)); + } + ThreadEvent::NewRequest => { + *completion_count_clone.lock() += 1; + } + _ => {} + }) + }); + + // First attempt + thread.update(cx, |thread, cx| { + thread.send_to_model(model.clone(), CompletionIntent::UserPrompt, None, cx); + }); + cx.run_until_parked(); + + // Should have 3 retries scheduled immediately due to errors + let mut total_delay = Duration::from_millis(0); + for i in 0..3 { + total_delay += Duration::from_secs(5 * 2u64.pow(i)); + cx.executor().advance_clock(total_delay); + cx.run_until_parked(); + } + + // Check exponential backoff + let events = retry_events.lock(); + assert_eq!(events.len(), 3, "Should have scheduled three retries"); + assert_eq!( + events[0], + (1, Duration::from_secs(5)), + "First retry: 5 seconds" + ); + assert_eq!( + events[1], + (2, Duration::from_secs(10)), + "Second retry: 10 seconds" + ); + assert_eq!( + events[2], + (3, Duration::from_secs(20)), + "Third retry: 20 seconds" + ); + } + + #[gpui::test] + async fn test_max_retries_exceeded(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; + + // 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); + }); + + // Track events + let error_shown = Arc::new(Mutex::new(false)); + let error_shown_clone = error_shown.clone(); + let retry_count = Arc::new(Mutex::new(0)); + let retry_count_clone = retry_count.clone(); + + let _subscription = thread.update(cx, |_, cx| { + cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| match event { + ThreadEvent::ShowError(_) => { + *error_shown_clone.lock() = true; + } + ThreadEvent::RetryScheduled { .. } => { + *retry_count_clone.lock() += 1; + } + _ => {} + }) + }); + + // Start initial completion + thread.update(cx, |thread, cx| { + thread.send_to_model(model.clone(), CompletionIntent::UserPrompt, None, cx); + }); + cx.run_until_parked(); + + // Advance through all retries + let delays = [5, 10, 20]; + for delay in delays { + cx.executor().advance_clock(Duration::from_secs(delay)); + cx.run_until_parked(); + } + + // After max retries (3), should show error + assert_eq!(*retry_count.lock(), 3, "Should have attempted 3 retries"); + assert!( + *error_shown.lock(), + "Should show error after max retries exceeded" + ); + + // Retry state should be cleared + thread.read_with(cx, |thread, _| { + assert!( + thread.retry_state.is_none(), + "Retry state should be cleared after max retries" + ); + }); + } + + #[gpui::test] + async fn test_retry_message_removed_on_retry(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; + + // We'll use a wrapper to switch behavior after first failure + struct RetryTestModel { + inner: Arc, + failed_once: Arc>, + } + + impl LanguageModel for RetryTestModel { + fn id(&self) -> LanguageModelId { + self.inner.id() + } + + fn name(&self) -> LanguageModelName { + self.inner.name() + } + + fn provider_id(&self) -> LanguageModelProviderId { + self.inner.provider_id() + } + + fn provider_name(&self) -> LanguageModelProviderName { + self.inner.provider_name() + } + + fn supports_tools(&self) -> bool { + self.inner.supports_tools() + } + + fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool { + self.inner.supports_tool_choice(choice) + } + + fn supports_images(&self) -> bool { + self.inner.supports_images() + } + + fn telemetry_id(&self) -> String { + self.inner.telemetry_id() + } + + fn max_token_count(&self) -> u64 { + self.inner.max_token_count() + } + + fn count_tokens( + &self, + request: LanguageModelRequest, + cx: &App, + ) -> BoxFuture<'static, Result> { + self.inner.count_tokens(request, cx) + } + + fn stream_completion( + &self, + request: LanguageModelRequest, + cx: &AsyncApp, + ) -> BoxFuture< + 'static, + Result< + BoxStream< + 'static, + Result, + >, + LanguageModelCompletionError, + >, + > { + if !*self.failed_once.lock() { + *self.failed_once.lock() = true; + // Return error on first attempt + let stream = futures::stream::once(async move { + Err(LanguageModelCompletionError::Other(anyhow::anyhow!( + "Language model provider's API is currently overloaded" + ))) + }); + async move { Ok(stream.boxed()) }.boxed() + } else { + // Succeed on retry + self.inner.stream_completion(request, cx) + } + } + + fn as_fake(&self) -> &FakeLanguageModel { + &self.inner + } + } + + let model = Arc::new(RetryTestModel { + inner: Arc::new(FakeLanguageModel::default()), + failed_once: Arc::new(Mutex::new(false)), + }); + + // Insert a user message + thread.update(cx, |thread, cx| { + thread.insert_user_message("Hello!", ContextLoadResult::default(), None, vec![], cx); + }); + + // Track message deletions + let deleted_messages = Arc::new(Mutex::new(Vec::new())); + let deleted_messages_clone = deleted_messages.clone(); + + let _subscription = thread.update(cx, |_, cx| { + cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| { + if let ThreadEvent::MessageDeleted(id) = event { + deleted_messages_clone.lock().push(*id); + } + }) + }); + + // Start completion + thread.update(cx, |thread, cx| { + thread.send_to_model(model.clone(), CompletionIntent::UserPrompt, None, cx); + }); + cx.run_until_parked(); + + // Get the retry message ID + let retry_message_id = thread.read_with(cx, |thread, _| { + thread + .messages() + .find(|msg| msg.role == Role::System && msg.is_hidden) + .map(|msg| msg.id) + .expect("Should have a retry message") + }); + + // Wait for retry + cx.executor().advance_clock(Duration::from_secs(5)); + cx.run_until_parked(); + + // Stream some successful content + let fake_model = model.as_fake(); + fake_model.stream_completion_response(&fake_model.pending_completions()[0], "Success!"); + fake_model.end_completion_stream(&fake_model.pending_completions()[0]); + cx.run_until_parked(); + + // Check that the retry message was deleted + assert!( + deleted_messages.lock().contains(&retry_message_id), + "Retry message should be deleted when retry starts" + ); + + // Message should no longer exist + thread.read_with(cx, |thread, _| { + assert!( + thread.message(retry_message_id).is_none(), + "Retry message should be removed from thread" + ); + }); + } + + #[gpui::test] + async fn test_successful_completion_clears_retry_state(cx: &mut TestAppContext) { + init_test_settings(cx); + + let project = create_test_project(cx, json!({})).await; + let (_, _, thread, _, model) = setup_test_environment(cx, project.clone()).await; + + // 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(); + + // Fail first attempt + let fake_model = model.as_fake(); + fake_model.end_completion_stream(&fake_model.pending_completions()[0]); + cx.executor().advance_clock(Duration::from_millis(100)); + cx.run_until_parked(); + + // Verify retry state exists + thread.read_with(cx, |thread, _| { + assert!( + thread.retry_state.is_some(), + "Should have retry state after failure" + ); + }); + + // Wait for retry + cx.executor().advance_clock(Duration::from_secs(5)); + cx.run_until_parked(); + + // Succeed on retry + fake_model.stream_completion_response(&fake_model.pending_completions()[0], "Hello there!"); + fake_model.end_completion_stream(&fake_model.pending_completions()[0]); + cx.run_until_parked(); + + // Retry state should be cleared on success + thread.read_with(cx, |thread, _| { + assert!( + thread.retry_state.is_none(), + "Retry state should be cleared on success" + ); + }); + } + fn test_summarize_error( model: &Arc, thread: &Entity, From 53b1868a39d0841ea57dc73472e76c64344b8cdb Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Mon, 23 Jun 2025 14:23:06 -0400 Subject: [PATCH 05/32] Use retry_after when available --- crates/agent/src/thread.rs | 243 +++++++++++++++++++++++++-- crates/agent_ui/src/active_thread.rs | 20 ++- 2 files changed, 244 insertions(+), 19 deletions(-) diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index 82b3bae61b4360..cac4cfd4be5899 100644 --- a/crates/agent/src/thread.rs +++ b/crates/agent/src/thread.rs @@ -1888,9 +1888,9 @@ impl Thread { provider_name.0.as_ref() ); - if !thread.handle_retryable_error_with_delay( + if !thread.handle_rate_limit_error( &error_message, - Some(*retry_after), + *retry_after, model.clone(), window, cx, @@ -2068,6 +2068,68 @@ impl Thread { }); } + fn handle_rate_limit_error( + &mut self, + error_message: &str, + retry_after: Duration, + model: Arc, + window: Option, + cx: &mut Context, + ) -> bool { + // For rate limit errors, we only retry once with the specified duration + // Clear any existing retry state since this is a special case + self.retry_state = None; + + let delay_secs = retry_after.as_secs(); + + // Add a transient message to inform the user (no attempt count for single retry) + let retry_message = format!("{}. Retrying in {} seconds...", error_message, delay_secs); + + let message_id = self.insert_message( + Role::System, + vec![MessageSegment::Text(retry_message)], + LoadedContext::default(), + Vec::new(), + true, // is_hidden - won't be persisted + cx, + ); + + cx.emit(ThreadEvent::RetryScheduled { + message_id, + delay: retry_after, + attempt: 1, + max_attempts: 1, + provider_name: model.provider_name(), + }); + + // Schedule the retry + let thread_handle = cx.entity().downgrade(); + let intent = self + .current_completion_intent + .clone() + .unwrap_or(CompletionIntent::UserPrompt); + + cx.spawn(async move |_thread, cx| { + cx.background_executor().timer(retry_after).await; + + thread_handle + .update(cx, |thread, cx| { + // Remove the retry notification message + if let Some(index) = thread.messages.iter().position(|m| m.id == message_id) { + thread.messages.remove(index); + cx.emit(ThreadEvent::MessageDeleted(message_id)); + } + + // Retry the completion + thread.send_to_model(model, intent, window, cx); + }) + .log_err(); + }) + .detach(); + + true + } + fn handle_retryable_error( &mut self, error_message: &str, @@ -3214,8 +3276,8 @@ mod tests { use icons::IconName; use language_model::fake_provider::{FakeLanguageModel, FakeLanguageModelProvider}; use language_model::{ - LanguageModelCompletionError, LanguageModelKnownError, LanguageModelName, - LanguageModelProviderId, LanguageModelToolChoice, + LanguageModelCompletionError, LanguageModelName, LanguageModelProviderId, + LanguageModelToolChoice, }; use parking_lot::Mutex; use project::{FakeFs, Project}; @@ -4078,13 +4140,13 @@ fn main() {{ >, > { let error = match self.error_type { - TestError::Overloaded => LanguageModelKnownError::Overloaded, - TestError::InternalServerError => LanguageModelKnownError::ApiInternalServerError, + TestError::Overloaded => LanguageModelCompletionError::Overloaded, + TestError::InternalServerError => { + LanguageModelCompletionError::ApiInternalServerError + } }; async move { - let stream = futures::stream::once(async move { - Err(LanguageModelCompletionError::Other(anyhow::anyhow!(error))) - }); + let stream = futures::stream::once(async move { Err(error) }); Ok(stream.boxed()) } .boxed() @@ -4431,9 +4493,7 @@ fn main() {{ *self.failed_once.lock() = true; // Return error on first attempt let stream = futures::stream::once(async move { - Err(LanguageModelCompletionError::Other(anyhow::anyhow!( - "Language model provider's API is currently overloaded" - ))) + Err(LanguageModelCompletionError::Overloaded) }); async move { Ok(stream.boxed()) }.boxed() } else { @@ -4559,6 +4619,165 @@ fn main() {{ }); } + #[gpui::test] + async fn test_rate_limit_retry_single_attempt(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; + + // Create a model that returns rate limit error with retry_after + struct RateLimitModel { + inner: Arc, + } + + impl LanguageModel for RateLimitModel { + fn id(&self) -> LanguageModelId { + self.inner.id() + } + + fn name(&self) -> LanguageModelName { + self.inner.name() + } + + fn provider_id(&self) -> LanguageModelProviderId { + self.inner.provider_id() + } + + fn provider_name(&self) -> LanguageModelProviderName { + self.inner.provider_name() + } + + fn supports_tools(&self) -> bool { + self.inner.supports_tools() + } + + fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool { + self.inner.supports_tool_choice(choice) + } + + fn supports_images(&self) -> bool { + self.inner.supports_images() + } + + fn telemetry_id(&self) -> String { + self.inner.telemetry_id() + } + + fn max_token_count(&self) -> u64 { + self.inner.max_token_count() + } + + fn count_tokens( + &self, + request: LanguageModelRequest, + cx: &App, + ) -> BoxFuture<'static, Result> { + self.inner.count_tokens(request, cx) + } + + fn stream_completion( + &self, + _request: LanguageModelRequest, + _cx: &AsyncApp, + ) -> BoxFuture< + 'static, + Result< + BoxStream< + 'static, + Result, + >, + LanguageModelCompletionError, + >, + > { + async move { + let stream = futures::stream::once(async move { + Err(LanguageModelCompletionError::RateLimitExceeded { + retry_after: Duration::from_secs(30), + }) + }); + Ok(stream.boxed()) + } + .boxed() + } + + fn as_fake(&self) -> &FakeLanguageModel { + &self.inner + } + } + + let model = Arc::new(RateLimitModel { + inner: Arc::new(FakeLanguageModel::default()), + }); + + // Insert a user message + thread.update(cx, |thread, cx| { + thread.insert_user_message("Hello!", ContextLoadResult::default(), None, vec![], cx); + }); + + // Track retry events + let retry_events = Arc::new(Mutex::new(Vec::new())); + let retry_events_clone = retry_events.clone(); + + let _subscription = thread.update(cx, |_, cx| { + cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| { + if let ThreadEvent::RetryScheduled { + attempt, + delay, + max_attempts, + .. + } = event + { + retry_events_clone + .lock() + .push((*attempt, *delay, *max_attempts)); + } + }) + }); + + // Start completion + thread.update(cx, |thread, cx| { + thread.send_to_model(model.clone(), CompletionIntent::UserPrompt, None, cx); + }); + + cx.run_until_parked(); + + // Check that a retry was scheduled with the specified delay + let events = retry_events.lock(); + assert_eq!(events.len(), 1, "Should have scheduled exactly one retry"); + assert_eq!(events[0].0, 1, "Should be first attempt"); + assert_eq!( + events[0].1, + Duration::from_secs(30), + "Should use the retry_after duration" + ); + assert_eq!( + events[0].2, 1, + "Should only have max 1 attempt for rate limit" + ); + + // Check that retry message doesn't include attempt count + thread.read_with(cx, |thread, _| { + let mut messages = thread.messages(); + assert!( + messages.any(|msg| { + msg.role == Role::System + && msg.is_hidden + && msg.segments.iter().any(|seg| { + if let MessageSegment::Text(text) = seg { + text.contains("rate limit exceeded") + && text.contains("Retrying in 30 seconds...") + && !text.contains("attempt") + } else { + false + } + }) + }), + "Should have added a system retry message without attempt count" + ); + }); + } + fn test_summarize_error( model: &Arc, thread: &Entity, diff --git a/crates/agent_ui/src/active_thread.rs b/crates/agent_ui/src/active_thread.rs index f718c6549091af..ab1a6b143e03ac 100644 --- a/crates/agent_ui/src/active_thread.rs +++ b/crates/agent_ui/src/active_thread.rs @@ -1160,13 +1160,19 @@ impl ActiveThread { } // Show a notification to the user - let notification_text = format!( - "{} retry in {}s (attempt {}/{})", - provider_name.0.as_ref(), - delay.as_secs(), - attempt, - max_attempts - ); + let notification_text = if *max_attempts == 1 { + // For single retry (e.g., rate limit with retry_after) + format!("{} retry in {}s", provider_name.0.as_ref(), delay.as_secs()) + } else { + // For multiple retry attempts + format!( + "{} retry in {}s (attempt {}/{})", + provider_name.0.as_ref(), + delay.as_secs(), + attempt, + max_attempts + ) + }; self.show_notification(¬ification_text, IconName::ArrowCircle, window, cx); cx.notify(); } From 6df2c1132f6cd10349e2b69733a4c7ce016c3bab Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Mon, 23 Jun 2025 14:23:40 -0400 Subject: [PATCH 06/32] Make the git diff a little cleaner --- crates/agent_ui/src/agent_diff.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/agent_ui/src/agent_diff.rs b/crates/agent_ui/src/agent_diff.rs index 15e0b7726cd900..d1255ea3dfc80b 100644 --- a/crates/agent_ui/src/agent_diff.rs +++ b/crates/agent_ui/src/agent_diff.rs @@ -1380,8 +1380,8 @@ impl AgentDiff { | ThreadEvent::ToolConfirmationNeeded | ThreadEvent::ToolUseLimitReached | ThreadEvent::CancelEditing - | ThreadEvent::ProfileChanged - | ThreadEvent::RetryScheduled { .. } => {} + | ThreadEvent::RetryScheduled { .. } + | ThreadEvent::ProfileChanged => {} } } From e5a5048d1ee90ced502d6a80e1a0f22e022f8870 Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Mon, 23 Jun 2025 14:25:49 -0400 Subject: [PATCH 07/32] We only need u8 retries, not u32 --- crates/agent/src/thread.rs | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index cac4cfd4be5899..2b0ebe3929568d 100644 --- a/crates/agent/src/thread.rs +++ b/crates/agent/src/thread.rs @@ -384,8 +384,8 @@ pub struct Thread { #[derive(Clone, Debug)] struct RetryState { - attempt: u32, - max_attempts: u32, + attempt: u8, + max_attempts: u8, } #[derive(Clone, Debug, PartialEq, Eq)] @@ -2148,7 +2148,7 @@ impl Thread { window: Option, cx: &mut Context, ) -> bool { - const MAX_RETRY_ATTEMPTS: u32 = 3; + const MAX_RETRY_ATTEMPTS: u8 = 3; const BASE_RETRY_DELAY_SECS: u64 = 5; // Check if we should create or update retry state @@ -2172,7 +2172,8 @@ impl Thread { let delay = if let Some(custom_delay) = custom_delay { custom_delay } else { - let delay_secs = BASE_RETRY_DELAY_SECS * 2u64.pow(attempt - 1); + let delay_secs = + BASE_RETRY_DELAY_SECS * 2u64.pow((attempt.saturating_sub(1)) as u32); Duration::from_secs(delay_secs) }; let delay_secs = delay.as_secs(); @@ -3168,8 +3169,8 @@ pub enum ThreadEvent { RetryScheduled { message_id: MessageId, delay: Duration, - attempt: u32, - max_attempts: u32, + attempt: u8, + max_attempts: u8, provider_name: LanguageModelProviderName, }, } @@ -4329,7 +4330,7 @@ fn main() {{ // Should have 3 retries scheduled immediately due to errors let mut total_delay = Duration::from_millis(0); - for i in 0..3 { + for i in 0..3u32 { total_delay += Duration::from_secs(5 * 2u64.pow(i)); cx.executor().advance_clock(total_delay); cx.run_until_parked(); @@ -4395,7 +4396,7 @@ fn main() {{ cx.run_until_parked(); // Advance through all retries - let delays = [5, 10, 20]; + let delays = [5u64, 10, 20]; for delay in delays { cx.executor().advance_clock(Duration::from_secs(delay)); cx.run_until_parked(); From 4edb91437018b0677c6d8baef4befc37030588bd Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Mon, 23 Jun 2025 14:27:33 -0400 Subject: [PATCH 08/32] Use string interpolation more --- crates/agent_ui/src/active_thread.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/crates/agent_ui/src/active_thread.rs b/crates/agent_ui/src/active_thread.rs index ab1a6b143e03ac..e0e1c1442e15ad 100644 --- a/crates/agent_ui/src/active_thread.rs +++ b/crates/agent_ui/src/active_thread.rs @@ -1166,11 +1166,9 @@ impl ActiveThread { } else { // For multiple retry attempts format!( - "{} retry in {}s (attempt {}/{})", + "{} retry in {}s (attempt {attempt}/{max_attempts})", provider_name.0.as_ref(), delay.as_secs(), - attempt, - max_attempts ) }; self.show_notification(¬ification_text, IconName::ArrowCircle, window, cx); From 75aecb0ec3c4f868be0697d6907dd3cc91749030 Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Mon, 23 Jun 2025 14:39:34 -0400 Subject: [PATCH 09/32] Introduce ui-only messages for showing retries --- crates/agent/src/thread.rs | 151 +++++++++++++++++++++++++++++++++---- 1 file changed, 135 insertions(+), 16 deletions(-) diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index 2b0ebe3929568d..860dc7187dc878 100644 --- a/crates/agent/src/thread.rs +++ b/crates/agent/src/thread.rs @@ -123,6 +123,7 @@ pub struct Message { pub loaded_context: LoadedContext, pub creases: Vec, pub is_hidden: bool, + pub ui_only: bool, } impl Message { @@ -574,6 +575,7 @@ impl Thread { }) .collect(), is_hidden: message.is_hidden, + ui_only: false, // UI-only messages are not persisted }) .collect(), next_message_id, @@ -1058,6 +1060,7 @@ impl Thread { loaded_context, creases, is_hidden, + ui_only: false, }); self.touch_updated_at(); cx.emit(ThreadEvent::MessageAdded(id)); @@ -1147,6 +1150,7 @@ impl Thread { updated_at: this.updated_at(), messages: this .messages() + .filter(|message| !message.ui_only) .map(|message| SerializedMessage { id: message.id, role: message.role, @@ -1323,6 +1327,11 @@ impl Thread { let mut message_ix_to_cache = None; for message in &self.messages { + // Skip UI-only messages - they're for UI only, not for the model + if message.ui_only { + continue; + } + let mut request_message = LanguageModelRequestMessage { role: message.role, content: Vec::new(), @@ -2085,14 +2094,19 @@ impl Thread { // Add a transient message to inform the user (no attempt count for single retry) let retry_message = format!("{}. Retrying in {} seconds...", error_message, delay_secs); - let message_id = self.insert_message( - Role::System, - vec![MessageSegment::Text(retry_message)], - LoadedContext::default(), - Vec::new(), - true, // is_hidden - won't be persisted - cx, - ); + // Add a UI-only message instead of a regular message + let id = self.next_message_id.post_inc(); + self.messages.push(Message { + id, + role: Role::System, + segments: vec![MessageSegment::Text(retry_message)], + loaded_context: LoadedContext::default(), + creases: Vec::new(), + is_hidden: true, + ui_only: true, + }); + cx.emit(ThreadEvent::MessageAdded(id)); + let message_id = id; cx.emit(ThreadEvent::RetryScheduled { message_id, @@ -2184,14 +2198,19 @@ impl Thread { error_message, attempt, max_attempts, delay_secs ); - let message_id = self.insert_message( - Role::System, - vec![MessageSegment::Text(retry_message)], - LoadedContext::default(), - Vec::new(), - true, // is_hidden - won't be persisted - cx, - ); + // Add a UI-only message instead of a regular message + let id = self.next_message_id.post_inc(); + self.messages.push(Message { + id, + role: Role::System, + segments: vec![MessageSegment::Text(retry_message)], + loaded_context: LoadedContext::default(), + creases: Vec::new(), + is_hidden: true, + ui_only: true, + }); + cx.emit(ThreadEvent::MessageAdded(id)); + let message_id = id; cx.emit(ThreadEvent::RetryScheduled { message_id, @@ -4779,6 +4798,106 @@ fn main() {{ }); } + #[gpui::test] + async fn test_ui_only_messages_not_sent_to_model(cx: &mut TestAppContext) { + init_test_settings(cx); + + let project = create_test_project(cx, json!({})).await; + let (_, _, thread, _, model) = setup_test_environment(cx, project.clone()).await; + + // Insert a regular user message + thread.update(cx, |thread, cx| { + thread.insert_user_message("Hello!", ContextLoadResult::default(), None, vec![], cx); + }); + + // Insert a UI-only message (like our retry notifications) + thread.update(cx, |thread, cx| { + let id = thread.next_message_id.post_inc(); + thread.messages.push(Message { + id, + role: Role::System, + segments: vec![MessageSegment::Text( + "This is a UI-only message that should not be sent to the model".to_string(), + )], + loaded_context: LoadedContext::default(), + creases: Vec::new(), + is_hidden: true, + ui_only: true, + }); + cx.emit(ThreadEvent::MessageAdded(id)); + }); + + // Insert another regular message + thread.update(cx, |thread, cx| { + thread.insert_user_message( + "How are you?", + ContextLoadResult::default(), + None, + vec![], + cx, + ); + }); + + // Generate the completion request + let request = thread.update(cx, |thread, cx| { + thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx) + }); + + // Verify that the request only contains non-UI-only messages + // Should have system prompt + 2 user messages, but not the UI-only message + let user_messages: Vec<_> = request + .messages + .iter() + .filter(|msg| msg.role == Role::User) + .collect(); + assert_eq!( + user_messages.len(), + 2, + "Should have exactly 2 user messages" + ); + + // Verify the UI-only content is not present anywhere in the request + let request_text = request + .messages + .iter() + .flat_map(|msg| &msg.content) + .filter_map(|content| match content { + MessageContent::Text(text) => Some(text.as_str()), + _ => None, + }) + .collect::(); + + assert!( + !request_text.contains("UI-only message"), + "UI-only message content should not be in the request" + ); + + // Verify the thread still has all 3 messages (including UI-only) + thread.read_with(cx, |thread, _| { + assert_eq!( + thread.messages().count(), + 3, + "Thread should have 3 messages" + ); + assert_eq!( + thread.messages().filter(|m| m.ui_only).count(), + 1, + "Thread should have 1 UI-only message" + ); + }); + + // Verify that UI-only messages are not serialized + let serialized = thread + .update(cx, |thread, cx| thread.serialize(cx)) + .await + .unwrap(); + assert_eq!( + serialized.messages.len(), + 2, + "Serialized thread should only have 2 messages (no UI-only)" + ); + } + fn test_summarize_error( model: &Arc, thread: &Entity, From f3adf53923dfca80aac30961073e06653f05d77b Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Mon, 23 Jun 2025 14:43:56 -0400 Subject: [PATCH 10/32] Drop some unnecessary comments --- crates/agent_ui/src/active_thread.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/crates/agent_ui/src/active_thread.rs b/crates/agent_ui/src/active_thread.rs index e0e1c1442e15ad..9bd2051af9a1aa 100644 --- a/crates/agent_ui/src/active_thread.rs +++ b/crates/agent_ui/src/active_thread.rs @@ -1161,10 +1161,8 @@ impl ActiveThread { // Show a notification to the user let notification_text = if *max_attempts == 1 { - // For single retry (e.g., rate limit with retry_after) format!("{} retry in {}s", provider_name.0.as_ref(), delay.as_secs()) } else { - // For multiple retry attempts format!( "{} retry in {}s (attempt {attempt}/{max_attempts})", provider_name.0.as_ref(), From 59bc8b042aca40ada13197d336e81f53cd69d84f Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Mon, 23 Jun 2025 14:45:19 -0400 Subject: [PATCH 11/32] Remove an unnecessary clone --- crates/agent/src/thread.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index 860dc7187dc878..cf8d81f68f3b9f 100644 --- a/crates/agent/src/thread.rs +++ b/crates/agent/src/thread.rs @@ -1248,7 +1248,7 @@ impl Thread { self.remaining_turns -= 1; // Store the intent for potential retries - self.current_completion_intent = Some(intent.clone()); + self.current_completion_intent = Some(intent); let request = self.to_completion_request(model.clone(), intent, cx); From ba2abc01feb3d1b9ebf35c8e6925f2d9b8855a0c Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Mon, 23 Jun 2025 14:45:54 -0400 Subject: [PATCH 12/32] Update a comment --- crates/agent/src/thread.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index cf8d81f68f3b9f..822f7a2d74cb96 100644 --- a/crates/agent/src/thread.rs +++ b/crates/agent/src/thread.rs @@ -1327,7 +1327,7 @@ impl Thread { let mut message_ix_to_cache = None; for message in &self.messages { - // Skip UI-only messages - they're for UI only, not for the model + // ui_only messages are for the UI only, not for the model if message.ui_only { continue; } From 3aae4277c39880a404f5cc137e69346b6c53d9ff Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Mon, 23 Jun 2025 14:47:23 -0400 Subject: [PATCH 13/32] Minor formatting cleanups --- crates/agent/src/thread.rs | 21 ++++++--------------- 1 file changed, 6 insertions(+), 15 deletions(-) diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index 822f7a2d74cb96..6c1681f841c314 100644 --- a/crates/agent/src/thread.rs +++ b/crates/agent/src/thread.rs @@ -1891,11 +1891,8 @@ impl Thread { cx.notify(); } LanguageModelKnownError::RateLimitExceeded { retry_after } => { - let provider_name = model.provider_name(); - let error_message = format!( - "{}'s API rate limit exceeded", - provider_name.0.as_ref() - ); + let provider = model.provider_name().0.as_ref(); + let error_message = format!("{provider}'s API rate limit exceeded"); if !thread.handle_rate_limit_error( &error_message, @@ -1908,11 +1905,8 @@ impl Thread { } } LanguageModelKnownError::Overloaded => { - let provider_name = model.provider_name(); - let error_message = format!( - "{}'s API servers are overloaded right now", - provider_name.0.as_ref() - ); + let provider_name = model.provider_name().0.as_ref(); + let error_message = format!("{provider}'s API servers are overloaded right now"); if !thread.handle_retryable_error( &error_message, @@ -1924,11 +1918,8 @@ impl Thread { } } LanguageModelKnownError::ApiInternalServerError => { - let provider_name = model.provider_name(); - let error_message = format!( - "{}'s API server reported an internal server error", - provider_name.0.as_ref() - ); + let provider = model.provider_name().0.as_ref(); + let error_message = format!("{provider}'s API server reported an internal server error"); if !thread.handle_retryable_error( &error_message, From de098616269f7adbf78baa4486949ae56aae91df Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Mon, 23 Jun 2025 15:32:41 -0400 Subject: [PATCH 14/32] Drop redundant retry state reset --- crates/agent/src/thread.rs | 25 +++++++++++++++---------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index 6c1681f841c314..3a6be34ec97b7a 100644 --- a/crates/agent/src/thread.rs +++ b/crates/agent/src/thread.rs @@ -1891,8 +1891,11 @@ impl Thread { cx.notify(); } LanguageModelKnownError::RateLimitExceeded { retry_after } => { - let provider = model.provider_name().0.as_ref(); - let error_message = format!("{provider}'s API rate limit exceeded"); + let provider_name = model.provider_name(); + let error_message = format!( + "{}'s API rate limit exceeded", + provider_name.0.as_ref() + ); if !thread.handle_rate_limit_error( &error_message, @@ -1905,8 +1908,11 @@ impl Thread { } } LanguageModelKnownError::Overloaded => { - let provider_name = model.provider_name().0.as_ref(); - let error_message = format!("{provider}'s API servers are overloaded right now"); + let provider_name = model.provider_name(); + let error_message = format!( + "{}'s API servers are overloaded right now", + provider_name.0.as_ref() + ); if !thread.handle_retryable_error( &error_message, @@ -1918,8 +1924,11 @@ impl Thread { } } LanguageModelKnownError::ApiInternalServerError => { - let provider = model.provider_name().0.as_ref(); - let error_message = format!("{provider}'s API server reported an internal server error"); + let provider_name = model.provider_name(); + let error_message = format!( + "{}'s API server reported an internal server error", + provider_name.0.as_ref() + ); if !thread.handle_retryable_error( &error_message, @@ -2076,10 +2085,6 @@ impl Thread { window: Option, cx: &mut Context, ) -> bool { - // For rate limit errors, we only retry once with the specified duration - // Clear any existing retry state since this is a special case - self.retry_state = None; - let delay_secs = retry_after.as_secs(); // Add a transient message to inform the user (no attempt count for single retry) From 31d292db2565d98a5ef9dea569c03f58461db1e6 Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Mon, 23 Jun 2025 15:56:45 -0400 Subject: [PATCH 15/32] Cancel pending retries when appropriate --- crates/agent/src/thread.rs | 104 +++++++++++++++++++++++++++++++------ 1 file changed, 87 insertions(+), 17 deletions(-) diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index 3a6be34ec97b7a..246d60a5dc1a24 100644 --- a/crates/agent/src/thread.rs +++ b/crates/agent/src/thread.rs @@ -1897,15 +1897,13 @@ impl Thread { provider_name.0.as_ref() ); - if !thread.handle_rate_limit_error( + thread.handle_rate_limit_error( &error_message, *retry_after, model.clone(), window, cx, - ) { - emit_generic_error(error, cx); - } + ); } LanguageModelKnownError::Overloaded => { let provider_name = model.provider_name(); @@ -2084,11 +2082,12 @@ impl Thread { model: Arc, window: Option, cx: &mut Context, - ) -> bool { - let delay_secs = retry_after.as_secs(); - - // Add a transient message to inform the user (no attempt count for single retry) - let retry_message = format!("{}. Retrying in {} seconds...", error_message, delay_secs); + ) { + // For rate limit errors, we only retry once with the specified duration + let retry_message = format!( + "{error_message}. Retrying in {} seconds…", + retry_after.as_secs() + ); // Add a UI-only message instead of a regular message let id = self.next_message_id.post_inc(); @@ -2136,8 +2135,6 @@ impl Thread { .log_err(); }) .detach(); - - true } fn handle_retryable_error( @@ -2182,13 +2179,12 @@ impl Thread { let delay = if let Some(custom_delay) = custom_delay { custom_delay } else { - let delay_secs = - BASE_RETRY_DELAY_SECS * 2u64.pow((attempt.saturating_sub(1)) as u32); + let delay_secs = BASE_RETRY_DELAY_SECS * 2u64.pow((attempt - 1) as u32); Duration::from_secs(delay_secs) }; - let delay_secs = delay.as_secs(); // Add a transient message to inform the user + let delay_secs = delay.as_secs(); let retry_message = format!( "{}. Retrying (attempt {} of {}) in {} seconds...", error_message, attempt, max_attempts, delay_secs @@ -2594,6 +2590,11 @@ impl Thread { ) -> bool { let mut canceled = self.pending_completions.pop().is_some(); + // Cancel any pending retry by clearing retry state + if self.retry_state.take().is_some() { + canceled = true; + } + for pending_tool_use in self.tool_use.cancel_pending() { canceled = true; self.tool_finished( @@ -4566,8 +4567,14 @@ fn main() {{ // Stream some successful content let fake_model = model.as_fake(); - fake_model.stream_completion_response(&fake_model.pending_completions()[0], "Success!"); - fake_model.end_completion_stream(&fake_model.pending_completions()[0]); + // After the retry, there should be a new pending completion + let pending = fake_model.pending_completions(); + assert!( + !pending.is_empty(), + "Should have a pending completion after retry" + ); + fake_model.stream_completion_response(&pending[0], "Success!"); + fake_model.end_completion_stream(&pending[0]); cx.run_until_parked(); // Check that the retry message was deleted @@ -4782,7 +4789,7 @@ fn main() {{ && msg.segments.iter().any(|seg| { if let MessageSegment::Text(text) = seg { text.contains("rate limit exceeded") - && text.contains("Retrying in 30 seconds...") + && text.contains("Retrying in 30 seconds…") && !text.contains("attempt") } else { false @@ -4894,6 +4901,69 @@ fn main() {{ ); } + #[gpui::test] + async fn test_retry_cancelled_on_stop(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; + + // 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); + }); + + // Track retry events + let retry_scheduled = Arc::new(Mutex::new(false)); + let retry_scheduled_clone = retry_scheduled.clone(); + + let _subscription = thread.update(cx, |_, cx| { + cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| { + if matches!(event, ThreadEvent::RetryScheduled { .. }) { + *retry_scheduled_clone.lock() = true; + } + }) + }); + + // Start completion + thread.update(cx, |thread, cx| { + thread.send_to_model(model.clone(), CompletionIntent::UserPrompt, None, cx); + }); + + cx.run_until_parked(); + + // Verify retry was scheduled + assert!(*retry_scheduled.lock(), "Should have scheduled a retry"); + + // Cancel the completion before the retry happens + thread.update(cx, |thread, cx| { + thread.cancel_last_completion(None, cx); + }); + + // Advance time past the retry delay + cx.executor().advance_clock(Duration::from_secs(10)); + cx.run_until_parked(); + + // The retry should not have happened - no pending completions + let fake_model = model.as_fake(); + assert_eq!( + fake_model.pending_completions().len(), + 0, + "Should have no pending completions after cancellation" + ); + + // Verify the retry was cancelled by checking retry state + thread.read_with(cx, |thread, _| { + assert!( + thread.retry_state.is_none(), + "retry_state should be cleared after cancellation" + ); + }); + } + fn test_summarize_error( model: &Arc, thread: &Entity, From d24644a10c626864fc6d535a1e4b8a08ff5f25e2 Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Mon, 23 Jun 2025 15:59:03 -0400 Subject: [PATCH 16/32] Drop redundant comments --- crates/agent_ui/src/active_thread.rs | 3 --- 1 file changed, 3 deletions(-) diff --git a/crates/agent_ui/src/active_thread.rs b/crates/agent_ui/src/active_thread.rs index 9bd2051af9a1aa..ba8ff1cb794f0b 100644 --- a/crates/agent_ui/src/active_thread.rs +++ b/crates/agent_ui/src/active_thread.rs @@ -1147,8 +1147,6 @@ impl ActiveThread { max_attempts, provider_name, } => { - // The retry message is already added to the thread as a system message - // We just need to ensure the UI updates let message_segments = self .thread .read(cx) @@ -1159,7 +1157,6 @@ impl ActiveThread { self.push_message(message_id, &message_segments, window, cx); } - // Show a notification to the user let notification_text = if *max_attempts == 1 { format!("{} retry in {}s", provider_name.0.as_ref(), delay.as_secs()) } else { From cdb2a34099bfb2acb2691b6cc0de08ad6b575bc9 Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Tue, 24 Jun 2025 13:00:01 -0400 Subject: [PATCH 17/32] Simplify retry_state management --- crates/agent/src/thread.rs | 91 ++++++++++++++++++-------------------- 1 file changed, 43 insertions(+), 48 deletions(-) diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index 246d60a5dc1a24..65f5df056b5632 100644 --- a/crates/agent/src/thread.rs +++ b/crates/agent/src/thread.rs @@ -1806,53 +1806,58 @@ impl Thread { .update(cx, |thread, cx| { thread.finalize_pending_checkpoint(cx); match result.as_ref() { - Ok(stop_reason) => match stop_reason { - StopReason::ToolUse => { - let tool_uses = thread.use_pending_tools(window, cx, model.clone()); - cx.emit(ThreadEvent::UsePendingTools { tool_uses }); - } - StopReason::EndTurn | StopReason::MaxTokens => { - thread.project.update(cx, |project, cx| { - project.set_agent_location(None, cx); - }); - } - StopReason::Refusal => { - thread.project.update(cx, |project, cx| { - project.set_agent_location(None, cx); - }); - - // Remove the turn that was refused. - // - // https://docs.anthropic.com/en/docs/test-and-evaluate/strengthen-guardrails/handle-streaming-refusals#reset-context-after-refusal - { - let mut messages_to_remove = Vec::new(); + Ok(stop_reason) => { + match stop_reason { + StopReason::ToolUse => { + let tool_uses = thread.use_pending_tools(window, cx, model.clone()); + cx.emit(ThreadEvent::UsePendingTools { tool_uses }); + } + StopReason::EndTurn | StopReason::MaxTokens => { + thread.project.update(cx, |project, cx| { + project.set_agent_location(None, cx); + }); + } + StopReason::Refusal => { + thread.project.update(cx, |project, cx| { + project.set_agent_location(None, cx); + }); - for (ix, message) in thread.messages.iter().enumerate().rev() { - messages_to_remove.push(message.id); + // Remove the turn that was refused. + // + // https://docs.anthropic.com/en/docs/test-and-evaluate/strengthen-guardrails/handle-streaming-refusals#reset-context-after-refusal + { + let mut messages_to_remove = Vec::new(); - if message.role == Role::User { - if ix == 0 { - break; - } + for (ix, message) in thread.messages.iter().enumerate().rev() { + messages_to_remove.push(message.id); - if let Some(prev_message) = thread.messages.get(ix - 1) { - if prev_message.role == Role::Assistant { + if message.role == Role::User { + if ix == 0 { break; } + + if let Some(prev_message) = thread.messages.get(ix - 1) { + if prev_message.role == Role::Assistant { + break; + } + } } } - } - for message_id in messages_to_remove { - thread.delete_message(message_id, cx); + for message_id in messages_to_remove { + thread.delete_message(message_id, cx); + } } - } - cx.emit(ThreadEvent::ShowError(ThreadError::Message { - header: "Language model refusal".into(), - message: "Model refused to generate content for safety reasons.".into(), - })); + cx.emit(ThreadEvent::ShowError(ThreadError::Message { + header: "Language model refusal".into(), + message: "Model refused to generate content for safety reasons.".into(), + })); + } } + + // We successfully completed, so cancel any remaining retries. + thread.retry_state = None; }, Err(error) => { thread.project.update(cx, |project, cx| { @@ -1948,17 +1953,10 @@ impl Thread { emit_generic_error(error, cx); } - // Clear retry state on any other error - thread.retry_state = None; thread.cancel_last_completion(window, cx); } } - // Clear retry state on successful completion - if result.is_ok() && thread.retry_state.is_some() { - thread.retry_state = None; - } - cx.emit(ThreadEvent::Stopped(result.map_err(Arc::new))); if let Some((request_callback, (request, response_events))) = thread @@ -2588,12 +2586,9 @@ impl Thread { window: Option, cx: &mut Context, ) -> bool { - let mut canceled = self.pending_completions.pop().is_some(); + let mut canceled = self.pending_completions.pop().is_some() || self.retry_state.is_some(); - // Cancel any pending retry by clearing retry state - if self.retry_state.take().is_some() { - canceled = true; - } + self.retry_state = None; for pending_tool_use in self.tool_use.cancel_pending() { canceled = true; From 1dad69b117ac8aacd40e91c06ec03e117c60bdbd Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Tue, 24 Jun 2025 14:32:03 -0400 Subject: [PATCH 18/32] Remove some segment cloning when rendering messages --- crates/agent_ui/src/active_thread.rs | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/crates/agent_ui/src/active_thread.rs b/crates/agent_ui/src/active_thread.rs index ba8ff1cb794f0b..7e90bbfb71377c 100644 --- a/crates/agent_ui/src/active_thread.rs +++ b/crates/agent_ui/src/active_thread.rs @@ -1147,14 +1147,16 @@ impl ActiveThread { max_attempts, provider_name, } => { - let message_segments = self - .thread - .read(cx) - .message(*message_id) - .map(|message| message.segments.clone()); - - if let Some(message_segments) = message_segments { - self.push_message(message_id, &message_segments, window, cx); + if let Some(rendered_message) = self.thread.update(cx, |thread, cx| { + thread.message(*message_id).map(|message| { + RenderedMessage::from_segments( + &message.segments, + self.language_registry.clone(), + cx, + ) + }) + }) { + self.push_rendered_message(*message_id, rendered_message); } let notification_text = if *max_attempts == 1 { From ed859348056ec523319150efe707c41902576e00 Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Tue, 24 Jun 2025 15:31:27 -0400 Subject: [PATCH 19/32] Don't clone an owned String to make a SharedString --- crates/agent_ui/src/active_thread.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/agent_ui/src/active_thread.rs b/crates/agent_ui/src/active_thread.rs index 7e90bbfb71377c..8ddf3bab7318df 100644 --- a/crates/agent_ui/src/active_thread.rs +++ b/crates/agent_ui/src/active_thread.rs @@ -1168,7 +1168,7 @@ impl ActiveThread { delay.as_secs(), ) }; - self.show_notification(¬ification_text, IconName::ArrowCircle, window, cx); + self.show_notification(notification_text, IconName::ArrowCircle, window, cx); cx.notify(); } } From e369f0c9fca0269498733c1f0728c7b0a6a4b061 Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Tue, 24 Jun 2025 16:12:03 -0400 Subject: [PATCH 20/32] Pass around intent rather than storing it --- Cargo.lock | 1 + crates/agent/Cargo.toml | 1 + crates/agent/src/thread.rs | 65 ++++++++++++---------------- crates/agent_ui/src/active_thread.rs | 8 +++- 4 files changed, 37 insertions(+), 38 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 0c832b83aa5983..735885d3362b6d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -77,6 +77,7 @@ dependencies = [ "language", "language_model", "log", + "parking_lot", "paths", "postage", "pretty_assertions", diff --git a/crates/agent/Cargo.toml b/crates/agent/Cargo.toml index f320e58d002f21..135363ab6552a9 100644 --- a/crates/agent/Cargo.toml +++ b/crates/agent/Cargo.toml @@ -72,6 +72,7 @@ gpui = { workspace = true, "features" = ["test-support"] } indoc.workspace = true language = { workspace = true, "features" = ["test-support"] } language_model = { workspace = true, "features" = ["test-support"] } +parking_lot.workspace = true pretty_assertions.workspace = true project = { workspace = true, features = ["test-support"] } workspace = { workspace = true, features = ["test-support"] } diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index 65f5df056b5632..6d12480b7941af 100644 --- a/crates/agent/src/thread.rs +++ b/crates/agent/src/thread.rs @@ -371,7 +371,6 @@ pub struct Thread { tool_use_limit_reached: bool, feedback: Option, retry_state: Option, - current_completion_intent: Option, message_feedback: HashMap, last_auto_capture_at: Option, last_received_chunk_at: Option, @@ -387,6 +386,7 @@ pub struct Thread { struct RetryState { attempt: u8, max_attempts: u8, + intent: CompletionIntent, } #[derive(Clone, Debug, PartialEq, Eq)] @@ -471,7 +471,6 @@ impl Thread { tool_use_limit_reached: false, feedback: None, retry_state: None, - current_completion_intent: None, message_feedback: HashMap::default(), last_auto_capture_at: None, last_received_chunk_at: None, @@ -539,7 +538,6 @@ impl Thread { detailed_summary_rx, completion_mode, retry_state: None, - current_completion_intent: None, messages: serialized .messages .into_iter() @@ -1247,12 +1245,9 @@ impl Thread { self.remaining_turns -= 1; - // Store the intent for potential retries - self.current_completion_intent = Some(intent); - let request = self.to_completion_request(model.clone(), intent, cx); - self.stream_completion(request, model, window, cx); + self.stream_completion(request, model, intent, window, cx); } pub fn used_tools_since_last_user_message(&self) -> bool { @@ -1486,6 +1481,7 @@ impl Thread { &mut self, request: LanguageModelRequest, model: Arc, + intent: CompletionIntent, window: Option, cx: &mut Context, ) { @@ -1906,6 +1902,7 @@ impl Thread { &error_message, *retry_after, model.clone(), + intent, window, cx, ); @@ -1920,6 +1917,7 @@ impl Thread { if !thread.handle_retryable_error( &error_message, model.clone(), + intent, window, cx, ) { @@ -1936,6 +1934,7 @@ impl Thread { if !thread.handle_retryable_error( &error_message, model.clone(), + intent, window, cx, ) { @@ -2078,6 +2077,7 @@ impl Thread { error_message: &str, retry_after: Duration, model: Arc, + intent: CompletionIntent, window: Option, cx: &mut Context, ) { @@ -2104,17 +2104,16 @@ impl Thread { cx.emit(ThreadEvent::RetryScheduled { message_id, delay: retry_after, - attempt: 1, - max_attempts: 1, + retry: RetryState { + attempt: 1, + max_attempts: 1, + intent, + }, provider_name: model.provider_name(), }); // Schedule the retry let thread_handle = cx.entity().downgrade(); - let intent = self - .current_completion_intent - .clone() - .unwrap_or(CompletionIntent::UserPrompt); cx.spawn(async move |_thread, cx| { cx.background_executor().timer(retry_after).await; @@ -2139,10 +2138,11 @@ impl Thread { &mut self, error_message: &str, model: Arc, + intent: CompletionIntent, window: Option, cx: &mut Context, ) -> bool { - self.handle_retryable_error_with_delay(error_message, None, model, window, cx) + self.handle_retryable_error_with_delay(error_message, None, model, intent, window, cx) } fn handle_retryable_error_with_delay( @@ -2150,28 +2150,25 @@ impl Thread { error_message: &str, custom_delay: Option, model: Arc, + intent: CompletionIntent, window: Option, cx: &mut Context, ) -> bool { const MAX_RETRY_ATTEMPTS: u8 = 3; const BASE_RETRY_DELAY_SECS: u64 = 5; - // Check if we should create or update retry state - let should_retry = if let Some(retry_state) = &mut self.retry_state { - retry_state.attempt += 1; - retry_state.attempt <= retry_state.max_attempts - } else { - self.retry_state = Some(RetryState { - attempt: 1, - max_attempts: MAX_RETRY_ATTEMPTS, - }); - true - }; + let retry_state = self.retry_state.get_or_insert(RetryState { + attempt: 0, // This will be immediately incremented + max_attempts: MAX_RETRY_ATTEMPTS, + intent, + }); + retry_state.attempt += 1; - if should_retry { + if retry_state.attempt <= retry_state.max_attempts { let retry_state = self.retry_state.as_ref().unwrap(); let attempt = retry_state.attempt; let max_attempts = retry_state.max_attempts; + let intent = retry_state.intent; // Use custom delay if provided (e.g., from rate limit), otherwise exponential backoff let delay = if let Some(custom_delay) = custom_delay { @@ -2205,17 +2202,12 @@ impl Thread { cx.emit(ThreadEvent::RetryScheduled { message_id, delay, - attempt, - max_attempts, + retry: retry_state, provider_name: model.provider_name(), }); // Schedule the retry let thread_handle = cx.entity().downgrade(); - let intent = self - .current_completion_intent - .clone() - .unwrap_or(CompletionIntent::UserPrompt); cx.spawn(async move |_thread, cx| { cx.background_executor().timer(delay).await; @@ -3180,8 +3172,7 @@ pub enum ThreadEvent { RetryScheduled { message_id: MessageId, delay: Duration, - attempt: u8, - max_attempts: u8, + retry: RetryState, provider_name: LanguageModelProviderName, }, } @@ -4210,7 +4201,7 @@ fn main() {{ assert_eq!( events[0].1, Duration::from_secs(5), - "Should have 5 second delay for first retry" + "Should have 5 second delay" ); // Check that a retry message was added @@ -4271,11 +4262,11 @@ fn main() {{ // Check that a retry was scheduled let events = retry_events.lock(); assert_eq!(events.len(), 1, "Should have scheduled one retry"); - assert_eq!(events[0].0, 1, "Should be first retry attempt"); + assert_eq!(events[0].0, 1, "Should be first attempt"); assert_eq!( events[0].1, Duration::from_secs(5), - "Should have 5 second delay for first retry" + "Should have 5 second delay" ); // Check that a retry message was added with provider name diff --git a/crates/agent_ui/src/active_thread.rs b/crates/agent_ui/src/active_thread.rs index 8ddf3bab7318df..98303f6f498c97 100644 --- a/crates/agent_ui/src/active_thread.rs +++ b/crates/agent_ui/src/active_thread.rs @@ -3796,7 +3796,13 @@ mod tests { thread.update(cx, |thread, cx| { let request = thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx); - thread.stream_completion(request, model, cx.active_window(), cx) + thread.stream_completion( + request, + model, + CompletionIntent::UserPrompt, + cx.active_window(), + cx, + ) }); // Follow the agent cx.update(|window, cx| { From 575a8d06ac6f0099e7b32ae0e607d37ca29bfe30 Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Tue, 24 Jun 2025 16:29:48 -0400 Subject: [PATCH 21/32] Simplify RetryScheduled --- crates/agent/src/thread.rs | 57 +++++++++++----------------- crates/agent_ui/src/active_thread.rs | 21 +--------- crates/eval/src/example.rs | 3 ++ 3 files changed, 27 insertions(+), 54 deletions(-) diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index 6d12480b7941af..bdc9a4fb2a8ade 100644 --- a/crates/agent/src/thread.rs +++ b/crates/agent/src/thread.rs @@ -21,13 +21,14 @@ use gpui::{ AnyWindowHandle, App, AppContext, AsyncApp, Context, Entity, EventEmitter, SharedString, Task, WeakEntity, Window, }; +use icons::IconName; use language_model::{ ConfiguredModel, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent, - LanguageModelId, LanguageModelKnownError, LanguageModelProviderName, LanguageModelRegistry, - LanguageModelRequest, LanguageModelRequestMessage, LanguageModelRequestTool, - LanguageModelToolResult, LanguageModelToolResultContent, LanguageModelToolUseId, - MessageContent, ModelRequestLimitReachedError, PaymentRequiredError, Role, SelectedModel, - StopReason, TokenUsage, + LanguageModelId, LanguageModelKnownError, LanguageModelRegistry, LanguageModelRequest, + LanguageModelRequestMessage, LanguageModelRequestTool, LanguageModelToolResult, + LanguageModelToolResultContent, LanguageModelToolUseId, MessageContent, + ModelRequestLimitReachedError, PaymentRequiredError, Role, SelectedModel, StopReason, + TokenUsage, }; use postage::stream::Stream as _; use project::{ @@ -2101,16 +2102,7 @@ impl Thread { cx.emit(ThreadEvent::MessageAdded(id)); let message_id = id; - cx.emit(ThreadEvent::RetryScheduled { - message_id, - delay: retry_after, - retry: RetryState { - attempt: 1, - max_attempts: 1, - intent, - }, - provider_name: model.provider_name(), - }); + cx.emit(ThreadEvent::RetryScheduled { message_id }); // Schedule the retry let thread_handle = cx.entity().downgrade(); @@ -2199,12 +2191,7 @@ impl Thread { cx.emit(ThreadEvent::MessageAdded(id)); let message_id = id; - cx.emit(ThreadEvent::RetryScheduled { - message_id, - delay, - retry: retry_state, - provider_name: model.provider_name(), - }); + cx.emit(ThreadEvent::RetryScheduled { message_id }); // Schedule the retry let thread_handle = cx.entity().downgrade(); @@ -2230,8 +2217,17 @@ impl Thread { true } else { - // Max retries exceeded, clear retry state + // Max retries exceeded; clear retry state and notify user that the thread stopped. self.retry_state = None; + + let notification_text = if retry_state.max_attempts == 1 { + format!("Failed after retrying."); + } else { + format!("Failed after retrying {} times.", retry_state.max_attempts); + }; + + self.show_notification(notification_text, IconName::ArrowCircle, window, cx); + false } } @@ -3171,9 +3167,6 @@ pub enum ThreadEvent { ProfileChanged, RetryScheduled { message_id: MessageId, - delay: Duration, - retry: RetryState, - provider_name: LanguageModelProviderName, }, } @@ -4181,7 +4174,7 @@ fn main() {{ let _subscription = thread.update(cx, |_, cx| { cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| { - if let ThreadEvent::RetryScheduled { attempt, delay, .. } = event { + if let ThreadEvent::RetryScheduled { .. } = event { retry_events_clone.lock().push((*attempt, *delay)); } }) @@ -4246,7 +4239,7 @@ fn main() {{ let _subscription = thread.update(cx, |_, cx| { cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| { - if let ThreadEvent::RetryScheduled { attempt, delay, .. } = event { + if let ThreadEvent::RetryScheduled { .. } = event { retry_events_clone.lock().push((*attempt, *delay)); } }) @@ -4314,7 +4307,7 @@ fn main() {{ let _subscription = thread.update(cx, |_, cx| { cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| match event { - ThreadEvent::RetryScheduled { attempt, delay, .. } => { + ThreadEvent::RetryScheduled { .. } => { retry_events_clone.lock().push((*attempt, *delay)); } ThreadEvent::NewRequest => { @@ -4730,13 +4723,7 @@ fn main() {{ let _subscription = thread.update(cx, |_, cx| { cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| { - if let ThreadEvent::RetryScheduled { - attempt, - delay, - max_attempts, - .. - } = event - { + if let ThreadEvent::RetryScheduled { .. } = event { retry_events_clone .lock() .push((*attempt, *delay, *max_attempts)); diff --git a/crates/agent_ui/src/active_thread.rs b/crates/agent_ui/src/active_thread.rs index 98303f6f498c97..5aacef9ebef862 100644 --- a/crates/agent_ui/src/active_thread.rs +++ b/crates/agent_ui/src/active_thread.rs @@ -1140,13 +1140,7 @@ impl ActiveThread { self.save_thread(cx); cx.notify(); } - ThreadEvent::RetryScheduled { - message_id, - delay, - attempt, - max_attempts, - provider_name, - } => { + ThreadEvent::RetryScheduled { message_id } => { if let Some(rendered_message) = self.thread.update(cx, |thread, cx| { thread.message(*message_id).map(|message| { RenderedMessage::from_segments( @@ -1157,19 +1151,8 @@ impl ActiveThread { }) }) { self.push_rendered_message(*message_id, rendered_message); + cx.notify(); } - - let notification_text = if *max_attempts == 1 { - format!("{} retry in {}s", provider_name.0.as_ref(), delay.as_secs()) - } else { - format!( - "{} retry in {}s (attempt {attempt}/{max_attempts})", - provider_name.0.as_ref(), - delay.as_secs(), - ) - }; - self.show_notification(notification_text, IconName::ArrowCircle, window, cx); - cx.notify(); } } } diff --git a/crates/eval/src/example.rs b/crates/eval/src/example.rs index 09770364cb6b46..fae7fd030873e8 100644 --- a/crates/eval/src/example.rs +++ b/crates/eval/src/example.rs @@ -221,6 +221,9 @@ impl ExampleContext { ThreadEvent::ShowError(thread_error) => { tx.try_send(Err(anyhow!(thread_error.clone()))).ok(); } + ThreadEvent::RetryScheduled { .. } => { + // Ignore retry events + } ThreadEvent::Stopped(reason) => match reason { Ok(StopReason::EndTurn) => { tx.close_channel(); From e05d5b07aedc8e0cf7e909712eac4b75b141c35a Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Tue, 24 Jun 2025 16:35:54 -0400 Subject: [PATCH 22/32] Fix when to show notification --- crates/agent/src/thread.rs | 17 ++++++++++++----- crates/agent_ui/src/active_thread.rs | 3 +++ crates/agent_ui/src/agent_diff.rs | 1 + 3 files changed, 16 insertions(+), 5 deletions(-) diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index bdc9a4fb2a8ade..bddab3f957dc74 100644 --- a/crates/agent/src/thread.rs +++ b/crates/agent/src/thread.rs @@ -21,7 +21,6 @@ use gpui::{ AnyWindowHandle, App, AppContext, AsyncApp, Context, Entity, EventEmitter, SharedString, Task, WeakEntity, Window, }; -use icons::IconName; use language_model::{ ConfiguredModel, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelId, LanguageModelKnownError, LanguageModelRegistry, LanguageModelRequest, @@ -2218,15 +2217,19 @@ impl Thread { true } else { // Max retries exceeded; clear retry state and notify user that the thread stopped. + let max_attempts = self.retry_state.as_ref().unwrap().max_attempts; self.retry_state = None; - let notification_text = if retry_state.max_attempts == 1 { - format!("Failed after retrying."); + let notification_text = if max_attempts == 1 { + "Failed after retrying.".into() } else { - format!("Failed after retrying {} times.", retry_state.max_attempts); + format!("Failed after retrying {} times.", max_attempts).into() }; - self.show_notification(notification_text, IconName::ArrowCircle, window, cx); + cx.emit(ThreadEvent::RetriesFailed { + message: notification_text, + max_attempts, + }); false } @@ -3168,6 +3171,10 @@ pub enum ThreadEvent { RetryScheduled { message_id: MessageId, }, + RetriesFailed { + message: SharedString, + max_attempts: u8, + }, } impl EventEmitter for Thread {} diff --git a/crates/agent_ui/src/active_thread.rs b/crates/agent_ui/src/active_thread.rs index 5aacef9ebef862..002a49e2a77a22 100644 --- a/crates/agent_ui/src/active_thread.rs +++ b/crates/agent_ui/src/active_thread.rs @@ -1154,6 +1154,9 @@ impl ActiveThread { cx.notify(); } } + ThreadEvent::RetriesFailed { message, .. } => { + self.show_notification(message.clone(), ui::IconName::Warning, window, cx); + } } } diff --git a/crates/agent_ui/src/agent_diff.rs b/crates/agent_ui/src/agent_diff.rs index d1255ea3dfc80b..2803fb53b5d6c3 100644 --- a/crates/agent_ui/src/agent_diff.rs +++ b/crates/agent_ui/src/agent_diff.rs @@ -1381,6 +1381,7 @@ impl AgentDiff { | ThreadEvent::ToolUseLimitReached | ThreadEvent::CancelEditing | ThreadEvent::RetryScheduled { .. } + | ThreadEvent::RetriesFailed { .. } => {} | ThreadEvent::ProfileChanged => {} } } From 42374954a0887d40c5402ff0344f1e7be9730779 Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Wed, 25 Jun 2025 14:39:32 -0400 Subject: [PATCH 23/32] Fix UI notifications for retries --- crates/agent/src/thread.rs | 58 ++--- crates/agent_ui/src/active_thread.rs | 342 +++++++++++++++------------ 2 files changed, 217 insertions(+), 183 deletions(-) diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index bddab3f957dc74..11368fba83be4c 100644 --- a/crates/agent/src/thread.rs +++ b/crates/agent/src/thread.rs @@ -1485,8 +1485,6 @@ impl Thread { window: Option, cx: &mut Context, ) { - // Clear any previous retry state when starting a new completion - self.retry_state = None; self.tool_use_limit_reached = false; let pending_completion_id = post_inc(&mut self.completion_count); @@ -1797,6 +1795,7 @@ impl Thread { }; let result = stream_completion.await; + let mut retry_scheduled = false; thread .update(cx, |thread, cx| { @@ -1906,6 +1905,7 @@ impl Thread { window, cx, ); + retry_scheduled = true; } LanguageModelKnownError::Overloaded => { let provider_name = model.provider_name(); @@ -1914,13 +1914,14 @@ impl Thread { provider_name.0.as_ref() ); - if !thread.handle_retryable_error( + retry_scheduled = thread.handle_retryable_error( &error_message, model.clone(), intent, window, cx, - ) { + ); + if !retry_scheduled { emit_generic_error(error, cx); } } @@ -1931,13 +1932,14 @@ impl Thread { provider_name.0.as_ref() ); - if !thread.handle_retryable_error( + retry_scheduled = thread.handle_retryable_error( &error_message, model.clone(), intent, window, cx, - ) { + ); + if !retry_scheduled { emit_generic_error(error, cx); } } @@ -1952,11 +1954,15 @@ impl Thread { emit_generic_error(error, cx); } - thread.cancel_last_completion(window, cx); + if !retry_scheduled { + thread.cancel_last_completion(window, cx); + } } } - cx.emit(ThreadEvent::Stopped(result.map_err(Arc::new))); + if !retry_scheduled { + cx.emit(ThreadEvent::Stopped(result.map_err(Arc::new))); + } if let Some((request_callback, (request, response_events))) = thread .request_callback @@ -2095,7 +2101,7 @@ impl Thread { segments: vec![MessageSegment::Text(retry_message)], loaded_context: LoadedContext::default(), creases: Vec::new(), - is_hidden: true, + is_hidden: false, ui_only: true, }); cx.emit(ThreadEvent::MessageAdded(id)); @@ -2111,12 +2117,6 @@ impl Thread { thread_handle .update(cx, |thread, cx| { - // Remove the retry notification message - if let Some(index) = thread.messages.iter().position(|m| m.id == message_id) { - thread.messages.remove(index); - cx.emit(ThreadEvent::MessageDeleted(message_id)); - } - // Retry the completion thread.send_to_model(model, intent, window, cx); }) @@ -2149,18 +2149,17 @@ impl Thread { const BASE_RETRY_DELAY_SECS: u64 = 5; let retry_state = self.retry_state.get_or_insert(RetryState { - attempt: 0, // This will be immediately incremented + attempt: 0, max_attempts: MAX_RETRY_ATTEMPTS, intent, }); - retry_state.attempt += 1; - if retry_state.attempt <= retry_state.max_attempts { - let retry_state = self.retry_state.as_ref().unwrap(); - let attempt = retry_state.attempt; - let max_attempts = retry_state.max_attempts; - let intent = retry_state.intent; + retry_state.attempt += 1; + let attempt = retry_state.attempt; + let max_attempts = retry_state.max_attempts; + let intent = retry_state.intent; + if attempt <= max_attempts { // Use custom delay if provided (e.g., from rate limit), otherwise exponential backoff let delay = if let Some(custom_delay) = custom_delay { custom_delay @@ -2184,7 +2183,7 @@ impl Thread { segments: vec![MessageSegment::Text(retry_message)], loaded_context: LoadedContext::default(), creases: Vec::new(), - is_hidden: true, + is_hidden: false, ui_only: true, }); cx.emit(ThreadEvent::MessageAdded(id)); @@ -2200,13 +2199,6 @@ impl Thread { thread_handle .update(cx, |thread, cx| { - // Remove the retry notification message - if let Some(index) = thread.messages.iter().position(|m| m.id == message_id) - { - thread.messages.remove(index); - cx.emit(ThreadEvent::MessageDeleted(message_id)); - } - // Retry the completion thread.send_to_model(model, intent, window, cx); }) @@ -2216,8 +2208,7 @@ impl Thread { true } else { - // Max retries exceeded; clear retry state and notify user that the thread stopped. - let max_attempts = self.retry_state.as_ref().unwrap().max_attempts; + // Max retries exceeded self.retry_state = None; let notification_text = if max_attempts == 1 { @@ -2226,6 +2217,9 @@ impl Thread { format!("Failed after retrying {} times.", max_attempts).into() }; + // Stop generating since we're giving up on retrying. + self.pending_completions.clear(); + cx.emit(ThreadEvent::RetriesFailed { message: notification_text, max_attempts, diff --git a/crates/agent_ui/src/active_thread.rs b/crates/agent_ui/src/active_thread.rs index 002a49e2a77a22..107ca4d4c6bde7 100644 --- a/crates/agent_ui/src/active_thread.rs +++ b/crates/agent_ui/src/active_thread.rs @@ -1140,22 +1140,11 @@ impl ActiveThread { self.save_thread(cx); cx.notify(); } - ThreadEvent::RetryScheduled { message_id } => { - if let Some(rendered_message) = self.thread.update(cx, |thread, cx| { - thread.message(*message_id).map(|message| { - RenderedMessage::from_segments( - &message.segments, - self.language_registry.clone(), - cx, - ) - }) - }) { - self.push_rendered_message(*message_id, rendered_message); - cx.notify(); - } + ThreadEvent::RetryScheduled { .. } => { + // Message already handled by MessageAdded event } ThreadEvent::RetriesFailed { message, .. } => { - self.show_notification(message.clone(), ui::IconName::Warning, window, cx); + self.show_notification(message, ui::IconName::Warning, window, cx); } } } @@ -1845,16 +1834,16 @@ impl ActiveThread { let tool_uses = thread.tool_uses_for_message(message_id, cx); let has_tool_uses = !tool_uses.is_empty(); - let editing_message_state = self .editing_message .as_ref() .filter(|(id, _)| *id == message_id) .map(|(_, state)| state); - let colors = cx.theme().colors(); - let editor_bg_color = colors.editor_background; - let panel_bg = colors.panel_background; + let (editor_bg_color, panel_bg) = { + let colors = cx.theme().colors(); + (colors.editor_background, colors.panel_background) + }; let open_as_markdown = IconButton::new(("open-as-markdown", ix), IconName::DocumentText) .icon_size(IconSize::XSmall) @@ -2042,152 +2031,167 @@ impl ActiveThread { } }); - let styled_message = match message.role { - Role::User => v_flex() - .id(("message-container", ix)) - .pt_2() - .pl_2() - .pr_2p5() - .pb_4() - .child( + let styled_message = if message.ui_only { + self.render_ui_notification(message_content, ix, cx) + } else { + match message.role { + Role::User => { + let colors = cx.theme().colors(); v_flex() - .id(("user-message", ix)) - .bg(editor_bg_color) - .rounded_lg() - .shadow_md() - .border_1() - .border_color(colors.border) - .hover(|hover| hover.border_color(colors.text_accent.opacity(0.5))) + .id(("message-container", ix)) + .pt_2() + .pl_2() + .pr_2p5() + .pb_4() .child( v_flex() - .p_2p5() - .gap_1() - .children(message_content) - .when_some(editing_message_state, |this, state| { - let focus_handle = state.editor.focus_handle(cx).clone(); - - this.child( - h_flex() - .w_full() - .gap_1() - .justify_between() - .flex_wrap() - .child( + .id(("user-message", ix)) + .bg(editor_bg_color) + .rounded_lg() + .shadow_md() + .border_1() + .border_color(colors.border) + .hover(|hover| hover.border_color(colors.text_accent.opacity(0.5))) + .child( + v_flex() + .p_2p5() + .gap_1() + .children(message_content) + .when_some(editing_message_state, |this, state| { + let focus_handle = state.editor.focus_handle(cx).clone(); + + this.child( h_flex() - .gap_1p5() + .w_full() + .gap_1() + .justify_between() + .flex_wrap() .child( - div() - .opacity(0.8) + h_flex() + .gap_1p5() .child( - Icon::new(IconName::Warning) - .size(IconSize::Indicator) - .color(Color::Warning) + div() + .opacity(0.8) + .child( + Icon::new(IconName::Warning) + .size(IconSize::Indicator) + .color(Color::Warning) + ), + ) + .child( + Label::new("Editing will restart the thread from this point.") + .color(Color::Muted) + .size(LabelSize::XSmall), ), ) .child( - Label::new("Editing will restart the thread from this point.") - .color(Color::Muted) - .size(LabelSize::XSmall), - ), - ) - .child( - h_flex() - .gap_0p5() - .child( - IconButton::new( - "cancel-edit-message", - IconName::Close, - ) - .shape(ui::IconButtonShape::Square) - .icon_color(Color::Error) - .icon_size(IconSize::Small) - .tooltip({ - let focus_handle = focus_handle.clone(); - move |window, cx| { - Tooltip::for_action_in( - "Cancel Edit", - &menu::Cancel, - &focus_handle, - window, - cx, + h_flex() + .gap_0p5() + .child( + IconButton::new( + "cancel-edit-message", + IconName::Close, ) - } - }) - .on_click(cx.listener(Self::handle_cancel_click)), - ) - .child( - IconButton::new( - "confirm-edit-message", - IconName::Return, - ) - .disabled(state.editor.read(cx).is_empty(cx)) - .shape(ui::IconButtonShape::Square) - .icon_color(Color::Muted) - .icon_size(IconSize::Small) - .tooltip({ - let focus_handle = focus_handle.clone(); - move |window, cx| { - Tooltip::for_action_in( - "Regenerate", - &menu::Confirm, - &focus_handle, - window, - cx, + .shape(ui::IconButtonShape::Square) + .icon_color(Color::Error) + .icon_size(IconSize::Small) + .tooltip({ + let focus_handle = focus_handle.clone(); + move |window, cx| { + Tooltip::for_action_in( + "Cancel Edit", + &menu::Cancel, + &focus_handle, + window, + cx, + ) + } + }) + .on_click(cx.listener(Self::handle_cancel_click)), + ) + .child( + IconButton::new( + "confirm-edit-message", + IconName::Return, ) - } - }) - .on_click( - cx.listener(Self::handle_regenerate_click), - ), - ), + .disabled(state.editor.read(cx).is_empty(cx)) + .shape(ui::IconButtonShape::Square) + .icon_color(Color::Muted) + .icon_size(IconSize::Small) + .tooltip({ + let focus_handle = focus_handle.clone(); + move |window, cx| { + Tooltip::for_action_in( + "Regenerate", + &menu::Confirm, + &focus_handle, + window, + cx, + ) + } + }) + .on_click( + cx.listener(Self::handle_regenerate_click), + ), + ), + ) ) - ) - }), + }), + ) + .on_click(cx.listener({ + let message_creases = message.creases.clone(); + move |this, _, window, cx| { + if let Some(message_text) = + this.thread.read(cx).message(message_id).and_then(|message| { + message.segments.first().and_then(|segment| { + match segment { + MessageSegment::Text(message_text) => { + Some(Into::>::into(message_text.as_str())) + } + _ => { + None + } + } + }) + }) + { + this.start_editing_message( + message_id, + message_text, + &message_creases, + window, + cx, + ); + } + } + })), ) - .on_click(cx.listener({ - let message_creases = message.creases.clone(); - move |this, _, window, cx| { - if let Some(message_text) = - this.thread.read(cx).message(message_id).and_then(|message| { - message.segments.first().and_then(|segment| { - match segment { - MessageSegment::Text(message_text) => { - Some(Into::>::into(message_text.as_str())) - } - _ => { - None - } - } - }) - }) - { - this.start_editing_message( - message_id, - message_text, - &message_creases, - window, - cx, - ); - } - } - })), - ), - Role::Assistant => v_flex() - .id(("message-container", ix)) - .px(RESPONSE_PADDING_X) - .gap_2() - .children(message_content) - .when(has_tool_uses, |parent| { - parent.children(tool_uses.into_iter().map(|tool_use| { - self.render_tool_use(tool_use, window, workspace.clone(), cx) - })) - }), - Role::System => div().id(("message-container", ix)).py_1().px_2().child( - v_flex() - .bg(colors.editor_background) - .rounded_sm() - .child(div().p_4().children(message_content)), - ), + } + Role::Assistant => v_flex() + .id(("message-container", ix)) + .px(RESPONSE_PADDING_X) + .gap_2() + .children(message_content) + .when(has_tool_uses, |parent| { + parent.children(tool_uses.into_iter().map(|tool_use| { + self.render_tool_use(tool_use, window, workspace.clone(), cx) + })) + }), + Role::System => { + let colors = cx.theme().colors(); + div().id(("message-container", ix)).py_1().px_2().child( + v_flex() + .w_full() + .bg(colors.editor_background) + .rounded_sm() + .child( + h_flex() + .p_4() + .child(v_flex().flex_1().children(message_content)), + ), + ) + } + } }; let after_editing_message = self @@ -2526,6 +2530,42 @@ impl ActiveThread { .blend(cx.theme().colors().editor_foreground.opacity(0.025)) } + fn render_ui_notification( + &self, + message_content: impl IntoIterator, + ix: usize, + cx: &mut Context, + ) -> Stateful
{ + let colors = cx.theme().colors(); + div().id(("message-container", ix)).py_1().px_2().child( + v_flex() + .w_full() + .bg(colors.editor_background) + .rounded_sm() + .child( + h_flex() + .w_full() + .p_2() + .gap_2() + .child( + div().flex_none().child( + Icon::new(IconName::Warning) + .size(IconSize::Small) + .color(Color::Warning), + ), + ) + .child( + v_flex() + .flex_1() + .min_w_0() + .text_size(TextSize::Small.rems(cx)) + .text_color(cx.theme().colors().text_muted) + .children(message_content), + ), + ), + ) + } + fn render_message_thinking_segment( &self, message_id: MessageId, From bacc2559d03a47dafe05abd437c9f70f13060f95 Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Wed, 25 Jun 2025 16:12:18 -0400 Subject: [PATCH 24/32] Fix retry logic tests --- crates/agent/src/thread.rs | 474 +++++++++++++++++++++++++++---------- crates/eval/src/example.rs | 3 + 2 files changed, 349 insertions(+), 128 deletions(-) diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index 11368fba83be4c..f2b34d34ee5d91 100644 --- a/crates/agent/src/thread.rs +++ b/crates/agent/src/thread.rs @@ -50,6 +50,10 @@ use util::{ResultExt as _, post_inc}; use uuid::Uuid; use zed_llm_client::{CompletionIntent, CompletionRequestStatus, UsageLimit}; +// Retry constants +const MAX_RETRY_ATTEMPTS: u8 = 3; +const BASE_RETRY_DELAY_SECS: u64 = 5; + #[derive( Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Serialize, Deserialize, JsonSchema, )] @@ -2145,9 +2149,6 @@ impl Thread { window: Option, cx: &mut Context, ) -> bool { - const MAX_RETRY_ATTEMPTS: u8 = 3; - const BASE_RETRY_DELAY_SECS: u64 = 5; - let retry_state = self.retry_state.get_or_insert(RetryState { attempt: 0, max_attempts: MAX_RETRY_ATTEMPTS, @@ -3264,6 +3265,9 @@ mod tests { use crate::{ context::load_context, context_store::ContextStore, thread_store, thread_store::ThreadStore, }; + + // Test-specific constants + const TEST_RATE_LIMIT_RETRY_SECS: u64 = 30; use agent_settings::{AgentProfileId, AgentSettings, LanguageModelParameters}; use assistant_tool::ToolRegistry; use futures::StreamExt; @@ -3274,7 +3278,7 @@ mod tests { use language_model::fake_provider::{FakeLanguageModel, FakeLanguageModelProvider}; use language_model::{ LanguageModelCompletionError, LanguageModelName, LanguageModelProviderId, - LanguageModelToolChoice, + LanguageModelProviderName, LanguageModelToolChoice, }; use parking_lot::Mutex; use project::{FakeFs, Project}; @@ -4170,13 +4174,13 @@ fn main() {{ }); // Track retry events - let retry_events = Arc::new(Mutex::new(Vec::new())); - let retry_events_clone = retry_events.clone(); + let retry_count = Arc::new(Mutex::new(0)); + let retry_count_clone = retry_count.clone(); let _subscription = thread.update(cx, |_, cx| { cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| { if let ThreadEvent::RetryScheduled { .. } = event { - retry_events_clone.lock().push((*attempt, *delay)); + *retry_count_clone.lock() += 1; } }) }); @@ -4189,14 +4193,17 @@ fn main() {{ cx.run_until_parked(); // Check that a retry was scheduled - let events = retry_events.lock(); - assert_eq!(events.len(), 1, "Should have scheduled one retry"); - assert_eq!(events[0].0, 1, "Should be first retry attempt"); - assert_eq!( - events[0].1, - Duration::from_secs(5), - "Should have 5 second delay" - ); + assert_eq!(*retry_count.lock(), 1, "Should have scheduled one retry"); + + thread.read_with(cx, |thread, _| { + assert!(thread.retry_state.is_some(), "Should have retry state"); + let retry_state = thread.retry_state.as_ref().unwrap(); + assert_eq!(retry_state.attempt, 1, "Should be first retry attempt"); + assert_eq!( + retry_state.max_attempts, MAX_RETRY_ATTEMPTS, + "Should have default max attempts" + ); + }); // Check that a retry message was added thread.read_with(cx, |thread, _| { @@ -4204,11 +4211,12 @@ fn main() {{ assert!( messages.any(|msg| { msg.role == Role::System - && msg.is_hidden + && msg.ui_only && msg.segments.iter().any(|seg| { if let MessageSegment::Text(text) = seg { text.contains("overloaded") - && text.contains("Retrying (attempt 1 of 3)") + && text + .contains(&format!("attempt 1 of {}", MAX_RETRY_ATTEMPTS)) } else { false } @@ -4235,13 +4243,13 @@ fn main() {{ }); // Track retry events - let retry_events = Arc::new(Mutex::new(Vec::new())); - let retry_events_clone = retry_events.clone(); + let retry_count = Arc::new(Mutex::new(0)); + let retry_count_clone = retry_count.clone(); let _subscription = thread.update(cx, |_, cx| { cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| { if let ThreadEvent::RetryScheduled { .. } = event { - retry_events_clone.lock().push((*attempt, *delay)); + *retry_count_clone.lock() += 1; } }) }); @@ -4254,14 +4262,18 @@ fn main() {{ cx.run_until_parked(); // Check that a retry was scheduled - let events = retry_events.lock(); - assert_eq!(events.len(), 1, "Should have scheduled one retry"); - assert_eq!(events[0].0, 1, "Should be first attempt"); - assert_eq!( - events[0].1, - Duration::from_secs(5), - "Should have 5 second delay" - ); + assert_eq!(*retry_count.lock(), 1, "Should have scheduled one retry"); + + // Check retry state on thread + thread.read_with(cx, |thread, _| { + assert!(thread.retry_state.is_some(), "Should have retry state"); + let retry_state = thread.retry_state.as_ref().unwrap(); + assert_eq!(retry_state.attempt, 1, "Should be first retry attempt"); + assert_eq!( + retry_state.max_attempts, MAX_RETRY_ATTEMPTS, + "Should have correct max attempts" + ); + }); // Check that a retry message was added with provider name thread.read_with(cx, |thread, _| { @@ -4269,12 +4281,13 @@ fn main() {{ assert!( messages.any(|msg| { msg.role == Role::System - && msg.is_hidden + && msg.ui_only && msg.segments.iter().any(|seg| { if let MessageSegment::Text(text) = seg { - text.contains("internal server error") && - text.contains("Fake") && // Provider name from fake_provider.rs - text.contains("Retrying (attempt 1 of 3)") + text.contains("internal") + && text.contains("Fake") + && text + .contains(&format!("attempt 1 of {}", MAX_RETRY_ATTEMPTS)) } else { false } @@ -4301,15 +4314,15 @@ fn main() {{ }); // Track retry events and completion count - let retry_events = Arc::new(Mutex::new(Vec::new())); - let retry_events_clone = retry_events.clone(); + let retry_count = Arc::new(Mutex::new(0)); + let retry_count_clone = retry_count.clone(); let completion_count = Arc::new(Mutex::new(0)); let completion_count_clone = completion_count.clone(); let _subscription = thread.update(cx, |_, cx| { cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| match event { ThreadEvent::RetryScheduled { .. } => { - retry_events_clone.lock().push((*attempt, *delay)); + *retry_count_clone.lock() += 1; } ThreadEvent::NewRequest => { *completion_count_clone.lock() += 1; @@ -4324,31 +4337,77 @@ fn main() {{ }); cx.run_until_parked(); - // Should have 3 retries scheduled immediately due to errors - let mut total_delay = Duration::from_millis(0); - for i in 0..3u32 { - total_delay += Duration::from_secs(5 * 2u64.pow(i)); - cx.executor().advance_clock(total_delay); - cx.run_until_parked(); - } + // Should have scheduled first retry + assert_eq!(*retry_count.lock(), 1, "Should have scheduled first retry"); + + // Check retry state + thread.read_with(cx, |thread, _| { + assert!(thread.retry_state.is_some(), "Should have retry state"); + let retry_state = thread.retry_state.as_ref().unwrap(); + assert_eq!(retry_state.attempt, 1, "Should be first retry attempt"); + }); + + // Advance clock for first retry + cx.executor() + .advance_clock(Duration::from_secs(BASE_RETRY_DELAY_SECS)); + cx.run_until_parked(); + + // Should have scheduled second retry + assert_eq!(*retry_count.lock(), 2, "Should have scheduled second retry"); + + // Check retry state updated + thread.read_with(cx, |thread, _| { + assert!(thread.retry_state.is_some(), "Should have retry state"); + let retry_state = thread.retry_state.as_ref().unwrap(); + assert_eq!(retry_state.attempt, 2, "Should be second retry attempt"); + assert_eq!( + retry_state.max_attempts, MAX_RETRY_ATTEMPTS, + "Should have correct max attempts" + ); + }); - // Check exponential backoff - let events = retry_events.lock(); - assert_eq!(events.len(), 3, "Should have scheduled three retries"); + // Advance clock for second retry (exponential backoff) + cx.executor() + .advance_clock(Duration::from_secs(BASE_RETRY_DELAY_SECS * 2)); + cx.run_until_parked(); + + // Should have scheduled third retry assert_eq!( - events[0], - (1, Duration::from_secs(5)), - "First retry: 5 seconds" + *retry_count.lock(), + MAX_RETRY_ATTEMPTS as usize, + "Should have scheduled third retry" ); + + // Check retry state updated + thread.read_with(cx, |thread, _| { + assert!(thread.retry_state.is_some(), "Should have retry state"); + let retry_state = thread.retry_state.as_ref().unwrap(); + assert_eq!( + retry_state.attempt, MAX_RETRY_ATTEMPTS, + "Should be at max retry attempt" + ); + assert_eq!( + retry_state.max_attempts, MAX_RETRY_ATTEMPTS, + "Should have correct max attempts" + ); + }); + + // Advance clock for third retry (exponential backoff) + cx.executor() + .advance_clock(Duration::from_secs(BASE_RETRY_DELAY_SECS * 4)); + cx.run_until_parked(); + + // No more retries should be scheduled assert_eq!( - events[1], - (2, Duration::from_secs(10)), - "Second retry: 10 seconds" + *retry_count.lock(), + MAX_RETRY_ATTEMPTS as usize, + "Should not exceed max retries" ); + // Final completion count should be initial + max retries assert_eq!( - events[2], - (3, Duration::from_secs(20)), - "Third retry: 20 seconds" + *completion_count.lock(), + (MAX_RETRY_ATTEMPTS + 1) as usize, + "Should have made initial + max retry attempts" ); } @@ -4368,15 +4427,15 @@ fn main() {{ }); // Track events - let error_shown = Arc::new(Mutex::new(false)); - let error_shown_clone = error_shown.clone(); + let retries_failed = Arc::new(Mutex::new(false)); + let retries_failed_clone = retries_failed.clone(); let retry_count = Arc::new(Mutex::new(0)); let retry_count_clone = retry_count.clone(); let _subscription = thread.update(cx, |_, cx| { cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| match event { - ThreadEvent::ShowError(_) => { - *error_shown_clone.lock() = true; + ThreadEvent::RetriesFailed { .. } => { + *retries_failed_clone.lock() = true; } ThreadEvent::RetryScheduled { .. } => { *retry_count_clone.lock() += 1; @@ -4392,17 +4451,32 @@ fn main() {{ cx.run_until_parked(); // Advance through all retries - let delays = [5u64, 10, 20]; - for delay in delays { + for i in 0..MAX_RETRY_ATTEMPTS { + let delay = if i == 0 { + BASE_RETRY_DELAY_SECS + } else { + BASE_RETRY_DELAY_SECS * 2u64.pow(i as u32 - 1) + }; cx.executor().advance_clock(Duration::from_secs(delay)); cx.run_until_parked(); } - // After max retries (3), should show error - assert_eq!(*retry_count.lock(), 3, "Should have attempted 3 retries"); + // After the 3rd retry is scheduled, we need to wait for it to execute and fail + // The 3rd retry has a delay of BASE_RETRY_DELAY_SECS * 4 (20 seconds) + let final_delay = BASE_RETRY_DELAY_SECS * 2u64.pow((MAX_RETRY_ATTEMPTS - 1) as u32); + cx.executor() + .advance_clock(Duration::from_secs(final_delay)); + cx.run_until_parked(); + + // After max retries, should emit RetriesFailed event + assert_eq!( + *retry_count.lock(), + MAX_RETRY_ATTEMPTS as usize, + "Should have attempted max retries" + ); assert!( - *error_shown.lock(), - "Should show error after max retries exceeded" + *retries_failed.lock(), + "Should emit RetriesFailed event after max retries exceeded" ); // Retry state should be cleared @@ -4411,6 +4485,17 @@ fn main() {{ thread.retry_state.is_none(), "Retry state should be cleared after max retries" ); + + // Verify we have the expected number of retry messages + let retry_messages = thread + .messages + .iter() + .filter(|msg| msg.ui_only && msg.role == Role::System) + .count(); + assert_eq!( + retry_messages, MAX_RETRY_ATTEMPTS as usize, + "Should have one retry message per attempt" + ); }); } @@ -4515,13 +4600,14 @@ fn main() {{ }); // Track message deletions - let deleted_messages = Arc::new(Mutex::new(Vec::new())); - let deleted_messages_clone = deleted_messages.clone(); + // Track when retry completes successfully + let retry_completed = Arc::new(Mutex::new(false)); + let retry_completed_clone = retry_completed.clone(); let _subscription = thread.update(cx, |_, cx| { cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| { - if let ThreadEvent::MessageDeleted(id) = event { - deleted_messages_clone.lock().push(*id); + if let ThreadEvent::StreamedCompletion = event { + *retry_completed_clone.lock() = true; } }) }); @@ -4536,13 +4622,14 @@ fn main() {{ let retry_message_id = thread.read_with(cx, |thread, _| { thread .messages() - .find(|msg| msg.role == Role::System && msg.is_hidden) + .find(|msg| msg.role == Role::System && msg.ui_only) .map(|msg| msg.id) .expect("Should have a retry message") }); // Wait for retry - cx.executor().advance_clock(Duration::from_secs(5)); + cx.executor() + .advance_clock(Duration::from_secs(BASE_RETRY_DELAY_SECS)); cx.run_until_parked(); // Stream some successful content @@ -4557,17 +4644,22 @@ fn main() {{ fake_model.end_completion_stream(&pending[0]); cx.run_until_parked(); - // Check that the retry message was deleted + // Check that the retry completed successfully assert!( - deleted_messages.lock().contains(&retry_message_id), - "Retry message should be deleted when retry starts" + *retry_completed.lock(), + "Retry should have completed successfully" ); - // Message should no longer exist + // Retry message should still exist but be marked as ui_only thread.read_with(cx, |thread, _| { - assert!( - thread.message(retry_message_id).is_none(), - "Retry message should be removed from thread" + let retry_msg = thread + .message(retry_message_id) + .expect("Retry message should still exist"); + assert!(retry_msg.ui_only, "Retry message should be ui_only"); + assert_eq!( + retry_msg.role, + Role::System, + "Retry message should have System role" ); }); } @@ -4577,26 +4669,116 @@ fn main() {{ init_test_settings(cx); let project = create_test_project(cx, json!({})).await; - let (_, _, thread, _, model) = setup_test_environment(cx, project.clone()).await; + let (_, _, thread, _, _base_model) = setup_test_environment(cx, project.clone()).await; + + // Create a model that fails once then succeeds + struct FailOnceModel { + inner: Arc, + failed_once: Arc>, + } + + impl LanguageModel for FailOnceModel { + fn id(&self) -> LanguageModelId { + self.inner.id() + } + + fn name(&self) -> LanguageModelName { + self.inner.name() + } + + fn provider_id(&self) -> LanguageModelProviderId { + self.inner.provider_id() + } + + fn provider_name(&self) -> LanguageModelProviderName { + self.inner.provider_name() + } + + fn supports_tools(&self) -> bool { + self.inner.supports_tools() + } + + fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool { + self.inner.supports_tool_choice(choice) + } + + fn supports_images(&self) -> bool { + self.inner.supports_images() + } + + fn telemetry_id(&self) -> String { + self.inner.telemetry_id() + } + + fn max_token_count(&self) -> u64 { + self.inner.max_token_count() + } + + fn count_tokens( + &self, + request: LanguageModelRequest, + cx: &App, + ) -> BoxFuture<'static, Result> { + self.inner.count_tokens(request, cx) + } + + fn stream_completion( + &self, + request: LanguageModelRequest, + cx: &AsyncApp, + ) -> BoxFuture< + 'static, + Result< + BoxStream< + 'static, + Result, + >, + LanguageModelCompletionError, + >, + > { + if !*self.failed_once.lock() { + *self.failed_once.lock() = true; + // Return error on first attempt + let stream = futures::stream::once(async move { + Err(LanguageModelCompletionError::Overloaded) + }); + async move { Ok(stream.boxed()) }.boxed() + } else { + // Succeed on retry + self.inner.stream_completion(request, cx) + } + } + } + + let fail_once_model = Arc::new(FailOnceModel { + inner: Arc::new(FakeLanguageModel::default()), + failed_once: Arc::new(Mutex::new(false)), + }); // Insert a user message thread.update(cx, |thread, cx| { - thread.insert_user_message("Hello!", ContextLoadResult::default(), None, vec![], cx); + thread.insert_user_message( + "Test message", + ContextLoadResult::default(), + None, + vec![], + cx, + ); }); - // Start completion + // Start completion with fail-once model thread.update(cx, |thread, cx| { - thread.send_to_model(model.clone(), CompletionIntent::UserPrompt, None, cx); + thread.send_to_model( + fail_once_model.clone(), + CompletionIntent::UserPrompt, + None, + cx, + ); }); - cx.run_until_parked(); - // Fail first attempt - let fake_model = model.as_fake(); - fake_model.end_completion_stream(&fake_model.pending_completions()[0]); - cx.executor().advance_clock(Duration::from_millis(100)); cx.run_until_parked(); - // Verify retry state exists + // Verify retry state exists after first failure thread.read_with(cx, |thread, _| { assert!( thread.retry_state.is_some(), @@ -4604,20 +4786,38 @@ fn main() {{ ); }); - // Wait for retry - cx.executor().advance_clock(Duration::from_secs(5)); + // Wait for retry delay + cx.executor() + .advance_clock(Duration::from_secs(BASE_RETRY_DELAY_SECS)); + cx.run_until_parked(); + + // The retry should now use our FailOnceModel which should succeed + // We need to help the FakeLanguageModel complete the stream + let inner_fake = fail_once_model.inner.clone(); + + // Wait a bit for the retry to start cx.run_until_parked(); - // Succeed on retry - fake_model.stream_completion_response(&fake_model.pending_completions()[0], "Hello there!"); - fake_model.end_completion_stream(&fake_model.pending_completions()[0]); + // Check for pending completions and complete them + if let Some(pending) = inner_fake.pending_completions().first() { + inner_fake.stream_completion_response(pending, "Success!"); + inner_fake.end_completion_stream(pending); + } cx.run_until_parked(); - // Retry state should be cleared on success thread.read_with(cx, |thread, _| { assert!( thread.retry_state.is_none(), - "Retry state should be cleared on success" + "Retry state should be cleared after successful completion" + ); + + let has_assistant_message = thread + .messages + .iter() + .any(|msg| msg.role == Role::Assistant && !msg.ui_only); + assert!( + has_assistant_message, + "Should have an assistant message after successful retry" ); }); } @@ -4696,7 +4896,7 @@ fn main() {{ async move { let stream = futures::stream::once(async move { Err(LanguageModelCompletionError::RateLimitExceeded { - retry_after: Duration::from_secs(30), + retry_after: Duration::from_secs(TEST_RATE_LIMIT_RETRY_SECS), }) }); Ok(stream.boxed()) @@ -4719,15 +4919,13 @@ fn main() {{ }); // Track retry events - let retry_events = Arc::new(Mutex::new(Vec::new())); - let retry_events_clone = retry_events.clone(); + let retry_count = Arc::new(Mutex::new(0)); + let retry_count_clone = retry_count.clone(); let _subscription = thread.update(cx, |_, cx| { cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| { if let ThreadEvent::RetryScheduled { .. } = event { - retry_events_clone - .lock() - .push((*attempt, *delay, *max_attempts)); + *retry_count_clone.lock() += 1; } }) }); @@ -4739,40 +4937,60 @@ fn main() {{ cx.run_until_parked(); - // Check that a retry was scheduled with the specified delay - let events = retry_events.lock(); - assert_eq!(events.len(), 1, "Should have scheduled exactly one retry"); - assert_eq!(events[0].0, 1, "Should be first attempt"); - assert_eq!( - events[0].1, - Duration::from_secs(30), - "Should use the retry_after duration" - ); - assert_eq!( - events[0].2, 1, - "Should only have max 1 attempt for rate limit" - ); + assert_eq!(*retry_count.lock(), 1, "Should have scheduled one retry"); - // Check that retry message doesn't include attempt count thread.read_with(cx, |thread, _| { - let mut messages = thread.messages(); assert!( - messages.any(|msg| { - msg.role == Role::System - && msg.is_hidden + thread.retry_state.is_none(), + "Rate limit errors should not set retry_state" + ); + }); + + // Verify we have one retry message + thread.read_with(cx, |thread, _| { + let retry_messages = thread + .messages + .iter() + .filter(|msg| { + msg.ui_only && msg.segments.iter().any(|seg| { if let MessageSegment::Text(text) = seg { text.contains("rate limit exceeded") - && text.contains("Retrying in 30 seconds…") - && !text.contains("attempt") } else { false } }) - }), - "Should have added a system retry message without attempt count" + }) + .count(); + assert_eq!( + retry_messages, 1, + "Should have one rate limit retry message" ); }); + + // Check that retry message doesn't include attempt count + thread.read_with(cx, |thread, _| { + let retry_message = thread + .messages + .iter() + .find(|msg| msg.role == Role::System && msg.ui_only) + .expect("Should have a retry message"); + + // Check that the message doesn't contain attempt count + if let Some(MessageSegment::Text(text)) = retry_message.segments.first() { + assert!( + !text.contains("attempt"), + "Rate limit retry message should not contain attempt count" + ); + assert!( + text.contains(&format!( + "Retrying in {} seconds", + TEST_RATE_LIMIT_RETRY_SECS + )), + "Rate limit retry message should contain retry delay" + ); + } + }); } #[gpui::test] @@ -4917,8 +5135,6 @@ fn main() {{ thread.cancel_last_completion(None, cx); }); - // Advance time past the retry delay - cx.executor().advance_clock(Duration::from_secs(10)); cx.run_until_parked(); // The retry should not have happened - no pending completions @@ -4931,10 +5147,12 @@ fn main() {{ // Verify the retry was cancelled by checking retry state thread.read_with(cx, |thread, _| { - assert!( - thread.retry_state.is_none(), - "retry_state should be cleared after cancellation" - ); + if let Some(retry_state) = &thread.retry_state { + panic!( + "retry_state should be cleared after cancellation, but found: attempt={}, max_attempts={}, intent={:?}", + retry_state.attempt, retry_state.max_attempts, retry_state.intent + ); + } }); } diff --git a/crates/eval/src/example.rs b/crates/eval/src/example.rs index fae7fd030873e8..c2813ac1050b77 100644 --- a/crates/eval/src/example.rs +++ b/crates/eval/src/example.rs @@ -224,6 +224,9 @@ impl ExampleContext { ThreadEvent::RetryScheduled { .. } => { // Ignore retry events } + ThreadEvent::RetriesFailed { .. } => { + // Ignore retries failed events + } ThreadEvent::Stopped(reason) => match reason { Ok(StopReason::EndTurn) => { tx.close_channel(); From a8bfc93ac9f15ffdbfd8dd62e3f921a5d0af0870 Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Wed, 25 Jun 2025 19:48:53 -0400 Subject: [PATCH 25/32] Simplify a branch --- crates/agent_ui/src/agent_diff.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/agent_ui/src/agent_diff.rs b/crates/agent_ui/src/agent_diff.rs index 2803fb53b5d6c3..192ca6b52a7933 100644 --- a/crates/agent_ui/src/agent_diff.rs +++ b/crates/agent_ui/src/agent_diff.rs @@ -1381,7 +1381,7 @@ impl AgentDiff { | ThreadEvent::ToolUseLimitReached | ThreadEvent::CancelEditing | ThreadEvent::RetryScheduled { .. } - | ThreadEvent::RetriesFailed { .. } => {} + | ThreadEvent::RetriesFailed { .. } | ThreadEvent::ProfileChanged => {} } } From 538d6bc182e4950fd2b30456258fe8e6608c0d65 Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Wed, 25 Jun 2025 20:16:48 -0400 Subject: [PATCH 26/32] Remove obsolete RetryScheduled event --- crates/agent/src/thread.rs | 260 +++++++++++++++++---------- crates/agent_ui/src/active_thread.rs | 3 - crates/agent_ui/src/agent_diff.rs | 2 - crates/eval/src/example.rs | 3 - 4 files changed, 166 insertions(+), 102 deletions(-) diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index 370aa5c20f5a5d..1ecd04de63a8e3 100644 --- a/crates/agent/src/thread.rs +++ b/crates/agent/src/thread.rs @@ -2109,10 +2109,6 @@ impl Thread { ui_only: true, }); cx.emit(ThreadEvent::MessageAdded(id)); - let message_id = id; - - cx.emit(ThreadEvent::RetryScheduled { message_id }); - // Schedule the retry let thread_handle = cx.entity().downgrade(); @@ -2188,9 +2184,6 @@ impl Thread { ui_only: true, }); cx.emit(ThreadEvent::MessageAdded(id)); - let message_id = id; - - cx.emit(ThreadEvent::RetryScheduled { message_id }); // Schedule the retry let thread_handle = cx.entity().downgrade(); @@ -3173,9 +3166,6 @@ pub enum ThreadEvent { CancelEditing, CompletionCanceled, ProfileChanged, - RetryScheduled { - message_id: MessageId, - }, RetriesFailed { message: SharedString, max_attempts: u8, @@ -4183,18 +4173,6 @@ fn main() {{ thread.insert_user_message("Hello!", ContextLoadResult::default(), None, vec![], cx); }); - // Track retry events - let retry_count = Arc::new(Mutex::new(0)); - let retry_count_clone = retry_count.clone(); - - let _subscription = thread.update(cx, |_, cx| { - cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| { - if let ThreadEvent::RetryScheduled { .. } = event { - *retry_count_clone.lock() += 1; - } - }) - }); - // Start completion thread.update(cx, |thread, cx| { thread.send_to_model(model.clone(), CompletionIntent::UserPrompt, None, cx); @@ -4202,9 +4180,6 @@ fn main() {{ cx.run_until_parked(); - // Check that a retry was scheduled - assert_eq!(*retry_count.lock(), 1, "Should have scheduled one retry"); - thread.read_with(cx, |thread, _| { assert!(thread.retry_state.is_some(), "Should have retry state"); let retry_state = thread.retry_state.as_ref().unwrap(); @@ -4235,6 +4210,25 @@ fn main() {{ "Should have added a system retry message" ); }); + + let retry_count = thread.update(cx, |thread, _| { + thread + .messages + .iter() + .filter(|m| { + m.ui_only + && m.segments.iter().any(|s| { + if let MessageSegment::Text(text) = s { + text.contains("Retrying") && text.contains("seconds") + } else { + false + } + }) + }) + .count() + }); + + assert_eq!(retry_count, 1, "Should have one retry message"); } #[gpui::test] @@ -4252,18 +4246,6 @@ fn main() {{ thread.insert_user_message("Hello!", ContextLoadResult::default(), None, vec![], cx); }); - // Track retry events - let retry_count = Arc::new(Mutex::new(0)); - let retry_count_clone = retry_count.clone(); - - let _subscription = thread.update(cx, |_, cx| { - cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| { - if let ThreadEvent::RetryScheduled { .. } = event { - *retry_count_clone.lock() += 1; - } - }) - }); - // Start completion thread.update(cx, |thread, cx| { thread.send_to_model(model.clone(), CompletionIntent::UserPrompt, None, cx); @@ -4271,9 +4253,6 @@ fn main() {{ cx.run_until_parked(); - // Check that a retry was scheduled - assert_eq!(*retry_count.lock(), 1, "Should have scheduled one retry"); - // Check retry state on thread thread.read_with(cx, |thread, _| { assert!(thread.retry_state.is_some(), "Should have retry state"); @@ -4306,6 +4285,26 @@ fn main() {{ "Should have added a system retry message with provider name" ); }); + + // Count retry messages + let retry_count = thread.update(cx, |thread, _| { + thread + .messages + .iter() + .filter(|m| { + m.ui_only + && m.segments.iter().any(|s| { + if let MessageSegment::Text(text) = s { + text.contains("Retrying") && text.contains("seconds") + } else { + false + } + }) + }) + .count() + }); + + assert_eq!(retry_count, 1, "Should have one retry message"); } #[gpui::test] @@ -4324,20 +4323,15 @@ fn main() {{ }); // Track retry events and completion count - let retry_count = Arc::new(Mutex::new(0)); - let retry_count_clone = retry_count.clone(); + // Track completion events let completion_count = Arc::new(Mutex::new(0)); let completion_count_clone = completion_count.clone(); let _subscription = thread.update(cx, |_, cx| { - cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| match event { - ThreadEvent::RetryScheduled { .. } => { - *retry_count_clone.lock() += 1; - } - ThreadEvent::NewRequest => { + cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| { + if let ThreadEvent::NewRequest = event { *completion_count_clone.lock() += 1; } - _ => {} }) }); @@ -4347,8 +4341,24 @@ fn main() {{ }); cx.run_until_parked(); - // Should have scheduled first retry - assert_eq!(*retry_count.lock(), 1, "Should have scheduled first retry"); + // Should have scheduled first retry - count retry messages + let retry_count = thread.update(cx, |thread, _| { + thread + .messages + .iter() + .filter(|m| { + m.ui_only + && m.segments.iter().any(|s| { + if let MessageSegment::Text(text) = s { + text.contains("Retrying") && text.contains("seconds") + } else { + false + } + }) + }) + .count() + }); + assert_eq!(retry_count, 1, "Should have scheduled first retry"); // Check retry state thread.read_with(cx, |thread, _| { @@ -4362,8 +4372,24 @@ fn main() {{ .advance_clock(Duration::from_secs(BASE_RETRY_DELAY_SECS)); cx.run_until_parked(); - // Should have scheduled second retry - assert_eq!(*retry_count.lock(), 2, "Should have scheduled second retry"); + // Should have scheduled second retry - count retry messages + let retry_count = thread.update(cx, |thread, _| { + thread + .messages + .iter() + .filter(|m| { + m.ui_only + && m.segments.iter().any(|s| { + if let MessageSegment::Text(text) = s { + text.contains("Retrying") && text.contains("seconds") + } else { + false + } + }) + }) + .count() + }); + assert_eq!(retry_count, 2, "Should have scheduled second retry"); // Check retry state updated thread.read_with(cx, |thread, _| { @@ -4382,9 +4408,25 @@ fn main() {{ cx.run_until_parked(); // Should have scheduled third retry + // Count all retry messages now + let retry_count = thread.update(cx, |thread, _| { + thread + .messages + .iter() + .filter(|m| { + m.ui_only + && m.segments.iter().any(|s| { + if let MessageSegment::Text(text) = s { + text.contains("Retrying") && text.contains("seconds") + } else { + false + } + }) + }) + .count() + }); assert_eq!( - *retry_count.lock(), - MAX_RETRY_ATTEMPTS as usize, + retry_count, MAX_RETRY_ATTEMPTS as usize, "Should have scheduled third retry" ); @@ -4407,12 +4449,28 @@ fn main() {{ .advance_clock(Duration::from_secs(BASE_RETRY_DELAY_SECS * 4)); cx.run_until_parked(); - // No more retries should be scheduled + // No more retries should be scheduled after clock was advanced. + let retry_count = thread.update(cx, |thread, _| { + thread + .messages + .iter() + .filter(|m| { + m.ui_only + && m.segments.iter().any(|s| { + if let MessageSegment::Text(text) = s { + text.contains("Retrying") && text.contains("seconds") + } else { + false + } + }) + }) + .count() + }); assert_eq!( - *retry_count.lock(), - MAX_RETRY_ATTEMPTS as usize, + retry_count, MAX_RETRY_ATTEMPTS as usize, "Should not exceed max retries" ); + // Final completion count should be initial + max retries assert_eq!( *completion_count.lock(), @@ -4439,18 +4497,12 @@ fn main() {{ // Track events let retries_failed = Arc::new(Mutex::new(false)); let retries_failed_clone = retries_failed.clone(); - let retry_count = Arc::new(Mutex::new(0)); - let retry_count_clone = retry_count.clone(); let _subscription = thread.update(cx, |_, cx| { - cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| match event { - ThreadEvent::RetriesFailed { .. } => { + cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| { + if let ThreadEvent::RetriesFailed { .. } = event { *retries_failed_clone.lock() = true; } - ThreadEvent::RetryScheduled { .. } => { - *retry_count_clone.lock() += 1; - } - _ => {} }) }); @@ -4478,10 +4530,26 @@ fn main() {{ .advance_clock(Duration::from_secs(final_delay)); cx.run_until_parked(); + let retry_count = thread.update(cx, |thread, _| { + thread + .messages + .iter() + .filter(|m| { + m.ui_only + && m.segments.iter().any(|s| { + if let MessageSegment::Text(text) = s { + text.contains("Retrying") && text.contains("seconds") + } else { + false + } + }) + }) + .count() + }); + // After max retries, should emit RetriesFailed event assert_eq!( - *retry_count.lock(), - MAX_RETRY_ATTEMPTS as usize, + retry_count, MAX_RETRY_ATTEMPTS as usize, "Should have attempted max retries" ); assert!( @@ -4928,18 +4996,6 @@ fn main() {{ thread.insert_user_message("Hello!", ContextLoadResult::default(), None, vec![], cx); }); - // Track retry events - let retry_count = Arc::new(Mutex::new(0)); - let retry_count_clone = retry_count.clone(); - - let _subscription = thread.update(cx, |_, cx| { - cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| { - if let ThreadEvent::RetryScheduled { .. } = event { - *retry_count_clone.lock() += 1; - } - }) - }); - // Start completion thread.update(cx, |thread, cx| { thread.send_to_model(model.clone(), CompletionIntent::UserPrompt, None, cx); @@ -4947,7 +5003,23 @@ fn main() {{ cx.run_until_parked(); - assert_eq!(*retry_count.lock(), 1, "Should have scheduled one retry"); + let retry_count = thread.update(cx, |thread, _| { + thread + .messages + .iter() + .filter(|m| { + m.ui_only + && m.segments.iter().any(|s| { + if let MessageSegment::Text(text) = s { + text.contains("rate limit exceeded") + } else { + false + } + }) + }) + .count() + }); + assert_eq!(retry_count, 1, "Should have scheduled one retry"); thread.read_with(cx, |thread, _| { assert!( @@ -5118,18 +5190,6 @@ fn main() {{ thread.insert_user_message("Hello!", ContextLoadResult::default(), None, vec![], cx); }); - // Track retry events - let retry_scheduled = Arc::new(Mutex::new(false)); - let retry_scheduled_clone = retry_scheduled.clone(); - - let _subscription = thread.update(cx, |_, cx| { - cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| { - if matches!(event, ThreadEvent::RetryScheduled { .. }) { - *retry_scheduled_clone.lock() = true; - } - }) - }); - // Start completion thread.update(cx, |thread, cx| { thread.send_to_model(model.clone(), CompletionIntent::UserPrompt, None, cx); @@ -5137,8 +5197,20 @@ fn main() {{ cx.run_until_parked(); - // Verify retry was scheduled - assert!(*retry_scheduled.lock(), "Should have scheduled a retry"); + // Verify retry was scheduled by checking for retry message + let has_retry_message = thread.read_with(cx, |thread, _| { + thread.messages.iter().any(|m| { + m.ui_only + && m.segments.iter().any(|s| { + if let MessageSegment::Text(text) = s { + text.contains("Retrying") && text.contains("seconds") + } else { + false + } + }) + }) + }); + assert!(has_retry_message, "Should have scheduled a retry"); // Cancel the completion before the retry happens thread.update(cx, |thread, cx| { diff --git a/crates/agent_ui/src/active_thread.rs b/crates/agent_ui/src/active_thread.rs index 9643ee11d3167f..384a691efecf7f 100644 --- a/crates/agent_ui/src/active_thread.rs +++ b/crates/agent_ui/src/active_thread.rs @@ -1140,9 +1140,6 @@ impl ActiveThread { self.save_thread(cx); cx.notify(); } - ThreadEvent::RetryScheduled { .. } => { - // Message already handled by MessageAdded event - } ThreadEvent::RetriesFailed { message, .. } => { self.show_notification(message, ui::IconName::Warning, window, cx); } diff --git a/crates/agent_ui/src/agent_diff.rs b/crates/agent_ui/src/agent_diff.rs index 192ca6b52a7933..aa5e49551b4e1c 100644 --- a/crates/agent_ui/src/agent_diff.rs +++ b/crates/agent_ui/src/agent_diff.rs @@ -1380,8 +1380,6 @@ impl AgentDiff { | ThreadEvent::ToolConfirmationNeeded | ThreadEvent::ToolUseLimitReached | ThreadEvent::CancelEditing - | ThreadEvent::RetryScheduled { .. } - | ThreadEvent::RetriesFailed { .. } | ThreadEvent::ProfileChanged => {} } } diff --git a/crates/eval/src/example.rs b/crates/eval/src/example.rs index c2813ac1050b77..904eca83e609dc 100644 --- a/crates/eval/src/example.rs +++ b/crates/eval/src/example.rs @@ -221,9 +221,6 @@ impl ExampleContext { ThreadEvent::ShowError(thread_error) => { tx.try_send(Err(anyhow!(thread_error.clone()))).ok(); } - ThreadEvent::RetryScheduled { .. } => { - // Ignore retry events - } ThreadEvent::RetriesFailed { .. } => { // Ignore retries failed events } From bc7dc882f51dd75a663cc0cbfdb0035d1f305331 Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Wed, 25 Jun 2025 20:18:44 -0400 Subject: [PATCH 27/32] Share code on intent --- crates/agent_ui/src/active_thread.rs | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/crates/agent_ui/src/active_thread.rs b/crates/agent_ui/src/active_thread.rs index 384a691efecf7f..520310fa2a22f9 100644 --- a/crates/agent_ui/src/active_thread.rs +++ b/crates/agent_ui/src/active_thread.rs @@ -3812,15 +3812,9 @@ mod tests { // Stream response to user message thread.update(cx, |thread, cx| { - let request = - thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx); - thread.stream_completion( - request, - model, - CompletionIntent::UserPrompt, - cx.active_window(), - cx, - ) + let intent = CompletionIntent::UserPrompt; + let request = thread.to_completion_request(model.clone(), intent, cx); + thread.stream_completion(request, model, intent, cx.active_window(), cx) }); // Follow the agent cx.update(|window, cx| { From 353056267520b7e8a79444c3466d62f29f945afe Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Wed, 25 Jun 2025 17:12:45 -0400 Subject: [PATCH 28/32] Simulate Overloaded error --- crates/agent/src/thread.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index 1ecd04de63a8e3..0465fd047c3278 100644 --- a/crates/agent/src/thread.rs +++ b/crates/agent/src/thread.rs @@ -1531,6 +1531,10 @@ impl Thread { } thread.update(cx, |thread, cx| { + // Simulate Overloaded error for testing retry logic + eprintln!("Ignoring event {event:?} and returning Overloaded"); + return Err(anyhow::Error::new(LanguageModelKnownError::Overloaded)); + let event = match event { Ok(event) => event, Err(error) => { From f96bf6fe2b490e2bc8871c69c9364fe9d05195a8 Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Wed, 25 Jun 2025 20:19:26 -0400 Subject: [PATCH 29/32] Revert "Simulate Overloaded error" This reverts commit 353056267520b7e8a79444c3466d62f29f945afe. --- crates/agent/src/thread.rs | 4 ---- 1 file changed, 4 deletions(-) diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index 0465fd047c3278..1ecd04de63a8e3 100644 --- a/crates/agent/src/thread.rs +++ b/crates/agent/src/thread.rs @@ -1531,10 +1531,6 @@ impl Thread { } thread.update(cx, |thread, cx| { - // Simulate Overloaded error for testing retry logic - eprintln!("Ignoring event {event:?} and returning Overloaded"); - return Err(anyhow::Error::new(LanguageModelKnownError::Overloaded)); - let event = match event { Ok(event) => event, Err(error) => { From 7e51a36bbd8cd5eb0d4258e3276e0690f3fef81e Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Wed, 25 Jun 2025 20:21:40 -0400 Subject: [PATCH 30/32] Re-add whitespace --- crates/agent_ui/src/active_thread.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/crates/agent_ui/src/active_thread.rs b/crates/agent_ui/src/active_thread.rs index 520310fa2a22f9..ad0901153ad2aa 100644 --- a/crates/agent_ui/src/active_thread.rs +++ b/crates/agent_ui/src/active_thread.rs @@ -1831,6 +1831,7 @@ impl ActiveThread { let tool_uses = thread.tool_uses_for_message(message_id, cx); let has_tool_uses = !tool_uses.is_empty(); + let editing_message_state = self .editing_message .as_ref() From c281f9256043b006357503e1607f9f19fdb3f79c Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Wed, 25 Jun 2025 20:22:56 -0400 Subject: [PATCH 31/32] Drop unused field in RetriesFailed --- crates/agent/src/thread.rs | 2 -- crates/agent_ui/src/active_thread.rs | 2 +- crates/agent_ui/src/agent_diff.rs | 1 + 3 files changed, 2 insertions(+), 3 deletions(-) diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index 1ecd04de63a8e3..5155b6bc7f9740 100644 --- a/crates/agent/src/thread.rs +++ b/crates/agent/src/thread.rs @@ -2216,7 +2216,6 @@ impl Thread { cx.emit(ThreadEvent::RetriesFailed { message: notification_text, - max_attempts, }); false @@ -3168,7 +3167,6 @@ pub enum ThreadEvent { ProfileChanged, RetriesFailed { message: SharedString, - max_attempts: u8, }, } diff --git a/crates/agent_ui/src/active_thread.rs b/crates/agent_ui/src/active_thread.rs index ad0901153ad2aa..4da959d36e9f77 100644 --- a/crates/agent_ui/src/active_thread.rs +++ b/crates/agent_ui/src/active_thread.rs @@ -1140,7 +1140,7 @@ impl ActiveThread { self.save_thread(cx); cx.notify(); } - ThreadEvent::RetriesFailed { message, .. } => { + ThreadEvent::RetriesFailed { message } => { self.show_notification(message, ui::IconName::Warning, window, cx); } } diff --git a/crates/agent_ui/src/agent_diff.rs b/crates/agent_ui/src/agent_diff.rs index aa5e49551b4e1c..b8e67512e2b069 100644 --- a/crates/agent_ui/src/agent_diff.rs +++ b/crates/agent_ui/src/agent_diff.rs @@ -1380,6 +1380,7 @@ impl AgentDiff { | ThreadEvent::ToolConfirmationNeeded | ThreadEvent::ToolUseLimitReached | ThreadEvent::CancelEditing + | ThreadEvent::RetriesFailed { .. } | ThreadEvent::ProfileChanged => {} } } From b409cd11940b133e5288fca3d28bc699d2707b12 Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Wed, 25 Jun 2025 20:23:32 -0400 Subject: [PATCH 32/32] Delete unnecessary comment --- crates/agent/src/thread.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index 5155b6bc7f9740..68624d7c3b9d30 100644 --- a/crates/agent/src/thread.rs +++ b/crates/agent/src/thread.rs @@ -50,7 +50,6 @@ use util::{ResultExt as _, post_inc}; use uuid::Uuid; use zed_llm_client::{CompletionIntent, CompletionRequestStatus, UsageLimit}; -// Retry constants const MAX_RETRY_ATTEMPTS: u8 = 3; const BASE_RETRY_DELAY_SECS: u64 = 5;