diff --git a/crates/agent/src/db.rs b/crates/agent/src/db.rs index 61768833506..9d0541475bb 100644 --- a/crates/agent/src/db.rs +++ b/crates/agent/src/db.rs @@ -668,31 +668,48 @@ impl ThreadsDatabase { let connection = self.connection.clone(); self.executor.spawn(async move { - let sandboxed_terminal_temp_dir = { + let sandboxed_terminal_temp_dirs = { let connection = connection.lock(); + let mut select_children = + connection.select_bound::, Arc>(indoc! {" + SELECT id FROM threads WHERE parent_id = ? + "})?; + + // Collect target thread together with all of its transitive + // subagent threads + let mut ids_to_delete = vec![id.0.clone()]; + let mut frontier = vec![id.0.clone()]; + while let Some(parent) = frontier.pop() { + for child in select_children(parent)? { + ids_to_delete.push(child.clone()); + frontier.push(child); + } + } + let mut select = connection.select_bound::, (DataType, Vec)>(indoc! {" SELECT data_type, data FROM threads WHERE id = ? LIMIT 1 "})?; - let sandboxed_terminal_temp_dir = select(id.0.clone())? - .into_iter() - .next() - .and_then(|(data_type, data)| { - Self::sandboxed_terminal_temp_dir(data_type, data) - }); - let mut delete = connection.exec_bound::>(indoc! {" DELETE FROM threads WHERE id = ? "})?; - delete(id.0)?; + let mut sandboxed_terminal_temp_dirs = Vec::new(); + for thread_id in ids_to_delete { + if let Some(temp_dir) = select(thread_id.clone())?.into_iter().next().and_then( + |(data_type, data)| Self::sandboxed_terminal_temp_dir(data_type, data), + ) { + sandboxed_terminal_temp_dirs.push(temp_dir); + } + delete(thread_id)?; + } - sandboxed_terminal_temp_dir + sandboxed_terminal_temp_dirs }; - if let Some(temp_dir) = sandboxed_terminal_temp_dir { + for temp_dir in sandboxed_terminal_temp_dirs { Self::remove_sandboxed_terminal_temp_dir(temp_dir); } @@ -1009,6 +1026,62 @@ mod tests { assert!(!temp_dir.exists()); } + #[gpui::test] + async fn test_delete_thread_deletes_subagent_threads(cx: &mut TestAppContext) { + let database = ThreadsDatabase::new(cx.executor()).unwrap(); + + let parent_id = session_id("parent-thread"); + let child_id = session_id("child-thread"); + let grandchild_id = session_id("grandchild-thread"); + let unrelated_id = session_id("unrelated-thread"); + + let parent_thread = make_thread( + "Parent Thread", + Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(), + ); + + let mut child_thread = make_thread( + "Child Subagent Thread", + Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(), + ); + child_thread.subagent_context = Some(crate::SubagentContext { + parent_thread_id: parent_id.clone(), + depth: 1, + }); + + let mut grandchild_thread = make_thread( + "Grandchild Subagent Thread", + Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(), + ); + grandchild_thread.subagent_context = Some(crate::SubagentContext { + parent_thread_id: child_id.clone(), + depth: 2, + }); + + let unrelated_thread = make_thread( + "Unrelated Thread", + Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(), + ); + + for (id, thread) in [ + (parent_id.clone(), parent_thread), + (child_id.clone(), child_thread), + (grandchild_id.clone(), grandchild_thread), + (unrelated_id.clone(), unrelated_thread), + ] { + database + .save_thread(id, thread, PathList::default()) + .await + .unwrap(); + } + + database.delete_thread(parent_id.clone()).await.unwrap(); + + let remaining = database.list_threads().await.unwrap(); + let remaining_ids: Vec<_> = remaining.iter().map(|thread| thread.id.clone()).collect(); + assert_eq!(remaining_ids, vec![unrelated_id]); + } + #[gpui::test] async fn test_subagent_context_roundtrips_through_save_load(cx: &mut TestAppContext) { let database = ThreadsDatabase::new(cx.executor()).unwrap();