From 04cd8dd0f2f88f73bb0416f67c1b7f11cf08166f Mon Sep 17 00:00:00 2001
From: Marshall Bowers <elliott.codes@gmail.com>
Date: Mon, 29 Apr 2024 19:25:58 -0400
Subject: [PATCH] assistant2: Add the ability to collapse chat messages
 (#11194)

This PR adds the ability to collapse/uncollapse chat messages.

I think the spacing might be a little off with the collapsed
calculations, so we'll need to look into that.


https://github.com/zed-industries/zed/assets/1486634/4009c831-b44e-4b30-85ed-0266cb5b8a26

Release Notes:

- N/A
---
 Cargo.lock                               |  1 +
 crates/assistant2/Cargo.toml             |  1 +
 crates/assistant2/src/assistant2.rs      | 32 ++++++++++++++++++++----
 crates/assistant2/src/ui/chat_message.rs | 10 ++++----
 4 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index adcc694c9c..f4ff1ee348 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -379,6 +379,7 @@ dependencies = [
  "assets",
  "assistant_tooling",
  "client",
+ "collections",
  "editor",
  "env_logger",
  "feature_flags",
diff --git a/crates/assistant2/Cargo.toml b/crates/assistant2/Cargo.toml
index f79ff8d97a..f6384c716b 100644
--- a/crates/assistant2/Cargo.toml
+++ b/crates/assistant2/Cargo.toml
@@ -12,6 +12,7 @@ path = "src/assistant2.rs"
 anyhow.workspace = true
 assistant_tooling.workspace = true
 client.workspace = true
+collections.workspace = true
 editor.workspace = true
 feature_flags.workspace = true
 fs.workspace = true
diff --git a/crates/assistant2/src/assistant2.rs b/crates/assistant2/src/assistant2.rs
index 5e61097d9e..3ee03f6d29 100644
--- a/crates/assistant2/src/assistant2.rs
+++ b/crates/assistant2/src/assistant2.rs
@@ -7,6 +7,7 @@ use ::ui::{div, prelude::*, Color, ViewContext};
 use anyhow::{Context, Result};
 use assistant_tooling::{ToolFunctionCall, ToolRegistry};
 use client::{proto, Client, UserStore};
+use collections::HashMap;
 use completion_provider::*;
 use editor::Editor;
 use feature_flags::FeatureFlagAppExt as _;
@@ -214,6 +215,7 @@ struct AssistantChat {
     composer_editor: View<Editor>,
     user_store: Model<UserStore>,
     next_message_id: MessageId,
+    collapsed_messages: HashMap<MessageId, bool>,
     pending_completion: Option<Task<()>>,
     tool_registry: Arc<ToolRegistry>,
 }
@@ -250,6 +252,7 @@ impl AssistantChat {
             user_store,
             language_registry,
             next_message_id: MessageId(0),
+            collapsed_messages: HashMap::default(),
             pending_completion: None,
             tool_registry,
         }
@@ -496,6 +499,15 @@ impl AssistantChat {
         }
     }
 
+    fn is_message_collapsed(&self, id: &MessageId) -> bool {
+        self.collapsed_messages.get(id).copied().unwrap_or_default()
+    }
+
+    fn toggle_message_collapsed(&mut self, id: MessageId) {
+        let entry = self.collapsed_messages.entry(id).or_insert(false);
+        *entry = !*entry;
+    }
+
     fn render_error(
         &self,
         error: Option<SharedString>,
@@ -531,8 +543,13 @@ impl AssistantChat {
                     *id,
                     UserOrAssistant::User(self.user_store.read(cx).current_user()),
                     body.clone().into_any_element(),
-                    false,
-                    Box::new(|_, _| {}),
+                    self.is_message_collapsed(id),
+                    Box::new(cx.listener({
+                        let id = *id;
+                        move |assistant_chat, _event, _cx| {
+                            assistant_chat.toggle_message_collapsed(id)
+                        }
+                    })),
                 ))
                 .into_any(),
             ChatMessage::Assistant(AssistantMessage {
@@ -554,8 +571,13 @@ impl AssistantChat {
                         *id,
                         UserOrAssistant::Assistant,
                         assistant_body.into_any_element(),
-                        false,
-                        Box::new(|_, _| {}),
+                        self.is_message_collapsed(id),
+                        Box::new(cx.listener({
+                            let id = *id;
+                            move |assistant_chat, _event, _cx| {
+                                assistant_chat.toggle_message_collapsed(id)
+                            }
+                        })),
                     ))
                     // TODO: Should the errors and tool calls get passed into `ChatMessage`?
                     .child(self.render_error(error.clone(), ix, cx))
@@ -665,7 +687,7 @@ impl Render for AssistantChat {
     }
 }
 
-#[derive(Debug, Copy, Clone, Eq, PartialEq)]
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy)]
 struct MessageId(usize);
 
 impl MessageId {
diff --git a/crates/assistant2/src/ui/chat_message.rs b/crates/assistant2/src/ui/chat_message.rs
index 8dde0bb500..d51f62aba5 100644
--- a/crates/assistant2/src/ui/chat_message.rs
+++ b/crates/assistant2/src/ui/chat_message.rs
@@ -1,7 +1,7 @@
 use std::sync::Arc;
 
 use client::User;
-use gpui::AnyElement;
+use gpui::{AnyElement, ClickEvent};
 use ui::{prelude::*, Avatar};
 
 use crate::MessageId;
@@ -17,7 +17,7 @@ pub struct ChatMessage {
     player: UserOrAssistant,
     message: AnyElement,
     collapsed: bool,
-    on_collapse: Box<dyn Fn(bool, &mut WindowContext) + 'static>,
+    on_collapse_handle_click: Box<dyn Fn(&ClickEvent, &mut WindowContext) + 'static>,
 }
 
 impl ChatMessage {
@@ -26,14 +26,14 @@ impl ChatMessage {
         player: UserOrAssistant,
         message: AnyElement,
         collapsed: bool,
-        on_collapse: Box<dyn Fn(bool, &mut WindowContext) + 'static>,
+        on_collapse_handle_click: Box<dyn Fn(&ClickEvent, &mut WindowContext) + 'static>,
     ) -> Self {
         Self {
             id,
             player,
             message,
             collapsed,
-            on_collapse,
+            on_collapse_handle_click,
         }
     }
 }
@@ -53,7 +53,7 @@ impl RenderOnce for ChatMessage {
             .w_1()
             .mx_2()
             .h_full()
-            .on_click(move |_event, cx| (self.on_collapse)(!self.collapsed, cx))
+            .on_click(self.on_collapse_handle_click)
             .child(
                 div()
                     .w_px()