From a5d12c2e7dfa1094ca9e1e5eba05fd910c5d4ea0 Mon Sep 17 00:00:00 2001
From: Taylor Ho <taylorkmho@gmail.com>
Date: Thu, 16 Apr 2026 13:42:45 -1000
Subject: [PATCH] feat: truncate backend conversation on edit/retry via
 _meta.truncate_before_message_id

Uses timestamp bridge: truncate thread_messages by message_id to get
the created_timestamp, then truncate messages table by that timestamp.
This handles the different ID schemes (tmsg_ vs msg_) across tables.
---
 crates/goose-acp/src/server.rs             | 37 ++++++++++++++++++++++
 crates/goose/src/session/thread_manager.rs | 37 ++++++++++++++++++++++
 2 files changed, 74 insertions(+)

diff --git a/crates/goose-acp/src/server.rs b/crates/goose-acp/src/server.rs
index bf902a9848..e1ae80a1e4 100644
--- a/crates/goose-acp/src/server.rs
+++ b/crates/goose-acp/src/server.rs
@@ -1691,6 +1691,43 @@ impl GooseAcpAgent {
             .get_session_agent(&thread_id, Some(cancel_token.clone()))
             .await?;
 
+        // Handle edit/retry truncation — if the client requests it via _meta,
+        // truncate both storage layers before appending the new message.
+        // The two tables use different ID schemes (tmsg_ vs msg_), so we use
+        // a timestamp bridge: truncate thread_messages by message_id to get
+        // the timestamp, then truncate messages by that timestamp.
+        if let Some(truncate_id) = args
+            .meta
+            .as_ref()
+            .and_then(|m| m.get("truncate_before_message_id"))
+            .and_then(|v| v.as_str())
+        {
+            let (rows_deleted, created_ts) = self
+                .thread_manager
+                .truncate_from_message(&thread_id, truncate_id)
+                .await
+                .map_err(|e| {
+                    sacp::Error::internal_error().data(format!("Failed to truncate thread: {}", e))
+                })?;
+
+            if created_ts > 0 {
+                self.session_manager
+                    .truncate_conversation(&internal_session_id, created_ts)
+                    .await
+                    .map_err(|e| {
+                        sacp::Error::internal_error()
+                            .data(format!("Failed to truncate session: {}", e))
+                    })?;
+            }
+
+            tracing::info!(
+                thread_id = %thread_id,
+                truncate_id = %truncate_id,
+                rows_deleted = rows_deleted,
+                "Truncated conversation for edit/retry"
+            );
+        }
+
         let user_message = self.convert_acp_prompt_to_message(args.prompt);
 
         self.thread_manager
diff --git a/crates/goose/src/session/thread_manager.rs b/crates/goose/src/session/thread_manager.rs
index 3a96cc69ff..a66e76bda4 100644
--- a/crates/goose/src/session/thread_manager.rs
+++ b/crates/goose/src/session/thread_manager.rs
@@ -378,6 +378,43 @@ impl ThreadManager {
         self.get_thread(&new_id).await
     }
 
+    /// Truncate a thread from a specific message onward (inclusive).
+    /// Returns (rows_deleted, created_timestamp) so callers can bridge
+    /// into session-level truncation which uses timestamps.
+    pub async fn truncate_from_message(
+        &self,
+        thread_id: &str,
+        message_id: &str,
+    ) -> Result<(u64, i64)> {
+        let pool = self.storage.pool().await?;
+
+        // Find the target message's autoincrement id and timestamp
+        let row = sqlx::query_as::<_, (i64, i64)>(
+            "SELECT id, created_timestamp FROM thread_messages WHERE thread_id = ? AND message_id = ? LIMIT 1",
+        )
+        .bind(thread_id)
+        .bind(message_id)
+        .fetch_optional(pool)
+        .await?;
+
+        let Some((row_id, created_ts)) = row else {
+            return Ok((0, 0));
+        };
+
+        let result = sqlx::query("DELETE FROM thread_messages WHERE thread_id = ? AND id >= ?")
+            .bind(thread_id)
+            .bind(row_id)
+            .execute(pool)
+            .await?;
+
+        sqlx::query("UPDATE threads SET updated_at = CURRENT_TIMESTAMP WHERE id = ?")
+            .bind(thread_id)
+            .execute(pool)
+            .await?;
+
+        Ok((result.rows_affected(), created_ts))
+    }
+
     pub async fn list_messages(&self, thread_id: &str) -> Result<Vec<Message>> {
         let pool = self.storage.pool().await?;
         let rows = sqlx::query_as::<_, (Option<String>, String, Option<String>, String, i64, String)>(