Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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

Browse source 
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
Download patch file Download diff file Expand all files Collapse all files

50
crates/assistant2/src/context_picker/directory_context_picker.rs
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))
}),
)
}
}

18
crates/assistant2/src/context_picker/fetch_context_picker.rs
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))
}),
)
}
}

111
crates/assistant2/src/context_picker/file_context_picker.rs
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,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(
h_flex()
.gap_2()
.child(Label::new(file_name))
.children(directory.map(|directory| {
Label::new(directory)
.size(LabelSize::Small)
.color(Color::Muted)
})),
),
ListItem::new(ix)
.inset(true)
.toggle_state(selected)
.child(
h_flex()
.gap_2()
.child(Label::new(file_name))
.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
}

22
crates/assistant2/src/context_picker/thread_context_picker.rs
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))
}),
)
}
}