feat(goose2): inline edit & retry

crates/goose-acp/src/server.rs

@@ -1523,12 +1523,17 @@ impl GooseAcpAgent {
                 continue;
             }
 
+            let msg_id = message.id.clone();
+
             for content_item in &message.content {
                 match content_item {
                     MessageContent::Text(text) => {
-                        let chunk = ContentChunk::new(ContentBlock::Text(TextContent::new(
+                        let mut chunk = ContentChunk::new(ContentBlock::Text(TextContent::new(
                             text.text.clone(),
                         )));
+                        if let Some(ref id) = msg_id {
+                            chunk = chunk.message_id(id.clone());
+                        }
                         let update = match message.role {
                             Role::User => SessionUpdate::UserMessageChunk(chunk),
                             Role::Assistant => SessionUpdate::AgentMessageChunk(chunk),
@@ -1691,7 +1696,49 @@ impl GooseAcpAgent {
             .get_session_agent(&thread_id, Some(cancel_token.clone()))
             .await?;
 
-        let user_message = self.convert_acp_prompt_to_message(args.prompt);
+        // 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::debug!(
+                thread_id = %thread_id,
+                truncate_id = %truncate_id,
+                rows_deleted = rows_deleted,
+                "Truncated conversation for edit/retry"
+            );
+        }
+
+        let mut user_message = self.convert_acp_prompt_to_message(args.prompt);
+        // Use the client-provided message_id so the frontend and backend
+        // share the same ID — required for edit/retry truncation to work.
+        if let Some(client_msg_id) = args.message_id {
+            user_message.id = Some(client_msg_id);
+        }
 
         self.thread_manager
             .append_message(&thread_id, Some(&internal_session_id), &user_message)

crates/goose/src/session/thread_manager.rs

@@ -378,6 +378,43 @@
         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)>(
@@ -431,3 +468,201 @@
     }
     Ok(serde_json::to_string(&items)?)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::conversation::message::{Message, MessageContent};
+    use rmcp::model::Role;
+    use std::sync::Arc;
+    use tempfile::TempDir;
+
+    /// Create a ThreadManager backed by a temporary SQLite database.
+    async fn setup() -> (ThreadManager, TempDir) {
+        let tmp = TempDir::new().unwrap();
+        let storage = Arc::new(SessionStorage::new(tmp.path().to_path_buf()));
+        // Force schema initialisation by touching the pool once.
+        let _ = storage.pool().await.unwrap();
+        (ThreadManager::new(storage), tmp)
+    }
+
+    /// Helper: build a simple text message with a given id, alternating role,
+    /// and a distinct `created` timestamp so every insert produces its own row.
+    fn make_msg(id: &str, role: Role, ts: i64) -> Message {
+        Message::new(role, ts, vec![MessageContent::text(format!("body-{id}"))]).with_id(id)
+    }
+
+    /// Insert `n` messages into `thread_id`, alternating User/Assistant so the
+    /// coalescing logic never merges them.  Returns the stored messages (with
+    /// their final ids).
+    async fn insert_messages(
+        tm: &ThreadManager,
+        thread_id: &str,
+        n: usize,
+    ) -> Vec<Message> {
+        let mut stored = Vec::new();
+        for i in 0..n {
+            let role = if i % 2 == 0 { Role::User } else { Role::Assistant };
+            let msg = make_msg(&format!("msg{i}"), role, (100 + i) as i64);
+            let s = tm.append_message(thread_id, None, &msg).await.unwrap();
+            stored.push(s);
+        }
+        stored
+    }
+
+    // ── Basic truncation from the middle ────────────────────────────────
+
+    #[tokio::test]
+    async fn truncate_from_middle_deletes_tail() {
+        let (tm, _tmp) = setup().await;
+        let thread = tm.create_thread(None, None, None).await.unwrap();
+        let msgs = insert_messages(&tm, &thread.id, 5).await;
+
+        // Truncate from msg2 (index 2) → should delete msgs 2, 3, 4
+        let (deleted, ts) = tm
+            .truncate_from_message(&thread.id, msgs[2].id.as_ref().unwrap())
+            .await
+            .unwrap();
+
+        assert_eq!(deleted, 3, "should delete msg2, msg3, msg4");
+        assert_eq!(ts, msgs[2].created, "returned timestamp must match msg2");
+
+        let remaining = tm.list_messages(&thread.id).await.unwrap();
+        assert_eq!(remaining.len(), 2);
+        assert_eq!(remaining[0].id, msgs[0].id);
+        assert_eq!(remaining[1].id, msgs[1].id);
+    }
+
+    // ── Truncation from the first message (delete everything) ───────────
+
+    #[tokio::test]
+    async fn truncate_from_first_deletes_all() {
+        let (tm, _tmp) = setup().await;
+        let thread = tm.create_thread(None, None, None).await.unwrap();
+        let msgs = insert_messages(&tm, &thread.id, 4).await;
+
+        let (deleted, ts) = tm
+            .truncate_from_message(&thread.id, msgs[0].id.as_ref().unwrap())
+            .await
+            .unwrap();
+
+        assert_eq!(deleted, 4);
+        assert_eq!(ts, msgs[0].created);
+
+        let remaining = tm.list_messages(&thread.id).await.unwrap();
+        assert!(remaining.is_empty(), "all messages should be deleted");
+    }
+
+    // ── Truncation from the last message (delete only one) ──────────────
+
+    #[tokio::test]
+    async fn truncate_from_last_deletes_only_one() {
+        let (tm, _tmp) = setup().await;
+        let thread = tm.create_thread(None, None, None).await.unwrap();
+        let msgs = insert_messages(&tm, &thread.id, 4).await;
+
+        let (deleted, ts) = tm
+            .truncate_from_message(&thread.id, msgs[3].id.as_ref().unwrap())
+            .await
+            .unwrap();
+
+        assert_eq!(deleted, 1);
+        assert_eq!(ts, msgs[3].created);
+
+        let remaining = tm.list_messages(&thread.id).await.unwrap();
+        assert_eq!(remaining.len(), 3);
+    }
+
+    // ── Message not found → (0, 0) ─────────────────────────────────────
+
+    #[tokio::test]
+    async fn truncate_message_not_found_returns_zero() {
+        let (tm, _tmp) = setup().await;
+        let thread = tm.create_thread(None, None, None).await.unwrap();
+        let _msgs = insert_messages(&tm, &thread.id, 3).await;
+
+        let (deleted, ts) = tm
+            .truncate_from_message(&thread.id, "nonexistent-id")
+            .await
+            .unwrap();
+
+        assert_eq!(deleted, 0);
+        assert_eq!(ts, 0);
+
+        // Nothing should have been removed.
+        let remaining = tm.list_messages(&thread.id).await.unwrap();
+        assert_eq!(remaining.len(), 3);
+    }
+
+    // ── Cross-thread isolation ──────────────────────────────────────────
+
+    #[tokio::test]
+    async fn truncate_does_not_affect_other_threads() {
+        let (tm, _tmp) = setup().await;
+        let thread_a = tm.create_thread(None, None, None).await.unwrap();
+        let thread_b = tm.create_thread(None, None, None).await.unwrap();
+
+        let msgs_a = insert_messages(&tm, &thread_a.id, 4).await;
+        let msgs_b = insert_messages(&tm, &thread_b.id, 3).await;
+
+        // Truncate thread_a from its first message (delete all of A).
+        let (deleted, _) = tm
+            .truncate_from_message(&thread_a.id, msgs_a[0].id.as_ref().unwrap())
+            .await
+            .unwrap();
+        assert_eq!(deleted, 4);
+
+        // Thread B must be completely untouched.
+        let remaining_b = tm.list_messages(&thread_b.id).await.unwrap();
+        assert_eq!(remaining_b.len(), 3);
+        for (i, m) in remaining_b.iter().enumerate() {
+            assert_eq!(m.id, msgs_b[i].id);
+        }
+    }
+
+    // ── Return-value correctness (rows_deleted + created_timestamp) ─────
+
+    #[tokio::test]
+    async fn return_value_matches_expectations() {
+        let (tm, _tmp) = setup().await;
+        let thread = tm.create_thread(None, None, None).await.unwrap();
+        let msgs = insert_messages(&tm, &thread.id, 6).await;
+
+        // Truncate from index 4 → should delete 2 messages (index 4, 5).
+        let (deleted, ts) = tm
+            .truncate_from_message(&thread.id, msgs[4].id.as_ref().unwrap())
+            .await
+            .unwrap();
+
+        assert_eq!(deleted, 2);
+        assert_eq!(ts, msgs[4].created);
+        assert_eq!(ts, 104); // 100 + 4
+
+        let remaining = tm.list_messages(&thread.id).await.unwrap();
+        assert_eq!(remaining.len(), 4);
+    }
+
+    // ── Wrong thread_id with valid message_id → not found ───────────────
+
+    #[tokio::test]
+    async fn truncate_wrong_thread_returns_zero() {
+        let (tm, _tmp) = setup().await;
+        let thread_a = tm.create_thread(None, None, None).await.unwrap();
+        let thread_b = tm.create_thread(None, None, None).await.unwrap();
+
+        let msgs_a = insert_messages(&tm, &thread_a.id, 3).await;
+
+        // Use thread_b's id but a message_id that belongs to thread_a.
+        let (deleted, ts) = tm
+            .truncate_from_message(&thread_b.id, msgs_a[1].id.as_ref().unwrap())
+            .await
+            .unwrap();
+
+        assert_eq!(deleted, 0);
+        assert_eq!(ts, 0);
+
+        // thread_a should be untouched.
+        let remaining = tm.list_messages(&thread_a.id).await.unwrap();
+        assert_eq!(remaining.len(), 3);
+    }
+}