tenant.openagents/omega

1mod prompts;
2pub mod rules_to_skills_migration;
3
4use anyhow::{Result, anyhow};
5use chrono::{DateTime, Utc};
6use collections::HashMap;
7use futures::FutureExt as _;
8use futures::future::Shared;
9
10use gpui::{App, AppContext, Entity, Global, ReadGlobal, SharedString, Task};
11use heed::{
12    Database, RoTxn,
13    types::{SerdeBincode, SerdeJson, Str},
14};
15use parking_lot::RwLock;
16pub use prompts::*;
17
18use serde::{Deserialize, Serialize};
19use std::{future::Future, path::PathBuf, sync::Arc};
20use strum::{EnumIter, IntoEnumIterator as _};
21use text::LineEnding;
22use util::ResultExt;
23use uuid::Uuid;
24
25/// Init starts loading the PromptStore in the background and assigns
26/// a shared future to a global.
27pub fn init(cx: &mut App) {
28    let db_path = paths::prompts_dir().join("prompts-library-db.0.mdb");
29    let prompt_store_task = PromptStore::new(db_path, cx);
30    let prompt_store_entity_task = cx
31        .spawn(async move |cx| {
32            prompt_store_task
33                .await
34                .map(|prompt_store| cx.new(|_cx| prompt_store))
35                .map_err(Arc::new)
36        })
37        .shared();
38    cx.set_global(GlobalPromptStore(prompt_store_entity_task))
39}
40
41#[derive(Clone, Debug, Serialize, Deserialize)]
42pub struct PromptMetadata {
43    pub id: PromptId,
44    pub title: Option<SharedString>,
45    pub default: bool,
46    pub saved_at: DateTime<Utc>,
47}
48
49impl PromptMetadata {
50    fn builtin(builtin: BuiltInPrompt) -> Self {
51        Self {
52            id: PromptId::BuiltIn(builtin),
53            title: Some(builtin.title().into()),
54            default: false,
55            saved_at: DateTime::default(),
56        }
57    }
58}
59
60/// Built-in prompts that have default content and can be customized by users.
61#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize, EnumIter)]
62pub enum BuiltInPrompt {
63    CommitMessage,
64}
65
66impl BuiltInPrompt {
67    pub fn title(&self) -> &'static str {
68        match self {
69            Self::CommitMessage => "Commit message",
70        }
71    }
72
73    /// Returns the default content for this built-in prompt.
74    pub fn default_content(&self) -> &'static str {
75        match self {
76            Self::CommitMessage => include_str!("../../git_ui/src/commit_message_prompt.txt"),
77        }
78    }
79}
80
81impl std::fmt::Display for BuiltInPrompt {
82    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
83        match self {
84            Self::CommitMessage => write!(f, "Commit message"),
85        }
86    }
87}
88
89#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
90#[serde(tag = "kind")]
91pub enum PromptId {
92    User { uuid: UserPromptId },
93    BuiltIn(BuiltInPrompt),
94}
95
96impl PromptId {
97    pub fn new() -> PromptId {
98        UserPromptId::new().into()
99    }
100
101    pub fn as_user(&self) -> Option<UserPromptId> {
102        match self {
103            Self::User { uuid } => Some(*uuid),
104            Self::BuiltIn { .. } => None,
105        }
106    }
107
108    pub fn as_built_in(&self) -> Option<BuiltInPrompt> {
109        match self {
110            Self::User { .. } => None,
111            Self::BuiltIn(builtin) => Some(*builtin),
112        }
113    }
114
115    pub fn is_built_in(&self) -> bool {
116        matches!(self, Self::BuiltIn { .. })
117    }
118}
119
120impl From<BuiltInPrompt> for PromptId {
121    fn from(builtin: BuiltInPrompt) -> Self {
122        PromptId::BuiltIn(builtin)
123    }
124}
125
126impl From<UserPromptId> for PromptId {
127    fn from(uuid: UserPromptId) -> Self {
128        PromptId::User { uuid }
129    }
130}
131
132#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
133#[serde(transparent)]
134pub struct UserPromptId(pub Uuid);
135
136impl UserPromptId {
137    pub fn new() -> UserPromptId {
138        UserPromptId(Uuid::new_v4())
139    }
140}
141
142impl From<Uuid> for UserPromptId {
143    fn from(uuid: Uuid) -> Self {
144        UserPromptId(uuid)
145    }
146}
147
148impl std::fmt::Display for PromptId {
149    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
150        match self {
151            PromptId::User { uuid } => write!(f, "{}", uuid.0),
152            PromptId::BuiltIn(builtin) => write!(f, "{}", builtin),
153        }
154    }
155}
156
157pub struct PromptStore {
158    env: heed::Env,
159    metadata_cache: RwLock<MetadataCache>,
160    bodies: Database<SerdeJson<PromptId>, Str>,
161}
162
163#[derive(Default)]
164struct MetadataCache {
165    metadata: Vec<PromptMetadata>,
166    metadata_by_id: HashMap<PromptId, PromptMetadata>,
167}
168
169impl MetadataCache {
170    fn from_db(
171        db: Database<SerdeJson<PromptId>, SerdeJson<PromptMetadata>>,
172        txn: &RoTxn,
173    ) -> Result<Self> {
174        let mut cache = MetadataCache::default();
175        for result in db.iter(txn)? {
176            // Fail-open: skip records that can't be decoded (e.g. from a different branch)
177            // rather than failing the entire prompt store initialization.
178            let Ok((prompt_id, metadata)) = result else {
179                log::warn!(
180                    "Skipping unreadable prompt record in database: {:?}",
181                    result.err()
182                );
183                continue;
184            };
185            cache.metadata.push(metadata.clone());
186            cache.metadata_by_id.insert(prompt_id, metadata);
187        }
188
189        // Insert all the built-in prompts that were not customized by the user
190        for builtin in BuiltInPrompt::iter() {
191            let builtin_id = PromptId::BuiltIn(builtin);
192            if !cache.metadata_by_id.contains_key(&builtin_id) {
193                let metadata = PromptMetadata::builtin(builtin);
194                cache.metadata.push(metadata.clone());
195                cache.metadata_by_id.insert(builtin_id, metadata);
196            }
197        }
198        cache.sort();
199        Ok(cache)
200    }
201
202    fn sort(&mut self) {
203        self.metadata.sort_unstable_by(|a, b| {
204            a.title
205                .cmp(&b.title)
206                .then_with(|| b.saved_at.cmp(&a.saved_at))
207        });
208    }
209}
210
211impl PromptStore {
212    pub fn global(cx: &App) -> impl Future<Output = Result<Entity<Self>>> + use<> {
213        let store = GlobalPromptStore::global(cx).0.clone();
214        async move { store.await.map_err(|err| anyhow!(err)) }
215    }
216
217    pub fn new(db_path: PathBuf, cx: &App) -> Task<Result<Self>> {
218        cx.background_spawn(async move {
219            std::fs::create_dir_all(&db_path)?;
220
221            let db_env = unsafe {
222                heed::EnvOpenOptions::new()
223                    .map_size(1024 * 1024 * 1024) // 1GB
224                    .max_dbs(4) // Metadata and bodies (possibly v1 of both as well)
225                    .open(db_path)?
226            };
227
228            let mut txn = db_env.write_txn()?;
229            let metadata = db_env.create_database(&mut txn, Some("metadata.v2"))?;
230            let bodies = db_env.create_database(&mut txn, Some("bodies.v2"))?;
231            txn.commit()?;
232
233            Self::upgrade_dbs(&db_env, metadata, bodies).log_err();
234
235            let txn = db_env.read_txn()?;
236            let metadata_cache = MetadataCache::from_db(metadata, &txn)?;
237            txn.commit()?;
238
239            Ok(PromptStore {
240                env: db_env,
241                metadata_cache: RwLock::new(metadata_cache),
242                bodies,
243            })
244        })
245    }
246
247    fn upgrade_dbs(
248        env: &heed::Env,
249        metadata_db: heed::Database<SerdeJson<PromptId>, SerdeJson<PromptMetadata>>,
250        bodies_db: heed::Database<SerdeJson<PromptId>, Str>,
251    ) -> Result<()> {
252        let mut txn = env.write_txn()?;
253        let Some(bodies_v1_db) = env
254            .open_database::<SerdeBincode<PromptIdV1>, SerdeBincode<String>>(
255                &txn,
256                Some("bodies"),
257            )?
258        else {
259            return Ok(());
260        };
261        let mut bodies_v1 = bodies_v1_db
262            .iter(&txn)?
263            .collect::<heed::Result<HashMap<_, _>>>()?;
264
265        let Some(metadata_v1_db) = env
266            .open_database::<SerdeBincode<PromptIdV1>, SerdeBincode<PromptMetadataV1>>(
267                &txn,
268                Some("metadata"),
269            )?
270        else {
271            return Ok(());
272        };
273        let metadata_v1 = metadata_v1_db
274            .iter(&txn)?
275            .collect::<heed::Result<HashMap<_, _>>>()?;
276
277        for (prompt_id_v1, metadata_v1) in metadata_v1 {
278            let prompt_id_v2 = UserPromptId(prompt_id_v1.0).into();
279            let Some(body_v1) = bodies_v1.remove(&prompt_id_v1) else {
280                continue;
281            };
282
283            if metadata_db
284                .get(&txn, &prompt_id_v2)?
285                .is_none_or(|metadata_v2| metadata_v1.saved_at > metadata_v2.saved_at)
286            {
287                metadata_db.put(
288                    &mut txn,
289                    &prompt_id_v2,
290                    &PromptMetadata {
291                        id: prompt_id_v2,
292                        title: metadata_v1.title.clone(),
293                        default: metadata_v1.default,
294                        saved_at: metadata_v1.saved_at,
295                    },
296                )?;
297                bodies_db.put(&mut txn, &prompt_id_v2, &body_v1)?;
298            }
299        }
300
301        txn.commit()?;
302
303        Ok(())
304    }
305
306    pub fn load(&self, id: PromptId, cx: &App) -> Task<Result<String>> {
307        let env = self.env.clone();
308        let bodies = self.bodies;
309        cx.background_spawn(async move {
310            let txn = env.read_txn()?;
311            let mut prompt: String = match bodies.get(&txn, &id)? {
312                Some(body) => body.into(),
313                None => {
314                    if let Some(built_in) = id.as_built_in() {
315                        built_in.default_content().into()
316                    } else {
317                        anyhow::bail!("prompt not found")
318                    }
319                }
320            };
321            LineEnding::normalize(&mut prompt);
322            Ok(prompt)
323        })
324    }
325
326    pub fn all_prompt_metadata(&self) -> Vec<PromptMetadata> {
327        self.metadata_cache.read().metadata.clone()
328    }
329}
330
331/// Deprecated: Legacy V1 prompt ID format, used only for migrating data from old databases. Use `PromptId` instead.
332#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize, Hash)]
333struct PromptIdV1(Uuid);
334
335impl From<UserPromptId> for PromptIdV1 {
336    fn from(id: UserPromptId) -> Self {
337        PromptIdV1(id.0)
338    }
339}
340
341/// Deprecated: Legacy V1 prompt metadata format, used only for migrating data from old databases. Use `PromptMetadata` instead.
342#[derive(Clone, Debug, Serialize, Deserialize)]
343struct PromptMetadataV1 {
344    id: PromptIdV1,
345    title: Option<SharedString>,
346    default: bool,
347    saved_at: DateTime<Utc>,
348}
349
350/// Wraps a shared future to a prompt store so it can be assigned as a context global.
351pub struct GlobalPromptStore(Shared<Task<Result<Entity<PromptStore>, Arc<anyhow::Error>>>>);
352
353impl Global for GlobalPromptStore {}
354
355#[cfg(test)]
356mod tests {
357    use super::*;
358    use gpui::TestAppContext;
359
360    #[gpui::test]
361    async fn test_built_in_prompt_load(cx: &mut TestAppContext) {
362        cx.executor().allow_parking();
363
364        let temp_dir = tempfile::tempdir().unwrap();
365        let db_path = temp_dir.path().join("prompts-db");
366
367        let store = cx.update(|cx| PromptStore::new(db_path, cx)).await.unwrap();
368        let store = cx.new(|_cx| store);
369
370        let commit_message_id = PromptId::BuiltIn(BuiltInPrompt::CommitMessage);
371
372        let loaded_content = store
373            .update(cx, |store, cx| store.load(commit_message_id, cx))
374            .await
375            .unwrap();
376
377        let mut expected_content = BuiltInPrompt::CommitMessage.default_content().to_string();
378        LineEnding::normalize(&mut expected_content);
379        assert_eq!(
380            loaded_content.trim(),
381            expected_content.trim(),
382            "Loading a built-in prompt not in DB should return default content"
383        );
384
385        assert!(
386            store.read_with(cx, |store, _| {
387                store
388                    .all_prompt_metadata()
389                    .iter()
390                    .any(|metadata| metadata.id == commit_message_id)
391            }),
392            "Built-in prompt should always be in cache"
393        );
394    }
395}
396