Add ability to attach rules as context (#29109)

Release Notes:

- agent: Added support for adding rules as context.

crates/agent/src/context_store.rs

@@ -9,6 +9,7 @@ use futures::{self, Future, FutureExt, future};
 use gpui::{App, AppContext as _, Context, Entity, SharedString, Task, WeakEntity};
 use language::{Buffer, File};
 use project::{Project, ProjectItem, ProjectPath, Worktree};
+use prompt_store::UserPromptId;
 use rope::{Point, Rope};
 use text::{Anchor, BufferId, OffsetRangeExt};
 use util::{ResultExt as _, maybe};
@@ -16,7 +17,7 @@ use util::{ResultExt as _, maybe};
 use crate::ThreadStore;
 use crate::context::{
     AssistantContext, ContextBuffer, ContextId, ContextSymbol, ContextSymbolId, DirectoryContext,
-    ExcerptContext, FetchedUrlContext, FileContext, SymbolContext, ThreadContext,
+    ExcerptContext, FetchedUrlContext, FileContext, RulesContext, SymbolContext, ThreadContext,
 };
 use crate::context_strip::SuggestedContext;
 use crate::thread::{Thread, ThreadId};
@@ -25,7 +26,6 @@ pub struct ContextStore {
     project: WeakEntity<Project>,
     context: Vec<AssistantContext>,
     thread_store: Option<WeakEntity<ThreadStore>>,
-    // TODO: If an EntityId is used for all context types (like BufferId), can remove ContextId.
     next_context_id: ContextId,
     files: BTreeMap<BufferId, ContextId>,
     directories: HashMap<ProjectPath, ContextId>,
@@ -35,6 +35,7 @@ pub struct ContextStore {
     threads: HashMap<ThreadId, ContextId>,
     thread_summary_tasks: Vec<Task<()>>,
     fetched_urls: HashMap<String, ContextId>,
+    user_rules: HashMap<UserPromptId, ContextId>,
 }
 
 impl ContextStore {
@@ -55,6 +56,7 @@ impl ContextStore {
             threads: HashMap::default(),
             thread_summary_tasks: Vec::new(),
             fetched_urls: HashMap::default(),
+            user_rules: HashMap::default(),
         }
     }
 
@@ -72,6 +74,7 @@ impl ContextStore {
         self.directories.clear();
         self.threads.clear();
         self.fetched_urls.clear();
+        self.user_rules.clear();
     }
 
     pub fn add_file_from_path(
@@ -390,6 +393,42 @@ impl ContextStore {
         cx.notify();
     }
 
+    pub fn add_rules(
+        &mut self,
+        prompt_id: UserPromptId,
+        title: impl Into<SharedString>,
+        text: impl Into<SharedString>,
+        remove_if_exists: bool,
+        cx: &mut Context<ContextStore>,
+    ) {
+        if let Some(context_id) = self.includes_user_rules(&prompt_id) {
+            if remove_if_exists {
+                self.remove_context(context_id, cx);
+            }
+        } else {
+            self.insert_user_rules(prompt_id, title, text, cx);
+        }
+    }
+
+    pub fn insert_user_rules(
+        &mut self,
+        prompt_id: UserPromptId,
+        title: impl Into<SharedString>,
+        text: impl Into<SharedString>,
+        cx: &mut Context<ContextStore>,
+    ) {
+        let id = self.next_context_id.post_inc();
+
+        self.user_rules.insert(prompt_id, id);
+        self.context.push(AssistantContext::Rules(RulesContext {
+            id,
+            prompt_id,
+            title: title.into(),
+            text: text.into(),
+        }));
+        cx.notify();
+    }
+
     pub fn add_fetched_url(
         &mut self,
         url: String,
@@ -518,6 +557,9 @@ impl ContextStore {
             AssistantContext::Thread(_) => {
                 self.threads.retain(|_, context_id| *context_id != id);
             }
+            AssistantContext::Rules(RulesContext { prompt_id, .. }) => {
+                self.user_rules.remove(&prompt_id);
+            }
         }
 
         cx.notify();
@@ -614,6 +656,10 @@ impl ContextStore {
         self.threads.get(thread_id).copied()
     }
 
+    pub fn includes_user_rules(&self, prompt_id: &UserPromptId) -> Option<ContextId> {
+        self.user_rules.get(prompt_id).copied()
+    }
+
     pub fn includes_url(&self, url: &str) -> Option<ContextId> {
         self.fetched_urls.get(url).copied()
     }
@@ -641,7 +687,8 @@ impl ContextStore {
                 | AssistantContext::Symbol(_)
                 | AssistantContext::Excerpt(_)
                 | AssistantContext::FetchedUrl(_)
-                | AssistantContext::Thread(_) => None,
+                | AssistantContext::Thread(_)
+                | AssistantContext::Rules(_) => None,
             })
             .collect()
     }
@@ -876,6 +923,10 @@ pub fn refresh_context_store_text(
                 // and doing the caching properly could be tricky (unless it's already handled by
                 // the HttpClient?).
                 AssistantContext::FetchedUrl(_) => {}
+                AssistantContext::Rules(user_rules_context) => {
+                    let context_store = context_store.clone();
+                    return Some(refresh_user_rules(context_store, user_rules_context, cx));
+                }
             }
 
             None
@@ -1026,6 +1077,45 @@ fn refresh_thread_text(
     })
 }
 
+fn refresh_user_rules(
+    context_store: Entity<ContextStore>,
+    user_rules_context: &RulesContext,
+    cx: &App,
+) -> Task<()> {
+    let id = user_rules_context.id;
+    let prompt_id = user_rules_context.prompt_id;
+    let Some(thread_store) = context_store.read(cx).thread_store.as_ref() else {
+        return Task::ready(());
+    };
+    let Ok(load_task) = thread_store.read_with(cx, |thread_store, cx| {
+        thread_store.load_rules(prompt_id, cx)
+    }) else {
+        return Task::ready(());
+    };
+    cx.spawn(async move |cx| {
+        if let Ok((metadata, text)) = load_task.await {
+            if let Some(title) = metadata.title.clone() {
+                context_store
+                    .update(cx, |context_store, _cx| {
+                        context_store.replace_context(AssistantContext::Rules(RulesContext {
+                            id,
+                            prompt_id,
+                            title,
+                            text: text.into(),
+                        }));
+                    })
+                    .ok();
+                return;
+            }
+        }
+        context_store
+            .update(cx, |context_store, cx| {
+                context_store.remove_context(id, cx);
+            })
+            .ok();
+    })
+}
+
 fn refresh_context_buffer(
     context_buffer: &ContextBuffer,
     cx: &App,