assistant: Refesh message headers only for dirty messages (#16881)

Piotr Osiewicz created

We've noticed performance issues in long conversations with assistants;
the profiles pointed to slowiness in WrapMap (and indeed there were some
low hanging fruits that we picked up in
https://github.com/zed-industries/zed/pull/16761). That however did not
fully resolve the issue, as WrapMap still cracked through in profiles;
basically, the speedup I've landed has just moved the post elsewhere.

The higher level issue is that we were trying to refresh message headers
for all messages, irrespective of whether they've actually needed to be
updated. This PR fixes that by using `replace_blocks` API where
possible.

Release Notes:

- Improved performance of Assistant Panel with long conversations.

Change summary

crates/assistant/src/assistant_panel.rs | 346 ++++++++++++++------------
crates/assistant/src/context.rs         |  13 
2 files changed, 203 insertions(+), 156 deletions(-)

Detailed changes

crates/assistant/src/assistant_panel.rs 🔗

@@ -13,9 +13,10 @@ use crate::{
     terminal_inline_assistant::TerminalInlineAssistant,
     Assist, CacheStatus, ConfirmCommand, Context, ContextEvent, ContextId, ContextStore,
     CycleMessageRole, DeployHistory, DeployPromptLibrary, InlineAssistId, InlineAssistant,
-    InsertIntoEditor, MessageStatus, ModelSelector, PendingSlashCommand, PendingSlashCommandStatus,
-    QuoteSelection, RemoteContextMetadata, SavedContextMetadata, Split, ToggleFocus,
-    ToggleModelSelector, WorkflowStepResolution, WorkflowStepView,
+    InsertIntoEditor, Message, MessageId, MessageMetadata, MessageStatus, ModelSelector,
+    PendingSlashCommand, PendingSlashCommandStatus, QuoteSelection, RemoteContextMetadata,
+    SavedContextMetadata, Split, ToggleFocus, ToggleModelSelector, WorkflowStepResolution,
+    WorkflowStepView,
 };
 use crate::{ContextStoreEvent, ModelPickerDelegate};
 use anyhow::{anyhow, Result};
@@ -1702,6 +1703,8 @@ struct WorkflowAssist {
     assist_ids: Vec<InlineAssistId>,
 }
 
+type MessageHeader = MessageMetadata;
+
 pub struct ContextEditor {
     context: Model<Context>,
     fs: Arc<dyn Fs>,
@@ -1709,7 +1712,7 @@ pub struct ContextEditor {
     project: Model<Project>,
     lsp_adapter_delegate: Option<Arc<dyn LspAdapterDelegate>>,
     editor: View<Editor>,
-    blocks: HashSet<CustomBlockId>,
+    blocks: HashMap<MessageId, (MessageHeader, CustomBlockId)>,
     image_blocks: HashSet<CustomBlockId>,
     scroll_position: Option<ScrollPosition>,
     remote_id: Option<workspace::ViewId>,
@@ -3036,176 +3039,209 @@ impl ContextEditor {
     fn update_message_headers(&mut self, cx: &mut ViewContext<Self>) {
         self.editor.update(cx, |editor, cx| {
             let buffer = editor.buffer().read(cx).snapshot(cx);
-            let excerpt_id = *buffer.as_singleton().unwrap().0;
-            let old_blocks = std::mem::take(&mut self.blocks);
-            let new_blocks = self
-                .context
-                .read(cx)
-                .messages(cx)
-                .map(|message| BlockProperties {
-                    position: buffer
-                        .anchor_in_excerpt(excerpt_id, message.anchor)
-                        .unwrap(),
-                    height: 2,
-                    style: BlockStyle::Sticky,
-                    render: Box::new({
-                        let context = self.context.clone();
-                        move |cx| {
-                            let message_id = message.id;
-                            let show_spinner = message.role == Role::Assistant
-                                && message.status == MessageStatus::Pending;
 
-                            let label = match message.role {
-                                Role::User => {
-                                    Label::new("You").color(Color::Default).into_any_element()
-                                }
-                                Role::Assistant => {
-                                    let label = Label::new("Assistant").color(Color::Info);
-                                    if show_spinner {
-                                        label
-                                            .with_animation(
-                                                "pulsating-label",
-                                                Animation::new(Duration::from_secs(2))
-                                                    .repeat()
-                                                    .with_easing(pulsating_between(0.4, 0.8)),
-                                                |label, delta| label.alpha(delta),
-                                            )
-                                            .into_any_element()
-                                    } else {
-                                        label.into_any_element()
-                                    }
+            let excerpt_id = *buffer.as_singleton().unwrap().0;
+            let mut old_blocks = std::mem::take(&mut self.blocks);
+            let mut blocks_to_remove: HashMap<_, _> = old_blocks
+                .iter()
+                .map(|(message_id, (_, block_id))| (*message_id, *block_id))
+                .collect();
+            let mut blocks_to_replace: HashMap<_, RenderBlock> = Default::default();
+
+            let render_block = |message: MessageMetadata| -> RenderBlock {
+                Box::new({
+                    let context = self.context.clone();
+                    move |cx| {
+                        let message_id = MessageId(message.timestamp);
+                        let show_spinner = message.role == Role::Assistant
+                            && message.status == MessageStatus::Pending;
+
+                        let label = match message.role {
+                            Role::User => {
+                                Label::new("You").color(Color::Default).into_any_element()
+                            }
+                            Role::Assistant => {
+                                let label = Label::new("Assistant").color(Color::Info);
+                                if show_spinner {
+                                    label
+                                        .with_animation(
+                                            "pulsating-label",
+                                            Animation::new(Duration::from_secs(2))
+                                                .repeat()
+                                                .with_easing(pulsating_between(0.4, 0.8)),
+                                            |label, delta| label.alpha(delta),
+                                        )
+                                        .into_any_element()
+                                } else {
+                                    label.into_any_element()
                                 }
+                            }
 
-                                Role::System => Label::new("System")
-                                    .color(Color::Warning)
-                                    .into_any_element(),
-                            };
+                            Role::System => Label::new("System")
+                                .color(Color::Warning)
+                                .into_any_element(),
+                        };
 
-                            let sender = ButtonLike::new("role")
-                                .style(ButtonStyle::Filled)
-                                .child(label)
-                                .tooltip(|cx| {
-                                    Tooltip::with_meta(
-                                        "Toggle message role",
-                                        None,
-                                        "Available roles: You (User), Assistant, System",
-                                        cx,
-                                    )
-                                })
-                                .on_click({
-                                    let context = context.clone();
-                                    move |_, cx| {
-                                        context.update(cx, |context, cx| {
-                                            context.cycle_message_roles(
-                                                HashSet::from_iter(Some(message_id)),
-                                                cx,
-                                            )
-                                        })
-                                    }
-                                });
+                        let sender = ButtonLike::new("role")
+                            .style(ButtonStyle::Filled)
+                            .child(label)
+                            .tooltip(|cx| {
+                                Tooltip::with_meta(
+                                    "Toggle message role",
+                                    None,
+                                    "Available roles: You (User), Assistant, System",
+                                    cx,
+                                )
+                            })
+                            .on_click({
+                                let context = context.clone();
+                                move |_, cx| {
+                                    context.update(cx, |context, cx| {
+                                        context.cycle_message_roles(
+                                            HashSet::from_iter(Some(message_id)),
+                                            cx,
+                                        )
+                                    })
+                                }
+                            });
 
-                            h_flex()
-                                .id(("message_header", message_id.as_u64()))
-                                .pl(cx.gutter_dimensions.full_width())
-                                .h_11()
-                                .w_full()
-                                .relative()
-                                .gap_1()
-                                .child(sender)
-                                .children(match &message.cache {
-                                    Some(cache) if cache.is_final_anchor => match cache.status {
-                                        CacheStatus::Cached => Some(
-                                            div()
-                                                .id("cached")
-                                                .child(
-                                                    Icon::new(IconName::DatabaseZap)
-                                                        .size(IconSize::XSmall)
-                                                        .color(Color::Hint),
-                                                )
-                                                .tooltip(|cx| {
-                                                    Tooltip::with_meta(
-                                                        "Context cached",
-                                                        None,
-                                                        "Large messages cached to optimize performance",
-                                                        cx,
-                                                    )
-                                                }).into_any_element()
-                                        ),
-                                        CacheStatus::Pending => Some(
-                                            div()
-                                                .child(
-                                                    Icon::new(IconName::Ellipsis)
-                                                        .size(IconSize::XSmall)
-                                                        .color(Color::Hint),
-                                                ).into_any_element()
-                                        ),
-                                    },
-                                    _ => None,
-                                })
-                                .children(match &message.status {
-                                    MessageStatus::Error(error) => Some(
-                                        Button::new("show-error", "Error")
-                                            .color(Color::Error)
-                                            .selected_label_color(Color::Error)
-                                            .selected_icon_color(Color::Error)
-                                            .icon(IconName::XCircle)
-                                            .icon_color(Color::Error)
-                                            .icon_size(IconSize::Small)
-                                            .icon_position(IconPosition::Start)
-                                            .tooltip(move |cx| {
+                        h_flex()
+                            .id(("message_header", message_id.as_u64()))
+                            .pl(cx.gutter_dimensions.full_width())
+                            .h_11()
+                            .w_full()
+                            .relative()
+                            .gap_1()
+                            .child(sender)
+                            .children(match &message.cache {
+                                Some(cache) if cache.is_final_anchor => match cache.status {
+                                    CacheStatus::Cached => Some(
+                                        div()
+                                            .id("cached")
+                                            .child(
+                                                Icon::new(IconName::DatabaseZap)
+                                                    .size(IconSize::XSmall)
+                                                    .color(Color::Hint),
+                                            )
+                                            .tooltip(|cx| {
                                                 Tooltip::with_meta(
-                                                    "Error interacting with language model",
+                                                    "Context cached",
                                                     None,
-                                                    "Click for more details",
+                                                    "Large messages cached to optimize performance",
                                                     cx,
                                                 )
                                             })
-                                            .on_click({
-                                                let context = context.clone();
-                                                let error = error.clone();
-                                                move |_, cx| {
-                                                    context.update(cx, |_, cx| {
-                                                        cx.emit(ContextEvent::ShowAssistError(
-                                                            error.clone(),
-                                                        ));
-                                                    });
-                                                }
-                                            })
                                             .into_any_element(),
                                     ),
-                                    MessageStatus::Canceled => Some(
-                                        ButtonLike::new("canceled")
-                                            .child(
-                                                Icon::new(IconName::XCircle).color(Color::Disabled),
-                                            )
+                                    CacheStatus::Pending => Some(
+                                        div()
                                             .child(
-                                                Label::new("Canceled")
-                                                    .size(LabelSize::Small)
-                                                    .color(Color::Disabled),
+                                                Icon::new(IconName::Ellipsis)
+                                                    .size(IconSize::XSmall)
+                                                    .color(Color::Hint),
                                             )
-                                            .tooltip(move |cx| {
-                                                Tooltip::with_meta(
-                                                    "Canceled",
-                                                    None,
-                                                    "Interaction with the assistant was canceled",
-                                                    cx,
-                                                )
-                                            })
                                             .into_any_element(),
                                     ),
-                                    _ => None,
-                                })
-                                .into_any_element()
-                        }
-                    }),
-                    disposition: BlockDisposition::Above,
-                    priority: usize::MAX,
+                                },
+                                _ => None,
+                            })
+                            .children(match &message.status {
+                                MessageStatus::Error(error) => Some(
+                                    Button::new("show-error", "Error")
+                                        .color(Color::Error)
+                                        .selected_label_color(Color::Error)
+                                        .selected_icon_color(Color::Error)
+                                        .icon(IconName::XCircle)
+                                        .icon_color(Color::Error)
+                                        .icon_size(IconSize::Small)
+                                        .icon_position(IconPosition::Start)
+                                        .tooltip(move |cx| {
+                                            Tooltip::with_meta(
+                                                "Error interacting with language model",
+                                                None,
+                                                "Click for more details",
+                                                cx,
+                                            )
+                                        })
+                                        .on_click({
+                                            let context = context.clone();
+                                            let error = error.clone();
+                                            move |_, cx| {
+                                                context.update(cx, |_, cx| {
+                                                    cx.emit(ContextEvent::ShowAssistError(
+                                                        error.clone(),
+                                                    ));
+                                                });
+                                            }
+                                        })
+                                        .into_any_element(),
+                                ),
+                                MessageStatus::Canceled => Some(
+                                    ButtonLike::new("canceled")
+                                        .child(Icon::new(IconName::XCircle).color(Color::Disabled))
+                                        .child(
+                                            Label::new("Canceled")
+                                                .size(LabelSize::Small)
+                                                .color(Color::Disabled),
+                                        )
+                                        .tooltip(move |cx| {
+                                            Tooltip::with_meta(
+                                                "Canceled",
+                                                None,
+                                                "Interaction with the assistant was canceled",
+                                                cx,
+                                            )
+                                        })
+                                        .into_any_element(),
+                                ),
+                                _ => None,
+                            })
+                            .into_any_element()
+                    }
                 })
-                .collect::<Vec<_>>();
+            };
+            let create_block_properties = |message: &Message| BlockProperties {
+                position: buffer
+                    .anchor_in_excerpt(excerpt_id, message.anchor)
+                    .unwrap(),
+                height: 2,
+                style: BlockStyle::Sticky,
+                disposition: BlockDisposition::Above,
+                priority: usize::MAX,
+                render: render_block(MessageMetadata::from(message)),
+            };
+            let mut new_blocks = vec![];
+            let mut block_index_to_message = vec![];
+            for message in self.context.read(cx).messages(cx) {
+                if let Some(_) = blocks_to_remove.remove(&message.id) {
+                    // This is an old message that we might modify.
+                    let Some((meta, block_id)) = old_blocks.get_mut(&message.id) else {
+                        debug_assert!(
+                            false,
+                            "old_blocks should contain a message_id we've just removed."
+                        );
+                        continue;
+                    };
+                    // Should we modify it?
+                    let message_meta = MessageMetadata::from(&message);
+                    if meta != &message_meta {
+                        blocks_to_replace.insert(*block_id, render_block(message_meta.clone()));
+                        *meta = message_meta;
+                    }
+                } else {
+                    // This is a new message.
+                    new_blocks.push(create_block_properties(&message));
+                    block_index_to_message.push((message.id, MessageMetadata::from(&message)));
+                }
+            }
+            editor.replace_blocks(blocks_to_replace, None, cx);
+            editor.remove_blocks(blocks_to_remove.into_values().collect(), None, cx);
 
-            editor.remove_blocks(old_blocks, None, cx);
             let ids = editor.insert_blocks(new_blocks, None, cx);
-            self.blocks = HashSet::from_iter(ids);
+            old_blocks.extend(ids.into_iter().zip(block_index_to_message).map(
+                |(block_id, (message_id, message_meta))| (message_id, (message_meta, block_id)),
+            ));
+            self.blocks = old_blocks;
         });
     }
 

crates/assistant/src/context.rs 🔗

@@ -330,11 +330,22 @@ pub struct MessageCacheMetadata {
 pub struct MessageMetadata {
     pub role: Role,
     pub status: MessageStatus,
-    timestamp: clock::Lamport,
+    pub(crate) timestamp: clock::Lamport,
     #[serde(skip)]
     pub cache: Option<MessageCacheMetadata>,
 }
 
+impl From<&Message> for MessageMetadata {
+    fn from(message: &Message) -> Self {
+        Self {
+            role: message.role,
+            status: message.status.clone(),
+            timestamp: message.id.0,
+            cache: message.cache.clone(),
+        }
+    }
+}
+
 impl MessageMetadata {
     pub fn is_cache_valid(&self, buffer: &BufferSnapshot, range: &Range<usize>) -> bool {
         let result = match &self.cache {