test: add unit tests for ThreadManager.truncate_from_message

Five test cases covering:
- truncate target + subsequent messages
- truncate from first message clears all
- nonexistent message ID is a no-op
- truncate last message only
- truncation is scoped to the target thread

Signed-off-by: Taylor Ho <taylorkmho@gmail.com>

crates/goose/src/session/thread_manager.rs

@@ -468,3 +468,201 @@ fn append_text_json(content_json: &str, new_text: &str) -> anyhow::Result<String
     }
     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);
+    }
+}

crates/goose/tests/truncate_from_message_test.rs

@@ -0,0 +1,122 @@
+use goose::conversation::message::Message;
+use goose::session::session_manager::SessionStorage;
+use goose::session::thread_manager::ThreadManager;
+use std::sync::Arc;
+use tempfile::TempDir;
+
+async fn setup() -> (ThreadManager, TempDir) {
+    let tmp = TempDir::new().unwrap();
+    let storage = SessionStorage::create(tmp.path()).await.unwrap();
+    let tm = ThreadManager::new(Arc::new(storage));
+    (tm, tmp)
+}
+
+/// Helper: append a user message with a known ID and return that ID.
+async fn append_user(tm: &ThreadManager, thread_id: &str, text: &str) -> String {
+    let mut msg = Message::user().with_text(text);
+    let id = format!("test_{}", uuid::Uuid::new_v4());
+    msg.id = Some(id.clone());
+    tm.append_message(thread_id, Some("s1"), &msg)
+        .await
+        .unwrap();
+    id
+}
+
+/// Helper: append an assistant message and return its ID.
+async fn append_assistant(tm: &ThreadManager, thread_id: &str, text: &str) -> String {
+    let mut msg = Message::assistant().with_text(text);
+    let id = format!("test_{}", uuid::Uuid::new_v4());
+    msg.id = Some(id.clone());
+    tm.append_message(thread_id, Some("s1"), &msg)
+        .await
+        .unwrap();
+    id
+}
+
+#[tokio::test]
+async fn truncate_removes_target_and_subsequent_messages() {
+    let (tm, _tmp) = setup().await;
+    let thread = tm.create_thread(None, None, None).await.unwrap();
+
+    let id1 = append_user(&tm, &thread.id, "msg 1").await;
+    let _id2 = append_assistant(&tm, &thread.id, "reply 1").await;
+    let id3 = append_user(&tm, &thread.id, "msg 2").await;
+    let _id4 = append_assistant(&tm, &thread.id, "reply 2").await;
+
+    // Truncate from msg 2 onward
+    let (deleted, ts) = tm.truncate_from_message(&thread.id, &id3).await.unwrap();
+    assert_eq!(deleted, 2, "should delete msg 2 + reply 2");
+    assert!(ts > 0, "should return a valid timestamp");
+
+    let remaining = tm.list_messages(&thread.id).await.unwrap();
+    assert_eq!(remaining.len(), 2);
+    assert_eq!(remaining[0].id.as_deref(), Some(id1.as_str()));
+}
+
+#[tokio::test]
+async fn truncate_from_first_message_clears_all() {
+    let (tm, _tmp) = setup().await;
+    let thread = tm.create_thread(None, None, None).await.unwrap();
+
+    let id1 = append_user(&tm, &thread.id, "first").await;
+    let _id2 = append_assistant(&tm, &thread.id, "reply").await;
+
+    let (deleted, _ts) = tm.truncate_from_message(&thread.id, &id1).await.unwrap();
+    assert_eq!(deleted, 2);
+
+    let remaining = tm.list_messages(&thread.id).await.unwrap();
+    assert!(remaining.is_empty());
+}
+
+#[tokio::test]
+async fn truncate_nonexistent_message_is_noop() {
+    let (tm, _tmp) = setup().await;
+    let thread = tm.create_thread(None, None, None).await.unwrap();
+
+    append_user(&tm, &thread.id, "keep me").await;
+
+    let (deleted, ts) = tm
+        .truncate_from_message(&thread.id, "does_not_exist")
+        .await
+        .unwrap();
+    assert_eq!(deleted, 0);
+    assert_eq!(ts, 0);
+
+    let remaining = tm.list_messages(&thread.id).await.unwrap();
+    assert_eq!(remaining.len(), 1);
+}
+
+#[tokio::test]
+async fn truncate_last_message_only() {
+    let (tm, _tmp) = setup().await;
+    let thread = tm.create_thread(None, None, None).await.unwrap();
+
+    let _id1 = append_user(&tm, &thread.id, "keep").await;
+    let _id2 = append_assistant(&tm, &thread.id, "keep too").await;
+    let id3 = append_user(&tm, &thread.id, "delete me").await;
+
+    let (deleted, _ts) = tm.truncate_from_message(&thread.id, &id3).await.unwrap();
+    assert_eq!(deleted, 1);
+
+    let remaining = tm.list_messages(&thread.id).await.unwrap();
+    assert_eq!(remaining.len(), 2);
+}
+
+#[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 id_a1 = append_user(&tm, &thread_a.id, "thread A msg").await;
+    append_user(&tm, &thread_b.id, "thread B msg").await;
+
+    tm.truncate_from_message(&thread_a.id, &id_a1)
+        .await
+        .unwrap();
+
+    let remaining_a = tm.list_messages(&thread_a.id).await.unwrap();
+    let remaining_b = tm.list_messages(&thread_b.id).await.unwrap();
+    assert!(remaining_a.is_empty());
+    assert_eq!(remaining_b.len(), 1, "thread B should be untouched");
+}