1use crate::{DbThread, DbThreadMetadata, ThreadsDatabase};
2use agent_client_protocol as acp;
3use anyhow::{Result, anyhow};
4use gpui::{App, Context, Entity, Global, Task, prelude::*};
5use util::path_list::PathList;
6
7struct GlobalThreadStore(Entity<ThreadStore>);
8
9impl Global for GlobalThreadStore {}
10
11pub struct ThreadStore {
12 threads: Vec<DbThreadMetadata>,
13}
14
15impl ThreadStore {
16 pub fn init_global(cx: &mut App) {
17 let thread_store = cx.new(|cx| Self::new(cx));
18 cx.set_global(GlobalThreadStore(thread_store));
19 }
20
21 pub fn global(cx: &App) -> Entity<Self> {
22 cx.global::<GlobalThreadStore>().0.clone()
23 }
24
25 pub fn new(cx: &mut Context<Self>) -> Self {
26 let this = Self {
27 threads: Vec::new(),
28 };
29 this.reload(cx);
30 this
31 }
32
33 pub fn thread_from_session_id(&self, session_id: &acp::SessionId) -> Option<&DbThreadMetadata> {
34 self.threads.iter().find(|thread| &thread.id == session_id)
35 }
36
37 pub fn load_thread(
38 &mut self,
39 id: acp::SessionId,
40 cx: &mut Context<Self>,
41 ) -> Task<Result<Option<DbThread>>> {
42 let database_future = ThreadsDatabase::connect(cx);
43 cx.background_spawn(async move {
44 let database = database_future.await.map_err(|err| anyhow!(err))?;
45 database.load_thread(id).await
46 })
47 }
48
49 pub fn save_thread(
50 &mut self,
51 id: acp::SessionId,
52 thread: crate::DbThread,
53 folder_paths: PathList,
54 cx: &mut Context<Self>,
55 ) -> Task<Result<()>> {
56 let database_future = ThreadsDatabase::connect(cx);
57 cx.spawn(async move |this, cx| {
58 let database = database_future.await.map_err(|err| anyhow!(err))?;
59 database.save_thread(id, thread, folder_paths).await?;
60 this.update(cx, |this, cx| this.reload(cx))
61 })
62 }
63
64 pub fn delete_thread(
65 &mut self,
66 id: acp::SessionId,
67 cx: &mut Context<Self>,
68 ) -> Task<Result<()>> {
69 let database_future = ThreadsDatabase::connect(cx);
70 cx.spawn(async move |this, cx| {
71 let database = database_future.await.map_err(|err| anyhow!(err))?;
72 database.delete_thread(id.clone()).await?;
73 this.update(cx, |this, cx| this.reload(cx))
74 })
75 }
76
77 pub fn delete_threads(&mut self, cx: &mut Context<Self>) -> Task<Result<()>> {
78 let database_future = ThreadsDatabase::connect(cx);
79 cx.spawn(async move |this, cx| {
80 let database = database_future.await.map_err(|err| anyhow!(err))?;
81 database.delete_threads().await?;
82 this.update(cx, |this, cx| this.reload(cx))
83 })
84 }
85
86 pub fn reload(&self, cx: &mut Context<Self>) {
87 let database_connection = ThreadsDatabase::connect(cx);
88 cx.spawn(async move |this, cx| {
89 let database = database_connection.await.map_err(|err| anyhow!(err))?;
90 let threads = database
91 .list_threads()
92 .await?
93 .into_iter()
94 .filter(|thread| thread.parent_session_id.is_none())
95 .collect::<Vec<_>>();
96 this.update(cx, |this, cx| {
97 this.threads = threads;
98 cx.notify();
99 })
100 })
101 .detach_and_log_err(cx);
102 }
103
104 pub fn is_empty(&self) -> bool {
105 self.threads.is_empty()
106 }
107
108 pub fn entries(&self) -> impl Iterator<Item = DbThreadMetadata> + '_ {
109 self.threads.iter().cloned()
110 }
111
112 /// Returns threads whose folder_paths match the given paths exactly.
113 pub fn threads_for_paths(&self, paths: &PathList) -> impl Iterator<Item = &DbThreadMetadata> {
114 self.threads
115 .iter()
116 .filter(move |thread| &thread.folder_paths == paths)
117 }
118}
119
120#[cfg(test)]
121mod tests {
122 use super::*;
123 use chrono::{DateTime, TimeZone, Utc};
124 use collections::HashMap;
125 use gpui::TestAppContext;
126 use std::sync::Arc;
127
128 fn session_id(value: &str) -> acp::SessionId {
129 acp::SessionId::new(Arc::<str>::from(value))
130 }
131
132 fn make_thread(title: &str, updated_at: DateTime<Utc>) -> DbThread {
133 DbThread {
134 title: title.to_string().into(),
135 messages: Vec::new(),
136 updated_at,
137 detailed_summary: None,
138 initial_project_snapshot: None,
139 cumulative_token_usage: Default::default(),
140 request_token_usage: HashMap::default(),
141 model: None,
142 profile: None,
143 imported: false,
144 subagent_context: None,
145 speed: None,
146 thinking_enabled: false,
147 thinking_effort: None,
148 draft_prompt: None,
149 }
150 }
151
152 #[gpui::test]
153 async fn test_entries_are_sorted_by_updated_at(cx: &mut TestAppContext) {
154 let thread_store = cx.new(|cx| ThreadStore::new(cx));
155 cx.run_until_parked();
156
157 let older_id = session_id("thread-a");
158 let newer_id = session_id("thread-b");
159
160 let older_thread = make_thread(
161 "Thread A",
162 Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(),
163 );
164 let newer_thread = make_thread(
165 "Thread B",
166 Utc.with_ymd_and_hms(2024, 1, 2, 0, 0, 0).unwrap(),
167 );
168
169 let save_older = thread_store.update(cx, |store, cx| {
170 store.save_thread(older_id.clone(), older_thread, PathList::default(), cx)
171 });
172 save_older.await.unwrap();
173
174 let save_newer = thread_store.update(cx, |store, cx| {
175 store.save_thread(newer_id.clone(), newer_thread, PathList::default(), cx)
176 });
177 save_newer.await.unwrap();
178
179 cx.run_until_parked();
180
181 let entries: Vec<_> = thread_store.read_with(cx, |store, _cx| store.entries().collect());
182 assert_eq!(entries.len(), 2);
183 assert_eq!(entries[0].id, newer_id);
184 assert_eq!(entries[1].id, older_id);
185 }
186
187 #[gpui::test]
188 async fn test_delete_threads_clears_entries(cx: &mut TestAppContext) {
189 let thread_store = cx.new(|cx| ThreadStore::new(cx));
190 cx.run_until_parked();
191
192 let thread_id = session_id("thread-a");
193 let thread = make_thread(
194 "Thread A",
195 Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(),
196 );
197
198 let save_task = thread_store.update(cx, |store, cx| {
199 store.save_thread(thread_id, thread, PathList::default(), cx)
200 });
201 save_task.await.unwrap();
202
203 cx.run_until_parked();
204 assert!(!thread_store.read_with(cx, |store, _cx| store.is_empty()));
205
206 let delete_task = thread_store.update(cx, |store, cx| store.delete_threads(cx));
207 delete_task.await.unwrap();
208 cx.run_until_parked();
209
210 assert!(thread_store.read_with(cx, |store, _cx| store.is_empty()));
211 }
212
213 #[gpui::test]
214 async fn test_delete_thread_removes_only_target(cx: &mut TestAppContext) {
215 let thread_store = cx.new(|cx| ThreadStore::new(cx));
216 cx.run_until_parked();
217
218 let first_id = session_id("thread-a");
219 let second_id = session_id("thread-b");
220
221 let first_thread = make_thread(
222 "Thread A",
223 Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(),
224 );
225 let second_thread = make_thread(
226 "Thread B",
227 Utc.with_ymd_and_hms(2024, 1, 2, 0, 0, 0).unwrap(),
228 );
229
230 let save_first = thread_store.update(cx, |store, cx| {
231 store.save_thread(first_id.clone(), first_thread, PathList::default(), cx)
232 });
233 save_first.await.unwrap();
234 let save_second = thread_store.update(cx, |store, cx| {
235 store.save_thread(second_id.clone(), second_thread, PathList::default(), cx)
236 });
237 save_second.await.unwrap();
238 cx.run_until_parked();
239
240 let delete_task =
241 thread_store.update(cx, |store, cx| store.delete_thread(first_id.clone(), cx));
242 delete_task.await.unwrap();
243 cx.run_until_parked();
244
245 let entries: Vec<_> = thread_store.read_with(cx, |store, _cx| store.entries().collect());
246 assert_eq!(entries.len(), 1);
247 assert_eq!(entries[0].id, second_id);
248 }
249
250 #[gpui::test]
251 async fn test_save_thread_refreshes_ordering(cx: &mut TestAppContext) {
252 let thread_store = cx.new(|cx| ThreadStore::new(cx));
253 cx.run_until_parked();
254
255 let first_id = session_id("thread-a");
256 let second_id = session_id("thread-b");
257
258 let first_thread = make_thread(
259 "Thread A",
260 Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(),
261 );
262 let second_thread = make_thread(
263 "Thread B",
264 Utc.with_ymd_and_hms(2024, 1, 2, 0, 0, 0).unwrap(),
265 );
266
267 let save_first = thread_store.update(cx, |store, cx| {
268 store.save_thread(first_id.clone(), first_thread, PathList::default(), cx)
269 });
270 save_first.await.unwrap();
271 let save_second = thread_store.update(cx, |store, cx| {
272 store.save_thread(second_id.clone(), second_thread, PathList::default(), cx)
273 });
274 save_second.await.unwrap();
275 cx.run_until_parked();
276
277 let updated_first = make_thread(
278 "Thread A",
279 Utc.with_ymd_and_hms(2024, 1, 3, 0, 0, 0).unwrap(),
280 );
281 let update_task = thread_store.update(cx, |store, cx| {
282 store.save_thread(first_id.clone(), updated_first, PathList::default(), cx)
283 });
284 update_task.await.unwrap();
285 cx.run_until_parked();
286
287 let entries: Vec<_> = thread_store.read_with(cx, |store, _cx| store.entries().collect());
288 assert_eq!(entries.len(), 2);
289 assert_eq!(entries[0].id, first_id);
290 assert_eq!(entries[1].id, second_id);
291 }
292
293 #[gpui::test]
294 async fn test_threads_for_paths_filters_correctly(cx: &mut TestAppContext) {
295 let thread_store = cx.new(|cx| ThreadStore::new(cx));
296 cx.run_until_parked();
297
298 let project_a_paths = PathList::new(&[std::path::PathBuf::from("/home/user/project-a")]);
299 let project_b_paths = PathList::new(&[std::path::PathBuf::from("/home/user/project-b")]);
300
301 let thread_a = make_thread(
302 "Thread in A",
303 Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(),
304 );
305 let thread_b = make_thread(
306 "Thread in B",
307 Utc.with_ymd_and_hms(2024, 1, 2, 0, 0, 0).unwrap(),
308 );
309 let thread_a_id = session_id("thread-a");
310 let thread_b_id = session_id("thread-b");
311
312 let save_a = thread_store.update(cx, |store, cx| {
313 store.save_thread(thread_a_id.clone(), thread_a, project_a_paths.clone(), cx)
314 });
315 save_a.await.unwrap();
316
317 let save_b = thread_store.update(cx, |store, cx| {
318 store.save_thread(thread_b_id.clone(), thread_b, project_b_paths.clone(), cx)
319 });
320 save_b.await.unwrap();
321
322 cx.run_until_parked();
323
324 thread_store.read_with(cx, |store, _cx| {
325 let a_threads: Vec<_> = store.threads_for_paths(&project_a_paths).collect();
326 assert_eq!(a_threads.len(), 1);
327 assert_eq!(a_threads[0].id, thread_a_id);
328
329 let b_threads: Vec<_> = store.threads_for_paths(&project_b_paths).collect();
330 assert_eq!(b_threads.len(), 1);
331 assert_eq!(b_threads[0].id, thread_b_id);
332
333 let nonexistent = PathList::new(&[std::path::PathBuf::from("/nonexistent")]);
334 let no_threads: Vec<_> = store.threads_for_paths(&nonexistent).collect();
335 assert!(no_threads.is_empty());
336 });
337 }
338}