Rewrite test

Co-authored-by: Ben Brandt <benjamin.j.brandt@gmail.com>

Cargo.lock

@@ -64,6 +64,7 @@ dependencies = [
  "buffer_diff",
  "chrono",
  "client",
+ "clock",
  "collections",
  "component",
  "context_server",

crates/agent/Cargo.toml

@@ -33,6 +33,7 @@ collections.workspace = true
 component.workspace = true
 context_server.workspace = true
 convert_case.workspace = true
+clock.workspace = true
 db.workspace = true
 editor.workspace = true
 extension.workspace = true

crates/agent/src/thread.rs

@@ -1,3 +1,4 @@
+use std::collections::BTreeMap;
 use std::fmt::Write as _;
 use std::io::Write;
 use std::ops::Range;
@@ -18,6 +19,7 @@ use gpui::{
     AnyWindowHandle, App, AppContext, AsyncApp, Context, Entity, EventEmitter, SharedString, Task,
     WeakEntity,
 };
+use language::Buffer;
 use language_model::{
     ConfiguredModel, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent,
     LanguageModelId, LanguageModelKnownError, LanguageModelRegistry, LanguageModelRequest,
@@ -344,6 +346,7 @@ pub struct Thread {
     tools: Entity<ToolWorkingSet>,
     tool_use: ToolUseState,
     action_log: Entity<ActionLog>,
+    last_buffer_notifications: BTreeMap<Entity<Buffer>, clock::Global>,
     last_restore_checkpoint: Option<LastRestoreCheckpoint>,
     pending_checkpoint: Option<ThreadCheckpoint>,
     initial_project_snapshot: Shared<Task<Option<Arc<ProjectSnapshot>>>>,
@@ -434,6 +437,7 @@ impl Thread {
             pending_checkpoint: None,
             tool_use: ToolUseState::new(tools.clone()),
             action_log: cx.new(|_| ActionLog::new(project.clone())),
+            last_buffer_notifications: BTreeMap::new(),
             initial_project_snapshot: {
                 let project_snapshot = Self::project_snapshot(project, cx);
                 cx.foreground_executor()
@@ -564,6 +568,7 @@ impl Thread {
             tools: tools.clone(),
             tool_use,
             action_log: cx.new(|_| ActionLog::new(project)),
+            last_buffer_notifications: BTreeMap::new(),
             initial_project_snapshot: Task::ready(serialized.initial_project_snapshot).shared(),
             request_token_usage: serialized.request_token_usage,
             cumulative_token_usage: serialized.cumulative_token_usage,
@@ -1045,6 +1050,33 @@ impl Thread {
         id
     }
 
+    pub fn insert_message_at(
+        &mut self,
+        index: usize,
+        role: Role,
+        segments: Vec<MessageSegment>,
+        loaded_context: LoadedContext,
+        creases: Vec<MessageCrease>,
+        is_hidden: bool,
+        cx: &mut Context<Self>,
+    ) -> MessageId {
+        let id = self.next_message_id.post_inc();
+        self.messages.insert(
+            index,
+            Message {
+                id,
+                role,
+                segments,
+                loaded_context,
+                creases,
+                is_hidden,
+            },
+        );
+        self.touch_updated_at();
+        cx.emit(ThreadEvent::MessageAdded(id));
+        id
+    }
+
     pub fn edit_message(
         &mut self,
         id: MessageId,
@@ -1224,6 +1256,19 @@ impl Thread {
 
         self.remaining_turns -= 1;
 
+        match intent {
+            CompletionIntent::UserPrompt | CompletionIntent::ToolResults => {
+                self.attach_tracked_files_state(cx);
+            }
+            CompletionIntent::ThreadSummarization
+            | CompletionIntent::ThreadContextSummarization
+            | CompletionIntent::CreateFile
+            | CompletionIntent::EditFile
+            | CompletionIntent::InlineAssist
+            | CompletionIntent::TerminalInlineAssist
+            | CompletionIntent::GenerateGitCommitMessage => {}
+        };
+
         let request = self.to_completion_request(model.clone(), intent, cx);
 
         self.stream_completion(request, model, window, cx);
@@ -1389,8 +1434,6 @@ impl Thread {
             request.messages[message_ix_to_cache].cache = true;
         }
 
-        self.attach_tracked_files_state(&mut request.messages, cx);
-
         request.tools = available_tools;
         request.mode = if model.supports_max_mode() {
             Some(self.completion_mode.into())
@@ -1453,18 +1496,18 @@ impl Thread {
         request
     }
 
-    fn attach_tracked_files_state(
-        &self,
-        messages: &mut Vec<LanguageModelRequestMessage>,
-        cx: &App,
-    ) {
+    pub fn attach_tracked_files_state(&mut self, cx: &mut Context<Self>) {
+        let action_log = self.action_log.read(cx);
         let mut stale_files = String::new();
 
-        let action_log = self.action_log.read(cx);
-
-        for stale_file in action_log.stale_buffers(cx) {
-            if let Some(file) = stale_file.read(cx).file() {
-                writeln!(&mut stale_files, "- {}", file.path().display()).ok();
+        for stale_buffer in action_log.stale_buffers(cx) {
+            let version = stale_buffer.read(cx).version();
+            if self.last_buffer_notifications.get(&stale_buffer) != Some(&version) {
+                if let Some(file) = stale_buffer.read(cx).file() {
+                    self.last_buffer_notifications
+                        .insert(stale_buffer.clone(), version);
+                    writeln!(&mut stale_files, "- {}", file.path().display()).ok();
+                }
             }
         }
 
@@ -1474,37 +1517,27 @@ impl Thread {
 
         // NOTE: Changes to this prompt require a symmetric update in the LLM Worker
         const STALE_FILES_HEADER: &str = include_str!("./prompts/stale_files_prompt_header.txt");
-        let content = MessageContent::Text(
-            format!("{STALE_FILES_HEADER}{stale_files}").replace("\r\n", "\n"),
-        );
+        let content = format!("{STALE_FILES_HEADER}{stale_files}").replace("\r\n", "\n");
 
         // Insert our message before the last Assistant message.
         // Inserting it to the tail distracts the agent too much
-        let insert_position = messages
+        let insert_position = self
+            .messages
             .iter()
             .enumerate()
             .rfind(|(_, message)| message.role == Role::Assistant)
-            .map_or(messages.len(), |(i, _)| i);
+            .map_or(self.messages.len(), |(i, _)| i);
 
-        let request_message = LanguageModelRequestMessage {
-            role: Role::User,
-            content: vec![content],
-            cache: false,
-        };
-
-        messages.insert(insert_position, request_message);
-
-        // It makes no sense to cache messages after this one because
-        // the cache is invalidated when this message is gone.
-        // Move the cache marker before this message.
-        let has_cached_messages_after = messages
-            .iter()
-            .skip(insert_position + 1)
-            .any(|message| message.cache);
-
-        if has_cached_messages_after {
-            messages[insert_position - 1].cache = true;
-        }
+        let is_hidden = true;
+        self.insert_message_at(
+            insert_position,
+            Role::User,
+            vec![MessageSegment::Text(content)],
+            LoadedContext::default(),
+            Vec::new(),
+            is_hidden,
+            cx,
+        );
     }
 
     pub fn stream_completion(
@@ -3239,7 +3272,7 @@ fn main() {{
         )
         .await;
 
-        let (_workspace, _thread_store, thread, context_store, model) =
+        let (_workspace, _thread_store, thread, context_store, _model) =
             setup_test_environment(cx, project.clone()).await;
 
         // Open buffer and add it to context
@@ -3258,24 +3291,14 @@ fn main() {{
             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
-        let initial_request = thread.update(cx, |thread, cx| {
-            thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
+        // Initially, no messages should have stale buffer notification
+        thread.read_with(cx, |thread, _| {
+            assert_eq!(thread.messages.len(), 1);
+            assert!(!thread.messages[0].is_hidden);
         });
 
-        // Make sure we don't have a stale file warning yet
-        let has_stale_warning = initial_request.messages.iter().any(|msg| {
-            msg.string_contents()
-                .contains("These files changed since last read:")
-        });
-        assert!(
-            !has_stale_warning,
-            "Should not have stale buffer warning before buffer is modified"
-        );
-
-        // Modify the buffer
+        // Modify the buffer to make it stale
         buffer.update(cx, |buffer, cx| {
-            // Find a position at the end of line 1
             buffer.edit(
                 [(1..1, "\n    println!(\"Added a new line\");\n")],
                 None,
@@ -3283,50 +3306,63 @@ fn main() {{
             );
         });
 
-        // Insert another user message without context
+        // Check that the stale file notification was added
         thread.update(cx, |thread, cx| {
-            thread.insert_user_message(
-                "What does the code do now?",
-                ContextLoadResult::default(),
-                None,
-                Vec::new(),
+            thread.attach_tracked_files_state(cx);
+        });
+        thread.read_with(cx, |thread, _| {
+            // Should have 2 messages now: original user message + hidden stale notification
+            assert_eq!(thread.messages.len(), 2);
+
+            let stale_msg = &thread.messages[1];
+            assert!(stale_msg.is_hidden, "Stale notification should be hidden");
+            assert_eq!(stale_msg.role, Role::User);
+
+            let expected_content = "[The following is an auto-generated notification; do not reply]\n\nThese files have changed since the last read:\n- code.rs\n";
+            assert_eq!(
+                stale_msg.to_string(),
+                expected_content,
+                "Stale buffer notification should have the correct format"
+            );
+        });
+
+        // Test that calling attach_tracked_files_state again doesn't add duplicate notifications
+        thread.update(cx, |thread, cx| {
+            thread.attach_tracked_files_state(cx);
+        });
+        thread.read_with(cx, |thread, _| {
+            assert_eq!(thread.messages.len(), 2);
+        });
+
+        // Test with assistant message - notification should be inserted before it
+        thread.update(cx, |thread, cx| {
+            thread.insert_assistant_message(
+                vec![MessageSegment::Text("Here's an explanation...".into())],
                 cx,
-            )
+            );
         });
 
-        // Create a new request and check for the stale buffer warning
-        let new_request = thread.update(cx, |thread, cx| {
-            thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
+        // Modify buffer again to create a new version
+        buffer.update(cx, |buffer, cx| {
+            buffer.edit([(1..1, "\n    // Another change\n")], None, cx);
         });
 
-        // We should have a stale file warning as the last message
-        let last_message = new_request
-            .messages
-            .last()
-            .expect("Request should have messages");
+        thread.update(cx, |thread, cx| {
+            thread.attach_tracked_files_state(cx);
+        });
 
-        // The last message should be the stale buffer notification
-        assert_eq!(last_message.role, Role::User);
+        thread.read_with(cx, |thread, _| {
+            // Should have 4 messages: user, stale notification, new stale notification, assistant
+            assert_eq!(thread.messages.len(), 4);
 
-        // Check the exact content of the message
-        let expected_content = "[The following is an auto-generated notification; do not reply]
-
-These files have changed since the last read:
-- code.rs
-";
-        assert_eq!(
-            last_message.string_contents(),
-            expected_content,
-            "Last message should be exactly the stale buffer notification"
-        );
-
-        // The message before the notification should be cached
-        let index = new_request.messages.len() - 2;
-        let previous_message = new_request.messages.get(index).unwrap();
-        assert!(
-            previous_message.cache,
-            "Message before the stale buffer notification should be cached"
-        );
+            // The new stale notification should be inserted before the assistant message
+            let new_stale_msg = &thread.messages[2];
+            assert!(
+                new_stale_msg
+                    .to_string()
+                    .contains("These files have changed since the last read:")
+            );
+        });
     }
 
     #[gpui::test]