assistant2: Background load of context + prep for refresh + efficiency (

#22935)

* Now loads context on background threads.

- For file and directory context, buffer ropes can be shared between
threads as they are immutable. This allows for traversal and
accumulation of buffer text on a background thread.

- For url context, the request, parsing, and rendering is now done on a
background thread.

* Prepares for support of buffer reload by individually storing the text
of directory buffers.

* Avoids some string copying / redundant strings.

- When attaching message context, no longer builds a string for each
context type.

- For directory context, does not build a `SharedString` for the full
text, instead has a slice of `SharedString` chunks which are then
directly appended to the message context.

- Building a fenced codeblock for a buffer now computes a precise
capacity in advance.

Release Notes:

- N/A

crates/assistant2/src/context.rs

@@ -2,7 +2,6 @@ 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};
@@ -29,8 +28,8 @@ pub struct ContextSnapshot {
     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,
+    /// Concatenating these strings yields text to send to the model. Not refreshed by `snapshot`.
+    pub text: Box<[SharedString]>,
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -60,27 +59,18 @@ impl Context {
     }
 }
 
-// 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,
+    pub buffer: ContextBuffer,
 }
 
 #[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 buffers: Vec<ContextBuffer>,
     pub snapshot: ContextSnapshot,
 }
 
@@ -101,6 +91,19 @@ pub struct ThreadContext {
     pub text: SharedString,
 }
 
+// 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 ContextBuffer {
+    #[allow(unused)]
+    pub id: BufferId,
+    pub buffer: Model<Buffer>,
+    #[allow(unused)]
+    pub version: clock::Global,
+    pub text: SharedString,
+}
+
 impl Context {
     pub fn snapshot(&self, cx: &AppContext) -> Option<ContextSnapshot> {
         match &self {
@@ -114,7 +117,7 @@ impl Context {
 
 impl FileContext {
     pub fn path(&self, cx: &AppContext) -> Option<Arc<Path>> {
-        let buffer = self.buffer.read(cx);
+        let buffer = self.buffer.buffer.read(cx);
         if let Some(file) = buffer.file() {
             Some(file.path().clone())
         } else {
@@ -141,7 +144,7 @@ impl FileContext {
             parent,
             tooltip: Some(full_path),
             kind: ContextKind::File,
-            text: self.text.clone(),
+            text: Box::new([self.buffer.text.clone()]),
         })
     }
 }
@@ -160,7 +163,7 @@ impl FetchedUrlContext {
             parent: None,
             tooltip: None,
             kind: ContextKind::FetchedUrl,
-            text: self.text.clone(),
+            text: Box::new([self.text.clone()]),
         }
     }
 }
@@ -174,7 +177,7 @@ impl ThreadContext {
             parent: None,
             tooltip: None,
             kind: ContextKind::Thread,
-            text: self.text.clone(),
+            text: Box::new([self.text.clone()]),
         }
     }
 }
@@ -183,55 +186,64 @@ pub fn attach_context_to_message(
     message: &mut LanguageModelRequestMessage,
     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();
+    let mut file_context = Vec::new();
+    let mut directory_context = Vec::new();
+    let mut fetch_context = Vec::new();
+    let mut thread_context = Vec::new();
 
     for context in contexts {
         match context.kind {
-            ContextKind::File => {
-                file_context.push_str(&context.text);
-                file_context.push('\n');
-            }
-            ContextKind::Directory => {
-                directory_context.push_str(&context.text);
-                directory_context.push('\n');
-            }
-            ContextKind::FetchedUrl => {
-                fetch_context.push_str(&context.name);
-                fetch_context.push('\n');
-                fetch_context.push_str(&context.text);
-                fetch_context.push('\n');
-            }
-            ContextKind::Thread { .. } => {
-                thread_context.push_str(&context.name);
-                thread_context.push('\n');
-                thread_context.push_str(&context.text);
-                thread_context.push('\n');
-            }
+            ContextKind::File => file_context.push(context),
+            ContextKind::Directory => directory_context.push(context),
+            ContextKind::FetchedUrl => fetch_context.push(context),
+            ContextKind::Thread => thread_context.push(context),
         }
     }
 
     let mut context_text = String::new();
+
     if !file_context.is_empty() {
         context_text.push_str("The following files are available:\n");
-        context_text.push_str(&file_context);
+        for context in file_context {
+            for chunk in context.text {
+                context_text.push_str(&chunk);
+            }
+            context_text.push('\n');
+        }
     }
 
     if !directory_context.is_empty() {
         context_text.push_str("The following directories are available:\n");
-        context_text.push_str(&directory_context);
+        for context in directory_context {
+            for chunk in context.text {
+                context_text.push_str(&chunk);
+            }
+            context_text.push('\n');
+        }
     }
 
     if !fetch_context.is_empty() {
         context_text.push_str("The following fetched results are available\n");
-        context_text.push_str(&fetch_context);
+        for context in fetch_context {
+            context_text.push_str(&context.name);
+            context_text.push('\n');
+            for chunk in context.text {
+                context_text.push_str(&chunk);
+            }
+            context_text.push('\n');
+        }
     }
 
     if !thread_context.is_empty() {
         context_text.push_str("The following previous conversation threads are available\n");
-        context_text.push_str(&thread_context);
+        for context in thread_context {
+            context_text.push_str(&context.name);
+            context_text.push('\n');
+            for chunk in context.text {
+                context_text.push_str(&chunk);
+            }
+            context_text.push('\n');
+        }
     }
 
     if !context_text.is_empty() {

crates/assistant2/src/context_picker/fetch_context_picker.rs

@@ -81,13 +81,12 @@ impl FetchContextPickerDelegate {
         }
     }
 
-    async fn build_message(http_client: Arc<HttpClientWithUrl>, url: &str) -> Result<String> {
-        let prefixed_url = if !url.starts_with("https://") && !url.starts_with("http://") {
-            Some(format!("https://{url}"))
+    async fn build_message(http_client: Arc<HttpClientWithUrl>, url: String) -> Result<String> {
+        let url = if !url.starts_with("https://") && !url.starts_with("http://") {
+            format!("https://{url}")
         } else {
-            None
+            url
         };
-        let url = prefixed_url.as_deref().unwrap_or(url);
 
         let mut response = http_client.get(&url, AsyncBody::default(), true).await?;
 
@@ -196,7 +195,10 @@ impl PickerDelegate for FetchContextPickerDelegate {
         let url = self.url.clone();
         let confirm_behavior = self.confirm_behavior;
         cx.spawn(|this, mut cx| async move {
-            let text = Self::build_message(http_client, &url).await?;
+            let text = cx
+                .background_executor()
+                .spawn(Self::build_message(http_client, url.clone()))
+                .await?;
 
             this.update(&mut cx, |this, cx| {
                 this.delegate

crates/assistant2/src/context_picker/file_context_picker.rs

@@ -201,7 +201,7 @@ impl PickerDelegate for FileContextPickerDelegate {
         let Some(task) = self
             .context_store
             .update(cx, |context_store, cx| {
-                context_store.add_file(project_path, cx)
+                context_store.add_file_from_path(project_path, cx)
             })
             .ok()
         else {