diff --git a/crates/forge_domain/src/conversation.rs b/crates/forge_domain/src/conversation.rs index c0bde6e4e8..0e40b606ac 100644 --- a/crates/forge_domain/src/conversation.rs +++ b/crates/forge_domain/src/conversation.rs @@ -6,7 +6,9 @@ use derive_setters::Setters; use serde::{Deserialize, Serialize}; use uuid::Uuid; -use crate::{Context, Error, Metrics, Result, TokenCount}; +use crate::{ + ChatRequest, Context, ContextMessage, Error, Event, EventValue, Metrics, Result, TokenCount, +}; #[derive(Debug, Default, Display, Serialize, Deserialize, Clone, PartialEq, Eq, Hash)] #[serde(transparent)] @@ -106,13 +108,29 @@ impl Conversation { /// Returns a vector of user messages, selecting the first message from /// each consecutive sequence of user messages. - pub fn first_user_messages(&self) -> Vec<&crate::ContextMessage> { + pub fn first_user_messages(&self) -> Vec<&ContextMessage> { self.context .as_ref() .map(|ctx| ctx.first_user_messages()) .unwrap_or_default() } + /// Rebuilds the request for the most recent user turn. + /// + /// This uses the first message in the last consecutive user block so that + /// retries replay the original prompt instead of trailing attachment or + /// context messages. + pub fn retry_request(&self) -> Option { + let message = self.first_user_messages().into_iter().last()?; + let event = match message.as_value() { + Some(EventValue::Text(user_prompt)) => Event::new(user_prompt.as_str()), + Some(EventValue::Command(command)) => Event::from(command.clone()), + None => message.content().map(Event::new)?, + }; + + Some(ChatRequest::new(event, self.id)) + } + /// Returns the total token usage across all messages in the conversation. /// /// This is a convenience method that aggregates usage from the context, @@ -207,7 +225,9 @@ mod tests { use pretty_assertions::assert_eq; use super::*; - use crate::{Context, ContextMessage, ToolOutput, ToolResult, ToolValue}; + use crate::{ + Context, ContextMessage, EventValue, Role, TextMessage, ToolOutput, ToolResult, ToolValue, + }; #[test] fn test_related_conversation_ids_empty() { @@ -355,4 +375,44 @@ mod tests { assert_eq!(actual, expected); } + + #[test] + fn test_retry_request_uses_first_message_from_last_user_block() { + let context = Context::default() + .add_message(ContextMessage::user("Original task", None)) + .add_message(ContextMessage::user("Attachment context", None)) + .add_message(ContextMessage::assistant("Working on it", None, None, None)) + .add_message(ContextMessage::user("Second task", None)) + .add_message(ContextMessage::user("More context", None)); + + let conversation = Conversation::generate().context(context); + let request = conversation.retry_request().expect("retry request"); + + let prompt = request + .event + .value + .as_ref() + .and_then(|value| value.as_user_prompt()) + .map(|prompt| prompt.as_str()); + + assert_eq!(prompt, Some("Second task")); + } + + #[test] + fn test_retry_request_prefers_raw_content() { + let message = TextMessage::new(Role::User, "Rendered prompt") + .raw_content(EventValue::text("Original prompt")); + let context = Context::default().add_message(ContextMessage::Text(message)); + let conversation = Conversation::generate().context(context); + let request = conversation.retry_request().expect("retry request"); + + let prompt = request + .event + .value + .as_ref() + .and_then(|value| value.as_user_prompt()) + .map(|prompt| prompt.as_str()); + + assert_eq!(prompt, Some("Original prompt")); + } } diff --git a/crates/forge_main/src/model.rs b/crates/forge_main/src/model.rs index aea647d14e..0d8cde56fa 100644 --- a/crates/forge_main/src/model.rs +++ b/crates/forge_main/src/model.rs @@ -1525,6 +1525,13 @@ mod tests { assert_eq!(actual, expected); } + #[test] + fn test_parse_retry_command() { + let fixture = ForgeCommandManager::default(); + let actual = fixture.parse("/retry").unwrap(); + assert!(matches!(actual, AppCommand::Retry)); + } + #[test] fn test_parse_tool_command() { // Setup diff --git a/crates/forge_main/src/ui.rs b/crates/forge_main/src/ui.rs index a517907b24..e6daea3843 100644 --- a/crates/forge_main/src/ui.rs +++ b/crates/forge_main/src/ui.rs @@ -951,15 +951,19 @@ impl A + Send + Sync> UI self.on_conversation_delete(conversation_id).await?; } ConversationCommand::Retry { id } => { - self.validate_conversation_exists(&id).await?; - + let conversation = self.validate_conversation_exists(&id).await?; + let chat = conversation.retry_request().ok_or_else(|| { + anyhow::anyhow!("Conversation '{id}' has no user message to retry.") + })?; let original_id = self.state.conversation_id; - self.state.conversation_id = Some(id); - self.spinner.start(None)?; - self.on_message(None).await?; + self.writeln_title(TitleFormat::action("Retrying"))?; + self.spinner.start(Some("Retrying"))?; + self.state.conversation_id = Some(id); + let result = self.on_chat(chat).await; self.state.conversation_id = original_id; + result?; } ConversationCommand::Resume { id } => { self.validate_conversation_exists(&id).await?; @@ -2289,8 +2293,20 @@ impl A + Send + Sync> UI return self.handle_provider_logout(None).await; } AppCommand::Retry => { - self.spinner.start(None)?; - self.on_message(None).await?; + let conversation_id = self + .state + .conversation_id + .ok_or_else(|| anyhow::anyhow!("No conversation available to retry."))?; + let conversation = self.validate_conversation_exists(&conversation_id).await?; + let chat = conversation.retry_request().ok_or_else(|| { + anyhow::anyhow!( + "Conversation '{conversation_id}' has no user message to retry." + ) + })?; + + self.writeln_title(TitleFormat::action("Retrying"))?; + self.spinner.start(Some("Retrying"))?; + self.on_chat(chat).await?; } AppCommand::Index => { let working_dir = self.state.cwd.clone(); @@ -4000,7 +4016,14 @@ impl A + Send + Sync> UI } async fn on_chat(&mut self, chat: ChatRequest) -> Result<()> { - let mut stream = self.api.chat(chat).await?; + let mut stream = match self.api.chat(chat).await { + Ok(stream) => stream, + Err(err) => { + self.spinner.stop(None)?; + self.spinner.reset(); + return Err(err); + } + }; // Always use streaming content writer let mut writer = StreamingWriter::new(self.spinner.clone(), self.api.clone());