From fffa40f97303d6c93b135bbc058a664af4093811 Mon Sep 17 00:00:00 2001
From: Marshall Bowers <elliott.codes@gmail.com>
Date: Tue, 7 Jan 2025 14:16:30 -0500
Subject: [PATCH] assistant2: Make context persistent in the thread (#22789)

This PR makes it so the context is persistent in the thread, rather than
having to reattach it for each message.

This PR intentionally does not make an attempt to refresh the attached
context if it changes. That will come in a follow-up.

Release Notes:

- N/A
---
 crates/assistant2/src/context_store.rs  |  6 ------
 crates/assistant2/src/message_editor.rs |  4 +++-
 crates/assistant2/src/thread.rs         | 28 ++++++++++++++++++++-----
 3 files changed, 26 insertions(+), 12 deletions(-)

diff --git a/crates/assistant2/src/context_store.rs b/crates/assistant2/src/context_store.rs
index 6ce35f8971..f2659f53a8 100644
--- a/crates/assistant2/src/context_store.rs
+++ b/crates/assistant2/src/context_store.rs
@@ -36,12 +36,6 @@ impl ContextStore {
         &self.context
     }
 
-    pub fn drain(&mut self) -> Vec<Context> {
-        let context = self.context.drain(..).collect();
-        self.clear();
-        context
-    }
-
     pub fn clear(&mut self) {
         self.context.clear();
         self.files.clear();
diff --git a/crates/assistant2/src/message_editor.rs b/crates/assistant2/src/message_editor.rs
index bad17d7f10..684b5c5df0 100644
--- a/crates/assistant2/src/message_editor.rs
+++ b/crates/assistant2/src/message_editor.rs
@@ -142,7 +142,9 @@ impl MessageEditor {
             editor.clear(cx);
             text
         });
-        let context = self.context_store.update(cx, |this, _cx| this.drain());
+        let context = self
+            .context_store
+            .update(cx, |this, _cx| this.context().clone());
 
         self.thread.update(cx, |thread, cx| {
             thread.insert_user_message(user_message, context, cx);
diff --git a/crates/assistant2/src/thread.rs b/crates/assistant2/src/thread.rs
index 0816e541b4..f3de36f8a7 100644
--- a/crates/assistant2/src/thread.rs
+++ b/crates/assistant2/src/thread.rs
@@ -3,7 +3,7 @@ use std::sync::Arc;
 use anyhow::Result;
 use assistant_tool::ToolWorkingSet;
 use chrono::{DateTime, Utc};
-use collections::HashMap;
+use collections::{HashMap, HashSet};
 use futures::future::Shared;
 use futures::{FutureExt as _, StreamExt as _};
 use gpui::{AppContext, EventEmitter, ModelContext, SharedString, Task};
@@ -209,7 +209,13 @@ impl Thread {
             temperature: None,
         };
 
+        let mut referenced_context_ids = HashSet::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(),
@@ -224,10 +230,6 @@ impl Thread {
                 }
             }
 
-            if let Some(context) = self.context_for_message(message.id) {
-                attach_context_to_message(&mut request_message, context.clone());
-            }
-
             if !message.text.is_empty() {
                 request_message
                     .content
@@ -245,6 +247,22 @@ impl Thread {
             request.messages.push(request_message);
         }
 
+        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))
+                .cloned();
+            attach_context_to_message(&mut context_message, referenced_context);
+
+            request.messages.push(context_message);
+        }
+
         request
     }