diff --git a/crates/goose/src/session/thread_manager.rs b/crates/goose/src/session/thread_manager.rs index a66e76bda4..bb4e3180bf 100644 --- a/crates/goose/src/session/thread_manager.rs +++ b/crates/goose/src/session/thread_manager.rs @@ -468,3 +468,201 @@ fn append_text_json(content_json: &str, new_text: &str) -> anyhow::Result (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 { + 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); + } +} diff --git a/crates/goose/tests/truncate_from_message_test.rs b/crates/goose/tests/truncate_from_message_test.rs new file mode 100644 index 0000000000..040c5e0c83 --- /dev/null +++ b/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"); +}