Simplify context refresh and only refresh new context

crates/agent/src/context_store.rs

@@ -12,7 +12,7 @@ use project::{Project, ProjectEntryId, ProjectItem, ProjectPath, Worktree};
 use prompt_store::UserPromptId;
 use rope::{Point, Rope};
 use text::{Anchor, BufferId, OffsetRangeExt};
-use util::{ResultExt as _, maybe};
+use util::ResultExt as _;
 
 use crate::ThreadStore;
 use crate::context::{
@@ -833,91 +833,33 @@ fn collect_files_in_path(worktree: &Worktree, path: &Path) -> Vec<Arc<Path>> {
     files
 }
 
-pub fn refresh_context_store_text(
+pub fn refresh_context_text(
     context_store: Entity<ContextStore>,
-    changed_buffers: &HashSet<Entity<Buffer>>,
+    context: &AssistantContext,
     cx: &App,
-) -> impl Future<Output = Vec<ContextId>> + use<> {
-    let mut tasks = Vec::new();
-
-    for context in &context_store.read(cx).context {
-        let id = context.id();
-
-        let task = maybe!({
-            match context {
-                AssistantContext::File(file_context) => {
-                    // TODO: Should refresh if the path has changed, as it's in the text.
-                    if changed_buffers.is_empty()
-                        || changed_buffers.contains(&file_context.context_buffer.buffer)
-                    {
-                        let context_store = context_store.clone();
-                        return refresh_file_text(context_store, file_context, cx);
-                    }
-                }
-                AssistantContext::Directory(directory_context) => {
-                    let directory_path = directory_context.project_path(cx)?;
-                    let should_refresh = directory_path.path != directory_context.last_path
-                        || changed_buffers.is_empty()
-                        || changed_buffers.iter().any(|buffer| {
-                            let Some(buffer_path) = buffer.read(cx).project_path(cx) else {
-                                return false;
-                            };
-                            buffer_path.starts_with(&directory_path)
-                        });
-
-                    if should_refresh {
-                        let context_store = context_store.clone();
-                        return refresh_directory_text(
-                            context_store,
-                            directory_context,
-                            directory_path,
-                            cx,
-                        );
-                    }
-                }
-                AssistantContext::Symbol(symbol_context) => {
-                    // TODO: Should refresh if the path has changed, as it's in the text.
-                    if changed_buffers.is_empty()
-                        || changed_buffers.contains(&symbol_context.context_symbol.buffer)
-                    {
-                        let context_store = context_store.clone();
-                        return refresh_symbol_text(context_store, symbol_context, cx);
-                    }
-                }
-                AssistantContext::Excerpt(excerpt_context) => {
-                    // TODO: Should refresh if the path has changed, as it's in the text.
-                    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();
-                        return Some(refresh_thread_text(context_store, thread_context, cx));
-                    }
-                }
-                // Intentionally omit refreshing fetched URLs as it doesn't seem all that useful,
-                // and doing the caching properly could be tricky (unless it's already handled by
-                // the HttpClient?).
-                AssistantContext::FetchedUrl(_) => {}
-                AssistantContext::Rules(user_rules_context) => {
-                    let context_store = context_store.clone();
-                    return Some(refresh_user_rules(context_store, user_rules_context, cx));
-                }
-            }
-
-            None
-        });
-
-        if let Some(task) = task {
-            tasks.push(task.map(move |_| id));
+) -> Option<Task<()>> {
+    match context {
+        AssistantContext::File(file_context) => refresh_file_text(context_store, file_context, cx),
+        AssistantContext::Directory(directory_context) => {
+            refresh_directory_text(context_store, directory_context, cx)
+        }
+        AssistantContext::Symbol(symbol_context) => {
+            refresh_symbol_text(context_store, symbol_context, cx)
+        }
+        AssistantContext::Excerpt(excerpt_context) => {
+            refresh_excerpt_text(context_store, excerpt_context, cx)
+        }
+        AssistantContext::Thread(thread_context) => {
+            Some(refresh_thread_text(context_store, thread_context, cx))
+        }
+        // Intentionally omit refreshing fetched URLs as it doesn't seem all that useful,
+        // and doing the caching properly could be tricky (unless it's already handled by
+        // the HttpClient?).
+        AssistantContext::FetchedUrl(_) => None,
+        AssistantContext::Rules(user_rules_context) => {
+            Some(refresh_user_rules(context_store, user_rules_context, cx))
         }
     }
-
-    future::join_all(tasks)
 }
 
 fn refresh_file_text(
@@ -945,10 +887,11 @@ fn refresh_file_text(
 fn refresh_directory_text(
     context_store: Entity<ContextStore>,
     directory_context: &DirectoryContext,
-    directory_path: ProjectPath,
     cx: &App,
 ) -> Option<Task<()>> {
-    let mut stale = false;
+    let directory_path = directory_context.project_path(cx)?;
+    let mut stale = directory_path.path != directory_context.last_path;
+
     let futures = directory_context
         .context_buffers
         .iter()
@@ -1101,7 +1044,10 @@ fn refresh_user_rules(
 
 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 file = buffer.file()?;
+    if buffer.version.changed_since(&context_buffer.version)
+        || file.full_path(cx).as_path() != context_buffer.last_full_path.as_ref()
+    {
         load_context_buffer(context_buffer.buffer.clone(), cx).log_err()
     } else {
         None
@@ -1114,7 +1060,10 @@ fn refresh_context_excerpt(
     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 file = buffer.file()?;
+    if buffer.version.changed_since(&context_buffer.version)
+        || file.full_path(cx).as_path() != context_buffer.last_full_path.as_ref()
+    {
         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)))
@@ -1129,6 +1078,7 @@ fn refresh_context_symbol(
 ) -> Option<impl Future<Output = ContextSymbol> + use<>> {
     let buffer = context_symbol.buffer.read(cx);
     let project_path = buffer.project_path(cx)?;
+    // TODO: Should refresh text when path has changed.
     if buffer.version.changed_since(&context_symbol.buffer_version) {
         let (_line_range, context_buffer_task) = load_context_buffer_range(
             context_symbol.buffer.clone(),

crates/agent/src/message_editor.rs

@@ -13,6 +13,7 @@ use editor::{
 };
 use file_icons::FileIcons;
 use fs::Fs;
+use futures::future;
 use gpui::{
     Animation, AnimationExt, App, Entity, EventEmitter, Focusable, Subscription, Task, TextStyle,
     WeakEntity, linear_color_stop, linear_gradient, point, pulsating_between,
@@ -31,7 +32,7 @@ use workspace::Workspace;
 
 use crate::assistant_model_selector::AssistantModelSelector;
 use crate::context_picker::{ContextPicker, ContextPickerCompletionProvider};
-use crate::context_store::{ContextStore, refresh_context_store_text};
+use crate::context_store::{ContextStore, refresh_context_text};
 use crate::context_strip::{ContextStrip, ContextStripEvent, SuggestContextKind};
 use crate::profile_selector::ProfileSelector;
 use crate::thread::{Thread, TokenUsageRatio};
@@ -269,8 +270,22 @@ impl MessageEditor {
         self.last_estimated_token_count.take();
         cx.emit(MessageEditorEvent::EstimatedTokenCount);
 
-        let refresh_task =
-            refresh_context_store_text(self.context_store.clone(), &HashSet::default(), cx);
+        // Only need to refresh context that's added to the message.
+        let new_context = self
+            .thread
+            .read(cx)
+            .filter_new_context(self.context_store.read(cx).context().iter())
+            .collect::<Vec<_>>();
+        let refresh_task = future::join_all(
+            new_context
+                .iter()
+                .flat_map(|context| refresh_context_text(self.context_store.clone(), context, cx))
+                .collect::<Vec<_>>(),
+        );
+        let new_context_ids = new_context
+            .into_iter()
+            .map(|context| context.id())
+            .collect::<Vec<_>>();
 
         let thread = self.thread.clone();
         let context_store = self.context_store.clone();
@@ -283,8 +298,13 @@ impl MessageEditor {
 
             thread
                 .update(cx, |thread, cx| {
-                    let context = context_store.read(cx).context().clone();
-                    thread.insert_user_message(user_message, context, checkpoint, cx);
+                    let context_store_ref = context_store.read(cx);
+                    let new_context = new_context_ids
+                        .into_iter()
+                        .flat_map(|id| context_store_ref.context_for_id(id))
+                        .cloned()
+                        .collect::<Vec<_>>();
+                    thread.insert_user_message(user_message, new_context, checkpoint, cx);
                 })
                 .log_err();
 

crates/agent/src/thread.rs

@@ -717,17 +717,13 @@ impl Thread {
         &self,
         context: impl Iterator<Item = &'a AssistantContext>,
     ) -> impl Iterator<Item = &'a AssistantContext> {
-        context.filter(|ctx| self.is_context_new(ctx))
-    }
-
-    fn is_context_new(&self, context: &AssistantContext) -> bool {
-        !self.context.contains_key(&context.id())
+        context.filter(|ctx| !self.context.contains_key(&ctx.id()))
     }
 
     pub fn insert_user_message(
         &mut self,
         text: impl Into<String>,
-        context: Vec<AssistantContext>,
+        new_context: Vec<AssistantContext>,
         git_checkpoint: Option<GitStoreCheckpoint>,
         cx: &mut Context<Self>,
     ) -> MessageId {
@@ -735,11 +731,6 @@ impl Thread {
 
         let message_id = self.insert_message(Role::User, vec![MessageSegment::Text(text)], cx);
 
-        let new_context: Vec<_> = context
-            .into_iter()
-            .filter(|ctx| self.is_context_new(ctx))
-            .collect();
-
         if !new_context.is_empty() {
             if let Some(context_string) = format_context_as_string(new_context.iter(), cx) {
                 if let Some(message) = self.messages.iter_mut().find(|m| m.id == message_id) {