Michael Sloan , Marshall Bowers , and Marshall created
Release Notes:
- N/A
---------
Co-authored-by: Marshall Bowers <elliott.codes@gmail.com>
Co-authored-by: Marshall <marshall@zed.dev>Cargo.lock | 2
crates/assistant2/Cargo.toml | 2
crates/assistant2/src/active_thread.rs | 12
crates/assistant2/src/buffer_codegen.rs | 5
crates/assistant2/src/context.rs | 145 +
crates/assistant2/src/context_picker/directory_context_picker.rs | 2
crates/assistant2/src/context_picker/fetch_context_picker.rs | 14
crates/assistant2/src/context_picker/file_context_picker.rs | 15
crates/assistant2/src/context_picker/thread_context_picker.rs | 2
crates/assistant2/src/context_store.rs | 266 +
crates/assistant2/src/context_strip.rs | 41
crates/assistant2/src/message_editor.rs | 7
crates/assistant2/src/terminal_inline_assistant.rs | 8
crates/assistant2/src/thread.rs | 12
crates/assistant2/src/ui/context_pill.rs | 12
15 files changed, 361 insertions(+), 184 deletions(-)
@@ -462,6 +462,7 @@ dependencies = [
"async-watch",
"chrono",
"client",
+ "clock",
"collections",
"command_palette_hooks",
"context_server",
@@ -476,6 +477,7 @@ dependencies = [
"html_to_markdown",
"http_client",
"indoc",
+ "itertools 0.13.0",
"language",
"language_model",
"language_model_selector",
@@ -19,6 +19,7 @@ assets.workspace = true
assistant_tool.workspace = true
async-watch.workspace = true
client.workspace = true
+clock.workspace = true
chrono.workspace = true
collections.workspace = true
command_palette_hooks.workspace = true
@@ -33,6 +34,7 @@ gpui.workspace = true
handlebars.workspace = true
html_to_markdown.workspace = true
http_client.workspace = true
+itertools.workspace = true
language.workspace = true
language_model.workspace = true
language_model_selector.workspace = true
@@ -282,11 +282,13 @@ impl ActiveThread {
.child(div().p_2p5().text_ui(cx).child(markdown.clone()))
.when_some(context, |parent, context| {
if !context.is_empty() {
- parent.child(h_flex().flex_wrap().gap_1().px_1p5().pb_1p5().children(
- context.iter().map(|context| {
- ContextPill::new_added(context.clone(), false, None)
- }),
- ))
+ parent.child(
+ h_flex().flex_wrap().gap_1().px_1p5().pb_1p5().children(
+ context.into_iter().map(|context| {
+ ContextPill::new_added(context, false, None)
+ }),
+ ),
+ )
} else {
parent
}
@@ -421,8 +421,7 @@ impl CodegenAlternative {
};
if let Some(context_store) = &self.context_store {
- let context = context_store.update(cx, |this, _cx| this.context().clone());
- attach_context_to_message(&mut request_message, context);
+ attach_context_to_message(&mut request_message, context_store.read(cx).snapshot(cx));
}
request_message.content.push(prompt.into());
@@ -1053,7 +1052,7 @@ mod tests {
stream::{self},
Stream,
};
- use gpui::{Context, TestAppContext};
+ use gpui::TestAppContext;
use indoc::indoc;
use language::{
language_settings, tree_sitter_rust, Buffer, Language, LanguageConfig, LanguageMatcher,
@@ -1,8 +1,17 @@
-use gpui::SharedString;
+use std::path::Path;
+use std::rc::Rc;
+use std::sync::Arc;
+
+use collections::BTreeMap;
+use gpui::{AppContext, Model, SharedString};
+use language::Buffer;
use language_model::{LanguageModelRequestMessage, MessageContent};
use serde::{Deserialize, Serialize};
+use text::BufferId;
use util::post_inc;
+use crate::thread::Thread;
+
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, Serialize, Deserialize)]
pub struct ContextId(pub(crate) usize);
@@ -14,16 +23,17 @@ impl ContextId {
/// Some context attached to a message in a thread.
#[derive(Debug, Clone)]
-pub struct Context {
+pub struct ContextSnapshot {
pub id: ContextId,
pub name: SharedString,
pub parent: Option<SharedString>,
pub tooltip: Option<SharedString>,
pub kind: ContextKind,
+ /// Text to send to the model. This is not refreshed by `snapshot`.
pub text: SharedString,
}
-#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ContextKind {
File,
Directory,
@@ -31,18 +41,139 @@ pub enum ContextKind {
Thread,
}
+#[derive(Debug)]
+pub enum Context {
+ File(FileContext),
+ Directory(DirectoryContext),
+ FetchedUrl(FetchedUrlContext),
+ Thread(ThreadContext),
+}
+
+impl Context {
+ pub fn id(&self) -> ContextId {
+ match self {
+ Self::File(file) => file.id,
+ Self::Directory(directory) => directory.snapshot.id,
+ Self::FetchedUrl(url) => url.id,
+ Self::Thread(thread) => thread.id,
+ }
+ }
+}
+
+// TODO: Model<Buffer> holds onto the buffer even if the file is deleted and closed. Should remove
+// the context from the message editor in this case.
+
+#[derive(Debug)]
+pub struct FileContext {
+ pub id: ContextId,
+ pub buffer: Model<Buffer>,
+ #[allow(unused)]
+ pub version: clock::Global,
+ pub text: SharedString,
+}
+
+#[derive(Debug)]
+pub struct DirectoryContext {
+ #[allow(unused)]
+ pub path: Rc<Path>,
+ // TODO: The choice to make this a BTreeMap was a result of use in a version of
+ // ContextStore::will_include_buffer before I realized that the path logic should be used there
+ // too.
+ #[allow(unused)]
+ pub buffers: BTreeMap<BufferId, (Model<Buffer>, clock::Global)>,
+ pub snapshot: ContextSnapshot,
+}
+
+#[derive(Debug)]
+pub struct FetchedUrlContext {
+ pub id: ContextId,
+ pub url: SharedString,
+ pub text: SharedString,
+}
+
+// TODO: Model<Thread> holds onto the thread even if the thread is deleted. Can either handle this
+// explicitly or have a WeakModel<Thread> and remove during snapshot.
+
+#[derive(Debug)]
+pub struct ThreadContext {
+ pub id: ContextId,
+ pub thread: Model<Thread>,
+ pub text: SharedString,
+}
+
+impl Context {
+ pub fn snapshot(&self, cx: &AppContext) -> Option<ContextSnapshot> {
+ match &self {
+ Self::File(file_context) => {
+ let path = file_context.path(cx)?;
+ let full_path: SharedString = path.to_string_lossy().into_owned().into();
+ let name = match path.file_name() {
+ Some(name) => name.to_string_lossy().into_owned().into(),
+ None => full_path.clone(),
+ };
+ let parent = path
+ .parent()
+ .and_then(|p| p.file_name())
+ .map(|p| p.to_string_lossy().into_owned().into());
+
+ Some(ContextSnapshot {
+ id: self.id(),
+ name,
+ parent,
+ tooltip: Some(full_path),
+ kind: ContextKind::File,
+ text: file_context.text.clone(),
+ })
+ }
+ Self::Directory(DirectoryContext { snapshot, .. }) => Some(snapshot.clone()),
+ Self::FetchedUrl(FetchedUrlContext { url, text, id }) => Some(ContextSnapshot {
+ id: *id,
+ name: url.clone(),
+ parent: None,
+ tooltip: None,
+ kind: ContextKind::FetchedUrl,
+ text: text.clone(),
+ }),
+ Self::Thread(thread_context) => {
+ let thread = thread_context.thread.read(cx);
+
+ Some(ContextSnapshot {
+ id: self.id(),
+ name: thread.summary().unwrap_or("New thread".into()),
+ parent: None,
+ tooltip: None,
+ kind: ContextKind::Thread,
+ text: thread_context.text.clone(),
+ })
+ }
+ }
+ }
+}
+
+impl FileContext {
+ pub fn path(&self, cx: &AppContext) -> Option<Arc<Path>> {
+ let buffer = self.buffer.read(cx);
+ if let Some(file) = buffer.file() {
+ Some(file.path().clone())
+ } else {
+ log::error!("Buffer that had a path unexpectedly no longer has a path.");
+ None
+ }
+ }
+}
+
pub fn attach_context_to_message(
message: &mut LanguageModelRequestMessage,
- context: impl IntoIterator<Item = Context>,
+ contexts: impl Iterator<Item = ContextSnapshot>,
) {
let mut file_context = String::new();
let mut directory_context = String::new();
let mut fetch_context = String::new();
let mut thread_context = String::new();
- for context in context.into_iter() {
+ for context in contexts {
match context.kind {
- ContextKind::File { .. } => {
+ ContextKind::File => {
file_context.push_str(&context.text);
file_context.push('\n');
}
@@ -56,7 +187,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);
@@ -240,7 +240,7 @@ impl PickerDelegate for DirectoryContextPickerDelegate {
let added = self.context_store.upgrade().map_or(false, |context_store| {
context_store
.read(cx)
- .included_directory(&path_match.path)
+ .includes_directory(&path_match.path)
.is_some()
});
@@ -82,10 +82,12 @@ impl FetchContextPickerDelegate {
}
async fn build_message(http_client: Arc<HttpClientWithUrl>, url: &str) -> Result<String> {
- let mut url = url.to_owned();
- if !url.starts_with("https://") && !url.starts_with("http://") {
- url = format!("https://{url}");
- }
+ let prefixed_url = if !url.starts_with("https://") && !url.starts_with("http://") {
+ Some(format!("https://{url}"))
+ } else {
+ None
+ };
+ let url = prefixed_url.as_deref().unwrap_or(url);
let mut response = http_client.get(&url, AsyncBody::default(), true).await?;
@@ -200,7 +202,7 @@ impl PickerDelegate for FetchContextPickerDelegate {
this.delegate
.context_store
.update(cx, |context_store, _cx| {
- if context_store.included_url(&url).is_none() {
+ if context_store.includes_url(&url).is_none() {
context_store.insert_fetched_url(url, text);
}
})?;
@@ -234,7 +236,7 @@ impl PickerDelegate for FetchContextPickerDelegate {
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()
+ context_store.read(cx).includes_url(&self.url).is_some()
});
Some(
@@ -11,7 +11,7 @@ use util::ResultExt as _;
use workspace::Workspace;
use crate::context_picker::{ConfirmBehavior, ContextPicker};
-use crate::context_store::{ContextStore, IncludedFile};
+use crate::context_store::{ContextStore, FileInclusion};
pub struct FileContextPicker {
picker: View<Picker<FileContextPickerDelegate>>,
@@ -275,10 +275,11 @@ 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));
+ let added = self.context_store.upgrade().and_then(|context_store| {
+ context_store
+ .read(cx)
+ .will_include_file_path(&path_match.path, cx)
+ });
Some(
ListItem::new(ix)
@@ -295,7 +296,7 @@ impl PickerDelegate for FileContextPickerDelegate {
})),
)
.when_some(added, |el, added| match added {
- IncludedFile::Direct(_) => el.end_slot(
+ FileInclusion::Direct(_) => el.end_slot(
h_flex()
.gap_1()
.child(
@@ -305,7 +306,7 @@ impl PickerDelegate for FileContextPickerDelegate {
)
.child(Label::new("Added").size(LabelSize::Small)),
),
- IncludedFile::InDirectory(dir_name) => {
+ FileInclusion::InDirectory(dir_name) => {
let dir_name = dir_name.to_string_lossy().into_owned();
el.end_slot(
@@ -194,7 +194,7 @@ impl PickerDelegate for ThreadContextPickerDelegate {
let thread = &self.matches[ix];
let added = self.context_store.upgrade().map_or(false, |context_store| {
- context_store.read(cx).included_thread(&thread.id).is_some()
+ context_store.read(cx).includes_thread(&thread.id).is_some()
});
Some(
@@ -3,23 +3,25 @@ use std::path::{Path, PathBuf};
use std::sync::Arc;
use anyhow::{anyhow, bail, Result};
-use collections::{HashMap, HashSet};
-use gpui::{Model, ModelContext, SharedString, Task, WeakView};
+use collections::{BTreeMap, HashMap};
+use gpui::{AppContext, Model, ModelContext, SharedString, Task, WeakView};
use language::Buffer;
use project::{ProjectPath, Worktree};
+use text::BufferId;
use workspace::Workspace;
-use crate::thread::Thread;
-use crate::{
- context::{Context, ContextId, ContextKind},
- thread::ThreadId,
+use crate::context::{
+ Context, ContextId, ContextKind, ContextSnapshot, DirectoryContext, FetchedUrlContext,
+ FileContext, ThreadContext,
};
+use crate::thread::{Thread, ThreadId};
pub struct ContextStore {
workspace: WeakView<Workspace>,
context: Vec<Context>,
+ // TODO: If an EntityId is used for all context types (like BufferId), can remove ContextId.
next_context_id: ContextId,
- files: HashMap<PathBuf, ContextId>,
+ files: BTreeMap<BufferId, ContextId>,
directories: HashMap<PathBuf, ContextId>,
threads: HashMap<ThreadId, ContextId>,
fetched_urls: HashMap<String, ContextId>,
@@ -31,13 +33,22 @@ impl ContextStore {
workspace,
context: Vec::new(),
next_context_id: ContextId(0),
- files: HashMap::default(),
+ files: BTreeMap::default(),
directories: HashMap::default(),
threads: HashMap::default(),
fetched_urls: HashMap::default(),
}
}
+ pub fn snapshot<'a>(
+ &'a self,
+ cx: &'a AppContext,
+ ) -> impl Iterator<Item = ContextSnapshot> + 'a {
+ self.context()
+ .iter()
+ .flat_map(|context| context.snapshot(cx))
+ }
+
pub fn context(&self) -> &Vec<Context> {
&self.context
}
@@ -63,64 +74,54 @@ impl ContextStore {
return Task::ready(Err(anyhow!("failed to read project")));
};
- let already_included = match self.included_file(&project_path.path) {
- Some(IncludedFile::Direct(context_id)) => {
- self.remove_context(&context_id);
- true
- }
- Some(IncludedFile::InDirectory(_)) => true,
- None => false,
- };
- if already_included {
- return Task::ready(Ok(()));
- }
-
cx.spawn(|this, mut cx| async move {
- let open_buffer_task =
- project.update(&mut cx, |project, cx| project.open_buffer(project_path, cx))?;
+ let open_buffer_task = project.update(&mut cx, |project, cx| {
+ project.open_buffer(project_path.clone(), cx)
+ })?;
let buffer = open_buffer_task.await?;
+ let buffer_id = buffer.update(&mut cx, |buffer, _cx| buffer.remote_id())?;
+
+ let already_included = this.update(&mut cx, |this, _cx| {
+ match this.will_include_buffer(buffer_id, &project_path.path) {
+ Some(FileInclusion::Direct(context_id)) => {
+ this.remove_context(context_id);
+ true
+ }
+ Some(FileInclusion::InDirectory(_)) => true,
+ None => false,
+ }
+ })?;
+
+ if already_included {
+ return anyhow::Ok(());
+ }
+
this.update(&mut cx, |this, cx| {
- this.insert_file(buffer.read(cx));
+ this.insert_file(buffer, cx);
})?;
anyhow::Ok(())
})
}
- pub fn insert_file(&mut self, buffer: &Buffer) {
+ pub fn insert_file(&mut self, buffer_model: Model<Buffer>, cx: &AppContext) {
+ let buffer = buffer_model.read(cx);
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 full_path: SharedString = path.to_string_lossy().into_owned().into();
-
- let name = match path.file_name() {
- Some(name) => name.to_string_lossy().into_owned().into(),
- None => full_path.clone(),
- };
-
- let parent = path
- .parent()
- .and_then(|p| p.file_name())
- .map(|p| p.to_string_lossy().into_owned().into());
-
let mut text = String::new();
- push_fenced_codeblock(path, buffer.text(), &mut text);
+ push_fenced_codeblock(file.path(), buffer.text(), &mut text);
- self.context.push(Context {
+ let id = self.next_context_id.post_inc();
+ self.files.insert(buffer.remote_id(), id);
+ self.context.push(Context::File(FileContext {
id,
- name,
- parent,
- tooltip: Some(full_path),
- kind: ContextKind::File,
+ buffer: buffer_model,
+ version: buffer.version.clone(),
text: text.into(),
- });
+ }));
}
pub fn add_directory(
@@ -136,9 +137,9 @@ impl ContextStore {
return Task::ready(Err(anyhow!("failed to read project")));
};
- let already_included = if let Some(context_id) = self.included_directory(&project_path.path)
+ let already_included = if let Some(context_id) = self.includes_directory(&project_path.path)
{
- self.remove_context(&context_id);
+ self.remove_context(context_id);
true
} else {
false
@@ -178,23 +179,24 @@ impl ContextStore {
this.update(&mut cx, |this, cx| {
let mut text = String::new();
- let mut added_files = 0;
-
- for buffer in buffers.into_iter().flatten() {
- let buffer = buffer.read(cx);
+ let mut directory_buffers = BTreeMap::new();
+ for buffer_model in buffers {
+ let buffer_model = buffer_model?;
+ let buffer = buffer_model.read(cx);
let path = buffer.file().map_or(&project_path.path, |file| file.path());
push_fenced_codeblock(&path, buffer.text(), &mut text);
- added_files += 1;
+ directory_buffers
+ .insert(buffer.remote_id(), (buffer_model, buffer.version.clone()));
}
- if added_files == 0 {
+ if directory_buffers.is_empty() {
bail!(
"could not read any text files from {}",
&project_path.path.display()
);
}
- this.insert_directory(&project_path.path, text);
+ this.insert_directory(&project_path.path, directory_buffers, text);
anyhow::Ok(())
})??;
@@ -203,7 +205,12 @@ impl ContextStore {
})
}
- pub fn insert_directory(&mut self, path: &Path, text: impl Into<SharedString>) {
+ pub fn insert_directory(
+ &mut self,
+ path: &Path,
+ buffers: BTreeMap<BufferId, (Model<Buffer>, clock::Global)>,
+ text: impl Into<SharedString>,
+ ) {
let id = self.next_context_id.post_inc();
self.directories.insert(path.to_path_buf(), id);
@@ -219,78 +226,104 @@ impl ContextStore {
.and_then(|p| p.file_name())
.map(|p| p.to_string_lossy().into_owned().into());
- self.context.push(Context {
- id,
- name,
- parent,
- tooltip: Some(full_path),
- kind: ContextKind::Directory,
- text: text.into(),
- });
+ self.context.push(Context::Directory(DirectoryContext {
+ path: path.into(),
+ buffers,
+ snapshot: ContextSnapshot {
+ id,
+ name,
+ parent,
+ tooltip: Some(full_path),
+ kind: ContextKind::Directory,
+ text: text.into(),
+ },
+ }));
}
pub fn add_thread(&mut self, thread: Model<Thread>, cx: &mut ModelContext<Self>) {
- if let Some(context_id) = self.included_thread(&thread.read(cx).id()) {
- self.remove_context(&context_id);
+ if let Some(context_id) = self.includes_thread(&thread.read(cx).id()) {
+ self.remove_context(context_id);
} else {
- self.insert_thread(thread.read(cx));
+ self.insert_thread(thread, cx);
}
}
- 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()),
- parent: None,
- tooltip: None,
- kind: ContextKind::Thread,
- text: thread.text().into(),
- });
+ pub fn insert_thread(&mut self, thread: Model<Thread>, cx: &AppContext) {
+ let id = self.next_context_id.post_inc();
+ let thread_ref = thread.read(cx);
+ let text = thread_ref.text().into();
+
+ self.threads.insert(thread_ref.id().clone(), id);
+ self.context
+ .push(Context::Thread(ThreadContext { id, thread, text }));
}
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(),
- parent: None,
- tooltip: None,
- kind: ContextKind::FetchedUrl,
+ let id = self.next_context_id.post_inc();
+
+ self.fetched_urls.insert(url.clone(), id);
+ self.context.push(Context::FetchedUrl(FetchedUrlContext {
+ id,
+ url: url.into(),
text: text.into(),
- });
+ }));
}
- pub fn remove_context(&mut self, id: &ContextId) {
- let Some(ix) = self.context.iter().position(|context| context.id == *id) else {
+ pub fn remove_context(&mut self, id: ContextId) {
+ let Some(ix) = self.context.iter().position(|context| context.id() == id) else {
return;
};
- match self.context.remove(ix).kind {
- ContextKind::File => {
- self.files.retain(|_, context_id| context_id != id);
+ match self.context.remove(ix) {
+ Context::File(_) => {
+ self.files.retain(|_, context_id| *context_id != id);
}
- ContextKind::Directory => {
- self.directories.retain(|_, context_id| context_id != id);
+ Context::Directory(_) => {
+ self.directories.retain(|_, context_id| *context_id != id);
}
- ContextKind::FetchedUrl => {
- self.fetched_urls.retain(|_, context_id| context_id != id);
+ Context::FetchedUrl(_) => {
+ self.fetched_urls.retain(|_, context_id| *context_id != id);
}
- ContextKind::Thread => {
- self.threads.retain(|_, context_id| context_id != id);
+ Context::Thread(_) => {
+ self.threads.retain(|_, context_id| *context_id != id);
}
}
}
- pub fn included_file(&self, path: &Path) -> Option<IncludedFile> {
- if let Some(id) = self.files.get(path) {
- return Some(IncludedFile::Direct(*id));
+ /// Returns whether the buffer is already included directly in the context, or if it will be
+ /// included in the context via a directory. Directory inclusion is based on paths rather than
+ /// buffer IDs as the directory will be re-scanned.
+ pub fn will_include_buffer(&self, buffer_id: BufferId, path: &Path) -> Option<FileInclusion> {
+ if let Some(context_id) = self.files.get(&buffer_id) {
+ return Some(FileInclusion::Direct(*context_id));
+ }
+
+ self.will_include_file_path_via_directory(path)
+ }
+
+ /// Returns whether this file path is already included directly in the context, or if it will be
+ /// included in the context via a directory.
+ pub fn will_include_file_path(&self, path: &Path, cx: &AppContext) -> Option<FileInclusion> {
+ if !self.files.is_empty() {
+ let found_file_context = self.context.iter().find(|context| match &context {
+ Context::File(file_context) => {
+ if let Some(file_path) = file_context.path(cx) {
+ *file_path == *path
+ } else {
+ false
+ }
+ }
+ _ => false,
+ });
+ if let Some(context) = found_file_context {
+ return Some(FileInclusion::Direct(context.id()));
+ }
}
+ self.will_include_file_path_via_directory(path)
+ }
+
+ fn will_include_file_path_via_directory(&self, path: &Path) -> Option<FileInclusion> {
if self.directories.is_empty() {
return None;
}
@@ -299,40 +332,27 @@ impl ContextStore {
while buf.pop() {
if let Some(_) = self.directories.get(&buf) {
- return Some(IncludedFile::InDirectory(buf));
+ return Some(FileInclusion::InDirectory(buf));
}
}
None
}
- pub fn included_directory(&self, path: &Path) -> Option<ContextId> {
+ pub fn includes_directory(&self, path: &Path) -> Option<ContextId> {
self.directories.get(path).copied()
}
- pub fn included_thread(&self, thread_id: &ThreadId) -> Option<ContextId> {
+ pub fn includes_thread(&self, thread_id: &ThreadId) -> Option<ContextId> {
self.threads.get(thread_id).copied()
}
- pub fn included_url(&self, url: &str) -> Option<ContextId> {
+ pub fn includes_url(&self, url: &str) -> Option<ContextId> {
self.fetched_urls.get(url).copied()
}
-
- pub fn duplicated_names(&self) -> HashSet<SharedString> {
- let mut seen = HashSet::default();
- let mut dupes = HashSet::default();
-
- for context in self.context().iter() {
- if !seen.insert(&context.name) {
- dupes.insert(context.name.clone());
- }
- }
-
- dupes
- }
}
-pub enum IncludedFile {
+pub enum FileInclusion {
Direct(ContextId),
InDirectory(PathBuf),
}
@@ -1,10 +1,12 @@
use std::rc::Rc;
+use collections::HashSet;
use editor::Editor;
use gpui::{
AppContext, DismissEvent, EventEmitter, FocusHandle, Model, Subscription, View, WeakModel,
WeakView,
};
+use itertools::Itertools;
use language::Buffer;
use ui::{prelude::*, KeyBinding, PopoverMenu, PopoverMenuHandle, Tooltip};
use workspace::Workspace;
@@ -73,11 +75,17 @@ impl ContextStrip {
let active_item = workspace.read(cx).active_item(cx)?;
let editor = active_item.to_any().downcast::<Editor>().ok()?.read(cx);
- let active_buffer = editor.buffer().read(cx).as_singleton()?;
+ let active_buffer_model = editor.buffer().read(cx).as_singleton()?;
+ let active_buffer = active_buffer_model.read(cx);
- let path = active_buffer.read(cx).file()?.path();
+ let path = active_buffer.file()?.path();
- if self.context_store.read(cx).included_file(path).is_some() {
+ if self
+ .context_store
+ .read(cx)
+ .will_include_buffer(active_buffer.remote_id(), path)
+ .is_some()
+ {
return None;
}
@@ -88,7 +96,7 @@ impl ContextStrip {
Some(SuggestedContext::File {
name,
- buffer: active_buffer.downgrade(),
+ buffer: active_buffer_model.downgrade(),
})
}
@@ -106,7 +114,7 @@ impl ContextStrip {
if self
.context_store
.read(cx)
- .included_thread(active_thread.id())
+ .includes_thread(active_thread.id())
.is_some()
{
return None;
@@ -131,13 +139,24 @@ impl ContextStrip {
impl Render for ContextStrip {
fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
let context_store = self.context_store.read(cx);
- let context = context_store.context().clone();
+ let context = context_store
+ .context()
+ .iter()
+ .flat_map(|context| context.snapshot(cx))
+ .collect::<Vec<_>>();
let context_picker = self.context_picker.clone();
let focus_handle = self.focus_handle.clone();
let suggested_context = self.suggested_context(cx);
- let dupe_names = context_store.duplicated_names();
+ let dupe_names = context
+ .iter()
+ .map(|context| context.name.clone())
+ .sorted()
+ .tuple_windows()
+ .filter(|(a, b)| a == b)
+ .map(|(a, _)| a)
+ .collect::<HashSet<SharedString>>();
h_flex()
.flex_wrap()
@@ -194,11 +213,11 @@ impl Render for ContextStrip {
context.clone(),
dupe_names.contains(&context.name),
Some({
- let context = context.clone();
+ let id = context.id;
let context_store = self.context_store.clone();
Rc::new(cx.listener(move |_this, _event, cx| {
context_store.update(cx, |this, _cx| {
- this.remove_context(&context.id);
+ this.remove_context(id);
});
cx.notify();
}))
@@ -284,12 +303,12 @@ impl SuggestedContext {
match self {
Self::File { buffer, name: _ } => {
if let Some(buffer) = buffer.upgrade() {
- context_store.insert_file(buffer.read(cx));
+ context_store.insert_file(buffer, cx);
};
}
Self::Thread { thread, name: _ } => {
if let Some(thread) = thread.upgrade() {
- context_store.insert_thread(thread.read(cx));
+ context_store.insert_thread(thread, cx);
};
}
}
@@ -147,11 +147,10 @@ impl MessageEditor {
editor.clear(cx);
text
});
- let context = self
- .context_store
- .update(cx, |this, _cx| this.context().clone());
- self.thread.update(cx, |thread, cx| {
+ let thread = self.thread.clone();
+ thread.update(cx, |thread, cx| {
+ let context = self.context_store.read(cx).snapshot(cx).collect::<Vec<_>>();
thread.insert_user_message(user_message, context, cx);
let mut request = thread.to_completion_request(request_kind, cx);
@@ -245,10 +245,10 @@ impl TerminalInlineAssistant {
cache: false,
};
- let context = assist
- .context_store
- .update(cx, |this, _cx| this.context().clone());
- attach_context_to_message(&mut request_message, context);
+ attach_context_to_message(
+ &mut request_message,
+ assist.context_store.read(cx).snapshot(cx),
+ );
request_message.content.push(prompt.into());
@@ -3,7 +3,7 @@ use std::sync::Arc;
use anyhow::Result;
use assistant_tool::ToolWorkingSet;
use chrono::{DateTime, Utc};
-use collections::{HashMap, HashSet};
+use collections::{BTreeMap, HashMap, HashSet};
use futures::future::Shared;
use futures::{FutureExt as _, StreamExt as _};
use gpui::{AppContext, EventEmitter, ModelContext, SharedString, Task};
@@ -17,7 +17,7 @@ use serde::{Deserialize, Serialize};
use util::{post_inc, TryFutureExt as _};
use uuid::Uuid;
-use crate::context::{attach_context_to_message, Context, ContextId};
+use crate::context::{attach_context_to_message, ContextId, ContextSnapshot};
#[derive(Debug, Clone, Copy)]
pub enum RequestKind {
@@ -64,7 +64,7 @@ pub struct Thread {
pending_summary: Task<Option<()>>,
messages: Vec<Message>,
next_message_id: MessageId,
- context: HashMap<ContextId, Context>,
+ context: BTreeMap<ContextId, ContextSnapshot>,
context_by_message: HashMap<MessageId, Vec<ContextId>>,
completion_count: usize,
pending_completions: Vec<PendingCompletion>,
@@ -83,7 +83,7 @@ impl Thread {
pending_summary: Task::ready(None),
messages: Vec::new(),
next_message_id: MessageId(0),
- context: HashMap::default(),
+ context: BTreeMap::default(),
context_by_message: HashMap::default(),
completion_count: 0,
pending_completions: Vec::new(),
@@ -131,7 +131,7 @@ impl Thread {
&self.tools
}
- pub fn context_for_message(&self, id: MessageId) -> Option<Vec<Context>> {
+ pub fn context_for_message(&self, id: MessageId) -> Option<Vec<ContextSnapshot>> {
let context = self.context_by_message.get(&id)?;
Some(
context
@@ -149,7 +149,7 @@ impl Thread {
pub fn insert_user_message(
&mut self,
text: impl Into<String>,
- context: Vec<Context>,
+ context: Vec<ContextSnapshot>,
cx: &mut ModelContext<Self>,
) {
let message_id = self.insert_message(Role::User, text, cx);
@@ -3,12 +3,12 @@ use std::rc::Rc;
use gpui::ClickEvent;
use ui::{prelude::*, IconButtonShape, Tooltip};
-use crate::context::{Context, ContextKind};
+use crate::context::{ContextKind, ContextSnapshot};
#[derive(IntoElement)]
pub enum ContextPill {
Added {
- context: Context,
+ context: ContextSnapshot,
dupe_name: bool,
on_remove: Option<Rc<dyn Fn(&ClickEvent, &mut WindowContext)>>,
},
@@ -21,7 +21,7 @@ pub enum ContextPill {
impl ContextPill {
pub fn new_added(
- context: Context,
+ context: ContextSnapshot,
dupe_name: bool,
on_remove: Option<Rc<dyn Fn(&ClickEvent, &mut WindowContext)>>,
) -> Self {
@@ -49,10 +49,10 @@ impl ContextPill {
}
}
- pub fn kind(&self) -> &ContextKind {
+ pub fn kind(&self) -> ContextKind {
match self {
- Self::Added { context, .. } => &context.kind,
- Self::Suggested { kind, .. } => kind,
+ Self::Added { context, .. } => context.kind,
+ Self::Suggested { kind, .. } => *kind,
}
}
}