From 0d6a549950d02a33cd7580c056baee24f036ea18 Mon Sep 17 00:00:00 2001
From: Michael Sloan <michael@zed.dev>
Date: Thu, 9 Jan 2025 21:40:11 -0700
Subject: [PATCH] assistant2: More improvement to prompt building efficiency
 (#22941)

Release Notes:

- N/A
---
 crates/assistant2/src/context.rs       | 76 ++++++++++++++++----------
 crates/assistant2/src/context_store.rs | 13 ++---
 2 files changed, 54 insertions(+), 35 deletions(-)

diff --git a/crates/assistant2/src/context.rs b/crates/assistant2/src/context.rs
index cb33e6a9c1..aad1805b9a 100644
--- a/crates/assistant2/src/context.rs
+++ b/crates/assistant2/src/context.rs
@@ -31,7 +31,7 @@ pub struct ContextSnapshot {
     pub tooltip: Option<SharedString>,
     pub icon_path: Option<SharedString>,
     pub kind: ContextKind,
-    /// Concatenating these strings yields text to send to the model. Not refreshed by `snapshot`.
+    /// Joining these strings separated by \n yields text for model. Not refreshed by `snapshot`.
     pub text: Box<[SharedString]>,
 }
 
@@ -210,7 +210,9 @@ pub fn attach_context_to_message(
     let mut fetch_context = Vec::new();
     let mut thread_context = Vec::new();
 
+    let mut capacity = 0;
     for context in contexts {
+        capacity += context.text.len();
         match context.kind {
             ContextKind::File => file_context.push(context),
             ContextKind::Directory => directory_context.push(context),
@@ -218,54 +220,72 @@ pub fn attach_context_to_message(
             ContextKind::Thread => thread_context.push(context),
         }
     }
+    if !file_context.is_empty() {
+        capacity += 1;
+    }
+    if !directory_context.is_empty() {
+        capacity += 1;
+    }
+    if !fetch_context.is_empty() {
+        capacity += 1 + fetch_context.len();
+    }
+    if !thread_context.is_empty() {
+        capacity += 1 + thread_context.len();
+    }
+    if capacity == 0 {
+        return;
+    }
 
-    let mut context_text = String::new();
+    let mut context_chunks = Vec::with_capacity(capacity);
 
     if !file_context.is_empty() {
-        context_text.push_str("The following files are available:\n");
-        for context in file_context {
-            for chunk in context.text {
-                context_text.push_str(&chunk);
+        context_chunks.push("The following files are available:\n");
+        for context in &file_context {
+            for chunk in &context.text {
+                context_chunks.push(&chunk);
             }
-            context_text.push('\n');
         }
     }
 
     if !directory_context.is_empty() {
-        context_text.push_str("The following directories are available:\n");
-        for context in directory_context {
-            for chunk in context.text {
-                context_text.push_str(&chunk);
+        context_chunks.push("The following directories are available:\n");
+        for context in &directory_context {
+            for chunk in &context.text {
+                context_chunks.push(&chunk);
             }
-            context_text.push('\n');
         }
     }
 
     if !fetch_context.is_empty() {
-        context_text.push_str("The following fetched results are available\n");
-        for context in fetch_context {
-            context_text.push_str(&context.name);
-            context_text.push('\n');
-            for chunk in context.text {
-                context_text.push_str(&chunk);
+        context_chunks.push("The following fetched results are available:\n");
+        for context in &fetch_context {
+            context_chunks.push(&context.name);
+            for chunk in &context.text {
+                context_chunks.push(&chunk);
             }
-            context_text.push('\n');
         }
     }
 
     if !thread_context.is_empty() {
-        context_text.push_str("The following previous conversation threads are available\n");
-        for context in thread_context {
-            context_text.push_str(&context.name);
-            context_text.push('\n');
-            for chunk in context.text {
-                context_text.push_str(&chunk);
+        context_chunks.push("The following previous conversation threads are available:\n");
+        for context in &thread_context {
+            context_chunks.push(&context.name);
+            for chunk in &context.text {
+                context_chunks.push(&chunk);
             }
-            context_text.push('\n');
         }
     }
 
-    if !context_text.is_empty() {
-        message.content.push(MessageContent::Text(context_text));
+    debug_assert!(
+        context_chunks.len() == capacity,
+        "attach_context_message calculated capacity of {}, but length was {}",
+        capacity,
+        context_chunks.len()
+    );
+
+    if !context_chunks.is_empty() {
+        message
+            .content
+            .push(MessageContent::Text(context_chunks.join("\n")));
     }
 }
diff --git a/crates/assistant2/src/context_store.rs b/crates/assistant2/src/context_store.rs
index 49518a1296..72f66689c9 100644
--- a/crates/assistant2/src/context_store.rs
+++ b/crates/assistant2/src/context_store.rs
@@ -462,13 +462,12 @@ fn to_fenced_codeblock(path: &Path, content: Rope) -> SharedString {
 
     buffer.push_str("```\n");
 
-    if buffer.len() > capacity {
-        log::error!(
-            "to_fenced_codeblock calculated capacity {} but length was {}",
-            capacity,
-            buffer.len()
-        );
-    }
+    debug_assert!(
+        buffer.len() == capacity - 1 || buffer.len() == capacity,
+        "to_fenced_codeblock calculated capacity of {}, but length was {}",
+        capacity,
+        buffer.len(),
+    );
 
     buffer.into()
 }