1mod prompts;
2
3use anyhow::{Result, anyhow};
4use chrono::{DateTime, Utc};
5use collections::HashMap;
6use futures::FutureExt as _;
7use futures::future::Shared;
8use fuzzy::StringMatchCandidate;
9use gpui::{
10 App, AppContext, Context, Entity, EventEmitter, Global, ReadGlobal, SharedString, Task,
11};
12use heed::{
13 Database, RoTxn,
14 types::{SerdeBincode, SerdeJson, Str},
15};
16use parking_lot::RwLock;
17pub use prompts::*;
18use rope::Rope;
19use serde::{Deserialize, Serialize};
20use std::{
21 cmp::Reverse,
22 future::Future,
23 path::PathBuf,
24 sync::{Arc, atomic::AtomicBool},
25};
26use strum::{EnumIter, IntoEnumIterator as _};
27use text::LineEnding;
28use util::ResultExt;
29use uuid::Uuid;
30
31/// Init starts loading the PromptStore in the background and assigns
32/// a shared future to a global.
33pub fn init(cx: &mut App) {
34 let db_path = paths::prompts_dir().join("prompts-library-db.0.mdb");
35 let prompt_store_task = PromptStore::new(db_path, cx);
36 let prompt_store_entity_task = cx
37 .spawn(async move |cx| {
38 prompt_store_task
39 .await
40 .map(|prompt_store| cx.new(|_cx| prompt_store))
41 .map_err(Arc::new)
42 })
43 .shared();
44 cx.set_global(GlobalPromptStore(prompt_store_entity_task))
45}
46
47#[derive(Clone, Debug, Serialize, Deserialize)]
48pub struct PromptMetadata {
49 pub id: PromptId,
50 pub title: Option<SharedString>,
51 pub default: bool,
52 pub saved_at: DateTime<Utc>,
53}
54
55impl PromptMetadata {
56 fn builtin(builtin: BuiltInPrompt) -> Self {
57 Self {
58 id: PromptId::BuiltIn(builtin),
59 title: Some(builtin.title().into()),
60 default: false,
61 saved_at: DateTime::default(),
62 }
63 }
64}
65
66/// Built-in prompts that have default content and can be customized by users.
67#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize, EnumIter)]
68pub enum BuiltInPrompt {
69 CommitMessage,
70}
71
72impl BuiltInPrompt {
73 pub fn title(&self) -> &'static str {
74 match self {
75 Self::CommitMessage => "Commit message",
76 }
77 }
78
79 /// Returns the default content for this built-in prompt.
80 pub fn default_content(&self) -> &'static str {
81 match self {
82 Self::CommitMessage => include_str!("../../git_ui/src/commit_message_prompt.txt"),
83 }
84 }
85}
86
87impl std::fmt::Display for BuiltInPrompt {
88 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
89 match self {
90 Self::CommitMessage => write!(f, "Commit message"),
91 }
92 }
93}
94
95#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
96#[serde(tag = "kind")]
97pub enum PromptId {
98 User { uuid: UserPromptId },
99 BuiltIn(BuiltInPrompt),
100}
101
102impl PromptId {
103 pub fn new() -> PromptId {
104 UserPromptId::new().into()
105 }
106
107 pub fn as_user(&self) -> Option<UserPromptId> {
108 match self {
109 Self::User { uuid } => Some(*uuid),
110 Self::BuiltIn { .. } => None,
111 }
112 }
113
114 pub fn as_built_in(&self) -> Option<BuiltInPrompt> {
115 match self {
116 Self::User { .. } => None,
117 Self::BuiltIn(builtin) => Some(*builtin),
118 }
119 }
120
121 pub fn is_built_in(&self) -> bool {
122 matches!(self, Self::BuiltIn { .. })
123 }
124
125 pub fn can_edit(&self) -> bool {
126 match self {
127 Self::User { .. } => true,
128 Self::BuiltIn(builtin) => match builtin {
129 BuiltInPrompt::CommitMessage => true,
130 },
131 }
132 }
133}
134
135impl From<BuiltInPrompt> for PromptId {
136 fn from(builtin: BuiltInPrompt) -> Self {
137 PromptId::BuiltIn(builtin)
138 }
139}
140
141impl From<UserPromptId> for PromptId {
142 fn from(uuid: UserPromptId) -> Self {
143 PromptId::User { uuid }
144 }
145}
146
147#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
148#[serde(transparent)]
149pub struct UserPromptId(pub Uuid);
150
151impl UserPromptId {
152 pub fn new() -> UserPromptId {
153 UserPromptId(Uuid::new_v4())
154 }
155}
156
157impl From<Uuid> for UserPromptId {
158 fn from(uuid: Uuid) -> Self {
159 UserPromptId(uuid)
160 }
161}
162
163impl std::fmt::Display for PromptId {
164 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
165 match self {
166 PromptId::User { uuid } => write!(f, "{}", uuid.0),
167 PromptId::BuiltIn(builtin) => write!(f, "{}", builtin),
168 }
169 }
170}
171
172pub struct PromptStore {
173 env: heed::Env,
174 metadata_cache: RwLock<MetadataCache>,
175 metadata: Database<SerdeJson<PromptId>, SerdeJson<PromptMetadata>>,
176 bodies: Database<SerdeJson<PromptId>, Str>,
177}
178
179pub struct PromptsUpdatedEvent;
180
181impl EventEmitter<PromptsUpdatedEvent> for PromptStore {}
182
183#[derive(Default)]
184struct MetadataCache {
185 metadata: Vec<PromptMetadata>,
186 metadata_by_id: HashMap<PromptId, PromptMetadata>,
187}
188
189impl MetadataCache {
190 fn from_db(
191 db: Database<SerdeJson<PromptId>, SerdeJson<PromptMetadata>>,
192 txn: &RoTxn,
193 ) -> Result<Self> {
194 let mut cache = MetadataCache::default();
195 for result in db.iter(txn)? {
196 // Fail-open: skip records that can't be decoded (e.g. from a different branch)
197 // rather than failing the entire prompt store initialization.
198 let Ok((prompt_id, metadata)) = result else {
199 log::warn!(
200 "Skipping unreadable prompt record in database: {:?}",
201 result.err()
202 );
203 continue;
204 };
205 cache.metadata.push(metadata.clone());
206 cache.metadata_by_id.insert(prompt_id, metadata);
207 }
208
209 // Insert all the built-in prompts that were not customized by the user
210 for builtin in BuiltInPrompt::iter() {
211 let builtin_id = PromptId::BuiltIn(builtin);
212 if !cache.metadata_by_id.contains_key(&builtin_id) {
213 let metadata = PromptMetadata::builtin(builtin);
214 cache.metadata.push(metadata.clone());
215 cache.metadata_by_id.insert(builtin_id, metadata);
216 }
217 }
218 cache.sort();
219 Ok(cache)
220 }
221
222 fn insert(&mut self, metadata: PromptMetadata) {
223 self.metadata_by_id.insert(metadata.id, metadata.clone());
224 if let Some(old_metadata) = self.metadata.iter_mut().find(|m| m.id == metadata.id) {
225 *old_metadata = metadata;
226 } else {
227 self.metadata.push(metadata);
228 }
229 self.sort();
230 }
231
232 fn remove(&mut self, id: PromptId) {
233 self.metadata.retain(|metadata| metadata.id != id);
234 self.metadata_by_id.remove(&id);
235 }
236
237 fn sort(&mut self) {
238 self.metadata.sort_unstable_by(|a, b| {
239 a.title
240 .cmp(&b.title)
241 .then_with(|| b.saved_at.cmp(&a.saved_at))
242 });
243 }
244}
245
246impl PromptStore {
247 pub fn global(cx: &App) -> impl Future<Output = Result<Entity<Self>>> + use<> {
248 let store = GlobalPromptStore::global(cx).0.clone();
249 async move { store.await.map_err(|err| anyhow!(err)) }
250 }
251
252 pub fn new(db_path: PathBuf, cx: &App) -> Task<Result<Self>> {
253 cx.background_spawn(async move {
254 std::fs::create_dir_all(&db_path)?;
255
256 let db_env = unsafe {
257 heed::EnvOpenOptions::new()
258 .map_size(1024 * 1024 * 1024) // 1GB
259 .max_dbs(4) // Metadata and bodies (possibly v1 of both as well)
260 .open(db_path)?
261 };
262
263 let mut txn = db_env.write_txn()?;
264 let metadata = db_env.create_database(&mut txn, Some("metadata.v2"))?;
265 let bodies = db_env.create_database(&mut txn, Some("bodies.v2"))?;
266 txn.commit()?;
267
268 Self::upgrade_dbs(&db_env, metadata, bodies).log_err();
269
270 let txn = db_env.read_txn()?;
271 let metadata_cache = MetadataCache::from_db(metadata, &txn)?;
272 txn.commit()?;
273
274 Ok(PromptStore {
275 env: db_env,
276 metadata_cache: RwLock::new(metadata_cache),
277 metadata,
278 bodies,
279 })
280 })
281 }
282
283 fn upgrade_dbs(
284 env: &heed::Env,
285 metadata_db: heed::Database<SerdeJson<PromptId>, SerdeJson<PromptMetadata>>,
286 bodies_db: heed::Database<SerdeJson<PromptId>, Str>,
287 ) -> Result<()> {
288 let mut txn = env.write_txn()?;
289 let Some(bodies_v1_db) = env
290 .open_database::<SerdeBincode<PromptIdV1>, SerdeBincode<String>>(
291 &txn,
292 Some("bodies"),
293 )?
294 else {
295 return Ok(());
296 };
297 let mut bodies_v1 = bodies_v1_db
298 .iter(&txn)?
299 .collect::<heed::Result<HashMap<_, _>>>()?;
300
301 let Some(metadata_v1_db) = env
302 .open_database::<SerdeBincode<PromptIdV1>, SerdeBincode<PromptMetadataV1>>(
303 &txn,
304 Some("metadata"),
305 )?
306 else {
307 return Ok(());
308 };
309 let metadata_v1 = metadata_v1_db
310 .iter(&txn)?
311 .collect::<heed::Result<HashMap<_, _>>>()?;
312
313 for (prompt_id_v1, metadata_v1) in metadata_v1 {
314 let prompt_id_v2 = UserPromptId(prompt_id_v1.0).into();
315 let Some(body_v1) = bodies_v1.remove(&prompt_id_v1) else {
316 continue;
317 };
318
319 if metadata_db
320 .get(&txn, &prompt_id_v2)?
321 .is_none_or(|metadata_v2| metadata_v1.saved_at > metadata_v2.saved_at)
322 {
323 metadata_db.put(
324 &mut txn,
325 &prompt_id_v2,
326 &PromptMetadata {
327 id: prompt_id_v2,
328 title: metadata_v1.title.clone(),
329 default: metadata_v1.default,
330 saved_at: metadata_v1.saved_at,
331 },
332 )?;
333 bodies_db.put(&mut txn, &prompt_id_v2, &body_v1)?;
334 }
335 }
336
337 txn.commit()?;
338
339 Ok(())
340 }
341
342 pub fn load(&self, id: PromptId, cx: &App) -> Task<Result<String>> {
343 let env = self.env.clone();
344 let bodies = self.bodies;
345 cx.background_spawn(async move {
346 let txn = env.read_txn()?;
347 let mut prompt: String = match bodies.get(&txn, &id)? {
348 Some(body) => body.into(),
349 None => {
350 if let Some(built_in) = id.as_built_in() {
351 built_in.default_content().into()
352 } else {
353 anyhow::bail!("prompt not found")
354 }
355 }
356 };
357 LineEnding::normalize(&mut prompt);
358 Ok(prompt)
359 })
360 }
361
362 pub fn all_prompt_metadata(&self) -> Vec<PromptMetadata> {
363 self.metadata_cache.read().metadata.clone()
364 }
365
366 pub fn default_prompt_metadata(&self) -> Vec<PromptMetadata> {
367 return self
368 .metadata_cache
369 .read()
370 .metadata
371 .iter()
372 .filter(|metadata| metadata.default)
373 .cloned()
374 .collect::<Vec<_>>();
375 }
376
377 pub fn delete(&self, id: PromptId, cx: &Context<Self>) -> Task<Result<()>> {
378 self.metadata_cache.write().remove(id);
379
380 let db_connection = self.env.clone();
381 let bodies = self.bodies;
382 let metadata = self.metadata;
383
384 let task = cx.background_spawn(async move {
385 let mut txn = db_connection.write_txn()?;
386
387 metadata.delete(&mut txn, &id)?;
388 bodies.delete(&mut txn, &id)?;
389
390 if let PromptId::User { uuid } = id {
391 let prompt_id_v1 = PromptIdV1::from(uuid);
392
393 if let Some(metadata_v1_db) = db_connection
394 .open_database::<SerdeBincode<PromptIdV1>, SerdeBincode<()>>(
395 &txn,
396 Some("metadata"),
397 )?
398 {
399 metadata_v1_db.delete(&mut txn, &prompt_id_v1)?;
400 }
401
402 if let Some(bodies_v1_db) = db_connection
403 .open_database::<SerdeBincode<PromptIdV1>, SerdeBincode<()>>(
404 &txn,
405 Some("bodies"),
406 )?
407 {
408 bodies_v1_db.delete(&mut txn, &prompt_id_v1)?;
409 }
410 }
411
412 txn.commit()?;
413 anyhow::Ok(())
414 });
415
416 cx.spawn(async move |this, cx| {
417 task.await?;
418 this.update(cx, |_, cx| cx.emit(PromptsUpdatedEvent)).ok();
419 anyhow::Ok(())
420 })
421 }
422
423 pub fn metadata(&self, id: PromptId) -> Option<PromptMetadata> {
424 self.metadata_cache.read().metadata_by_id.get(&id).cloned()
425 }
426
427 pub fn first(&self) -> Option<PromptMetadata> {
428 self.metadata_cache.read().metadata.first().cloned()
429 }
430
431 pub fn id_for_title(&self, title: &str) -> Option<PromptId> {
432 let metadata_cache = self.metadata_cache.read();
433 let metadata = metadata_cache
434 .metadata
435 .iter()
436 .find(|metadata| metadata.title.as_ref().map(|title| &***title) == Some(title))?;
437 Some(metadata.id)
438 }
439
440 pub fn search(
441 &self,
442 query: String,
443 cancellation_flag: Arc<AtomicBool>,
444 cx: &App,
445 ) -> Task<Vec<PromptMetadata>> {
446 let cached_metadata = self.metadata_cache.read().metadata.clone();
447 let executor = cx.background_executor().clone();
448 cx.background_spawn(async move {
449 let mut matches = if query.is_empty() {
450 cached_metadata
451 } else {
452 let candidates = cached_metadata
453 .iter()
454 .enumerate()
455 .filter_map(|(ix, metadata)| {
456 Some(StringMatchCandidate::new(ix, metadata.title.as_ref()?))
457 })
458 .collect::<Vec<_>>();
459 let matches = fuzzy::match_strings(
460 &candidates,
461 &query,
462 false,
463 true,
464 100,
465 &cancellation_flag,
466 executor,
467 )
468 .await;
469 matches
470 .into_iter()
471 .map(|mat| cached_metadata[mat.candidate_id].clone())
472 .collect()
473 };
474 matches.sort_by_key(|metadata| Reverse(metadata.default));
475 matches
476 })
477 }
478
479 pub fn save(
480 &self,
481 id: PromptId,
482 title: Option<SharedString>,
483 default: bool,
484 body: Rope,
485 cx: &Context<Self>,
486 ) -> Task<Result<()>> {
487 if !id.can_edit() {
488 return Task::ready(Err(anyhow!("this prompt cannot be edited")));
489 }
490
491 let body = body.to_string();
492 let is_default_content = id
493 .as_built_in()
494 .is_some_and(|builtin| body.trim() == builtin.default_content().trim());
495
496 let metadata = if let Some(builtin) = id.as_built_in() {
497 PromptMetadata::builtin(builtin)
498 } else {
499 PromptMetadata {
500 id,
501 title,
502 default,
503 saved_at: Utc::now(),
504 }
505 };
506
507 self.metadata_cache.write().insert(metadata.clone());
508
509 let db_connection = self.env.clone();
510 let bodies = self.bodies;
511 let metadata_db = self.metadata;
512
513 let task = cx.background_spawn(async move {
514 let mut txn = db_connection.write_txn()?;
515
516 if is_default_content {
517 metadata_db.delete(&mut txn, &id)?;
518 bodies.delete(&mut txn, &id)?;
519 } else {
520 metadata_db.put(&mut txn, &id, &metadata)?;
521 bodies.put(&mut txn, &id, &body)?;
522 }
523
524 txn.commit()?;
525
526 anyhow::Ok(())
527 });
528
529 cx.spawn(async move |this, cx| {
530 task.await?;
531 this.update(cx, |_, cx| cx.emit(PromptsUpdatedEvent)).ok();
532 anyhow::Ok(())
533 })
534 }
535
536 pub fn save_metadata(
537 &self,
538 id: PromptId,
539 mut title: Option<SharedString>,
540 default: bool,
541 cx: &Context<Self>,
542 ) -> Task<Result<()>> {
543 let mut cache = self.metadata_cache.write();
544
545 if !id.can_edit() {
546 title = cache
547 .metadata_by_id
548 .get(&id)
549 .and_then(|metadata| metadata.title.clone());
550 }
551
552 let prompt_metadata = PromptMetadata {
553 id,
554 title,
555 default,
556 saved_at: Utc::now(),
557 };
558
559 cache.insert(prompt_metadata.clone());
560
561 let db_connection = self.env.clone();
562 let metadata = self.metadata;
563
564 let task = cx.background_spawn(async move {
565 let mut txn = db_connection.write_txn()?;
566 metadata.put(&mut txn, &id, &prompt_metadata)?;
567 txn.commit()?;
568
569 anyhow::Ok(())
570 });
571
572 cx.spawn(async move |this, cx| {
573 task.await?;
574 this.update(cx, |_, cx| cx.emit(PromptsUpdatedEvent)).ok();
575 anyhow::Ok(())
576 })
577 }
578}
579
580/// Deprecated: Legacy V1 prompt ID format, used only for migrating data from old databases. Use `PromptId` instead.
581#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize, Hash)]
582struct PromptIdV1(Uuid);
583
584impl From<UserPromptId> for PromptIdV1 {
585 fn from(id: UserPromptId) -> Self {
586 PromptIdV1(id.0)
587 }
588}
589
590/// Deprecated: Legacy V1 prompt metadata format, used only for migrating data from old databases. Use `PromptMetadata` instead.
591#[derive(Clone, Debug, Serialize, Deserialize)]
592struct PromptMetadataV1 {
593 id: PromptIdV1,
594 title: Option<SharedString>,
595 default: bool,
596 saved_at: DateTime<Utc>,
597}
598
599/// Wraps a shared future to a prompt store so it can be assigned as a context global.
600pub struct GlobalPromptStore(Shared<Task<Result<Entity<PromptStore>, Arc<anyhow::Error>>>>);
601
602impl Global for GlobalPromptStore {}
603
604#[cfg(test)]
605mod tests {
606 use super::*;
607 use gpui::TestAppContext;
608
609 #[gpui::test]
610 async fn test_built_in_prompt_load_save(cx: &mut TestAppContext) {
611 cx.executor().allow_parking();
612
613 let temp_dir = tempfile::tempdir().unwrap();
614 let db_path = temp_dir.path().join("prompts-db");
615
616 let store = cx.update(|cx| PromptStore::new(db_path, cx)).await.unwrap();
617 let store = cx.new(|_cx| store);
618
619 let commit_message_id = PromptId::BuiltIn(BuiltInPrompt::CommitMessage);
620
621 let loaded_content = store
622 .update(cx, |store, cx| store.load(commit_message_id, cx))
623 .await
624 .unwrap();
625
626 let mut expected_content = BuiltInPrompt::CommitMessage.default_content().to_string();
627 LineEnding::normalize(&mut expected_content);
628 assert_eq!(
629 loaded_content.trim(),
630 expected_content.trim(),
631 "Loading a built-in prompt not in DB should return default content"
632 );
633
634 let metadata = store.read_with(cx, |store, _| store.metadata(commit_message_id));
635 assert!(
636 metadata.is_some(),
637 "Built-in prompt should always have metadata"
638 );
639 assert!(
640 store.read_with(cx, |store, _| {
641 store
642 .metadata_cache
643 .read()
644 .metadata_by_id
645 .contains_key(&commit_message_id)
646 }),
647 "Built-in prompt should always be in cache"
648 );
649
650 let custom_content = "Custom commit message prompt";
651 store
652 .update(cx, |store, cx| {
653 store.save(
654 commit_message_id,
655 Some("Commit message".into()),
656 false,
657 Rope::from(custom_content),
658 cx,
659 )
660 })
661 .await
662 .unwrap();
663
664 let loaded_custom = store
665 .update(cx, |store, cx| store.load(commit_message_id, cx))
666 .await
667 .unwrap();
668 assert_eq!(
669 loaded_custom.trim(),
670 custom_content.trim(),
671 "Custom content should be loaded after saving"
672 );
673
674 assert!(
675 store
676 .read_with(cx, |store, _| store.metadata(commit_message_id))
677 .is_some(),
678 "Built-in prompt should have metadata after customization"
679 );
680
681 store
682 .update(cx, |store, cx| {
683 store.save(
684 commit_message_id,
685 Some("Commit message".into()),
686 false,
687 Rope::from(BuiltInPrompt::CommitMessage.default_content()),
688 cx,
689 )
690 })
691 .await
692 .unwrap();
693
694 let metadata_after_reset =
695 store.read_with(cx, |store, _| store.metadata(commit_message_id));
696 assert!(
697 metadata_after_reset.is_some(),
698 "Built-in prompt should still have metadata after reset"
699 );
700 assert_eq!(
701 metadata_after_reset
702 .as_ref()
703 .and_then(|m| m.title.as_ref().map(|t| t.as_ref())),
704 Some("Commit message"),
705 "Built-in prompt should have default title after reset"
706 );
707
708 let loaded_after_reset = store
709 .update(cx, |store, cx| store.load(commit_message_id, cx))
710 .await
711 .unwrap();
712 let mut expected_content_after_reset =
713 BuiltInPrompt::CommitMessage.default_content().to_string();
714 LineEnding::normalize(&mut expected_content_after_reset);
715 assert_eq!(
716 loaded_after_reset.trim(),
717 expected_content_after_reset.trim(),
718 "Content should be back to default after saving default content"
719 );
720 }
721
722 /// Test that the prompt store initializes successfully even when the database
723 /// contains records with incompatible/undecodable PromptId keys (e.g., from
724 /// a different branch that used a different serialization format).
725 ///
726 /// This is a regression test for the "fail-open" behavior: we should skip
727 /// bad records rather than failing the entire store initialization.
728 #[gpui::test]
729 async fn test_prompt_store_handles_incompatible_db_records(cx: &mut TestAppContext) {
730 cx.executor().allow_parking();
731
732 let temp_dir = tempfile::tempdir().unwrap();
733 let db_path = temp_dir.path().join("prompts-db-with-bad-records");
734 std::fs::create_dir_all(&db_path).unwrap();
735
736 // First, create the DB and write an incompatible record directly.
737 // We simulate a record written by a different branch that used
738 // `{"kind":"CommitMessage"}` instead of `{"kind":"BuiltIn", ...}`.
739 {
740 let db_env = unsafe {
741 heed::EnvOpenOptions::new()
742 .map_size(1024 * 1024 * 1024)
743 .max_dbs(4)
744 .open(&db_path)
745 .unwrap()
746 };
747
748 let mut txn = db_env.write_txn().unwrap();
749 // Create the metadata.v2 database with raw bytes so we can write
750 // an incompatible key format.
751 let metadata_db: Database<heed::types::Bytes, heed::types::Bytes> = db_env
752 .create_database(&mut txn, Some("metadata.v2"))
753 .unwrap();
754
755 // Write an incompatible PromptId key: `{"kind":"CommitMessage"}`
756 // This is the old/branch format that current code can't decode.
757 let bad_key = br#"{"kind":"CommitMessage"}"#;
758 let dummy_metadata = br#"{"id":{"kind":"CommitMessage"},"title":"Bad Record","default":false,"saved_at":"2024-01-01T00:00:00Z"}"#;
759 metadata_db.put(&mut txn, bad_key, dummy_metadata).unwrap();
760
761 // Also write a valid record to ensure we can still read good data.
762 let good_key = br#"{"kind":"User","uuid":"550e8400-e29b-41d4-a716-446655440000"}"#;
763 let good_metadata = br#"{"id":{"kind":"User","uuid":"550e8400-e29b-41d4-a716-446655440000"},"title":"Good Record","default":false,"saved_at":"2024-01-01T00:00:00Z"}"#;
764 metadata_db.put(&mut txn, good_key, good_metadata).unwrap();
765
766 txn.commit().unwrap();
767 }
768
769 // Now try to create a PromptStore from this DB.
770 // With fail-open behavior, this should succeed and skip the bad record.
771 // Without fail-open, this would return an error.
772 let store_result = cx.update(|cx| PromptStore::new(db_path, cx)).await;
773
774 assert!(
775 store_result.is_ok(),
776 "PromptStore should initialize successfully even with incompatible DB records. \
777 Got error: {:?}",
778 store_result.err()
779 );
780
781 let store = cx.new(|_cx| store_result.unwrap());
782
783 // Verify the good record was loaded.
784 let good_id = PromptId::User {
785 uuid: UserPromptId("550e8400-e29b-41d4-a716-446655440000".parse().unwrap()),
786 };
787 let metadata = store.read_with(cx, |store, _| store.metadata(good_id));
788 assert!(
789 metadata.is_some(),
790 "Valid records should still be loaded after skipping bad ones"
791 );
792 assert_eq!(
793 metadata
794 .as_ref()
795 .and_then(|m| m.title.as_ref().map(|t| t.as_ref())),
796 Some("Good Record"),
797 "Valid record should have correct title"
798 );
799 }
800
801 #[gpui::test]
802 async fn test_deleted_prompt_does_not_reappear_after_migration(cx: &mut TestAppContext) {
803 cx.executor().allow_parking();
804
805 let temp_dir = tempfile::tempdir().unwrap();
806 let db_path = temp_dir.path().join("prompts-db-v1-migration");
807 std::fs::create_dir_all(&db_path).unwrap();
808
809 let prompt_uuid: Uuid = "550e8400-e29b-41d4-a716-446655440001".parse().unwrap();
810 let prompt_id_v1 = PromptIdV1(prompt_uuid);
811 let prompt_id_v2 = PromptId::User {
812 uuid: UserPromptId(prompt_uuid),
813 };
814
815 // Create V1 database with a prompt
816 {
817 let db_env = unsafe {
818 heed::EnvOpenOptions::new()
819 .map_size(1024 * 1024 * 1024)
820 .max_dbs(4)
821 .open(&db_path)
822 .unwrap()
823 };
824
825 let mut txn = db_env.write_txn().unwrap();
826
827 let metadata_v1_db: Database<SerdeBincode<PromptIdV1>, SerdeBincode<PromptMetadataV1>> =
828 db_env.create_database(&mut txn, Some("metadata")).unwrap();
829
830 let bodies_v1_db: Database<SerdeBincode<PromptIdV1>, SerdeBincode<String>> =
831 db_env.create_database(&mut txn, Some("bodies")).unwrap();
832
833 let metadata_v1 = PromptMetadataV1 {
834 id: prompt_id_v1.clone(),
835 title: Some("V1 Prompt".into()),
836 default: false,
837 saved_at: Utc::now(),
838 };
839
840 metadata_v1_db
841 .put(&mut txn, &prompt_id_v1, &metadata_v1)
842 .unwrap();
843 bodies_v1_db
844 .put(&mut txn, &prompt_id_v1, &"V1 prompt body".to_string())
845 .unwrap();
846
847 txn.commit().unwrap();
848 }
849
850 // Migrate V1 to V2 by creating PromptStore
851 let store = cx
852 .update(|cx| PromptStore::new(db_path.clone(), cx))
853 .await
854 .unwrap();
855 let store = cx.new(|_cx| store);
856
857 // Verify the prompt was migrated
858 let metadata = store.read_with(cx, |store, _| store.metadata(prompt_id_v2));
859 assert!(metadata.is_some(), "V1 prompt should be migrated to V2");
860 assert_eq!(
861 metadata
862 .as_ref()
863 .and_then(|m| m.title.as_ref().map(|t| t.as_ref())),
864 Some("V1 Prompt"),
865 "Migrated prompt should have correct title"
866 );
867
868 // Delete the prompt
869 store
870 .update(cx, |store, cx| store.delete(prompt_id_v2, cx))
871 .await
872 .unwrap();
873
874 // Verify prompt is deleted
875 let metadata_after_delete = store.read_with(cx, |store, _| store.metadata(prompt_id_v2));
876 assert!(
877 metadata_after_delete.is_none(),
878 "Prompt should be deleted from V2"
879 );
880
881 drop(store);
882
883 // "Restart" by creating a new PromptStore from the same path
884 let store_after_restart = cx.update(|cx| PromptStore::new(db_path, cx)).await.unwrap();
885 let store_after_restart = cx.new(|_cx| store_after_restart);
886
887 // Test the prompt does not reappear
888 let metadata_after_restart =
889 store_after_restart.read_with(cx, |store, _| store.metadata(prompt_id_v2));
890 assert!(
891 metadata_after_restart.is_none(),
892 "Deleted prompt should NOT reappear after restart/migration"
893 );
894 }
895}