Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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

Browse source 
Release Notes:

- N/A
This commit is contained in:
Michael Sloan 2025-04-21 21:16:46 -06:00 • committed by GitHub
parent dfbd132d9f
commit 19ef56ba7c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 106 additions and 145 deletions
Download patch file Download diff file Expand all files Collapse all files

19
crates/agent/src/context.rs
View file

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

228
crates/agent/src/context_store.rs
View file

@ -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 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::<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>();
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<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> {
) -> Result<(Range<Point>, 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.text_for_range(range.clone()).collect::<Rope>();
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<Buffer>, cx: &App) -> Result<Task<ContextBuffer>> {
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<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

4
crates/agent/src/ui/context_pill.rs
View file

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