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

Support @-mentions in inline assists and when editing old agent panel messages (#29734)

Browse source 
Closes #ISSUE

Co-authored-by: Bennet <bennet@zed.dev>

Release Notes:

- Added support for context `@mentions` in the inline prompt editor and
when editing past messages in the agent panel.

---------

Co-authored-by: Bennet Bo Fenner <bennet@zed.dev>
Co-authored-by: Ben Brandt <benjamin.j.brandt@gmail.com>
This commit is contained in:
Cole Miller 2025-05-02 16:08:53 -04:00 • committed by GitHub
parent c918f6cde1
commit 9547d42b15
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
18 changed files with 499 additions and 156 deletions
Download patch file Download diff file Expand all files Collapse all files

14
crates/agent/src/active_thread.rs
View file

@ -2,9 +2,10 @@ use crate::context::{AgentContextHandle, RULES_ICON};
use crate::context_picker::{ContextPicker, MentionLink};
use crate::context_store::ContextStore;
use crate::context_strip::{ContextStrip, ContextStripEvent, SuggestContextKind};
use crate::message_editor::insert_message_creases;
use crate::thread::{
LastRestoreCheckpoint, MessageId, MessageSegment, Thread, ThreadError, ThreadEvent,
ThreadFeedback,
LastRestoreCheckpoint, MessageCrease, MessageId, MessageSegment, Thread, ThreadError,
ThreadEvent, ThreadFeedback,
};
use crate::thread_store::{RulesLoadingError, ThreadStore};
use crate::tool_use::{PendingToolUseStatus, ToolUse};
@ -1240,6 +1241,7 @@ impl ActiveThread {
&mut self,
message_id: MessageId,
message_segments: &[MessageSegment],
message_creases: &[MessageCrease],
window: &mut Window,
cx: &mut Context<Self>,
) {
@ -1258,6 +1260,7 @@ impl ActiveThread {
);
editor.update(cx, |editor, cx| {
editor.set_text(message_text.clone(), window, cx);
insert_message_creases(editor, message_creases, &self.context_store, window, cx);
editor.focus_handle(cx).focus(window);
editor.move_to_end(&editor::actions::MoveToEnd, window, cx);
});
@ -1705,6 +1708,7 @@ impl ActiveThread {
let Some(message) = self.thread.read(cx).message(message_id) else {
return Empty.into_any();
};
let message_creases = message.creases.clone();
let Some(rendered_message) = self.rendered_messages_by_id.get(&message_id) else {
return Empty.into_any();
@ -1900,6 +1904,7 @@ impl ActiveThread {
open_context(&context, workspace, window, cx);
cx.notify();
}
cx.stop_propagation();
}
})),
)
@ -1985,15 +1990,13 @@ impl ActiveThread {
)
}),
)
.when(editing_message_state.is_none(), |this| {
this.tooltip(Tooltip::text("Click To Edit"))
})
.on_click(cx.listener({
let message_segments = message.segments.clone();
move |this, _, window, cx| {
this.start_editing_message(
message_id,
&message_segments,
&message_creases,
window,
cx,
);
@ -2361,6 +2364,7 @@ impl ActiveThread {
let workspace = self.workspace.clone();
move |text, window, cx| {
open_markdown_link(text, workspace.clone(), window, cx);
cx.stop_propagation();
}
}))
.into_any_element()

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

@ -4,10 +4,12 @@ use std::path::PathBuf;
use std::{ops::Range, path::Path, sync::Arc};
use assistant_tool::outline;
use collections::HashSet;
use collections::{HashMap, HashSet};
use editor::display_map::CreaseId;
use editor::{Addon, Editor};
use futures::future;
use futures::{FutureExt, future::Shared};
use gpui::{App, AppContext as _, Entity, SharedString, Task};
use gpui::{App, AppContext as _, Entity, SharedString, Subscription, Task};
use language::{Buffer, ParseStatus};
use language_model::{LanguageModelImage, LanguageModelRequestMessage, MessageContent};
use project::{Project, ProjectEntryId, ProjectPath, Worktree};
@ -15,10 +17,11 @@ use prompt_store::{PromptStore, UserPromptId};
use ref_cast::RefCast;
use rope::Point;
use text::{Anchor, OffsetRangeExt as _};
use ui::{ElementId, IconName};
use ui::{Context, ElementId, IconName};
use util::markdown::MarkdownCodeBlock;
use util::{ResultExt as _, post_inc};
use crate::context_store::{ContextStore, ContextStoreEvent};
use crate::thread::Thread;
pub const RULES_ICON: IconName = IconName::Context;
@ -67,7 +70,7 @@ pub enum AgentContextHandle {
}
impl AgentContextHandle {
fn id(&self) -> ContextId {
pub fn id(&self) -> ContextId {
match self {
Self::File(context) => context.context_id,
Self::Directory(context) => context.context_id,
@ -1036,6 +1039,69 @@ impl Hash for AgentContextKey {
}
}
#[derive(Default)]
pub struct ContextCreasesAddon {
creases: HashMap<AgentContextKey, Vec<(CreaseId, SharedString)>>,
_subscription: Option<Subscription>,
}
impl Addon for ContextCreasesAddon {
fn to_any(&self) -> &dyn std::any::Any {
self
}
fn to_any_mut(&mut self) -> Option<&mut dyn std::any::Any> {
Some(self)
}
}
impl ContextCreasesAddon {
pub fn new() -> Self {
Self {
creases: HashMap::default(),
_subscription: None,
}
}
pub fn add_creases(
&mut self,
context_store: &Entity<ContextStore>,
key: AgentContextKey,
creases: impl IntoIterator<Item = (CreaseId, SharedString)>,
cx: &mut Context<Editor>,
) {
self.creases.entry(key).or_default().extend(creases);
self._subscription = Some(cx.subscribe(
&context_store,
|editor, _, event, cx| match event {
ContextStoreEvent::ContextRemoved(key) => {
let Some(this) = editor.addon_mut::<Self>() else {
return;
};
let (crease_ids, replacement_texts): (Vec<_>, Vec<_>) = this
.creases
.remove(key)
.unwrap_or_default()
.into_iter()
.unzip();
let ranges = editor
.remove_creases(crease_ids, cx)
.into_iter()
.map(|(_, range)| range)
.collect::<Vec<_>>();
editor.unfold_ranges(&ranges, false, false, cx);
editor.edit(ranges.into_iter().zip(replacement_texts), cx);
cx.notify();
}
},
))
}
pub fn into_inner(self) -> HashMap<AgentContextKey, Vec<(CreaseId, SharedString)>> {
self.creases
}
}
#[cfg(test)]
mod tests {
use super::*;

27
crates/agent/src/context_picker.rs
View file

@ -11,7 +11,7 @@ use std::sync::Arc;
use anyhow::{Result, anyhow};
pub use completion_provider::ContextPickerCompletionProvider;
use editor::display_map::{Crease, FoldId};
use editor::display_map::{Crease, CreaseId, CreaseMetadata, FoldId};
use editor::{Anchor, AnchorRangeExt as _, Editor, ExcerptId, FoldPlaceholder, ToOffset};
use fetch_context_picker::FetchContextPicker;
use file_context_picker::FileContextPicker;
@ -675,21 +675,20 @@ fn selection_ranges(
})
}
pub(crate) fn insert_fold_for_mention(
pub(crate) fn insert_crease_for_mention(
excerpt_id: ExcerptId,
crease_start: text::Anchor,
content_len: usize,
crease_label: SharedString,
crease_icon_path: SharedString,
editor_entity: Entity<Editor>,
window: &mut Window,
cx: &mut App,
) {
) -> Option<CreaseId> {
editor_entity.update(cx, |editor, cx| {
let snapshot = editor.buffer().read(cx).snapshot(cx);
let Some(start) = snapshot.anchor_in_excerpt(excerpt_id, crease_start) else {
return;
};
let start = snapshot.anchor_in_excerpt(excerpt_id, crease_start)?;
let start = start.bias_right(&snapshot);
let end = snapshot.anchor_before(start.to_offset(&snapshot) + content_len);
@ -701,10 +700,10 @@ pub(crate) fn insert_fold_for_mention(
editor_entity.downgrade(),
);
editor.display_map.update(cx, |display_map, cx| {
display_map.fold(vec![crease], cx);
});
});
let ids = editor.insert_creases(vec![crease.clone()], cx);
editor.fold_creases(vec![crease], false, window, cx);
Some(ids[0])
})
}
pub fn crease_for_mention(
@ -714,20 +713,20 @@ pub fn crease_for_mention(
editor_entity: WeakEntity<Editor>,
) -> Crease<Anchor> {
let placeholder = FoldPlaceholder {
render: render_fold_icon_button(icon_path, label, editor_entity),
render: render_fold_icon_button(icon_path.clone(), label.clone(), editor_entity),
merge_adjacent: false,
..Default::default()
};
let render_trailer = move |_row, _unfold, _window: &mut Window, _cx: &mut App| Empty.into_any();
let crease = Crease::inline(
Crease::inline(
range,
placeholder.clone(),
fold_toggle("mention"),
render_trailer,
);
crease
)
.with_metadata(CreaseMetadata { icon_path, label })
}
fn render_fold_icon_button(

127
crates/agent/src/context_picker/completion_provider.rs
View file

@ -19,9 +19,11 @@ use prompt_store::PromptStore;
use rope::Point;
use text::{Anchor, OffsetRangeExt, ToPoint};
use ui::prelude::*;
use util::ResultExt as _;
use workspace::Workspace;
use crate::context::RULES_ICON;
use crate::Thread;
use crate::context::{AgentContextHandle, AgentContextKey, ContextCreasesAddon, RULES_ICON};
use crate::context_store::ContextStore;
use crate::thread_store::ThreadStore;
@ -310,7 +312,7 @@ impl ContextPickerCompletionProvider {
let context_store = context_store.clone();
let selections = selections.clone();
let selection_infos = selection_infos.clone();
move |_, _: &mut Window, cx: &mut App| {
move |_, window: &mut Window, cx: &mut App| {
context_store.update(cx, |context_store, cx| {
for (buffer, range) in &selections {
context_store.add_selection(
@ -323,7 +325,7 @@ impl ContextPickerCompletionProvider {
let editor = editor.clone();
let selection_infos = selection_infos.clone();
cx.defer(move |cx| {
window.defer(cx, move |window, cx| {
let mut current_offset = 0;
for (file_name, link, line_range) in selection_infos.iter() {
let snapshot =
@ -354,9 +356,8 @@ impl ContextPickerCompletionProvider {
);
editor.update(cx, |editor, cx| {
editor.display_map.update(cx, |display_map, cx| {
display_map.fold(vec![crease], cx);
});
editor.insert_creases(vec![crease.clone()], cx);
editor.fold_creases(vec![crease], false, window, cx);
});
current_offset += text_len + 1;
@ -419,21 +420,26 @@ impl ContextPickerCompletionProvider {
source_range.start,
new_text_len,
editor.clone(),
context_store.clone(),
move |cx| {
let thread_id = thread_entry.id.clone();
let context_store = context_store.clone();
let thread_store = thread_store.clone();
cx.spawn(async move |cx| {
let thread = thread_store
cx.spawn::<_, Option<_>>(async move |cx| {
let thread: Entity<Thread> = thread_store
.update(cx, |thread_store, cx| {
thread_store.open_thread(&thread_id, cx)
})?
.await?;
context_store.update(cx, |context_store, cx| {
context_store.add_thread(thread, false, cx)
})
})
.ok()?
.await
.log_err()?;
let context = context_store
.update(cx, |context_store, cx| {
context_store.add_thread(thread, false, cx)
})
.ok()??;
Some(context)
})
.detach_and_log_err(cx);
},
)),
}
@ -463,11 +469,13 @@ impl ContextPickerCompletionProvider {
source_range.start,
new_text_len,
editor.clone(),
context_store.clone(),
move |cx| {
let user_prompt_id = rules.prompt_id;
context_store.update(cx, |context_store, cx| {
context_store.add_rules(user_prompt_id, false, cx);
let context = context_store.update(cx, |context_store, cx| {
context_store.add_rules(user_prompt_id, false, cx)
});
Task::ready(context)
},
)),
}
@ -498,27 +506,33 @@ impl ContextPickerCompletionProvider {
source_range.start,
new_text_len,
editor.clone(),
context_store.clone(),
move |cx| {
let context_store = context_store.clone();
let http_client = http_client.clone();
let url_to_fetch = url_to_fetch.clone();
cx.spawn(async move |cx| {
if context_store.update(cx, |context_store, _| {
context_store.includes_url(&url_to_fetch)
})? {
return Ok(());
if let Some(context) = context_store
.update(cx, |context_store, _| {
context_store.get_url_context(url_to_fetch.clone())
})
.ok()?
{
return Some(context);
}
let content = cx
.background_spawn(fetch_url_content(
http_client,
url_to_fetch.to_string(),
))
.await?;
context_store.update(cx, |context_store, cx| {
context_store.add_fetched_url(url_to_fetch.to_string(), content, cx)
})
.await
.log_err()?;
context_store
.update(cx, |context_store, cx| {
context_store.add_fetched_url(url_to_fetch.to_string(), content, cx)
})
.ok()
})
.detach_and_log_err(cx);
},
)),
}
@ -577,15 +591,23 @@ impl ContextPickerCompletionProvider {
source_range.start,
new_text_len,
editor,
context_store.clone(),
move |cx| {
context_store.update(cx, |context_store, cx| {
let task = if is_directory {
Task::ready(context_store.add_directory(&project_path, false, cx))
} else {
if is_directory {
Task::ready(
context_store
.update(cx, |context_store, cx| {
context_store.add_directory(&project_path, false, cx)
})
.log_err()
.flatten(),
)
} else {
let result = context_store.update(cx, |context_store, cx| {
context_store.add_file_from_path(project_path.clone(), false, cx)
};
task.detach_and_log_err(cx);
})
});
cx.spawn(async move |_| result.await.log_err().flatten())
}
},
)),
}
@ -640,18 +662,19 @@ impl ContextPickerCompletionProvider {
source_range.start,
new_text_len,
editor.clone(),
context_store.clone(),
move |cx| {
let symbol = symbol.clone();
let context_store = context_store.clone();
let workspace = workspace.clone();
super::symbol_context_picker::add_symbol(
let result = super::symbol_context_picker::add_symbol(
symbol.clone(),
false,
workspace.clone(),
context_store.downgrade(),
cx,
)
.detach_and_log_err(cx);
);
cx.spawn(async move |_| result.await.log_err()?.0)
},
)),
})
@ -873,24 +896,44 @@ fn confirm_completion_callback(
start: Anchor,
content_len: usize,
editor: Entity<Editor>,
add_context_fn: impl Fn(&mut App) -> () + Send + Sync + 'static,
context_store: Entity<ContextStore>,
add_context_fn: impl Fn(&mut App) -> Task<Option<AgentContextHandle>> + Send + Sync + 'static,
) -> Arc<dyn Fn(CompletionIntent, &mut Window, &mut App) -> bool + Send + Sync> {
Arc::new(move |_, _, cx| {
add_context_fn(cx);
Arc::new(move |_, window, cx| {
let context = add_context_fn(cx);
let crease_text = crease_text.clone();
let crease_icon_path = crease_icon_path.clone();
let editor = editor.clone();
cx.defer(move |cx| {
crate::context_picker::insert_fold_for_mention(
let context_store = context_store.clone();
window.defer(cx, move |window, cx| {
let crease_id = crate::context_picker::insert_crease_for_mention(
excerpt_id,
start,
content_len,
crease_text,
crease_text.clone(),
crease_icon_path,
editor,
editor.clone(),
window,
cx,
);
cx.spawn(async move |cx| {
let crease_id = crease_id?;
let context = context.await?;
editor
.update(cx, |editor, cx| {
if let Some(addon) = editor.addon_mut::<ContextCreasesAddon>() {
addon.add_creases(
&context_store,
AgentContextKey(context),
[(crease_id, crease_text)],
cx,
);
}
})
.ok()
})
.detach();
});
false
})

18
crates/agent/src/context_picker/file_context_picker.rs
View file

@ -130,21 +130,19 @@ impl PickerDelegate for FileContextPickerDelegate {
let is_directory = mat.is_dir;
let Some(task) = self
.context_store
self.context_store
.update(cx, |context_store, cx| {
if is_directory {
Task::ready(context_store.add_directory(&project_path, true, cx))
context_store
.add_directory(&project_path, true, cx)
.log_err();
} else {
context_store.add_file_from_path(project_path.clone(), true, cx)
context_store
.add_file_from_path(project_path.clone(), true, cx)
.detach_and_log_err(cx);
}
})
.ok()
else {
return;
};
task.detach_and_log_err(cx);
.ok();
}
fn dismissed(&mut self, _: &mut Window, cx: &mut Context<Picker<Self>>) {

5
crates/agent/src/context_picker/symbol_context_picker.rs
View file

@ -14,6 +14,7 @@ use ui::{ListItem, prelude::*};
use util::ResultExt as _;
use workspace::Workspace;
use crate::context::AgentContextHandle;
use crate::context_picker::ContextPicker;
use crate::context_store::ContextStore;
@ -143,7 +144,7 @@ impl PickerDelegate for SymbolContextPickerDelegate {
let selected_index = self.selected_index;
cx.spawn(async move |this, cx| {
let included = add_symbol_task.await?;
let (_, included) = add_symbol_task.await?;
this.update(cx, |this, _| {
if let Some(mat) = this.delegate.matches.get_mut(selected_index) {
mat.is_included = included;
@ -187,7 +188,7 @@ pub(crate) fn add_symbol(
workspace: Entity<Workspace>,
context_store: WeakEntity<ContextStore>,
cx: &mut App,
) -> Task<Result<bool>> {
) -> Task<Result<(Option<AgentContextHandle>, bool)>> {
let project = workspace.read(cx).project().clone();
let open_buffer_task = project.update(cx, |project, cx| {
project.open_buffer(symbol.path.clone(), cx)

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

@ -5,7 +5,7 @@ use std::sync::Arc;
use anyhow::{Result, anyhow};
use collections::{HashSet, IndexSet};
use futures::{self, FutureExt};
use gpui::{App, Context, Entity, Image, SharedString, Task, WeakEntity};
use gpui::{App, Context, Entity, EventEmitter, Image, SharedString, Task, WeakEntity};
use language::Buffer;
use language_model::LanguageModelImage;
use project::image_store::is_image_file;
@ -31,6 +31,12 @@ pub struct ContextStore {
context_thread_ids: HashSet<ThreadId>,
}
pub enum ContextStoreEvent {
ContextRemoved(AgentContextKey),
}
impl EventEmitter<ContextStoreEvent> for ContextStore {}
impl ContextStore {
pub fn new(
project: WeakEntity<Project>,
@ -82,7 +88,7 @@ impl ContextStore {
project_path: ProjectPath,
remove_if_exists: bool,
cx: &mut Context<Self>,
) -> Task<Result<()>> {
) -> Task<Result<Option<AgentContextHandle>>> {
let Some(project) = self.project.upgrade() else {
return Task::ready(Err(anyhow!("failed to read project")));
};
@ -108,21 +114,22 @@ impl ContextStore {
buffer: Entity<Buffer>,
remove_if_exists: bool,
cx: &mut Context<Self>,
) {
) -> Option<AgentContextHandle> {
let context_id = self.next_context_id.post_inc();
let context = AgentContextHandle::File(FileContextHandle { buffer, context_id });
let already_included = if self.has_context(&context) {
if let Some(key) = self.context_set.get(AgentContextKey::ref_cast(&context)) {
if remove_if_exists {
self.remove_context(&context, cx);
None
} else {
Some(key.as_ref().clone())
}
true
} else if self.path_included_in_directory(project_path, cx).is_some() {
None
} else {
self.path_included_in_directory(project_path, cx).is_some()
};
if !already_included {
self.insert_context(context, cx);
self.insert_context(context.clone(), cx);
Some(context)
}
}
@ -131,7 +138,7 @@ impl ContextStore {
project_path: &ProjectPath,
remove_if_exists: bool,
cx: &mut Context<Self>,
) -> Result<()> {
) -> Result<Option<AgentContextHandle>> {
let Some(project) = self.project.upgrade() else {
return Err(anyhow!("failed to read project"));
};
@ -150,15 +157,20 @@ impl ContextStore {
context_id,
});
if self.has_context(&context) {
if remove_if_exists {
self.remove_context(&context, cx);
}
} else if self.path_included_in_directory(project_path, cx).is_none() {
self.insert_context(context, cx);
}
let context =
if let Some(existing) = self.context_set.get(AgentContextKey::ref_cast(&context)) {
if remove_if_exists {
self.remove_context(&context, cx);
None
} else {
Some(existing.as_ref().clone())
}
} else {
self.insert_context(context.clone(), cx);
Some(context)
};
anyhow::Ok(())
anyhow::Ok(context)
}
pub fn add_symbol(
@ -169,7 +181,7 @@ impl ContextStore {
enclosing_range: Range<Anchor>,
remove_if_exists: bool,
cx: &mut Context<Self>,
) -> bool {
) -> (Option<AgentContextHandle>, bool) {
let context_id = self.next_context_id.post_inc();
let context = AgentContextHandle::Symbol(SymbolContextHandle {
buffer,
@ -179,14 +191,18 @@ impl ContextStore {
context_id,
});
if self.has_context(&context) {
if remove_if_exists {
if let Some(key) = self.context_set.get(AgentContextKey::ref_cast(&context)) {
let handle = if remove_if_exists {
self.remove_context(&context, cx);
}
return false;
None
} else {
Some(key.as_ref().clone())
};
return (handle, false);
}
self.insert_context(context, cx)
let included = self.insert_context(context.clone(), cx);
(Some(context), included)
}
pub fn add_thread(
@ -194,16 +210,20 @@ impl ContextStore {
thread: Entity<Thread>,
remove_if_exists: bool,
cx: &mut Context<Self>,
) {
) -> Option<AgentContextHandle> {
let context_id = self.next_context_id.post_inc();
let context = AgentContextHandle::Thread(ThreadContextHandle { thread, context_id });
if self.has_context(&context) {
if let Some(existing) = self.context_set.get(AgentContextKey::ref_cast(&context)) {
if remove_if_exists {
self.remove_context(&context, cx);
None
} else {
Some(existing.as_ref().clone())
}
} else {
self.insert_context(context, cx);
self.insert_context(context.clone(), cx);
Some(context)
}
}
@ -212,19 +232,23 @@ impl ContextStore {
prompt_id: UserPromptId,
remove_if_exists: bool,
cx: &mut Context<ContextStore>,
) {
) -> Option<AgentContextHandle> {
let context_id = self.next_context_id.post_inc();
let context = AgentContextHandle::Rules(RulesContextHandle {
prompt_id,
context_id,
});
if self.has_context(&context) {
if let Some(existing) = self.context_set.get(AgentContextKey::ref_cast(&context)) {
if remove_if_exists {
self.remove_context(&context, cx);
None
} else {
Some(existing.as_ref().clone())
}
} else {
self.insert_context(context, cx);
self.insert_context(context.clone(), cx);
Some(context)
}
}
@ -233,14 +257,15 @@ impl ContextStore {
url: String,
text: impl Into<SharedString>,
cx: &mut Context<ContextStore>,
) {
) -> AgentContextHandle {
let context = AgentContextHandle::FetchedUrl(FetchedUrlContext {
url: url.into(),
text: text.into(),
context_id: self.next_context_id.post_inc(),
});
self.insert_context(context, cx);
self.insert_context(context.clone(), cx);
context
}
pub fn add_image_from_path(
@ -248,7 +273,7 @@ impl ContextStore {
project_path: ProjectPath,
remove_if_exists: bool,
cx: &mut Context<ContextStore>,
) -> Task<Result<()>> {
) -> Task<Result<Option<AgentContextHandle>>> {
let project = self.project.clone();
cx.spawn(async move |this, cx| {
let open_image_task = project.update(cx, |project, cx| {
@ -262,7 +287,7 @@ impl ContextStore {
image,
remove_if_exists,
cx,
);
)
})
})
}
@ -277,7 +302,7 @@ impl ContextStore {
image: Arc<Image>,
remove_if_exists: bool,
cx: &mut Context<ContextStore>,
) {
) -> Option<AgentContextHandle> {
let image_task = LanguageModelImage::from_image(image.clone(), cx).shared();
let context = AgentContextHandle::Image(ImageContext {
project_path,
@ -288,11 +313,12 @@ impl ContextStore {
if self.has_context(&context) {
if remove_if_exists {
self.remove_context(&context, cx);
return;
return None;
}
}
self.insert_context(context, cx);
self.insert_context(context.clone(), cx);
Some(context)
}
pub fn add_selection(
@ -364,9 +390,9 @@ impl ContextStore {
}
pub fn remove_context(&mut self, context: &AgentContextHandle, cx: &mut Context<Self>) {
if self
if let Some((_, key)) = self
.context_set
.shift_remove(AgentContextKey::ref_cast(context))
.shift_remove_full(AgentContextKey::ref_cast(context))
{
match context {
AgentContextHandle::Thread(thread_context) => {
@ -375,6 +401,7 @@ impl ContextStore {
}
_ => {}
}
cx.emit(ContextStoreEvent::ContextRemoved(key));
cx.notify();
}
}
@ -451,6 +478,12 @@ impl ContextStore {
.contains(&FetchedUrlContext::lookup_key(url.into()))
}
pub fn get_url_context(&self, url: SharedString) -> Option<AgentContextHandle> {
self.context_set
.get(&FetchedUrlContext::lookup_key(url))
.map(|key| key.as_ref().clone())
}
pub fn file_paths(&self, cx: &App) -> HashSet<ProjectPath> {
self.context()
.filter_map(|context| match context {

1
crates/agent/src/inline_assistant.rs
View file

@ -1199,6 +1199,7 @@ impl InlineAssistant {
) -> Vec<InlineAssistId> {
let assist_group = self.assist_groups.get_mut(&assist_group_id).unwrap();
assist_group.linked = false;
for assist_id in &assist_group.assist_ids {
let assist = self.assists.get_mut(assist_id).unwrap();
if let Some(editor_decorations) = assist.decorations.as_ref() {

37
crates/agent/src/inline_prompt_editor.rs
View file

@ -1,8 +1,10 @@
use crate::assistant_model_selector::{AssistantModelSelector, ModelType};
use crate::buffer_codegen::BufferCodegen;
use crate::context_picker::ContextPicker;
use crate::context::ContextCreasesAddon;
use crate::context_picker::{ContextPicker, ContextPickerCompletionProvider};
use crate::context_store::ContextStore;
use crate::context_strip::{ContextStrip, ContextStripEvent, SuggestContextKind};
use crate::message_editor::{extract_message_creases, insert_message_creases};
use crate::terminal_codegen::TerminalCodegen;
use crate::thread_store::ThreadStore;
use crate::{CycleNextInlineAssist, CyclePreviousInlineAssist};
@ -245,13 +247,22 @@ impl<T: 'static> PromptEditor<T> {
pub fn unlink(&mut self, window: &mut Window, cx: &mut Context<Self>) {
let prompt = self.prompt(cx);
let existing_creases = self.editor.update(cx, extract_message_creases);
let focus = self.editor.focus_handle(cx).contains_focused(window, cx);
self.editor = cx.new(|cx| {
let mut editor = Editor::auto_height(Self::MAX_LINES as usize, window, cx);
editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
editor.set_placeholder_text(Self::placeholder_text(&self.mode, window, cx), cx);
editor.set_placeholder_text("Add a prompt…", cx);
editor.set_text(prompt, window, cx);
insert_message_creases(
&mut editor,
&existing_creases,
&self.context_store,
window,
cx,
);
if focus {
window.focus(&editor.focus_handle(cx));
}
@ -860,8 +871,19 @@ impl PromptEditor<BufferCodegen> {
// typing in one will make what you typed appear in all of them.
editor.set_show_cursor_when_unfocused(true, cx);
editor.set_placeholder_text(Self::placeholder_text(&mode, window, cx), cx);
editor.register_addon(ContextCreasesAddon::new());
editor
});
let prompt_editor_entity = prompt_editor.downgrade();
prompt_editor.update(cx, |editor, _| {
editor.set_completion_provider(Some(Box::new(ContextPickerCompletionProvider::new(
workspace.clone(),
context_store.downgrade(),
thread_store.clone(),
prompt_editor_entity,
))));
});
let context_picker_menu_handle = PopoverMenuHandle::default();
let model_selector_menu_handle = PopoverMenuHandle::default();
@ -1015,6 +1037,17 @@ impl PromptEditor<TerminalCodegen> {
editor.set_placeholder_text(Self::placeholder_text(&mode, window, cx), cx);
editor
});
let prompt_editor_entity = prompt_editor.downgrade();
prompt_editor.update(cx, |editor, _| {
editor.set_completion_provider(Some(Box::new(ContextPickerCompletionProvider::new(
workspace.clone(),
context_store.downgrade(),
thread_store.clone(),
prompt_editor_entity,
))));
});
let context_picker_menu_handle = PopoverMenuHandle::default();
let model_selector_menu_handle = PopoverMenuHandle::default();

110
crates/agent/src/message_editor.rs
View file

@ -2,15 +2,15 @@ use std::collections::BTreeMap;
use std::sync::Arc;
use crate::assistant_model_selector::{AssistantModelSelector, ModelType};
use crate::context::{ContextLoadResult, load_context};
use crate::context::{AgentContextKey, ContextCreasesAddon, ContextLoadResult, load_context};
use crate::tool_compatibility::{IncompatibleToolsState, IncompatibleToolsTooltip};
use crate::ui::{AgentPreview, AnimatedLabel};
use buffer_diff::BufferDiff;
use collections::HashSet;
use collections::{HashMap, HashSet};
use editor::actions::{MoveUp, Paste};
use editor::{
ContextMenuOptions, ContextMenuPlacement, Editor, EditorElement, EditorEvent, EditorMode,
EditorStyle, MultiBuffer,
AnchorRangeExt, ContextMenuOptions, ContextMenuPlacement, Editor, EditorElement, EditorEvent,
EditorMode, EditorStyle, MultiBuffer,
};
use feature_flags::{FeatureFlagAppExt, NewBillingFeatureFlag};
use file_icons::FileIcons;
@ -35,11 +35,11 @@ use util::ResultExt as _;
use workspace::Workspace;
use zed_llm_client::CompletionMode;
use crate::context_picker::{ContextPicker, ContextPickerCompletionProvider};
use crate::context_picker::{ContextPicker, ContextPickerCompletionProvider, crease_for_mention};
use crate::context_store::ContextStore;
use crate::context_strip::{ContextStrip, ContextStripEvent, SuggestContextKind};
use crate::profile_selector::ProfileSelector;
use crate::thread::{Thread, TokenUsageRatio};
use crate::thread::{MessageCrease, Thread, TokenUsageRatio};
use crate::thread_store::ThreadStore;
use crate::{
ActiveThread, AgentDiff, Chat, ExpandMessageEditor, NewThread, OpenAgentDiff, RemoveAllContext,
@ -105,6 +105,7 @@ pub(crate) fn create_editor(
max_entries_visible: 12,
placement: Some(ContextMenuPlacement::Above),
});
editor.register_addon(ContextCreasesAddon::new());
editor
});
@ -290,10 +291,11 @@ impl MessageEditor {
return;
}
let user_message = self.editor.update(cx, |editor, cx| {
let (user_message, user_message_creases) = self.editor.update(cx, |editor, cx| {
let creases = extract_message_creases(editor, cx);
let text = editor.text(cx);
editor.clear(window, cx);
text
(text, creases)
});
self.last_estimated_token_count.take();
@ -311,7 +313,13 @@ impl MessageEditor {
thread
.update(cx, |thread, cx| {
thread.insert_user_message(user_message, loaded_context, checkpoint.ok(), cx);
thread.insert_user_message(
user_message,
loaded_context,
checkpoint.ok(),
user_message_creases,
cx,
);
})
.log_err();
@ -1164,6 +1172,53 @@ impl MessageEditor {
}
}
pub fn extract_message_creases(
editor: &mut Editor,
cx: &mut Context<'_, Editor>,
) -> Vec<MessageCrease> {
let buffer_snapshot = editor.buffer().read(cx).snapshot(cx);
let mut contexts_by_crease_id = editor
.addon_mut::<ContextCreasesAddon>()
.map(std::mem::take)
.unwrap_or_default()
.into_inner()
.into_iter()
.flat_map(|(key, creases)| {
let context = key.0;
creases
.into_iter()
.map(move |(id, _)| (id, context.clone()))
})
.collect::<HashMap<_, _>>();
// Filter the addon's list of creases based on what the editor reports,
// since the addon might have removed creases in it.
let creases = editor.display_map.update(cx, |display_map, cx| {
display_map
.snapshot(cx)
.crease_snapshot
.creases()
.filter_map(|(id, crease)| {
Some((
id,
(
crease.range().to_offset(&buffer_snapshot),
crease.metadata()?.clone(),
),
))
})
.map(|(id, (range, metadata))| {
let context = contexts_by_crease_id.remove(&id);
MessageCrease {
range,
metadata,
context,
}
})
.collect()
});
creases
}
impl EventEmitter<MessageEditorEvent> for MessageEditor {}
pub enum MessageEditorEvent {
@ -1204,6 +1259,43 @@ impl Render for MessageEditor {
}
}
pub fn insert_message_creases(
editor: &mut Editor,
message_creases: &[MessageCrease],
context_store: &Entity<ContextStore>,
window: &mut Window,
cx: &mut Context<'_, Editor>,
) {
let buffer_snapshot = editor.buffer().read(cx).snapshot(cx);
let creases = message_creases
.iter()
.map(|crease| {
let start = buffer_snapshot.anchor_after(crease.range.start);
let end = buffer_snapshot.anchor_before(crease.range.end);
crease_for_mention(
crease.metadata.label.clone(),
crease.metadata.icon_path.clone(),
start..end,
cx.weak_entity(),
)
})
.collect::<Vec<_>>();
let ids = editor.insert_creases(creases.clone(), cx);
editor.fold_creases(creases, false, window, cx);
if let Some(addon) = editor.addon_mut::<ContextCreasesAddon>() {
for (crease, id) in message_creases.iter().zip(ids) {
if let Some(context) = crease.context.as_ref() {
let key = AgentContextKey(context.clone());
addon.add_creases(
context_store,
key,
vec![(id, crease.metadata.label.clone())],
cx,
);
}
}
}
}
impl Component for MessageEditor {
fn scope() -> ComponentScope {
ComponentScope::Agent

70
crates/agent/src/thread.rs
View file

@ -9,6 +9,7 @@ use assistant_settings::AssistantSettings;
use assistant_tool::{ActionLog, AnyToolCard, Tool, ToolWorkingSet};
use chrono::{DateTime, Utc};
use collections::HashMap;
use editor::display_map::CreaseMetadata;
use feature_flags::{self, FeatureFlagAppExt};
use futures::future::Shared;
use futures::{FutureExt, StreamExt as _};
@ -39,10 +40,10 @@ use uuid::Uuid;
use zed_llm_client::CompletionMode;
use crate::ThreadStore;
use crate::context::{AgentContext, ContextLoadResult, LoadedContext};
use crate::context::{AgentContext, AgentContextHandle, ContextLoadResult, LoadedContext};
use crate::thread_store::{
SerializedLanguageModel, SerializedMessage, SerializedMessageSegment, SerializedThread,
SerializedToolResult, SerializedToolUse, SharedProjectContext,
SerializedCrease, SerializedLanguageModel, SerializedMessage, SerializedMessageSegment,
SerializedThread, SerializedToolResult, SerializedToolUse, SharedProjectContext,
};
use crate::tool_use::{PendingToolUse, ToolUse, ToolUseMetadata, ToolUseState};
@ -96,6 +97,15 @@ impl MessageId {
}
}
/// Stored information that can be used to resurrect a context crease when creating an editor for a past message.
#[derive(Clone, Debug)]
pub struct MessageCrease {
pub range: Range<usize>,
pub metadata: CreaseMetadata,
/// None for a deserialized message, Some otherwise.
pub context: Option<AgentContextHandle>,
}
/// A message in a [`Thread`].
#[derive(Debug, Clone)]
pub struct Message {
@ -103,6 +113,7 @@ pub struct Message {
pub role: Role,
pub segments: Vec<MessageSegment>,
pub loaded_context: LoadedContext,
pub creases: Vec<MessageCrease>,
}
impl Message {
@ -473,6 +484,18 @@ impl Thread {
text: message.context,
images: Vec::new(),
},
creases: message
.creases
.into_iter()
.map(|crease| MessageCrease {
range: crease.start..crease.end,
metadata: CreaseMetadata {
icon_path: crease.icon_path,
label: crease.label,
},
context: None,
})
.collect(),
})
.collect(),
next_message_id,
@ -826,6 +849,7 @@ impl Thread {
text: impl Into<String>,
loaded_context: ContextLoadResult,
git_checkpoint: Option<GitStoreCheckpoint>,
creases: Vec<MessageCrease>,
cx: &mut Context<Self>,
) -> MessageId {
if !loaded_context.referenced_buffers.is_empty() {
@ -840,6 +864,7 @@ impl Thread {
Role::User,
vec![MessageSegment::Text(text.into())],
loaded_context.loaded_context,
creases,
cx,
);
@ -860,7 +885,13 @@ impl Thread {
segments: Vec<MessageSegment>,
cx: &mut Context<Self>,
) -> MessageId {
self.insert_message(Role::Assistant, segments, LoadedContext::default(), cx)
self.insert_message(
Role::Assistant,
segments,
LoadedContext::default(),
Vec::new(),
cx,
)
}
pub fn insert_message(
@ -868,6 +899,7 @@ impl Thread {
role: Role,
segments: Vec<MessageSegment>,
loaded_context: LoadedContext,
creases: Vec<MessageCrease>,
cx: &mut Context<Self>,
) -> MessageId {
let id = self.next_message_id.post_inc();
@ -876,6 +908,7 @@ impl Thread {
role,
segments,
loaded_context,
creases,
});
self.touch_updated_at();
cx.emit(ThreadEvent::MessageAdded(id));
@ -995,6 +1028,16 @@ impl Thread {
})
.collect(),
context: message.loaded_context.text.clone(),
creases: message
.creases
.iter()
.map(|crease| SerializedCrease {
start: crease.range.start,
end: crease.range.end,
icon_path: crease.metadata.icon_path.clone(),
label: crease.metadata.label.clone(),
})
.collect(),
})
.collect(),
initial_project_snapshot,
@ -2502,7 +2545,13 @@ mod tests {
// Insert user message with context
let message_id = thread.update(cx, |thread, cx| {
thread.insert_user_message("Please explain this code", loaded_context, None, cx)
thread.insert_user_message(
"Please explain this code",
loaded_context,
None,
Vec::new(),
cx,
)
});
// Check content and context in message object
@ -2578,7 +2627,7 @@ fn main() {{
.update(|cx| load_context(new_contexts, &project, &None, cx))
.await;
let message1_id = thread.update(cx, |thread, cx| {
thread.insert_user_message("Message 1", loaded_context, None, cx)
thread.insert_user_message("Message 1", loaded_context, None, Vec::new(), cx)
});
// Second message with contexts 1 and 2 (context 1 should be skipped as it's already included)
@ -2593,7 +2642,7 @@ fn main() {{
.update(|cx| load_context(new_contexts, &project, &None, cx))
.await;
let message2_id = thread.update(cx, |thread, cx| {
thread.insert_user_message("Message 2", loaded_context, None, cx)
thread.insert_user_message("Message 2", loaded_context, None, Vec::new(), cx)
});
// Third message with all three contexts (contexts 1 and 2 should be skipped)
@ -2609,7 +2658,7 @@ fn main() {{
.update(|cx| load_context(new_contexts, &project, &None, cx))
.await;
let message3_id = thread.update(cx, |thread, cx| {
thread.insert_user_message("Message 3", loaded_context, None, cx)
thread.insert_user_message("Message 3", loaded_context, None, Vec::new(), cx)
});
// Check what contexts are included in each message
@ -2723,6 +2772,7 @@ fn main() {{
"What is the best way to learn Rust?",
ContextLoadResult::default(),
None,
Vec::new(),
cx,
)
});
@ -2756,6 +2806,7 @@ fn main() {{
"Are there any good books?",
ContextLoadResult::default(),
None,
Vec::new(),
cx,
)
});
@ -2805,7 +2856,7 @@ fn main() {{
// Insert user message with the buffer as context
thread.update(cx, |thread, cx| {
thread.insert_user_message("Explain this code", loaded_context, None, cx)
thread.insert_user_message("Explain this code", loaded_context, None, Vec::new(), cx)
});
// Create a request and check that it doesn't have a stale buffer warning yet
@ -2839,6 +2890,7 @@ fn main() {{
"What does the code do now?",
ContextLoadResult::default(),
None,
Vec::new(),
cx,
)
});

11
crates/agent/src/thread_store.rs
View file

@ -733,6 +733,8 @@ pub struct SerializedMessage {
pub tool_results: Vec<SerializedToolResult>,
#[serde(default)]
pub context: String,
#[serde(default)]
pub creases: Vec<SerializedCrease>,
}
#[derive(Debug, Serialize, Deserialize)]
@ -813,10 +815,19 @@ impl LegacySerializedMessage {
tool_uses: self.tool_uses,
tool_results: self.tool_results,
context: String::new(),
creases: Vec::new(),
}
}
}
#[derive(Debug, Serialize, Deserialize)]
pub struct SerializedCrease {
pub start: usize,
pub end: usize,
pub icon_path: SharedString,
pub label: SharedString,
}
struct GlobalThreadsDatabase(
Shared<BoxFuture<'static, Result<Arc<ThreadsDatabase>, Arc<anyhow::Error>>>>,
);

4
crates/assistant/src/inline_assistant.rs
View file

@ -1815,10 +1815,6 @@ impl PromptEditor {
self.editor = cx.new(|cx| {
let mut editor = Editor::auto_height(Self::MAX_LINES as usize, window, cx);
editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
editor.set_placeholder_text(
Self::placeholder_text(self.codegen.read(cx), window, cx),
cx,
);
editor.set_placeholder_text("Add a prompt…", cx);
editor.set_text(prompt, window, cx);
if focus {

12
crates/assistant_context_editor/src/context_editor.rs
View file

@ -1054,7 +1054,7 @@ impl ContextEditor {
|_, _, _, _| Empty.into_any_element(),
)
.with_metadata(CreaseMetadata {
icon: IconName::Ai,
icon_path: SharedString::from(IconName::Ai.path()),
label: "Thinking Process".into(),
}),
);
@ -1097,7 +1097,7 @@ impl ContextEditor {
FoldPlaceholder {
render: render_fold_icon_button(
cx.entity().downgrade(),
section.icon,
section.icon.path().into(),
section.label.clone(),
),
merge_adjacent: false,
@ -1107,7 +1107,7 @@ impl ContextEditor {
|_, _, _, _| Empty.into_any_element(),
)
.with_metadata(CreaseMetadata {
icon: section.icon,
icon_path: section.icon.path().into(),
label: section.label,
}),
);
@ -2055,7 +2055,7 @@ impl ContextEditor {
FoldPlaceholder {
render: render_fold_icon_button(
weak_editor.clone(),
metadata.crease.icon,
metadata.crease.icon_path.clone(),
metadata.crease.label.clone(),
),
..Default::default()
@ -2851,7 +2851,7 @@ fn render_thought_process_fold_icon_button(
fn render_fold_icon_button(
editor: WeakEntity<Editor>,
icon: IconName,
icon_path: SharedString,
label: SharedString,
) -> Arc<dyn Send + Sync + Fn(FoldId, Range<Anchor>, &mut App) -> AnyElement> {
Arc::new(move |fold_id, fold_range, _cx| {
@ -2859,7 +2859,7 @@ fn render_fold_icon_button(
ButtonLike::new(fold_id)
.style(ButtonStyle::Filled)
.layer(ElevationIndex::ElevatedSurface)
.child(Icon::new(icon))
.child(Icon::from_path(icon_path.clone()))
.child(Label::new(label.clone()).single_line())
.on_click(move |_, window, cx| {
editor

2
crates/editor/src/display_map.rs
View file

@ -391,7 +391,7 @@ impl DisplayMap {
&mut self,
crease_ids: impl IntoIterator<Item = CreaseId>,
cx: &mut Context<Self>,
) {
) -> Vec<(CreaseId, Range<Anchor>)> {
let snapshot = self.buffer.read(cx).snapshot(cx);
self.crease_map.remove(crease_ids, &snapshot)
}

25
crates/editor/src/display_map/crease_map.rs
View file

@ -5,7 +5,7 @@ use serde::{Deserialize, Serialize};
use std::{cmp::Ordering, fmt::Debug, ops::Range, sync::Arc};
use sum_tree::{Bias, SeekTarget, SumTree};
use text::Point;
use ui::{App, IconName, SharedString, Window};
use ui::{App, SharedString, Window};
use crate::{BlockStyle, FoldPlaceholder, RenderBlock};
@ -40,6 +40,10 @@ impl CreaseSnapshot {
}
}
pub fn creases(&self) -> impl Iterator<Item = (CreaseId, &Crease<Anchor>)> {
self.creases.iter().map(|item| (item.id, &item.crease))
}
/// Returns the first Crease starting on the specified buffer row.
pub fn query_row<'a>(
&'a self,
@ -147,7 +151,7 @@ pub enum Crease<T> {
/// Metadata about a [`Crease`], that is used for serialization.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct CreaseMetadata {
pub icon: IconName,
pub icon_path: SharedString,
pub label: SharedString,
}
@ -237,6 +241,13 @@ impl<T> Crease<T> {
Crease::Block { range, .. } => range,
}
}
pub fn metadata(&self) -> Option<&CreaseMetadata> {
match self {
Self::Inline { metadata, .. } => metadata.as_ref(),
Self::Block { .. } => None,
}
}
}
impl<T> std::fmt::Debug for Crease<T>
@ -305,7 +316,7 @@ impl CreaseMap {
&mut self,
ids: impl IntoIterator<Item = CreaseId>,
snapshot: &MultiBufferSnapshot,
) {
) -> Vec<(CreaseId, Range<Anchor>)> {
let mut removals = Vec::new();
for id in ids {
if let Some(range) = self.id_to_range.remove(&id) {
@ -320,11 +331,11 @@ impl CreaseMap {
let mut new_creases = SumTree::new(snapshot);
let mut cursor = self.snapshot.creases.cursor::<ItemSummary>(snapshot);
for (id, range) in removals {
new_creases.append(cursor.slice(&range, Bias::Left, snapshot), snapshot);
for (id, range) in &removals {
new_creases.append(cursor.slice(range, Bias::Left, snapshot), snapshot);
while let Some(item) = cursor.item() {
cursor.next(snapshot);
if item.id == id {
if item.id == *id {
break;
} else {
new_creases.push(item.clone(), snapshot);
@ -335,6 +346,8 @@ impl CreaseMap {
new_creases.append(cursor.suffix(snapshot), snapshot);
new_creases
};
removals
}
}

4
crates/editor/src/editor.rs
View file

@ -16239,9 +16239,9 @@ impl Editor {
&mut self,
ids: impl IntoIterator<Item = CreaseId>,
cx: &mut Context<Self>,
) {
) -> Vec<(CreaseId, Range<Anchor>)> {
self.display_map
.update(cx, |map, cx| map.remove_creases(ids, cx));
.update(cx, |map, cx| map.remove_creases(ids, cx))
}
pub fn longest_row(&self, cx: &mut App) -> DisplayRow {

1
crates/eval/src/example.rs
View file

@ -119,6 +119,7 @@ impl ExampleContext {
text.to_string(),
ContextLoadResult::default(),
None,
Vec::new(),
cx,
);
})