agent: Include context with first message that introduced it (#27925)

We were including the context at the end which meant it never got
cached. We'll now include it with the first message that introduced it
so it's cached as long as it doesn't change.

This is an improvement, but we probably still need to think of ways to
optimize caching for cases where files in context change.

Release Notes:

- N/A

crates/agent/src/context.rs

@@ -2,7 +2,7 @@ use std::{ops::Range, sync::Arc};
 
 use gpui::{App, Entity, SharedString};
 use language::{Buffer, File};
-use language_model::{LanguageModelRequestMessage, MessageContent};
+use language_model::LanguageModelRequestMessage;
 use project::ProjectPath;
 use serde::{Deserialize, Serialize};
 use text::{Anchor, BufferId};
@@ -170,52 +170,61 @@ pub fn attach_context_to_message<'a>(
     let mut context_chunks = Vec::new();
 
     if !file_context.is_empty() {
-        context_chunks.push("The following files are available:\n");
+        context_chunks.push("<files>\n");
         for context in file_context {
             context_chunks.push(&context.context_buffer.text);
         }
+        context_chunks.push("\n</files>\n");
     }
 
     if !directory_context.is_empty() {
-        context_chunks.push("The following directories are available:\n");
+        context_chunks.push("<directories>\n");
         for context in directory_context {
             for context_buffer in &context.context_buffers {
                 context_chunks.push(&context_buffer.text);
             }
         }
+        context_chunks.push("\n</directories>\n");
     }
 
     if !symbol_context.is_empty() {
-        context_chunks.push("The following symbols are available:\n");
+        context_chunks.push("<symbols>\n");
         for context in symbol_context {
             context_chunks.push(&context.context_symbol.text);
         }
+        context_chunks.push("\n</symbols>\n");
     }
 
     if !fetch_context.is_empty() {
-        context_chunks.push("The following fetched results are available:\n");
+        context_chunks.push("<fetched_urls>\n");
         for context in &fetch_context {
             context_chunks.push(&context.url);
             context_chunks.push(&context.text);
         }
+        context_chunks.push("\n</fetched_urls>\n");
     }
 
     // Need to own the SharedString for summary so that it can be referenced.
     let mut thread_context_chunks = Vec::new();
     if !thread_context.is_empty() {
-        context_chunks.push("The following previous conversation threads are available:\n");
+        context_chunks.push("<conversation_threads>\n");
         for context in &thread_context {
             thread_context_chunks.push(context.summary(cx));
             thread_context_chunks.push(context.text.clone());
         }
+        context_chunks.push("\n</conversation_threads>\n");
     }
+
     for chunk in &thread_context_chunks {
         context_chunks.push(chunk);
     }
 
     if !context_chunks.is_empty() {
-        message
-            .content
-            .push(MessageContent::Text(context_chunks.join("\n")));
+        message.content.push(
+            "\n<context>\n\
+                The following items were attached by the user. You don't need to use other tools to read them.\n\n".into(),
+        );
+        message.content.push(context_chunks.join("\n").into());
+        message.content.push("\n</context>\n".into());
     }
 }

crates/agent/src/thread.rs

@@ -881,13 +881,9 @@ impl Thread {
             log::error!("system_prompt_context not set.")
         }
 
-        let mut referenced_context_ids = HashSet::default();
+        let mut added_context_ids = HashSet::<ContextId>::default();
 
         for message in &self.messages {
-            if let Some(context_ids) = self.context_by_message.get(&message.id) {
-                referenced_context_ids.extend(context_ids);
-            }
-
             let mut request_message = LanguageModelRequestMessage {
                 role: message.role,
                 content: Vec::new(),
@@ -907,6 +903,23 @@ impl Thread {
                 }
             }
 
+            // Attach context to this message if it's the first to reference it
+            if let Some(context_ids) = self.context_by_message.get(&message.id) {
+                let new_context_ids: Vec<_> = context_ids
+                    .iter()
+                    .filter(|id| !added_context_ids.contains(id))
+                    .collect();
+
+                if !new_context_ids.is_empty() {
+                    let referenced_context = new_context_ids
+                        .iter()
+                        .filter_map(|context_id| self.context.get(*context_id));
+
+                    attach_context_to_message(&mut request_message, referenced_context, cx);
+                    added_context_ids.extend(context_ids.iter());
+                }
+            }
+
             if !message.segments.is_empty() {
                 request_message
                     .content
@@ -933,21 +946,6 @@ impl Thread {
             message.cache = index == breakpoint_index;
         }
 
-        if !referenced_context_ids.is_empty() {
-            let mut context_message = LanguageModelRequestMessage {
-                role: Role::User,
-                content: Vec::new(),
-                cache: false,
-            };
-
-            let referenced_context = referenced_context_ids
-                .into_iter()
-                .filter_map(|context_id| self.context.get(context_id));
-            attach_context_to_message(&mut context_message, referenced_context, cx);
-
-            request.messages.push(context_message);
-        }
-
         self.attached_tracked_files_state(&mut request.messages, cx);
 
         request