use crate::{DbThread, DbThreadMetadata, ThreadsDatabase}; use agent_client_protocol as acp; use anyhow::{Result, anyhow}; use gpui::{App, Context, Entity, Global, Task, prelude::*}; use util::path_list::PathList; struct GlobalThreadStore(Entity); impl Global for GlobalThreadStore {} pub struct ThreadStore { threads: Vec, } impl ThreadStore { pub fn init_global(cx: &mut App) { let thread_store = cx.new(|cx| Self::new(cx)); cx.set_global(GlobalThreadStore(thread_store)); } pub fn global(cx: &App) -> Entity { cx.global::().0.clone() } pub fn new(cx: &mut Context) -> Self { let this = Self { threads: Vec::new(), }; this.reload(cx); this } pub fn thread_from_session_id(&self, session_id: &acp::SessionId) -> Option<&DbThreadMetadata> { self.threads.iter().find(|thread| &thread.id == session_id) } pub fn load_thread( &mut self, id: acp::SessionId, cx: &mut Context, ) -> Task>> { let database_future = ThreadsDatabase::connect(cx); cx.background_spawn(async move { let database = database_future.await.map_err(|err| anyhow!(err))?; database.load_thread(id).await }) } pub fn save_thread( &mut self, id: acp::SessionId, thread: crate::DbThread, folder_paths: PathList, cx: &mut Context, ) -> Task> { let database_future = ThreadsDatabase::connect(cx); cx.spawn(async move |this, cx| { let database = database_future.await.map_err(|err| anyhow!(err))?; database.save_thread(id, thread, folder_paths).await?; this.update(cx, |this, cx| this.reload(cx)) }) } pub fn delete_thread( &mut self, id: acp::SessionId, cx: &mut Context, ) -> Task> { let database_future = ThreadsDatabase::connect(cx); cx.spawn(async move |this, cx| { let database = database_future.await.map_err(|err| anyhow!(err))?; database.delete_thread(id.clone()).await?; this.update(cx, |this, cx| this.reload(cx)) }) } pub fn delete_threads(&mut self, cx: &mut Context) -> Task> { let database_future = ThreadsDatabase::connect(cx); cx.spawn(async move |this, cx| { let database = database_future.await.map_err(|err| anyhow!(err))?; database.delete_threads().await?; this.update(cx, |this, cx| this.reload(cx)) }) } pub fn reload(&self, cx: &mut Context) { let database_connection = ThreadsDatabase::connect(cx); cx.spawn(async move |this, cx| { let database = database_connection.await.map_err(|err| anyhow!(err))?; let threads = database .list_threads() .await? .into_iter() .filter(|thread| thread.parent_session_id.is_none()) .collect::>(); this.update(cx, |this, cx| { this.threads = threads; cx.notify(); }) }) .detach_and_log_err(cx); } pub fn is_empty(&self) -> bool { self.threads.is_empty() } pub fn entries(&self) -> impl Iterator + '_ { self.threads.iter().cloned() } /// Returns threads whose folder_paths match the given paths exactly. pub fn threads_for_paths(&self, paths: &PathList) -> impl Iterator { self.threads .iter() .filter(move |thread| &thread.folder_paths == paths) } } #[cfg(test)] mod tests { use super::*; use chrono::{DateTime, TimeZone, Utc}; use collections::HashMap; use gpui::TestAppContext; use std::sync::Arc; fn session_id(value: &str) -> acp::SessionId { acp::SessionId::new(Arc::::from(value)) } fn make_thread(title: &str, updated_at: DateTime) -> DbThread { DbThread { title: title.to_string().into(), messages: Vec::new(), updated_at, detailed_summary: None, initial_project_snapshot: None, cumulative_token_usage: Default::default(), request_token_usage: HashMap::default(), model: None, profile: None, imported: false, subagent_context: None, speed: None, thinking_enabled: false, thinking_effort: None, draft_prompt: None, ui_scroll_position: None, } } #[gpui::test] async fn test_entries_are_sorted_by_updated_at(cx: &mut TestAppContext) { let thread_store = cx.new(|cx| ThreadStore::new(cx)); cx.run_until_parked(); let older_id = session_id("thread-a"); let newer_id = session_id("thread-b"); let older_thread = make_thread( "Thread A", Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(), ); let newer_thread = make_thread( "Thread B", Utc.with_ymd_and_hms(2024, 1, 2, 0, 0, 0).unwrap(), ); let save_older = thread_store.update(cx, |store, cx| { store.save_thread(older_id.clone(), older_thread, PathList::default(), cx) }); save_older.await.unwrap(); let save_newer = thread_store.update(cx, |store, cx| { store.save_thread(newer_id.clone(), newer_thread, PathList::default(), cx) }); save_newer.await.unwrap(); cx.run_until_parked(); let entries: Vec<_> = thread_store.read_with(cx, |store, _cx| store.entries().collect()); assert_eq!(entries.len(), 2); assert_eq!(entries[0].id, newer_id); assert_eq!(entries[1].id, older_id); } #[gpui::test] async fn test_delete_threads_clears_entries(cx: &mut TestAppContext) { let thread_store = cx.new(|cx| ThreadStore::new(cx)); cx.run_until_parked(); let thread_id = session_id("thread-a"); let thread = make_thread( "Thread A", Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(), ); let save_task = thread_store.update(cx, |store, cx| { store.save_thread(thread_id, thread, PathList::default(), cx) }); save_task.await.unwrap(); cx.run_until_parked(); assert!(!thread_store.read_with(cx, |store, _cx| store.is_empty())); let delete_task = thread_store.update(cx, |store, cx| store.delete_threads(cx)); delete_task.await.unwrap(); cx.run_until_parked(); assert!(thread_store.read_with(cx, |store, _cx| store.is_empty())); } #[gpui::test] async fn test_delete_thread_removes_only_target(cx: &mut TestAppContext) { let thread_store = cx.new(|cx| ThreadStore::new(cx)); cx.run_until_parked(); let first_id = session_id("thread-a"); let second_id = session_id("thread-b"); let first_thread = make_thread( "Thread A", Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(), ); let second_thread = make_thread( "Thread B", Utc.with_ymd_and_hms(2024, 1, 2, 0, 0, 0).unwrap(), ); let save_first = thread_store.update(cx, |store, cx| { store.save_thread(first_id.clone(), first_thread, PathList::default(), cx) }); save_first.await.unwrap(); let save_second = thread_store.update(cx, |store, cx| { store.save_thread(second_id.clone(), second_thread, PathList::default(), cx) }); save_second.await.unwrap(); cx.run_until_parked(); let delete_task = thread_store.update(cx, |store, cx| store.delete_thread(first_id.clone(), cx)); delete_task.await.unwrap(); cx.run_until_parked(); let entries: Vec<_> = thread_store.read_with(cx, |store, _cx| store.entries().collect()); assert_eq!(entries.len(), 1); assert_eq!(entries[0].id, second_id); } #[gpui::test] async fn test_save_thread_refreshes_ordering(cx: &mut TestAppContext) { let thread_store = cx.new(|cx| ThreadStore::new(cx)); cx.run_until_parked(); let first_id = session_id("thread-a"); let second_id = session_id("thread-b"); let first_thread = make_thread( "Thread A", Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(), ); let second_thread = make_thread( "Thread B", Utc.with_ymd_and_hms(2024, 1, 2, 0, 0, 0).unwrap(), ); let save_first = thread_store.update(cx, |store, cx| { store.save_thread(first_id.clone(), first_thread, PathList::default(), cx) }); save_first.await.unwrap(); let save_second = thread_store.update(cx, |store, cx| { store.save_thread(second_id.clone(), second_thread, PathList::default(), cx) }); save_second.await.unwrap(); cx.run_until_parked(); let updated_first = make_thread( "Thread A", Utc.with_ymd_and_hms(2024, 1, 3, 0, 0, 0).unwrap(), ); let update_task = thread_store.update(cx, |store, cx| { store.save_thread(first_id.clone(), updated_first, PathList::default(), cx) }); update_task.await.unwrap(); cx.run_until_parked(); let entries: Vec<_> = thread_store.read_with(cx, |store, _cx| store.entries().collect()); assert_eq!(entries.len(), 2); assert_eq!(entries[0].id, first_id); assert_eq!(entries[1].id, second_id); } #[gpui::test] async fn test_threads_for_paths_filters_correctly(cx: &mut TestAppContext) { let thread_store = cx.new(|cx| ThreadStore::new(cx)); cx.run_until_parked(); let project_a_paths = PathList::new(&[std::path::PathBuf::from("/home/user/project-a")]); let project_b_paths = PathList::new(&[std::path::PathBuf::from("/home/user/project-b")]); let thread_a = make_thread( "Thread in A", Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(), ); let thread_b = make_thread( "Thread in B", Utc.with_ymd_and_hms(2024, 1, 2, 0, 0, 0).unwrap(), ); let thread_a_id = session_id("thread-a"); let thread_b_id = session_id("thread-b"); let save_a = thread_store.update(cx, |store, cx| { store.save_thread(thread_a_id.clone(), thread_a, project_a_paths.clone(), cx) }); save_a.await.unwrap(); let save_b = thread_store.update(cx, |store, cx| { store.save_thread(thread_b_id.clone(), thread_b, project_b_paths.clone(), cx) }); save_b.await.unwrap(); cx.run_until_parked(); thread_store.read_with(cx, |store, _cx| { let a_threads: Vec<_> = store.threads_for_paths(&project_a_paths).collect(); assert_eq!(a_threads.len(), 1); assert_eq!(a_threads[0].id, thread_a_id); let b_threads: Vec<_> = store.threads_for_paths(&project_b_paths).collect(); assert_eq!(b_threads.len(), 1); assert_eq!(b_threads[0].id, thread_b_id); let nonexistent = PathList::new(&[std::path::PathBuf::from("/nonexistent")]); let no_threads: Vec<_> = store.threads_for_paths(&nonexistent).collect(); assert!(no_threads.is_empty()); }); } }