From bc7e0632444cdc665a2bf715f0f317a24ec4a474 Mon Sep 17 00:00:00 2001
From: Douwe Osinga <douwe@block.xyz>
Date: Fri, 13 Mar 2026 04:51:28 -0400
Subject: [PATCH] Keep messages in sync (#7850)

Co-authored-by: Douwe Osinga <douwe@squareup.com>
---
 crates/goose/src/agents/agent.rs | 31 ++++++++++++++++++-------------
 1 file changed, 18 insertions(+), 13 deletions(-)

diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs
index 62b40d6ee6..f6f096a472 100644
--- a/crates/goose/src/agents/agent.rs
+++ b/crates/goose/src/agents/agent.rs
@@ -1180,7 +1180,6 @@ impl Agent {
                 ).await?;
 
                 let mut no_tools_called = true;
-                let mut messages_to_add = Conversation::default();
                 let mut tools_updated = false;
                 let mut did_recovery_compact_this_iteration = false;
                 let mut exit_chat = false;
@@ -1235,7 +1234,8 @@ impl Agent {
                                     if !text.is_empty() {
                                         last_assistant_text = text;
                                     }
-                                    messages_to_add.push(response.clone());
+                                    session_manager.add_message(&session_config.id, &response).await?;
+                                    conversation.push(response);
                                     continue;
                                 }
 
@@ -1431,7 +1431,8 @@ impl Agent {
                                         response.created,
                                         thinking_content,
                                     ).with_id(format!("msg_{}", Uuid::new_v4()));
-                                    messages_to_add.push(thinking_msg);
+                                    session_manager.add_message(&session_config.id, &thinking_msg).await?;
+                                    conversation.push(thinking_msg);
                                 }
 
                                 // Collect reasoning content to attach to tool request messages
@@ -1459,11 +1460,14 @@ impl Agent {
                                                 request.metadata.as_ref(),
                                                 request.tool_meta.clone(),
                                             );
-                                        messages_to_add.push(request_msg);
                                         let final_response = tool_response_messages[idx]
                                                                 .lock().await.clone();
-                                        yield AgentEvent::Message(final_response.clone());
-                                        messages_to_add.push(final_response);
+                                        // Persist the tool request and response as a pair
+                                        session_manager.add_message(&session_config.id, &request_msg).await?;
+                                        session_manager.add_message(&session_config.id, &final_response).await?;
+                                        conversation.push(request_msg);
+                                        conversation.push(final_response.clone());
+                                        yield AgentEvent::Message(final_response);
                                     } else {
                                         error!(
                                             "Tool call could not be parsed: {}",
@@ -1600,11 +1604,13 @@ impl Agent {
                         if final_output_tool.final_output.is_none() {
                             warn!("Final output tool has not been called yet. Continuing agent loop.");
                             let message = Message::user().with_text(FINAL_OUTPUT_CONTINUATION_MESSAGE);
-                            messages_to_add.push(message.clone());
+                            session_manager.add_message(&session_config.id, &message).await?;
+                            conversation.push(message.clone());
                             yield AgentEvent::Message(message);
                         } else {
                             let message = Message::assistant().with_text(final_output_tool.final_output.clone().unwrap());
-                            messages_to_add.push(message.clone());
+                            session_manager.add_message(&session_config.id, &message).await?;
+                            conversation.push(message.clone());
                             yield AgentEvent::Message(message);
                             exit_chat = true;
                         }
@@ -1615,6 +1621,8 @@ impl Agent {
                             Ok(should_retry) => {
                                 if should_retry {
                                     info!("Retry logic triggered, restarting agent loop");
+                                    session_manager.replace_conversation(&session_config.id, &conversation).await?;
+                                    yield AgentEvent::HistoryReplaced(conversation.clone());
                                 } else {
                                     exit_chat = true;
                                 }
@@ -1655,17 +1663,14 @@ impl Agent {
                             }).await?;
                         }
                         conversation = Conversation::new_unvalidated(updated_messages);
-                        messages_to_add.push(summary_msg);
+                        session_manager.add_message(&session_config.id, &summary_msg).await?;
+                        conversation.push(summary_msg);
                     } else {
                         warn!("Expected a tool request/reply pair, but found {} matching messages",
                             matching.len());
                     }
                 }
 
-                for msg in &messages_to_add {
-                    session_manager.add_message(&session_config.id, msg).await?;
-                }
-                conversation.extend(messages_to_add);
                 if exit_chat {
                     break;
                 }