assistant2: Do not allow a context entry to be added multiple times (#22712)

https://github.com/user-attachments/assets/81674c88-031b-4d55-b362-43819492b93d


Release Notes:

- N/A

crates/assistant2/src/context_store.rs

@@ -1,6 +1,11 @@
-use gpui::SharedString;
-use project::ProjectEntryId;
+use std::fmt::Write as _;
+use std::path::{Path, PathBuf};
 
+use collections::HashMap;
+use gpui::SharedString;
+use language::Buffer;
+
+use crate::thread::Thread;
 use crate::{
     context::{Context, ContextId, ContextKind},
     thread::ThreadId,
@@ -9,6 +14,10 @@ use crate::{
 pub struct ContextStore {
     context: Vec<Context>,
     next_context_id: ContextId,
+    files: HashMap<PathBuf, ContextId>,
+    directories: HashMap<PathBuf, ContextId>,
+    threads: HashMap<ThreadId, ContextId>,
+    fetched_urls: HashMap<String, ContextId>,
 }
 
 impl ContextStore {
@@ -16,6 +25,10 @@ impl ContextStore {
         Self {
             context: Vec::new(),
             next_context_id: ContextId(0),
+            files: HashMap::default(),
+            directories: HashMap::default(),
+            threads: HashMap::default(),
+            fetched_urls: HashMap::default(),
         }
     }
 
@@ -24,42 +37,154 @@ impl ContextStore {
     }
 
     pub fn drain(&mut self) -> Vec<Context> {
+        self.files.clear();
+        self.directories.clear();
         self.context.drain(..).collect()
     }
 
     pub fn clear(&mut self) {
         self.context.clear();
+        self.files.clear();
+        self.directories.clear();
     }
 
-    pub fn insert_context(
-        &mut self,
-        kind: ContextKind,
-        name: impl Into<SharedString>,
-        text: impl Into<SharedString>,
-    ) {
+    pub fn insert_file(&mut self, buffer: &Buffer) {
+        let Some(file) = buffer.file() else {
+            return;
+        };
+
+        let path = file.path();
+
+        let id = self.next_context_id.post_inc();
+        self.files.insert(path.to_path_buf(), id);
+
+        let name = path.to_string_lossy().into_owned().into();
+
+        let mut text = String::new();
+        push_fenced_codeblock(path, buffer.text(), &mut text);
+
         self.context.push(Context {
-            id: self.next_context_id.post_inc(),
-            name: name.into(),
-            kind,
+            id,
+            name,
+            kind: ContextKind::File,
+            text: text.into(),
+        });
+    }
+
+    pub fn insert_directory(&mut self, path: &Path, text: impl Into<SharedString>) {
+        let id = self.next_context_id.post_inc();
+        self.directories.insert(path.to_path_buf(), id);
+
+        let name = path.to_string_lossy().into_owned().into();
+
+        self.context.push(Context {
+            id,
+            name,
+            kind: ContextKind::Directory,
+            text: text.into(),
+        });
+    }
+
+    pub fn insert_thread(&mut self, thread: &Thread) {
+        let context_id = self.next_context_id.post_inc();
+        self.threads.insert(thread.id().clone(), context_id);
+
+        self.context.push(Context {
+            id: context_id,
+            name: thread.summary().unwrap_or("New thread".into()),
+            kind: ContextKind::Thread,
+            text: thread.text().into(),
+        });
+    }
+
+    pub fn insert_fetched_url(&mut self, url: String, text: impl Into<SharedString>) {
+        let context_id = self.next_context_id.post_inc();
+        self.fetched_urls.insert(url.clone(), context_id);
+
+        self.context.push(Context {
+            id: context_id,
+            name: url.into(),
+            kind: ContextKind::FetchedUrl,
             text: text.into(),
         });
     }
 
     pub fn remove_context(&mut self, id: &ContextId) {
-        self.context.retain(|context| context.id != *id);
+        let Some(ix) = self.context.iter().position(|c| c.id == *id) else {
+            return;
+        };
+
+        match self.context.remove(ix).kind {
+            ContextKind::File => {
+                self.files.retain(|_, p_id| p_id != id);
+            }
+            ContextKind::Directory => {
+                self.directories.retain(|_, p_id| p_id != id);
+            }
+            ContextKind::FetchedUrl => {
+                self.fetched_urls.retain(|_, p_id| p_id != id);
+            }
+            ContextKind::Thread => {
+                self.threads.retain(|_, p_id| p_id != id);
+            }
+        }
     }
 
-    pub fn contains_project_entry(&self, entry_id: ProjectEntryId) -> bool {
-        self.context.iter().any(|probe| match probe.kind {
-            ContextKind::File(probe_entry_id) => probe_entry_id == entry_id,
-            ContextKind::Directory | ContextKind::FetchedUrl | ContextKind::Thread(_) => false,
-        })
+    pub fn included_file(&self, path: &Path) -> Option<IncludedFile> {
+        if let Some(id) = self.files.get(path) {
+            return Some(IncludedFile::Direct(*id));
+        }
+
+        if self.directories.is_empty() {
+            return None;
+        }
+
+        let mut buf = path.to_path_buf();
+
+        while buf.pop() {
+            if let Some(_) = self.directories.get(&buf) {
+                return Some(IncludedFile::InDirectory(buf));
+            }
+        }
+
+        None
     }
 
-    pub fn contains_thread(&self, thread_id: &ThreadId) -> bool {
-        self.context.iter().any(|probe| match probe.kind {
-            ContextKind::Thread(ref probe_thread_id) => probe_thread_id == thread_id,
-            ContextKind::File(_) | ContextKind::Directory | ContextKind::FetchedUrl => false,
-        })
+    pub fn included_directory(&self, path: &Path) -> Option<ContextId> {
+        self.directories.get(path).copied()
+    }
+
+    pub fn included_thread(&self, thread_id: &ThreadId) -> Option<ContextId> {
+        self.threads.get(thread_id).copied()
+    }
+
+    pub fn included_url(&self, url: &str) -> Option<ContextId> {
+        self.fetched_urls.get(url).copied()
     }
 }
+
+pub enum IncludedFile {
+    Direct(ContextId),
+    InDirectory(PathBuf),
+}
+
+pub(crate) fn push_fenced_codeblock(path: &Path, content: String, buf: &mut String) {
+    buf.reserve(content.len() + 64);
+
+    write!(buf, "```").unwrap();
+
+    if let Some(extension) = path.extension().and_then(|ext| ext.to_str()) {
+        write!(buf, "{} ", extension).unwrap();
+    }
+
+    write!(buf, "{}", path.display()).unwrap();
+
+    buf.push('\n');
+    buf.push_str(&content);
+
+    if !buf.ends_with('\n') {
+        buf.push('\n');
+    }
+
+    buf.push_str("```\n");
+}