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.rs

@@ -1,11 +1,8 @@
 use gpui::SharedString;
 use language_model::{LanguageModelRequestMessage, MessageContent};
-use project::ProjectEntryId;
 use serde::{Deserialize, Serialize};
 use util::post_inc;
 
-use crate::thread::ThreadId;
-
 #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, Serialize, Deserialize)]
 pub struct ContextId(pub(crate) usize);
 
@@ -26,10 +23,10 @@ pub struct Context {
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub enum ContextKind {
-    File(ProjectEntryId),
+    File,
     Directory,
     FetchedUrl,
-    Thread(ThreadId),
+    Thread,
 }
 
 pub fn attach_context_to_message(
@@ -43,7 +40,7 @@ pub fn attach_context_to_message(
 
     for context in context.into_iter() {
         match context.kind {
-            ContextKind::File(_) => {
+            ContextKind::File => {
                 file_context.push_str(&context.text);
                 file_context.push('\n');
             }
@@ -57,7 +54,7 @@ pub fn attach_context_to_message(
                 fetch_context.push_str(&context.text);
                 fetch_context.push('\n');
             }
-            ContextKind::Thread(_) => {
+            ContextKind::Thread => {
                 thread_context.push_str(&context.name);
                 thread_context.push('\n');
                 thread_context.push_str(&context.text);

crates/assistant2/src/context_picker.rs

@@ -14,6 +14,7 @@ use ui::{prelude::*, ListItem, ListItemSpacing};
 use util::ResultExt;
 use workspace::Workspace;
 
+use crate::context::ContextKind;
 use crate::context_picker::directory_context_picker::DirectoryContextPicker;
 use crate::context_picker::fetch_context_picker::FetchContextPicker;
 use crate::context_picker::file_context_picker::FileContextPicker;
@@ -52,24 +53,24 @@ impl ContextPicker {
         let mut entries = Vec::new();
         entries.push(ContextPickerEntry {
             name: "File".into(),
-            kind: ContextPickerEntryKind::File,
+            kind: ContextKind::File,
             icon: IconName::File,
         });
         entries.push(ContextPickerEntry {
             name: "Folder".into(),
-            kind: ContextPickerEntryKind::Directory,
+            kind: ContextKind::Directory,
             icon: IconName::Folder,
         });
         entries.push(ContextPickerEntry {
             name: "Fetch".into(),
-            kind: ContextPickerEntryKind::FetchedUrl,
+            kind: ContextKind::FetchedUrl,
             icon: IconName::Globe,
         });
 
         if thread_store.is_some() {
             entries.push(ContextPickerEntry {
                 name: "Thread".into(),
-                kind: ContextPickerEntryKind::Thread,
+                kind: ContextKind::Thread,
                 icon: IconName::MessageCircle,
             });
         }
@@ -133,18 +134,10 @@ impl Render for ContextPicker {
 #[derive(Clone)]
 struct ContextPickerEntry {
     name: SharedString,
-    kind: ContextPickerEntryKind,
+    kind: ContextKind,
     icon: IconName,
 }
 
-#[derive(Debug, Clone)]
-enum ContextPickerEntryKind {
-    File,
-    Directory,
-    FetchedUrl,
-    Thread,
-}
-
 pub(crate) struct ContextPickerDelegate {
     context_picker: WeakView<ContextPicker>,
     workspace: WeakView<Workspace>,
@@ -184,7 +177,7 @@ impl PickerDelegate for ContextPickerDelegate {
             self.context_picker
                 .update(cx, |this, cx| {
                     match entry.kind {
-                        ContextPickerEntryKind::File => {
+                        ContextKind::File => {
                             this.mode = ContextPickerMode::File(cx.new_view(|cx| {
                                 FileContextPicker::new(
                                     self.context_picker.clone(),
@@ -195,7 +188,7 @@ impl PickerDelegate for ContextPickerDelegate {
                                 )
                             }));
                         }
-                        ContextPickerEntryKind::Directory => {
+                        ContextKind::Directory => {
                             this.mode = ContextPickerMode::Directory(cx.new_view(|cx| {
                                 DirectoryContextPicker::new(
                                     self.context_picker.clone(),
@@ -206,7 +199,7 @@ impl PickerDelegate for ContextPickerDelegate {
                                 )
                             }));
                         }
-                        ContextPickerEntryKind::FetchedUrl => {
+                        ContextKind::FetchedUrl => {
                             this.mode = ContextPickerMode::Fetch(cx.new_view(|cx| {
                                 FetchContextPicker::new(
                                     self.context_picker.clone(),
@@ -217,7 +210,7 @@ impl PickerDelegate for ContextPickerDelegate {
                                 )
                             }));
                         }
-                        ContextPickerEntryKind::Thread => {
+                        ContextKind::Thread => {
                             if let Some(thread_store) = self.thread_store.as_ref() {
                                 this.mode = ContextPickerMode::Thread(cx.new_view(|cx| {
                                     ThreadContextPicker::new(

crates/assistant2/src/context_picker/directory_context_picker.rs

@@ -11,10 +11,8 @@ use ui::{prelude::*, ListItem};
 use util::ResultExt as _;
 use workspace::Workspace;
 
-use crate::context::ContextKind;
-use crate::context_picker::file_context_picker::codeblock_fence_for_path;
 use crate::context_picker::{ConfirmBehavior, ContextPicker};
-use crate::context_store::ContextStore;
+use crate::context_store::{push_fenced_codeblock, ContextStore};
 
 pub struct DirectoryContextPicker {
     picker: View<Picker<DirectoryContextPickerDelegate>>,
@@ -189,6 +187,22 @@ impl PickerDelegate for DirectoryContextPickerDelegate {
             return;
         };
         let path = mat.path.clone();
+
+        if self
+            .context_store
+            .update(cx, |context_store, _cx| {
+                if let Some(context_id) = context_store.included_directory(&path) {
+                    context_store.remove_context(&context_id);
+                    true
+                } else {
+                    false
+                }
+            })
+            .unwrap_or(true)
+        {
+            return;
+        }
+
         let worktree_id = WorktreeId::from_usize(mat.worktree_id);
         let confirm_behavior = self.confirm_behavior;
         cx.spawn(|this, mut cx| async move {
@@ -235,23 +249,15 @@ impl PickerDelegate for DirectoryContextPickerDelegate {
                 let mut text = String::new();
 
                 for buffer in buffers {
-                    text.push_str(&codeblock_fence_for_path(Some(&path), None));
+                    let buffer = buffer.read(cx);
-                    text.push_str(&buffer.read(cx).text());
+                    let path = buffer.file().map_or(&path, |file| file.path());
-                    if !text.ends_with('\n') {
+                    push_fenced_codeblock(&path, buffer.text(), &mut text);
-                        text.push('\n');
-                    }
-
-                    text.push_str("```\n");
                 }
 
                 this.delegate
                     .context_store
                     .update(cx, |context_store, _cx| {
-                        context_store.insert_context(
+                        context_store.insert_directory(&path, text);
-                            ContextKind::Directory,
-                            path.to_string_lossy().to_string(),
-                            text,
-                        );
                     })?;
 
                 match confirm_behavior {
@@ -280,16 +286,26 @@ impl PickerDelegate for DirectoryContextPickerDelegate {
         &self,
         ix: usize,
         selected: bool,
-        _cx: &mut ViewContext<Picker<Self>>,
+        cx: &mut ViewContext<Picker<Self>>,
     ) -> Option<Self::ListItem> {
         let path_match = &self.matches[ix];
         let directory_name = path_match.path.to_string_lossy().to_string();
 
+        let added = self.context_store.upgrade().map_or(false, |context_store| {
+            context_store
+                .read(cx)
+                .included_directory(&path_match.path)
+                .is_some()
+        });
+
         Some(
             ListItem::new(ix)
                 .inset(true)
                 .toggle_state(selected)
-                .child(h_flex().gap_2().child(Label::new(directory_name))),
+                .child(h_flex().gap_2().child(Label::new(directory_name)))
+                .when(added, |el| {
+                    el.end_slot(Label::new("Added").size(LabelSize::XSmall))
+                }),
         )
     }
 }

crates/assistant2/src/context_picker/fetch_context_picker.rs

@@ -11,7 +11,6 @@ use picker::{Picker, PickerDelegate};
 use ui::{prelude::*, ListItem, ViewContext};
 use workspace::Workspace;
 
-use crate::context::ContextKind;
 use crate::context_picker::{ConfirmBehavior, ContextPicker};
 use crate::context_store::ContextStore;
 
@@ -201,7 +200,9 @@ impl PickerDelegate for FetchContextPickerDelegate {
                 this.delegate
                     .context_store
                     .update(cx, |context_store, _cx| {
-                        context_store.insert_context(ContextKind::FetchedUrl, url, text);
+                        if context_store.included_url(&url).is_none() {
+                            context_store.insert_fetched_url(url, text);
+                        }
                     })?;
 
                 match confirm_behavior {
@@ -230,13 +231,22 @@ impl PickerDelegate for FetchContextPickerDelegate {
         &self,
         ix: usize,
         selected: bool,
-        _cx: &mut ViewContext<Picker<Self>>,
+        cx: &mut ViewContext<Picker<Self>>,
     ) -> Option<Self::ListItem> {
+        let added = self.context_store.upgrade().map_or(false, |context_store| {
+            context_store.read(cx).included_url(&self.url).is_some()
+        });
+
         Some(
             ListItem::new(ix)
                 .inset(true)
                 .toggle_state(selected)
-                .child(Label::new(self.url.clone())),
+                .child(Label::new(self.url.clone()))
+                .when(added, |child| {
+                    child
+                        .disabled(true)
+                        .end_slot(Label::new("Added").size(LabelSize::XSmall))
+                }),
         )
     }
 }

crates/assistant2/src/context_picker/file_context_picker.rs

@@ -1,5 +1,3 @@
-use std::fmt::Write as _;
-use std::ops::RangeInclusive;
 use std::path::Path;
 use std::sync::atomic::AtomicBool;
 use std::sync::Arc;
@@ -8,13 +6,12 @@ use fuzzy::PathMatch;
 use gpui::{AppContext, DismissEvent, FocusHandle, FocusableView, Task, View, WeakModel, WeakView};
 use picker::{Picker, PickerDelegate};
 use project::{PathMatchCandidateSet, ProjectPath, WorktreeId};
-use ui::{prelude::*, ListItem};
+use ui::{prelude::*, ListItem, Tooltip};
 use util::ResultExt as _;
 use workspace::Workspace;
 
-use crate::context::ContextKind;
 use crate::context_picker::{ConfirmBehavior, ContextPicker};
-use crate::context_store::ContextStore;
+use crate::context_store::{ContextStore, IncludedFile};
 
 pub struct FileContextPicker {
     picker: View<Picker<FileContextPickerDelegate>>,
@@ -204,20 +201,37 @@ impl PickerDelegate for FileContextPickerDelegate {
             return;
         };
         let path = mat.path.clone();
+
+        if self
+            .context_store
+            .update(cx, |context_store, _cx| {
+                match context_store.included_file(&path) {
+                    Some(IncludedFile::Direct(context_id)) => {
+                        context_store.remove_context(&context_id);
+                        true
+                    }
+                    Some(IncludedFile::InDirectory(_)) => true,
+                    None => false,
+                }
+            })
+            .unwrap_or(true)
+        {
+            return;
+        }
+
         let worktree_id = WorktreeId::from_usize(mat.worktree_id);
         let confirm_behavior = self.confirm_behavior;
         cx.spawn(|this, mut cx| async move {
-            let Some((entry_id, open_buffer_task)) = project
+            let Some(open_buffer_task) = project
                 .update(&mut cx, |project, cx| {
                     let project_path = ProjectPath {
                         worktree_id,
                         path: path.clone(),
                     };
 
-                    let entry_id = project.entry_for_path(&project_path, cx)?.id;
                     let task = project.open_buffer(project_path, cx);
 
-                    Some((entry_id, task))
+                    Some(task)
                 })
                 .ok()
                 .flatten()
@@ -231,20 +245,7 @@ impl PickerDelegate for FileContextPickerDelegate {
                 this.delegate
                     .context_store
                     .update(cx, |context_store, cx| {
-                        let mut text = String::new();
+                        context_store.insert_file(buffer.read(cx));
-                        text.push_str(&codeblock_fence_for_path(Some(&path), None));
-                        text.push_str(&buffer.read(cx).text());
-                        if !text.ends_with('\n') {
-                            text.push('\n');
-                        }
-
-                        text.push_str("```\n");
-
-                        context_store.insert_context(
-                            ContextKind::File(entry_id),
-                            path.to_string_lossy().to_string(),
-                            text,
-                        );
                     })?;
 
                 match confirm_behavior {
@@ -273,7 +274,7 @@ impl PickerDelegate for FileContextPickerDelegate {
         &self,
         ix: usize,
         selected: bool,
-        _cx: &mut ViewContext<Picker<Self>>,
+        cx: &mut ViewContext<Picker<Self>>,
     ) -> Option<Self::ListItem> {
         let path_match = &self.matches[ix];
 
@@ -301,42 +302,36 @@ impl PickerDelegate for FileContextPickerDelegate {
             (file_name, Some(directory))
         };
 
+        let added = self
+            .context_store
+            .upgrade()
+            .and_then(|context_store| context_store.read(cx).included_file(&path_match.path));
+
         Some(
-            ListItem::new(ix).inset(true).toggle_state(selected).child(
+            ListItem::new(ix)
-                h_flex()
+                .inset(true)
-                    .gap_2()
+                .toggle_state(selected)
-                    .child(Label::new(file_name))
+                .child(
-                    .children(directory.map(|directory| {
+                    h_flex()
-                        Label::new(directory)
+                        .gap_2()
-                            .size(LabelSize::Small)
+                        .child(Label::new(file_name))
-                            .color(Color::Muted)
+                        .children(directory.map(|directory| {
-                    })),
+                            Label::new(directory)
-            ),
+                                .size(LabelSize::Small)
+                                .color(Color::Muted)
+                        })),
+                )
+                .when_some(added, |el, added| match added {
+                    IncludedFile::Direct(_) => {
+                        el.end_slot(Label::new("Added").size(LabelSize::XSmall))
+                    }
+                    IncludedFile::InDirectory(dir_name) => {
+                        let dir_name = dir_name.to_string_lossy().into_owned();
+
+                        el.end_slot(Label::new("Included").size(LabelSize::XSmall))
+                            .tooltip(move |cx| Tooltip::text(format!("in {dir_name}"), cx))
+                    }
+                }),
         )
     }
 }
-
-pub(crate) fn codeblock_fence_for_path(
-    path: Option<&Path>,
-    row_range: Option<RangeInclusive<u32>>,
-) -> String {
-    let mut text = String::new();
-    write!(text, "```").unwrap();
-
-    if let Some(path) = path {
-        if let Some(extension) = path.extension().and_then(|ext| ext.to_str()) {
-            write!(text, "{} ", extension).unwrap();
-        }
-
-        write!(text, "{}", path.display()).unwrap();
-    } else {
-        write!(text, "untitled").unwrap();
-    }
-
-    if let Some(row_range) = row_range {
-        write!(text, ":{}-{}", row_range.start() + 1, row_range.end() + 1).unwrap();
-    }
-
-    text.push('\n');
-    text
-}

crates/assistant2/src/context_picker/thread_context_picker.rs

@@ -5,7 +5,6 @@ use gpui::{AppContext, DismissEvent, FocusHandle, FocusableView, Task, View, Wea
 use picker::{Picker, PickerDelegate};
 use ui::{prelude::*, ListItem};
 
-use crate::context::ContextKind;
 use crate::context_picker::{ConfirmBehavior, ContextPicker};
 use crate::context_store;
 use crate::thread::ThreadId;
@@ -169,11 +168,11 @@ impl PickerDelegate for ThreadContextPickerDelegate {
 
         self.context_store
             .update(cx, |context_store, cx| {
-                context_store.insert_context(
+                if let Some(context_id) = context_store.included_thread(&entry.id) {
-                    ContextKind::Thread(thread.read(cx).id().clone()),
+                    context_store.remove_context(&context_id);
-                    entry.summary.clone(),
+                } else {
-                    thread.read(cx).text(),
+                    context_store.insert_thread(thread.read(cx));
-                );
+                }
             })
             .ok();
 
@@ -196,15 +195,22 @@ impl PickerDelegate for ThreadContextPickerDelegate {
         &self,
         ix: usize,
         selected: bool,
-        _cx: &mut ViewContext<Picker<Self>>,
+        cx: &mut ViewContext<Picker<Self>>,
     ) -> Option<Self::ListItem> {
         let thread = &self.matches[ix];
 
+        let added = self.context_store.upgrade().map_or(false, |ctx_store| {
+            ctx_store.read(cx).included_thread(&thread.id).is_some()
+        });
+
         Some(
             ListItem::new(ix)
                 .inset(true)
                 .toggle_state(selected)
-                .child(Label::new(thread.summary.clone())),
+                .child(Label::new(thread.summary.clone()))
+                .when(added, |el| {
+                    el.end_slot(Label::new("Added").size(LabelSize::XSmall))
+                }),
         )
     }
 }

crates/assistant2/src/context_store.rs

@@ -1,6 +1,11 @@
-use gpui::SharedString;
+use std::fmt::Write as _;
-use project::ProjectEntryId;
+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(
+    pub fn insert_file(&mut self, buffer: &Buffer) {
-        &mut self,
+        let Some(file) = buffer.file() else {
-        kind: ContextKind,
+            return;
-        name: impl Into<SharedString>,
+        };
-        text: impl Into<SharedString>,
+
-    ) {
+        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(),
+            id,
-            name: name.into(),
+            name,
-            kind,
+            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 {
+    pub fn included_file(&self, path: &Path) -> Option<IncludedFile> {
-        self.context.iter().any(|probe| match probe.kind {
+        if let Some(id) = self.files.get(path) {
-            ContextKind::File(probe_entry_id) => probe_entry_id == entry_id,
+            return Some(IncludedFile::Direct(*id));
-            ContextKind::Directory | ContextKind::FetchedUrl | ContextKind::Thread(_) => false,
+        }
-        })
+
+        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 {
+    pub fn included_directory(&self, path: &Path) -> Option<ContextId> {
-        self.context.iter().any(|probe| match probe.kind {
+        self.directories.get(path).copied()
-            ContextKind::Thread(ref probe_thread_id) => probe_thread_id == thread_id,
+    }
-            ContextKind::File(_) | ContextKind::Directory | ContextKind::FetchedUrl => false,
+
-        })
+    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");
+}

crates/assistant2/src/context_strip.rs

@@ -3,14 +3,12 @@ use std::rc::Rc;
 use editor::Editor;
 use gpui::{AppContext, FocusHandle, Model, View, WeakModel, WeakView};
 use language::Buffer;
-use project::ProjectEntryId;
 use ui::{prelude::*, KeyBinding, PopoverMenu, PopoverMenuHandle, Tooltip};
 use workspace::Workspace;
 
-use crate::context::ContextKind;
 use crate::context_picker::{ConfirmBehavior, ContextPicker};
 use crate::context_store::ContextStore;
-use crate::thread::{Thread, ThreadId};
+use crate::thread::Thread;
 use crate::thread_store::ThreadStore;
 use crate::ui::ContextPill;
 use crate::{AssistantPanel, ToggleContextPicker};
@@ -62,20 +60,19 @@ impl ContextStrip {
     fn suggested_file(&self, cx: &ViewContext<Self>) -> Option<SuggestedContext> {
         let workspace = self.workspace.upgrade()?;
         let active_item = workspace.read(cx).active_item(cx)?;
-        let entry_id = *active_item.project_entry_ids(cx).first()?;
-
-        if self.context_store.read(cx).contains_project_entry(entry_id) {
-            return None;
-        }
 
         let editor = active_item.to_any().downcast::<Editor>().ok()?.read(cx);
         let active_buffer = editor.buffer().read(cx).as_singleton()?;
 
-        let file = active_buffer.read(cx).file()?;
+        let path = active_buffer.read(cx).file()?.path();
-        let title = file.path().to_string_lossy().into_owned().into();
+
+        if self.context_store.read(cx).included_file(path).is_some() {
+            return None;
+        }
+
+        let title = path.to_string_lossy().into_owned().into();
 
         Some(SuggestedContext::File {
-            entry_id,
             title,
             buffer: active_buffer.downgrade(),
         })
@@ -95,13 +92,13 @@ impl ContextStrip {
         if self
             .context_store
             .read(cx)
-            .contains_thread(active_thread.id())
+            .included_thread(active_thread.id())
+            .is_some()
         {
             return None;
         }
 
         Some(SuggestedContext::Thread {
-            id: active_thread.id().clone(),
             title: active_thread.summary().unwrap_or("Active Thread".into()),
             thread: weak_active_thread,
         })
@@ -230,12 +227,10 @@ pub enum SuggestContextKind {
 #[derive(Clone)]
 pub enum SuggestedContext {
     File {
-        entry_id: ProjectEntryId,
         title: SharedString,
         buffer: WeakModel<Buffer>,
     },
     Thread {
-        id: ThreadId,
         title: SharedString,
         thread: WeakModel<Thread>,
     },
@@ -251,32 +246,15 @@ impl SuggestedContext {
 
     pub fn accept(&self, context_store: &mut ContextStore, cx: &mut AppContext) {
         match self {
-            Self::File {
+            Self::File { buffer, title: _ } => {
-                entry_id,
+                if let Some(buffer) = buffer.upgrade() {
-                title,
+                    context_store.insert_file(buffer.read(cx));
-                buffer,
-            } => {
-                let Some(buffer) = buffer.upgrade() else {
-                    return;
                 };
-                let text = buffer.read(cx).text();
-
-                context_store.insert_context(
-                    ContextKind::File(*entry_id),
-                    title.clone(),
-                    text.clone(),
-                );
             }
-            Self::Thread { id, title, thread } => {
+            Self::Thread { thread, title: _ } => {
-                let Some(thread) = thread.upgrade() else {
+                if let Some(thread) = thread.upgrade() {
-                    return;
+                    context_store.insert_thread(thread.read(cx));
                 };
-
-                context_store.insert_context(
-                    ContextKind::Thread(id.clone()),
-                    title.clone(),
-                    thread.read(cx).text(),
-                );
             }
         }
     }

crates/assistant2/src/ui/context_pill.rs

@@ -33,10 +33,10 @@ impl RenderOnce for ContextPill {
             px(4.)
         };
         let icon = match self.context.kind {
-            ContextKind::File(_) => IconName::File,
+            ContextKind::File => IconName::File,
             ContextKind::Directory => IconName::Folder,
             ContextKind::FetchedUrl => IconName::Globe,
-            ContextKind::Thread(_) => IconName::MessageCircle,
+            ContextKind::Thread => IconName::MessageCircle,
         };
 
         h_flex()