agent: Use default prompts from prompt library in system prompt (#28915)

Related to #28490.

- Default prompts from the prompt library are now included as "user
rules" in the system prompt.
- Presence of these user rules is shown at the beginning of the thread
in the UI.
_ Now uses an `Entity<PromptStore>` instead of an `Arc<PromptStore>`.
Motivation for this is emitting a `PromptsUpdatedEvent`.
- Now disallows concurrent reloading of the system prompt. Before this
change it was possible for reloads to race.

Release Notes:

- agent: Added support for including default prompts from the Prompt
Library as "user rules" in the system prompt.

---------

Co-authored-by: Danilo Leal <daniloleal09@gmail.com>

crates/prompt_store/src/prompt_store.rs

@@ -4,9 +4,11 @@ use anyhow::{Result, anyhow};
 use chrono::{DateTime, Utc};
 use collections::HashMap;
 use futures::FutureExt as _;
-use futures::future::{self, BoxFuture, Shared};
+use futures::future::Shared;
 use fuzzy::StringMatchCandidate;
-use gpui::{App, BackgroundExecutor, Global, ReadGlobal, SharedString, Task};
+use gpui::{
+    App, AppContext, Context, Entity, EventEmitter, Global, ReadGlobal, SharedString, Task,
+};
 use heed::{
     Database, RoTxn,
     types::{SerdeBincode, SerdeJson, Str},
@@ -29,11 +31,16 @@ use uuid::Uuid;
 /// a shared future to a global.
 pub fn init(cx: &mut App) {
     let db_path = paths::prompts_dir().join("prompts-library-db.0.mdb");
-    let prompt_store_future = PromptStore::new(db_path, cx.background_executor().clone())
-        .then(|result| future::ready(result.map(Arc::new).map_err(Arc::new)))
-        .boxed()
+    let prompt_store_task = PromptStore::new(db_path, cx);
+    let prompt_store_entity_task = cx
+        .spawn(async move |cx| {
+            prompt_store_task
+                .await
+                .and_then(|prompt_store| cx.new(|_cx| prompt_store))
+                .map_err(Arc::new)
+        })
         .shared();
-    cx.set_global(GlobalPromptStore(prompt_store_future))
+    cx.set_global(GlobalPromptStore(prompt_store_entity_task))
 }
 
 #[derive(Clone, Debug, Serialize, Deserialize)]
@@ -64,13 +71,16 @@ impl PromptId {
 }
 
 pub struct PromptStore {
-    executor: BackgroundExecutor,
     env: heed::Env,
     metadata_cache: RwLock<MetadataCache>,
     metadata: Database<SerdeJson<PromptId>, SerdeJson<PromptMetadata>>,
     bodies: Database<SerdeJson<PromptId>, Str>,
 }
 
+pub struct PromptsUpdatedEvent;
+
+impl EventEmitter<PromptsUpdatedEvent> for PromptStore {}
+
 #[derive(Default)]
 struct MetadataCache {
     metadata: Vec<PromptMetadata>,
@@ -117,49 +127,45 @@ impl MetadataCache {
 }
 
 impl PromptStore {
-    pub fn global(cx: &App) -> impl Future<Output = Result<Arc<Self>>> + use<> {
+    pub fn global(cx: &App) -> impl Future<Output = Result<Entity<Self>>> + use<> {
         let store = GlobalPromptStore::global(cx).0.clone();
         async move { store.await.map_err(|err| anyhow!(err)) }
     }
 
-    pub fn new(db_path: PathBuf, executor: BackgroundExecutor) -> Task<Result<Self>> {
-        executor.spawn({
-            let executor = executor.clone();
-            async move {
-                std::fs::create_dir_all(&db_path)?;
+    pub fn new(db_path: PathBuf, cx: &App) -> Task<Result<Self>> {
+        cx.background_spawn(async move {
+            std::fs::create_dir_all(&db_path)?;
 
-                let db_env = unsafe {
-                    heed::EnvOpenOptions::new()
-                        .map_size(1024 * 1024 * 1024) // 1GB
-                        .max_dbs(4) // Metadata and bodies (possibly v1 of both as well)
-                        .open(db_path)?
-                };
+            let db_env = unsafe {
+                heed::EnvOpenOptions::new()
+                    .map_size(1024 * 1024 * 1024) // 1GB
+                    .max_dbs(4) // Metadata and bodies (possibly v1 of both as well)
+                    .open(db_path)?
+            };
 
-                let mut txn = db_env.write_txn()?;
-                let metadata = db_env.create_database(&mut txn, Some("metadata.v2"))?;
-                let bodies = db_env.create_database(&mut txn, Some("bodies.v2"))?;
+            let mut txn = db_env.write_txn()?;
+            let metadata = db_env.create_database(&mut txn, Some("metadata.v2"))?;
+            let bodies = db_env.create_database(&mut txn, Some("bodies.v2"))?;
 
-                // Remove edit workflow prompt, as we decided to opt into it using
-                // a slash command instead.
-                metadata.delete(&mut txn, &PromptId::EditWorkflow).ok();
-                bodies.delete(&mut txn, &PromptId::EditWorkflow).ok();
+            // Remove edit workflow prompt, as we decided to opt into it using
+            // a slash command instead.
+            metadata.delete(&mut txn, &PromptId::EditWorkflow).ok();
+            bodies.delete(&mut txn, &PromptId::EditWorkflow).ok();
 
-                txn.commit()?;
+            txn.commit()?;
 
-                Self::upgrade_dbs(&db_env, metadata, bodies).log_err();
+            Self::upgrade_dbs(&db_env, metadata, bodies).log_err();
 
-                let txn = db_env.read_txn()?;
-                let metadata_cache = MetadataCache::from_db(metadata, &txn)?;
-                txn.commit()?;
+            let txn = db_env.read_txn()?;
+            let metadata_cache = MetadataCache::from_db(metadata, &txn)?;
+            txn.commit()?;
 
-                Ok(PromptStore {
-                    executor,
-                    env: db_env,
-                    metadata_cache: RwLock::new(metadata_cache),
-                    metadata,
-                    bodies,
-                })
-            }
+            Ok(PromptStore {
+                env: db_env,
+                metadata_cache: RwLock::new(metadata_cache),
+                metadata,
+                bodies,
+            })
         })
     }
 
@@ -237,10 +243,10 @@ impl PromptStore {
         Ok(())
     }
 
-    pub fn load(&self, id: PromptId) -> Task<Result<String>> {
+    pub fn load(&self, id: PromptId, cx: &App) -> Task<Result<String>> {
         let env = self.env.clone();
         let bodies = self.bodies;
-        self.executor.spawn(async move {
+        cx.background_spawn(async move {
             let txn = env.read_txn()?;
             let mut prompt = bodies
                 .get(&txn, &id)?
@@ -262,21 +268,27 @@ impl PromptStore {
             .collect::<Vec<_>>();
     }
 
-    pub fn delete(&self, id: PromptId) -> Task<Result<()>> {
+    pub fn delete(&self, id: PromptId, cx: &Context<Self>) -> Task<Result<()>> {
         self.metadata_cache.write().remove(id);
 
         let db_connection = self.env.clone();
         let bodies = self.bodies;
         let metadata = self.metadata;
 
-        self.executor.spawn(async move {
+        let task = cx.background_spawn(async move {
             let mut txn = db_connection.write_txn()?;
 
             metadata.delete(&mut txn, &id)?;
             bodies.delete(&mut txn, &id)?;
 
             txn.commit()?;
-            Ok(())
+            anyhow::Ok(())
+        });
+
+        cx.spawn(async move |this, cx| {
+            task.await?;
+            this.update(cx, |_, cx| cx.emit(PromptsUpdatedEvent)).ok();
+            anyhow::Ok(())
         })
     }
 
@@ -302,10 +314,10 @@ impl PromptStore {
         Some(metadata.id)
     }
 
-    pub fn search(&self, query: String) -> Task<Vec<PromptMetadata>> {
+    pub fn search(&self, query: String, cx: &App) -> Task<Vec<PromptMetadata>> {
         let cached_metadata = self.metadata_cache.read().metadata.clone();
-        let executor = self.executor.clone();
-        self.executor.spawn(async move {
+        let executor = cx.background_executor().clone();
+        cx.background_spawn(async move {
             let mut matches = if query.is_empty() {
                 cached_metadata
             } else {
@@ -341,6 +353,7 @@ impl PromptStore {
         title: Option<SharedString>,
         default: bool,
         body: Rope,
+        cx: &Context<Self>,
     ) -> Task<Result<()>> {
         if id.is_built_in() {
             return Task::ready(Err(anyhow!("built-in prompts cannot be saved")));
@@ -358,7 +371,7 @@ impl PromptStore {
         let bodies = self.bodies;
         let metadata = self.metadata;
 
-        self.executor.spawn(async move {
+        let task = cx.background_spawn(async move {
             let mut txn = db_connection.write_txn()?;
 
             metadata.put(&mut txn, &id, &prompt_metadata)?;
@@ -366,7 +379,13 @@ impl PromptStore {
 
             txn.commit()?;
 
-            Ok(())
+            anyhow::Ok(())
+        });
+
+        cx.spawn(async move |this, cx| {
+            task.await?;
+            this.update(cx, |_, cx| cx.emit(PromptsUpdatedEvent)).ok();
+            anyhow::Ok(())
         })
     }
 
@@ -375,6 +394,7 @@ impl PromptStore {
         id: PromptId,
         mut title: Option<SharedString>,
         default: bool,
+        cx: &Context<Self>,
     ) -> Task<Result<()>> {
         let mut cache = self.metadata_cache.write();
 
@@ -397,19 +417,23 @@ impl PromptStore {
         let db_connection = self.env.clone();
         let metadata = self.metadata;
 
-        self.executor.spawn(async move {
+        let task = cx.background_spawn(async move {
             let mut txn = db_connection.write_txn()?;
             metadata.put(&mut txn, &id, &prompt_metadata)?;
             txn.commit()?;
 
-            Ok(())
+            anyhow::Ok(())
+        });
+
+        cx.spawn(async move |this, cx| {
+            task.await?;
+            this.update(cx, |_, cx| cx.emit(PromptsUpdatedEvent)).ok();
+            anyhow::Ok(())
         })
     }
 }
 
 /// Wraps a shared future to a prompt store so it can be assigned as a context global.
-pub struct GlobalPromptStore(
-    Shared<BoxFuture<'static, Result<Arc<PromptStore>, Arc<anyhow::Error>>>>,
-);
+pub struct GlobalPromptStore(Shared<Task<Result<Entity<PromptStore>, Arc<anyhow::Error>>>>);
 
 impl Global for GlobalPromptStore {}

crates/prompt_store/src/prompts.rs

@@ -19,20 +19,29 @@ use util::{ResultExt, get_system_shell};
 #[derive(Debug, Clone, Serialize)]
 pub struct ProjectContext {
     pub worktrees: Vec<WorktreeContext>,
+    /// Whether any worktree has a rules_file. Provided as a field because handlebars can't do this.
     pub has_rules: bool,
+    pub default_user_rules: Vec<DefaultUserRulesContext>,
+    /// `!default_user_rules.is_empty()` - provided as a field because handlebars can't do this.
+    pub has_default_user_rules: bool,
     pub os: String,
     pub arch: String,
     pub shell: String,
 }
 
 impl ProjectContext {
-    pub fn new(worktrees: Vec<WorktreeContext>) -> Self {
+    pub fn new(
+        worktrees: Vec<WorktreeContext>,
+        default_user_rules: Vec<DefaultUserRulesContext>,
+    ) -> Self {
         let has_rules = worktrees
             .iter()
             .any(|worktree| worktree.rules_file.is_some());
         Self {
             worktrees,
             has_rules,
+            has_default_user_rules: !default_user_rules.is_empty(),
+            default_user_rules,
             os: std::env::consts::OS.to_string(),
             arch: std::env::consts::ARCH.to_string(),
             shell: get_system_shell(),
@@ -40,6 +49,12 @@ impl ProjectContext {
     }
 }
 
+#[derive(Debug, Clone, Serialize)]
+pub struct DefaultUserRulesContext {
+    pub title: Option<String>,
+    pub contents: String,
+}
+
 #[derive(Debug, Clone, Serialize)]
 pub struct WorktreeContext {
     pub root_name: String,
@@ -377,3 +392,30 @@ impl PromptBuilder {
         self.handlebars.lock().render("suggest_edits", &())
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_assistant_system_prompt_renders() {
+        let worktrees = vec![WorktreeContext {
+            root_name: "path".into(),
+            abs_path: Path::new("/some/path").into(),
+            rules_file: Some(RulesFileContext {
+                path_in_worktree: Path::new(".rules").into(),
+                abs_path: Path::new("/some/path/.rules").into(),
+                text: "".into(),
+            }),
+        }];
+        let default_user_rules = vec![DefaultUserRulesContext {
+            title: Some("Rules title".into()),
+            contents: "Rules contents".into(),
+        }];
+        let project_context = ProjectContext::new(worktrees, default_user_rules);
+        PromptBuilder::new(None)
+            .unwrap()
+            .generate_assistant_system_prompt(&project_context)
+            .unwrap();
+    }
+}