agent: Fix file context renames affecting display + simplify loading code (#29192)

Michael Sloan created

Release Notes:

- N/A

Change summary

crates/agent/src/context.rs         |  19 ++
crates/agent/src/context_store.rs   | 212 +++++++++++-------------------
crates/agent/src/ui/context_pill.rs |   4 
3 files changed, 98 insertions(+), 137 deletions(-)

Detailed changes

crates/agent/src/context.rs 🔗

@@ -1,7 +1,11 @@
-use std::{ops::Range, path::Path, sync::Arc};
+use std::{
+    ops::Range,
+    path::{Path, PathBuf},
+    sync::Arc,
+};
 
 use gpui::{App, Entity, SharedString};
-use language::{Buffer, File};
+use language::Buffer;
 use language_model::LanguageModelRequestMessage;
 use project::{ProjectEntryId, ProjectPath, Worktree};
 use prompt_store::UserPromptId;
@@ -142,11 +146,20 @@ pub struct ContextBuffer {
     // TODO: Entity<Buffer> holds onto the buffer even if the buffer is deleted. Should probably be
     // a WeakEntity and handle removal from the UI when it has dropped.
     pub buffer: Entity<Buffer>,
-    pub file: Arc<dyn File>,
+    pub last_full_path: Arc<Path>,
     pub version: clock::Global,
     pub text: SharedString,
 }
 
+impl ContextBuffer {
+    pub fn full_path(&self, cx: &App) -> PathBuf {
+        let file = self.buffer.read(cx).file();
+        // Note that in practice file can't be `None` because it is present when this is created and
+        // there's no way for buffers to go from having a file to not.
+        file.map_or(self.last_full_path.to_path_buf(), |file| file.full_path(cx))
+    }
+}
+
 impl std::fmt::Debug for ContextBuffer {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("ContextBuffer")

crates/agent/src/context_store.rs 🔗

@@ -7,7 +7,7 @@ use collections::{BTreeMap, HashMap, HashSet};
 use futures::future::join_all;
 use futures::{self, Future, FutureExt, future};
 use gpui::{App, AppContext as _, Context, Entity, SharedString, Task, WeakEntity};
-use language::{Buffer, File};
+use language::Buffer;
 use project::{Project, ProjectEntryId, ProjectItem, ProjectPath, Worktree};
 use prompt_store::UserPromptId;
 use rope::{Point, Rope};
@@ -112,13 +112,12 @@ impl ContextStore {
                 return anyhow::Ok(());
             }
 
-            let (buffer_info, text_task) =
-                this.update(cx, |_, cx| collect_buffer_info_and_text(buffer, cx))??;
-
-            let text = text_task.await;
+            let context_buffer = this
+                .update(cx, |_, cx| load_context_buffer(buffer, cx))??
+                .await;
 
             this.update(cx, |this, cx| {
-                this.insert_file(make_context_buffer(buffer_info, text), cx);
+                this.insert_file(context_buffer, cx);
             })?;
 
             anyhow::Ok(())
@@ -131,14 +130,11 @@ impl ContextStore {
         cx: &mut Context<Self>,
     ) -> Task<Result<()>> {
         cx.spawn(async move |this, cx| {
-            let (buffer_info, text_task) =
-                this.update(cx, |_, cx| collect_buffer_info_and_text(buffer, cx))??;
-
-            let text = text_task.await;
+            let context_buffer = this
+                .update(cx, |_, cx| load_context_buffer(buffer, cx))??
+                .await;
 
-            this.update(cx, |this, cx| {
-                this.insert_file(make_context_buffer(buffer_info, text), cx)
-            })?;
+            this.update(cx, |this, cx| this.insert_file(context_buffer, cx))?;
 
             anyhow::Ok(())
         })
@@ -211,27 +207,15 @@ impl ContextStore {
 
             let buffers = open_buffers_task.await;
 
-            let mut buffer_infos = Vec::new();
-            let mut text_tasks = Vec::new();
-            this.update(cx, |_, cx| {
-                // Skip all binary files and other non-UTF8 files
-                for buffer in buffers.into_iter().flatten() {
-                    if let Some((buffer_info, text_task)) =
-                        collect_buffer_info_and_text(buffer, cx).log_err()
-                    {
-                        buffer_infos.push(buffer_info);
-                        text_tasks.push(text_task);
-                    }
-                }
-                anyhow::Ok(())
-            })??;
+            let context_buffer_tasks = this.update(cx, |_, cx| {
+                buffers
+                    .into_iter()
+                    .flatten()
+                    .flat_map(move |buffer| load_context_buffer(buffer, cx).log_err())
+                    .collect::<Vec<_>>()
+            })?;
 
-            let buffer_texts = future::join_all(text_tasks).await;
-            let context_buffers = buffer_infos
-                .into_iter()
-                .zip(buffer_texts)
-                .map(|(info, text)| make_context_buffer(info, text))
-                .collect::<Vec<_>>();
+            let context_buffers = future::join_all(context_buffer_tasks).await;
 
             if context_buffers.is_empty() {
                 let full_path = cx.update(|cx| worktree.read(cx).full_path(&project_path.path))?;
@@ -303,27 +287,23 @@ impl ContextStore {
             }
         }
 
-        let (buffer_info, collect_content_task) = match collect_buffer_info_and_text_for_range(
-            buffer,
-            symbol_enclosing_range.clone(),
-            cx,
-        ) {
-            Ok((_, buffer_info, collect_context_task)) => (buffer_info, collect_context_task),
-            Err(err) => return Task::ready(Err(err)),
-        };
+        let context_buffer_task =
+            match load_context_buffer_range(buffer, symbol_enclosing_range.clone(), cx) {
+                Ok((_line_range, context_buffer_task)) => context_buffer_task,
+                Err(err) => return Task::ready(Err(err)),
+            };
 
         cx.spawn(async move |this, cx| {
-            let content = collect_content_task.await;
+            let context_buffer = context_buffer_task.await;
 
             this.update(cx, |this, cx| {
                 this.insert_symbol(
                     make_context_symbol(
-                        buffer_info,
+                        context_buffer,
                         project_path,
                         symbol_name,
                         symbol_range,
                         symbol_enclosing_range,
-                        content,
                     ),
                     cx,
                 )
@@ -475,19 +455,14 @@ impl ContextStore {
         cx: &mut Context<ContextStore>,
     ) -> Task<Result<()>> {
         cx.spawn(async move |this, cx| {
-            let (line_range, buffer_info, text_task) = this.update(cx, |_, cx| {
-                collect_buffer_info_and_text_for_range(buffer, range.clone(), cx)
+            let (line_range, context_buffer_task) = this.update(cx, |_, cx| {
+                load_context_buffer_range(buffer, range.clone(), cx)
             })??;
 
-            let text = text_task.await;
+            let context_buffer = context_buffer_task.await;
 
             this.update(cx, |this, cx| {
-                this.insert_excerpt(
-                    make_context_buffer(buffer_info, text),
-                    range,
-                    line_range,
-                    cx,
-                )
+                this.insert_excerpt(context_buffer, range, line_range, cx)
             })?;
 
             anyhow::Ok(())
@@ -713,92 +688,78 @@ pub enum FileInclusion {
     InDirectory(ProjectPath),
 }
 
-// ContextBuffer without text.
-struct BufferInfo {
-    id: BufferId,
-    buffer: Entity<Buffer>,
-    file: Arc<dyn File>,
-    version: clock::Global,
-}
-
-fn make_context_buffer(info: BufferInfo, text: SharedString) -> ContextBuffer {
-    ContextBuffer {
-        id: info.id,
-        buffer: info.buffer,
-        file: info.file,
-        version: info.version,
-        text,
-    }
-}
-
 fn make_context_symbol(
-    info: BufferInfo,
+    context_buffer: ContextBuffer,
     path: ProjectPath,
     name: SharedString,
     range: Range<Anchor>,
     enclosing_range: Range<Anchor>,
-    text: SharedString,
 ) -> ContextSymbol {
     ContextSymbol {
         id: ContextSymbolId { name, range, path },
-        buffer_version: info.version,
+        buffer_version: context_buffer.version,
         enclosing_range,
-        buffer: info.buffer,
-        text,
+        buffer: context_buffer.buffer,
+        text: context_buffer.text,
     }
 }
 
-fn collect_buffer_info_and_text_for_range(
+fn load_context_buffer_range(
     buffer: Entity<Buffer>,
     range: Range<Anchor>,
     cx: &App,
-) -> Result<(Range<Point>, BufferInfo, Task<SharedString>)> {
-    let content = buffer
-        .read(cx)
-        .text_for_range(range.clone())
-        .collect::<Rope>();
+) -> Result<(Range<Point>, Task<ContextBuffer>)> {
+    let buffer_ref = buffer.read(cx);
+    let id = buffer_ref.remote_id();
 
-    let line_range = range.to_point(&buffer.read(cx).snapshot());
+    let file = buffer_ref.file().context("context buffer missing path")?;
+    let full_path = file.full_path(cx);
 
-    let buffer_info = collect_buffer_info(buffer, cx)?;
-    let full_path = buffer_info.file.full_path(cx);
+    // Important to collect version at the same time as content so that staleness logic is correct.
+    let version = buffer_ref.version();
+    let content = buffer_ref.text_for_range(range.clone()).collect::<Rope>();
+    let line_range = range.to_point(&buffer_ref.snapshot());
 
-    let text_task = cx.background_spawn({
+    // Build the text on a background thread.
+    let task = cx.background_spawn({
         let line_range = line_range.clone();
-        async move { to_fenced_codeblock(&full_path, content, Some(line_range)) }
+        async move {
+            let text = to_fenced_codeblock(&full_path, content, Some(line_range));
+            ContextBuffer {
+                id,
+                buffer,
+                last_full_path: full_path.into(),
+                version,
+                text,
+            }
+        }
     });
 
-    Ok((line_range, buffer_info, text_task))
+    Ok((line_range, task))
 }
 
-fn collect_buffer_info_and_text(
-    buffer: Entity<Buffer>,
-    cx: &App,
-) -> Result<(BufferInfo, Task<SharedString>)> {
-    let content = buffer.read(cx).as_rope().clone();
-
-    let buffer_info = collect_buffer_info(buffer, cx)?;
-    let full_path = buffer_info.file.full_path(cx);
-
-    let text_task =
-        cx.background_spawn(async move { to_fenced_codeblock(&full_path, content, None) });
-
-    Ok((buffer_info, text_task))
-}
-
-fn collect_buffer_info(buffer: Entity<Buffer>, cx: &App) -> Result<BufferInfo> {
+fn load_context_buffer(buffer: Entity<Buffer>, cx: &App) -> Result<Task<ContextBuffer>> {
     let buffer_ref = buffer.read(cx);
-    let file = buffer_ref.file().context("file context must have a path")?;
+    let id = buffer_ref.remote_id();
+
+    let file = buffer_ref.file().context("context buffer missing path")?;
+    let full_path = file.full_path(cx);
 
     // Important to collect version at the same time as content so that staleness logic is correct.
     let version = buffer_ref.version();
+    let content = buffer_ref.as_rope().clone();
 
-    Ok(BufferInfo {
-        buffer,
-        id: buffer_ref.remote_id(),
-        file: file.clone(),
-        version,
-    })
+    // Build the text on a background thread.
+    Ok(cx.background_spawn(async move {
+        let text = to_fenced_codeblock(&full_path, content, None);
+        ContextBuffer {
+            id,
+            buffer,
+            last_full_path: full_path.into(),
+            version,
+            text,
+        }
+    }))
 }
 
 fn to_fenced_codeblock(
@@ -1138,15 +1099,10 @@ fn refresh_user_rules(
     })
 }
 
-fn refresh_context_buffer(
-    context_buffer: &ContextBuffer,
-    cx: &App,
-) -> Option<impl Future<Output = ContextBuffer> + use<>> {
+fn refresh_context_buffer(context_buffer: &ContextBuffer, cx: &App) -> Option<Task<ContextBuffer>> {
     let buffer = context_buffer.buffer.read(cx);
     if buffer.version.changed_since(&context_buffer.version) {
-        let (buffer_info, text_task) =
-            collect_buffer_info_and_text(context_buffer.buffer.clone(), cx).log_err()?;
-        Some(text_task.map(move |text| make_context_buffer(buffer_info, text)))
+        load_context_buffer(context_buffer.buffer.clone(), cx).log_err()
     } else {
         None
     }
@@ -1159,10 +1115,9 @@ fn refresh_context_excerpt(
 ) -> Option<impl Future<Output = (Range<Point>, ContextBuffer)> + use<>> {
     let buffer = context_buffer.buffer.read(cx);
     if buffer.version.changed_since(&context_buffer.version) {
-        let (line_range, buffer_info, text_task) =
-            collect_buffer_info_and_text_for_range(context_buffer.buffer.clone(), range, cx)
-                .log_err()?;
-        Some(text_task.map(move |text| (line_range, make_context_buffer(buffer_info, text))))
+        let (line_range, context_buffer_task) =
+            load_context_buffer_range(context_buffer.buffer.clone(), range, cx).log_err()?;
+        Some(context_buffer_task.map(move |context_buffer| (line_range, context_buffer)))
     } else {
         None
     }
@@ -1175,7 +1130,7 @@ fn refresh_context_symbol(
     let buffer = context_symbol.buffer.read(cx);
     let project_path = buffer.project_path(cx)?;
     if buffer.version.changed_since(&context_symbol.buffer_version) {
-        let (_, buffer_info, text_task) = collect_buffer_info_and_text_for_range(
+        let (_line_range, context_buffer_task) = load_context_buffer_range(
             context_symbol.buffer.clone(),
             context_symbol.enclosing_range.clone(),
             cx,
@@ -1184,15 +1139,8 @@ fn refresh_context_symbol(
         let name = context_symbol.id.name.clone();
         let range = context_symbol.id.range.clone();
         let enclosing_range = context_symbol.enclosing_range.clone();
-        Some(text_task.map(move |text| {
-            make_context_symbol(
-                buffer_info,
-                project_path,
-                name,
-                range,
-                enclosing_range,
-                text,
-            )
+        Some(context_buffer_task.map(move |context_buffer| {
+            make_context_symbol(context_buffer, project_path, name, range, enclosing_range)
         }))
     } else {
         None

crates/agent/src/ui/context_pill.rs 🔗

@@ -241,7 +241,7 @@ impl AddedContext {
     pub fn new(context: &AssistantContext, cx: &App) -> AddedContext {
         match context {
             AssistantContext::File(file_context) => {
-                let full_path = file_context.context_buffer.file.full_path(cx);
+                let full_path = file_context.context_buffer.full_path(cx);
                 let full_path_string: SharedString =
                     full_path.to_string_lossy().into_owned().into();
                 let name = full_path
@@ -304,7 +304,7 @@ impl AddedContext {
             },
 
             AssistantContext::Excerpt(excerpt_context) => {
-                let full_path = excerpt_context.context_buffer.file.full_path(cx);
+                let full_path = excerpt_context.context_buffer.full_path(cx);
                 let mut full_path_string = full_path.to_string_lossy().into_owned();
                 let mut name = full_path
                     .file_name()