assistant2: Add live context type and use in message editor (#22865)

Release Notes:

- N/A

---------

Co-authored-by: Marshall Bowers <elliott.codes@gmail.com>
Co-authored-by: Marshall <marshall@zed.dev>

crates/assistant2/Cargo.toml

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

crates/assistant2/src/active_thread.rs

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

crates/assistant2/src/buffer_codegen.rs

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

crates/assistant2/src/context.rs

@@ -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,35 +23,157 @@ 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,
     FetchedUrl,
     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);

crates/assistant2/src/context_picker/directory_context_picker.rs

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

crates/assistant2/src/context_picker/fetch_context_picker.rs

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

crates/assistant2/src/context_picker/file_context_picker.rs

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

crates/assistant2/src/context_picker/thread_context_picker.rs

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