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
This commit is contained in:
Agus Zubiaga 2025-01-06 15:55:20 -03:00 committed by GitHub
parent c74e5f5de2
commit 3a061a91e7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 293 additions and 173 deletions

View file

@ -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);

View file

@ -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(

View file

@ -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));
text.push_str(&buffer.read(cx).text());
if !text.ends_with('\n') {
text.push('\n');
}
text.push_str("```\n");
let buffer = buffer.read(cx);
let path = buffer.file().map_or(&path, |file| file.path());
push_fenced_codeblock(&path, buffer.text(), &mut text);
}
this.delegate
.context_store
.update(cx, |context_store, _cx| {
context_store.insert_context(
ContextKind::Directory,
path.to_string_lossy().to_string(),
text,
);
context_store.insert_directory(&path, 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))
}),
)
}
}

View file

@ -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))
}),
)
}
}

View file

@ -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();
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,
);
context_store.insert_file(buffer.read(cx));
})?;
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,8 +302,16 @@ 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)
.inset(true)
.toggle_state(selected)
.child(
h_flex()
.gap_2()
.child(Label::new(file_name))
@ -311,32 +320,18 @@ impl PickerDelegate for FileContextPickerDelegate {
.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
}

View file

@ -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(
ContextKind::Thread(thread.read(cx).id().clone()),
entry.summary.clone(),
thread.read(cx).text(),
);
if let Some(context_id) = context_store.included_thread(&entry.id) {
context_store.remove_context(&context_id);
} else {
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))
}),
)
}
}

View file

@ -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));
}
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,
})
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 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");
}

View file

@ -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 title = file.path().to_string_lossy().into_owned().into();
let path = active_buffer.read(cx).file()?.path();
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 {
entry_id,
title,
buffer,
} => {
let Some(buffer) = buffer.upgrade() else {
return;
Self::File { buffer, title: _ } => {
if let Some(buffer) = buffer.upgrade() {
context_store.insert_file(buffer.read(cx));
};
let text = buffer.read(cx).text();
context_store.insert_context(
ContextKind::File(*entry_id),
title.clone(),
text.clone(),
);
}
Self::Thread { id, title, thread } => {
let Some(thread) = thread.upgrade() else {
return;
Self::Thread { thread, title: _ } => {
if let Some(thread) = thread.upgrade() {
context_store.insert_thread(thread.read(cx));
};
context_store.insert_context(
ContextKind::Thread(id.clone()),
title.clone(),
thread.read(cx).text(),
);
}
}
}

View file

@ -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()