From 0e87863aeb7d86ce586cf351f11c9f4e4d160609 Mon Sep 17 00:00:00 2001 From: duongynhi000005-oss Date: Mon, 25 May 2026 21:07:13 +0000 Subject: [PATCH] feat: add first-class retry API for conversations Implements API::retry(conversation_id) and wires /retry + conversation retry to reuse last non-droppable user message. Co-Authored-By: ForgeCode --- crates/forge_api/src/api.rs | 6 +++++ crates/forge_api/src/forge_api.rs | 23 ++++++++++++++++++ crates/forge_domain/src/conversation.rs | 29 +++++++++++++++++++++++ crates/forge_main/src/ui.rs | 31 +++++++++++++++++++++---- 4 files changed, 85 insertions(+), 4 deletions(-) diff --git a/crates/forge_api/src/api.rs b/crates/forge_api/src/api.rs index 5a2a5217fe..7ba484861c 100644 --- a/crates/forge_api/src/api.rs +++ b/crates/forge_api/src/api.rs @@ -46,6 +46,12 @@ pub trait API: Sync + Send { /// Executes a chat request and returns a stream of responses async fn chat(&self, chat: ChatRequest) -> Result>>; + /// Retries the last non-droppable user message in the conversation. + async fn retry( + &self, + conversation_id: ConversationId, + ) -> Result>>; + /// Commits changes with an AI-generated commit message async fn commit( &self, diff --git a/crates/forge_api/src/forge_api.rs b/crates/forge_api/src/forge_api.rs index a056705761..273aebeeaf 100644 --- a/crates/forge_api/src/forge_api.rs +++ b/crates/forge_api/src/forge_api.rs @@ -147,6 +147,29 @@ impl< self.app().chat(agent_id, chat).await } + async fn retry( + &self, + conversation_id: ConversationId, + ) -> anyhow::Result>> { + let conversation = self + .services + .find_conversation(&conversation_id) + .await? + .ok_or_else(|| forge_domain::Error::ConversationNotFound(conversation_id))?; + + let user_message = conversation.last_user_message() + .ok_or_else(|| anyhow::anyhow!("No user message found to retry"))?; + + let event = user_message + .raw_content + .clone() + .map(Event::new) + .unwrap_or_else(|| Event::new(user_message.content.clone())); + + let chat = ChatRequest::new(event, conversation_id); + self.chat(chat).await + } + async fn upsert_conversation(&self, conversation: Conversation) -> anyhow::Result<()> { self.services.upsert_conversation(conversation).await } diff --git a/crates/forge_domain/src/conversation.rs b/crates/forge_domain/src/conversation.rs index c0bde6e4e8..4369673e8a 100644 --- a/crates/forge_domain/src/conversation.rs +++ b/crates/forge_domain/src/conversation.rs @@ -113,6 +113,20 @@ impl Conversation { .unwrap_or_default() } + /// Returns the last non-droppable user message in the conversation. + pub fn last_user_message(&self) -> Option<&crate::TextMessage> { + self.context.as_ref()?.messages.iter().rev().find_map(|message| { + match &**message { + crate::ContextMessage::Text(text) + if text.role == crate::Role::User && !text.droppable => + { + Some(text) + } + _ => None, + } + }) + } + /// Returns the total token usage across all messages in the conversation. /// /// This is a convenience method that aggregates usage from the context, @@ -281,6 +295,21 @@ mod tests { assert_eq!(actual, vec![agent_conv_id]); } + #[test] + fn test_last_user_message_returns_latest_non_droppable_message() { + let context = Context::default() + .add_message(ContextMessage::user("First", None)) + .add_message(ContextMessage::assistant("Reply", None, None, None)) + .add_message(ContextMessage::user("Retry me", None)) + .add_message(crate::TextMessage::new(crate::Role::User, "Ignored").droppable(true)); + + let conversation = Conversation::generate().context(context); + let actual = conversation.last_user_message().map(|message| message.content.as_str()); + let expected = Some("Retry me"); + + assert_eq!(actual, expected); + } + #[test] fn test_total_cost() { use crate::{MessageEntry, Usage}; diff --git a/crates/forge_main/src/ui.rs b/crates/forge_main/src/ui.rs index 0bfd3112df..16bd858765 100644 --- a/crates/forge_main/src/ui.rs +++ b/crates/forge_main/src/ui.rs @@ -905,8 +905,7 @@ impl A + Send + Sync> UI let original_id = self.state.conversation_id; self.state.conversation_id = Some(id); - self.spinner.start(None)?; - self.on_message(None).await?; + self.on_retry(id).await?; self.state.conversation_id = original_id; } @@ -2191,8 +2190,8 @@ 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.init_conversation().await?; + self.on_retry(conversation_id).await?; } AppCommand::Index => { let working_dir = self.state.cwd.clone(); @@ -3923,6 +3922,30 @@ impl A + Send + Sync> UI Ok(()) } + async fn on_retry(&mut self, conversation_id: ConversationId) -> Result<()> { + self.spinner.start(None)?; + self.writeln_title(TitleFormat::action("Retrying last prompt"))?; + let mut stream = self.api.retry(conversation_id).await?; + let mut writer = StreamingWriter::new(self.spinner.clone(), self.api.clone()); + + while let Some(message) = stream.next().await { + match message { + Ok(message) => self.handle_chat_response(message, &mut writer).await?, + Err(err) => { + writer.finish()?; + self.spinner.stop(None)?; + self.spinner.reset(); + return Err(err); + } + } + } + + writer.finish()?; + self.spinner.stop(None)?; + self.spinner.reset(); + Ok(()) + } + /// Fetches related conversations for a given conversation in parallel. /// /// Returns a vector of related conversations that could be successfully