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

agent: Support adding selection as context (#28964)

Browse source 
https://github.com/user-attachments/assets/42ebe911-3392-48f7-8583-caab285aca09

Release Notes:

- agent: Support adding selections via @selection or `assistant: Quote
selection` as context
This commit is contained in:
Bennet Bo Fenner 2025-04-17 16:55:15 +02:00 • committed by GitHub
parent f07695c4cd
commit 002235d0da
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 495 additions and 196 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -9,14 +9,14 @@ use futures::{self, Future, FutureExt, future};
use gpui::{App, AppContext as _, Context, Entity, SharedString, Task, WeakEntity};
use language::{Buffer, File};
use project::{Project, ProjectItem, ProjectPath, Worktree};
use rope::Rope;
use rope::{Point, Rope};
use text::{Anchor, BufferId, OffsetRangeExt};
use util::{ResultExt as _, maybe};
use crate::ThreadStore;
use crate::context::{
AssistantContext, ContextBuffer, ContextId, ContextSymbol, ContextSymbolId, DirectoryContext,
FetchedUrlContext, FileContext, SymbolContext, ThreadContext,
ExcerptContext, FetchedUrlContext, FileContext, SymbolContext, ThreadContext,
};
use crate::context_strip::SuggestedContext;
use crate::thread::{Thread, ThreadId};
@ -110,7 +110,7 @@ impl ContextStore {
}
let (buffer_info, text_task) =
this.update(cx, |_, cx| collect_buffer_info_and_text(buffer, None, cx))??;
this.update(cx, |_, cx| collect_buffer_info_and_text(buffer, cx))??;
let text = text_task.await;
@ -129,7 +129,7 @@ impl ContextStore {
) -> Task<Result<()>> {
cx.spawn(async move |this, cx| {
let (buffer_info, text_task) =
this.update(cx, |_, cx| collect_buffer_info_and_text(buffer, None, cx))??;
this.update(cx, |_, cx| collect_buffer_info_and_text(buffer, cx))??;
let text = text_task.await;
@ -206,7 +206,7 @@ impl ContextStore {
// 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, None, cx).log_err()
collect_buffer_info_and_text(buffer, cx).log_err()
{
buffer_infos.push(buffer_info);
text_tasks.push(text_task);
@ -290,11 +290,14 @@ impl ContextStore {
}
}
let (buffer_info, collect_content_task) =
match collect_buffer_info_and_text(buffer, Some(symbol_enclosing_range.clone()), cx) {
Ok((buffer_info, collect_context_task)) => (buffer_info, collect_context_task),
Err(err) => return Task::ready(Err(err)),
};
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)),
};
cx.spawn(async move |this, cx| {
let content = collect_content_task.await;
@ -416,6 +419,49 @@ impl ContextStore {
cx.notify();
}
pub fn add_excerpt(
&mut self,
range: Range<Anchor>,
buffer: Entity<Buffer>,
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 text = text_task.await;
this.update(cx, |this, cx| {
this.insert_excerpt(
make_context_buffer(buffer_info, text),
range,
line_range,
cx,
)
})?;
anyhow::Ok(())
})
}
fn insert_excerpt(
&mut self,
context_buffer: ContextBuffer,
range: Range<Anchor>,
line_range: Range<Point>,
cx: &mut Context<Self>,
) {
let id = self.next_context_id.post_inc();
self.context.push(AssistantContext::Excerpt(ExcerptContext {
id,
range,
line_range,
context_buffer,
}));
cx.notify();
}
pub fn accept_suggested_context(
&mut self,
suggested: &SuggestedContext,
@ -465,6 +511,7 @@ impl ContextStore {
self.symbol_buffers.remove(&symbol.context_symbol.id);
self.symbols.retain(|_, context_id| *context_id != id);
}
AssistantContext::Excerpt(_) => {}
AssistantContext::FetchedUrl(_) => {
self.fetched_urls.retain(|_, context_id| *context_id != id);
}
@ -592,6 +639,7 @@ impl ContextStore {
}
AssistantContext::Directory(_)
| AssistantContext::Symbol(_)
| AssistantContext::Excerpt(_)
| AssistantContext::FetchedUrl(_)
| AssistantContext::Thread(_) => None,
})
@ -643,41 +691,78 @@ fn make_context_symbol(
}
}
fn collect_buffer_info_and_text_for_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>,
range: Option<Range<Anchor>>,
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> {
let buffer_ref = buffer.read(cx);
let file = buffer_ref.file().context("file context must have a path")?;
// Important to collect version at the same time as content so that staleness logic is correct.
let version = buffer_ref.version();
let content = if let Some(range) = range {
buffer_ref.text_for_range(range).collect::<Rope>()
} else {
buffer_ref.as_rope().clone()
};
let buffer_info = BufferInfo {
Ok(BufferInfo {
buffer,
id: buffer_ref.remote_id(),
file: file.clone(),
version,
};
let full_path = file.full_path(cx);
let text_task = cx.background_spawn(async move { to_fenced_codeblock(&full_path, content) });
Ok((buffer_info, text_task))
})
}
fn to_fenced_codeblock(path: &Path, content: Rope) -> SharedString {
fn to_fenced_codeblock(
path: &Path,
content: Rope,
line_range: Option<Range<Point>>,
) -> SharedString {
let line_range_text = line_range.map(|range| {
if range.start.row == range.end.row {
format!(":{}", range.start.row + 1)
} else {
format!(":{}-{}", range.start.row + 1, range.end.row + 1)
}
});
let path_extension = path.extension().and_then(|ext| ext.to_str());
let path_string = path.to_string_lossy();
let capacity = 3
+ path_extension.map_or(0, |extension| extension.len() + 1)
+ path_string.len()
+ line_range_text.as_ref().map_or(0, |text| text.len())
+ 1
+ content.len()
+ 5;
@ -691,6 +776,10 @@ fn to_fenced_codeblock(path: &Path, content: Rope) -> SharedString {
}
buffer.push_str(&path_string);
if let Some(line_range_text) = line_range_text {
buffer.push_str(&line_range_text);
}
buffer.push('\n');
for chunk in content.chunks() {
buffer.push_str(&chunk);
@ -769,6 +858,14 @@ pub fn refresh_context_store_text(
return refresh_symbol_text(context_store, symbol_context, cx);
}
}
AssistantContext::Excerpt(excerpt_context) => {
if changed_buffers.is_empty()
|| changed_buffers.contains(&excerpt_context.context_buffer.buffer)
{
let context_store = context_store.clone();
return refresh_excerpt_text(context_store, excerpt_context, cx);
}
}
AssistantContext::Thread(thread_context) => {
if changed_buffers.is_empty() {
let context_store = context_store.clone();
@ -880,6 +977,34 @@ fn refresh_symbol_text(
}
}
fn refresh_excerpt_text(
context_store: Entity<ContextStore>,
excerpt_context: &ExcerptContext,
cx: &App,
) -> Option<Task<()>> {
let id = excerpt_context.id;
let range = excerpt_context.range.clone();
let task = refresh_context_excerpt(&excerpt_context.context_buffer, range.clone(), cx);
if let Some(task) = task {
Some(cx.spawn(async move |cx| {
let (line_range, context_buffer) = task.await;
context_store
.update(cx, |context_store, _| {
let new_excerpt_context = ExcerptContext {
id,
range,
line_range,
context_buffer,
};
context_store.replace_context(AssistantContext::Excerpt(new_excerpt_context));
})
.ok();
}))
} else {
None
}
}
fn refresh_thread_text(
context_store: Entity<ContextStore>,
thread_context: &ThreadContext,
@ -908,13 +1033,29 @@ fn refresh_context_buffer(
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(), None, cx).log_err()?;
collect_buffer_info_and_text(context_buffer.buffer.clone(), cx).log_err()?;
Some(text_task.map(move |text| make_context_buffer(buffer_info, text)))
} else {
None
}
}
fn refresh_context_excerpt(
context_buffer: &ContextBuffer,
range: Range<Anchor>,
cx: &App,
) -> 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))))
} else {
None
}
}
fn refresh_context_symbol(
context_symbol: &ContextSymbol,
cx: &App,
@ -922,9 +1063,9 @@ 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(
let (_, buffer_info, text_task) = collect_buffer_info_and_text_for_range(
context_symbol.buffer.clone(),
Some(context_symbol.enclosing_range.clone()),
context_symbol.enclosing_range.clone(),
cx,
)
.log_err()?;