From 19ef56ba7caa20e262e47fecdbe18c25a983ca23 Mon Sep 17 00:00:00 2001 From: Michael Sloan Date: Mon, 21 Apr 2025 21:16:46 -0600 Subject: [PATCH] agent: Fix file context renames affecting display + simplify loading code (#29192) Release Notes: - N/A --- crates/agent/src/context.rs | 19 ++- crates/agent/src/context_store.rs | 228 +++++++++++----------------- crates/agent/src/ui/context_pill.rs | 4 +- 3 files changed, 106 insertions(+), 145 deletions(-) diff --git a/crates/agent/src/context.rs b/crates/agent/src/context.rs index d74cfc2ad4..3caef5b7fe 100644 --- a/crates/agent/src/context.rs +++ b/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 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, - pub file: Arc, + pub last_full_path: Arc, 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") diff --git a/crates/agent/src/context_store.rs b/crates/agent/src/context_store.rs index 864c3b3f69..a6255a7901 100644 --- a/crates/agent/src/context_store.rs +++ b/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, ) -> Task> { cx.spawn(async move |this, cx| { - let (buffer_info, text_task) = - this.update(cx, |_, cx| collect_buffer_info_and_text(buffer, cx))??; + let context_buffer = this + .update(cx, |_, cx| load_context_buffer(buffer, cx))?? + .await; - let text = text_task.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::>() + })?; - 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::>(); + 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, ) -> Task> { 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, - file: Arc, - 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, enclosing_range: Range, - 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, range: Range, cx: &App, -) -> Result<(Range, BufferInfo, Task)> { - let content = buffer - .read(cx) - .text_for_range(range.clone()) - .collect::(); - - let line_range = range.to_point(&buffer.read(cx).snapshot()); - - let buffer_info = collect_buffer_info(buffer, cx)?; - let full_path = buffer_info.file.full_path(cx); - - let text_task = cx.background_spawn({ - let line_range = line_range.clone(); - async move { to_fenced_codeblock(&full_path, content, Some(line_range)) } - }); - - Ok((line_range, buffer_info, text_task)) -} - -fn collect_buffer_info_and_text( - buffer: Entity, - cx: &App, -) -> Result<(BufferInfo, Task)> { - 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, cx: &App) -> Result { +) -> Result<(Range, Task)> { 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.text_for_range(range.clone()).collect::(); + let line_range = range.to_point(&buffer_ref.snapshot()); - Ok(BufferInfo { - buffer, - id: buffer_ref.remote_id(), - file: file.clone(), - version, - }) + // Build the text on a background thread. + let task = cx.background_spawn({ + let line_range = line_range.clone(); + 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, task)) +} + +fn load_context_buffer(buffer: Entity, cx: &App) -> Result> { + let buffer_ref = buffer.read(cx); + 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(); + + // 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 + use<>> { +fn refresh_context_buffer(context_buffer: &ContextBuffer, cx: &App) -> Option> { 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, 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 diff --git a/crates/agent/src/ui/context_pill.rs b/crates/agent/src/ui/context_pill.rs index 4f6b3430e8..6f739a1d41 100644 --- a/crates/agent/src/ui/context_pill.rs +++ b/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()