tenant.openagents/omega

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