tenant.openagents/omega

1use anyhow::Context as _;
2use gpui::App;
3use sqlez_macros::sql;
4use util::ResultExt as _;
5
6use crate::{
7    query,
8    sqlez::{domain::Domain, thread_safe_connection::ThreadSafeConnection},
9    write_and_log,
10};
11
12pub struct KeyValueStore(crate::sqlez::thread_safe_connection::ThreadSafeConnection);
13
14impl KeyValueStore {
15    pub fn from_app_db(db: &crate::AppDatabase) -> Self {
16        Self(db.0.clone())
17    }
18}
19
20impl Domain for KeyValueStore {
21    const NAME: &str = stringify!(KeyValueStore);
22
23    const MIGRATIONS: &[&str] = &[
24        sql!(
25            CREATE TABLE IF NOT EXISTS kv_store(
26                key TEXT PRIMARY KEY,
27                value TEXT NOT NULL
28            ) STRICT;
29        ),
30        sql!(
31            CREATE TABLE IF NOT EXISTS scoped_kv_store(
32                namespace TEXT NOT NULL,
33                key TEXT NOT NULL,
34                value TEXT NOT NULL,
35                PRIMARY KEY(namespace, key)
36            ) STRICT;
37        ),
38    ];
39}
40
41crate::static_connection!(KeyValueStore, []);
42
43pub trait Dismissable {
44    const KEY: &'static str;
45
46    fn dismissed(cx: &App) -> bool {
47        KeyValueStore::global(cx)
48            .read_kvp(Self::KEY)
49            .log_err()
50            .is_some_and(|s| s.is_some())
51    }
52
53    fn set_dismissed(is_dismissed: bool, cx: &mut App) {
54        let db = KeyValueStore::global(cx);
55        write_and_log(cx, move || async move {
56            if is_dismissed {
57                db.write_kvp(Self::KEY.into(), "1".into()).await
58            } else {
59                db.delete_kvp(Self::KEY.into()).await
60            }
61        })
62    }
63}
64
65impl KeyValueStore {
66    query! {
67        pub fn read_kvp(key: &str) -> Result<Option<String>> {
68            SELECT value FROM kv_store WHERE key = (?)
69        }
70    }
71
72    pub async fn write_kvp(&self, key: String, value: String) -> anyhow::Result<()> {
73        log::debug!("Writing key-value pair for key {key}");
74        self.write_kvp_inner(key, value).await
75    }
76
77    query! {
78        async fn write_kvp_inner(key: String, value: String) -> Result<()> {
79            INSERT OR REPLACE INTO kv_store(key, value) VALUES ((?), (?))
80        }
81    }
82
83    query! {
84        pub async fn delete_kvp(key: String) -> Result<()> {
85            DELETE FROM kv_store WHERE key = (?)
86        }
87    }
88
89    pub fn scoped<'a>(&'a self, namespace: &'a str) -> ScopedKeyValueStore<'a> {
90        ScopedKeyValueStore {
91            store: self,
92            namespace,
93        }
94    }
95}
96
97pub struct ScopedKeyValueStore<'a> {
98    store: &'a KeyValueStore,
99    namespace: &'a str,
100}
101
102impl ScopedKeyValueStore<'_> {
103    pub fn read(&self, key: &str) -> anyhow::Result<Option<String>> {
104        self.store.select_row_bound::<(&str, &str), String>(
105            "SELECT value FROM scoped_kv_store WHERE namespace = (?) AND key = (?)",
106        )?((self.namespace, key))
107        .context("Failed to read from scoped_kv_store")
108    }
109
110    pub async fn write(&self, key: String, value: String) -> anyhow::Result<()> {
111        let namespace = self.namespace.to_owned();
112        self.store
113            .write(move |connection| {
114                connection.exec_bound::<(&str, &str, &str)>(
115                    "INSERT OR REPLACE INTO scoped_kv_store(namespace, key, value) VALUES ((?), (?), (?))",
116                )?((&namespace, &key, &value))
117                .context("Failed to write to scoped_kv_store")
118            })
119            .await
120    }
121
122    pub async fn delete(&self, key: String) -> anyhow::Result<()> {
123        let namespace = self.namespace.to_owned();
124        self.store
125            .write(move |connection| {
126                connection.exec_bound::<(&str, &str)>(
127                    "DELETE FROM scoped_kv_store WHERE namespace = (?) AND key = (?)",
128                )?((&namespace, &key))
129                .context("Failed to delete from scoped_kv_store")
130            })
131            .await
132    }
133
134    pub async fn delete_all(&self) -> anyhow::Result<()> {
135        let namespace = self.namespace.to_owned();
136        self.store
137            .write(move |connection| {
138                connection
139                    .exec_bound::<&str>("DELETE FROM scoped_kv_store WHERE namespace = (?)")?(
140                    &namespace,
141                )
142                .context("Failed to delete_all from scoped_kv_store")
143            })
144            .await
145    }
146}
147
148#[cfg(test)]
149mod tests {
150    use crate::kvp::KeyValueStore;
151
152    #[gpui::test]
153    async fn test_kvp() {
154        let db = KeyValueStore::open_test_db("test_kvp").await;
155
156        assert_eq!(db.read_kvp("key-1").unwrap(), None);
157
158        db.write_kvp("key-1".to_string(), "one".to_string())
159            .await
160            .unwrap();
161        assert_eq!(db.read_kvp("key-1").unwrap(), Some("one".to_string()));
162
163        db.write_kvp("key-1".to_string(), "one-2".to_string())
164            .await
165            .unwrap();
166        assert_eq!(db.read_kvp("key-1").unwrap(), Some("one-2".to_string()));
167
168        db.write_kvp("key-2".to_string(), "two".to_string())
169            .await
170            .unwrap();
171        assert_eq!(db.read_kvp("key-2").unwrap(), Some("two".to_string()));
172
173        db.delete_kvp("key-1".to_string()).await.unwrap();
174        assert_eq!(db.read_kvp("key-1").unwrap(), None);
175    }
176
177    #[gpui::test]
178    async fn test_scoped_kvp() {
179        let db = KeyValueStore::open_test_db("test_scoped_kvp").await;
180
181        let scope_a = db.scoped("namespace-a");
182        let scope_b = db.scoped("namespace-b");
183
184        // Reading a missing key returns None
185        assert_eq!(scope_a.read("key-1").unwrap(), None);
186
187        // Writing and reading back a key works
188        scope_a
189            .write("key-1".to_string(), "value-a1".to_string())
190            .await
191            .unwrap();
192        assert_eq!(scope_a.read("key-1").unwrap(), Some("value-a1".to_string()));
193
194        // Two namespaces with the same key don't collide
195        scope_b
196            .write("key-1".to_string(), "value-b1".to_string())
197            .await
198            .unwrap();
199        assert_eq!(scope_a.read("key-1").unwrap(), Some("value-a1".to_string()));
200        assert_eq!(scope_b.read("key-1").unwrap(), Some("value-b1".to_string()));
201
202        // delete removes a single key without affecting others in the namespace
203        scope_a
204            .write("key-2".to_string(), "value-a2".to_string())
205            .await
206            .unwrap();
207        scope_a.delete("key-1".to_string()).await.unwrap();
208        assert_eq!(scope_a.read("key-1").unwrap(), None);
209        assert_eq!(scope_a.read("key-2").unwrap(), Some("value-a2".to_string()));
210        assert_eq!(scope_b.read("key-1").unwrap(), Some("value-b1".to_string()));
211
212        // delete_all removes all keys in a namespace without affecting other namespaces
213        scope_a
214            .write("key-3".to_string(), "value-a3".to_string())
215            .await
216            .unwrap();
217        scope_a.delete_all().await.unwrap();
218        assert_eq!(scope_a.read("key-2").unwrap(), None);
219        assert_eq!(scope_a.read("key-3").unwrap(), None);
220        assert_eq!(scope_b.read("key-1").unwrap(), Some("value-b1".to_string()));
221    }
222}
223
224pub struct GlobalKeyValueStore(ThreadSafeConnection);
225
226impl Domain for GlobalKeyValueStore {
227    const NAME: &str = stringify!(GlobalKeyValueStore);
228    const MIGRATIONS: &[&str] = &[sql!(
229        CREATE TABLE IF NOT EXISTS kv_store(
230            key TEXT PRIMARY KEY,
231            value TEXT NOT NULL
232        ) STRICT;
233    )];
234}
235
236impl std::ops::Deref for GlobalKeyValueStore {
237    type Target = ThreadSafeConnection;
238    fn deref(&self) -> &Self::Target {
239        &self.0
240    }
241}
242
243static GLOBAL_KEY_VALUE_STORE: std::sync::LazyLock<GlobalKeyValueStore> =
244    std::sync::LazyLock::new(|| {
245        let db_dir = crate::database_dir();
246        GlobalKeyValueStore(gpui::block_on(crate::open_db::<GlobalKeyValueStore>(
247            db_dir,
248            crate::GlobalDbScope,
249        )))
250    });
251
252impl GlobalKeyValueStore {
253    pub fn global() -> &'static Self {
254        &GLOBAL_KEY_VALUE_STORE
255    }
256
257    query! {
258        pub fn read_kvp(key: &str) -> Result<Option<String>> {
259            SELECT value FROM kv_store WHERE key = (?)
260        }
261    }
262
263    pub async fn write_kvp(&self, key: String, value: String) -> anyhow::Result<()> {
264        log::debug!("Writing global key-value pair for key {key}");
265        self.write_kvp_inner(key, value).await
266    }
267
268    query! {
269        async fn write_kvp_inner(key: String, value: String) -> Result<()> {
270            INSERT OR REPLACE INTO kv_store(key, value) VALUES ((?), (?))
271        }
272    }
273
274    query! {
275        pub async fn delete_kvp(key: String) -> Result<()> {
276            DELETE FROM kv_store WHERE key = (?)
277        }
278    }
279}
280