From 8232f628eaae24b2f850371f823ec0a54d4c368e Mon Sep 17 00:00:00 2001
From: dal <dallin@buster.so>
Date: Mon, 3 Mar 2025 16:43:44 -0700
Subject: [PATCH] add in initialize

---
 api/libs/handlers/src/chats/post_chat_handler.rs  | 15 +++++++++++++--
 api/libs/handlers/src/chats/streaming_parser.rs   |  4 ++--
 api/libs/handlers/src/chats/types.rs              |  2 +-
 api/libs/handlers/src/messages/types.rs           |  6 +++---
 .../routes/ws/threads_and_messages/post_thread.rs |  3 +++
 5 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/api/libs/handlers/src/chats/post_chat_handler.rs b/api/libs/handlers/src/chats/post_chat_handler.rs
index 7fe2506d2..4e27ae52d 100644
--- a/api/libs/handlers/src/chats/post_chat_handler.rs
+++ b/api/libs/handlers/src/chats/post_chat_handler.rs
@@ -36,6 +36,7 @@ pub enum ThreadEvent {
     GeneratingResponseMessage,
     GeneratingReasoningMessage,
     GeneratingTitle,
+    InitializeChat,
 }
 
 #[derive(Debug, Deserialize, Clone)]
@@ -101,6 +102,15 @@ pub async fn post_chat_handler(
         created_by_avatar: None,
     };
 
+    // Send initial chat state to client
+    if let Some(tx) = tx.clone() {
+        tx.send(Ok((
+            BusterContainer::Chat(chat_with_messages.clone()),
+            ThreadEvent::InitializeChat,
+        )))
+        .await?;
+    }
+
     // Create database connection
     let mut conn = get_pg_pool().get().await?;
 
@@ -439,7 +449,7 @@ pub struct BusterReasoningText {
     #[serde(rename = "type")]
     pub reasoning_type: String,
     pub title: String,
-    pub secondary_title: Option<String>,
+    pub secondary_title: String,
     pub message: Option<String>,
     pub message_chunk: Option<String>,
     pub status: Option<String>,
@@ -491,6 +501,7 @@ pub struct BusterFileMetadata {
 #[derive(Debug, Serialize, Clone)]
 #[serde(untagged)]
 pub enum BusterContainer {
+    Chat(ChatWithMessages),
     ChatMessage(BusterChatMessageContainer),
     ReasoningMessage(BusterReasoningMessageContainer),
     GeneratingTitle(BusterGeneratingTitle),
@@ -1170,7 +1181,7 @@ fn assistant_create_plan(
                         id: text.id,
                         thought_type: "thought".to_string(),
                         title: text.title,
-                        secondary_title: text.secondary_title.unwrap_or_default(),
+                        secondary_title: text.secondary_title,
                         pill_containers: Some(vec![plan_container]),
                         status: status.to_string(),
                     })])
diff --git a/api/libs/handlers/src/chats/streaming_parser.rs b/api/libs/handlers/src/chats/streaming_parser.rs
index c1f1a4854..365d22c86 100644
--- a/api/libs/handlers/src/chats/streaming_parser.rs
+++ b/api/libs/handlers/src/chats/streaming_parser.rs
@@ -50,7 +50,7 @@ impl StreamingParser {
                     id,
                     reasoning_type: "text".to_string(),
                     title: "Creating a plan...".to_string(),
-                    secondary_title: None,
+                    secondary_title: String::from(""),
                     message: None,
                     message_chunk: Some(plan_markdown.to_string()),
                     status: Some("loading".to_string()),
@@ -85,7 +85,7 @@ impl StreamingParser {
                     id,
                     reasoning_type: "text".to_string(),
                     title: "Searching your data catalog...".to_string(),
-                    secondary_title: None,
+                    secondary_title: String::from(""),
                     message: None,
                     message_chunk: Some(search_requirements.to_string()),
                     status: Some("loading".to_string()),
diff --git a/api/libs/handlers/src/chats/types.rs b/api/libs/handlers/src/chats/types.rs
index b17ec1389..6d689342b 100644
--- a/api/libs/handlers/src/chats/types.rs
+++ b/api/libs/handlers/src/chats/types.rs
@@ -3,7 +3,7 @@ use uuid::Uuid;
 
 use crate::messages::types::ChatMessage;
 
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Serialize, Deserialize, Clone)]
 pub struct ChatWithMessages {
     pub id: Uuid,
     pub title: String,
diff --git a/api/libs/handlers/src/messages/types.rs b/api/libs/handlers/src/messages/types.rs
index 6978756a5..c11216055 100644
--- a/api/libs/handlers/src/messages/types.rs
+++ b/api/libs/handlers/src/messages/types.rs
@@ -2,7 +2,7 @@ use serde::{Deserialize, Serialize};
 use serde_json::Value;
 use uuid::Uuid;
 
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Serialize, Deserialize, Clone)]
 pub struct ChatMessage {
     pub id: Uuid,
     pub request_message: ChatUserMessage,
@@ -11,10 +11,10 @@ pub struct ChatMessage {
     pub created_at: String,
 }
 
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Serialize, Deserialize, Clone)]
 pub struct ChatUserMessage {
     pub request: String,
     pub sender_id: Uuid,
     pub sender_name: String,
     pub sender_avatar: Option<String>,
-}
\ No newline at end of file
+}
diff --git a/api/src/routes/ws/threads_and_messages/post_thread.rs b/api/src/routes/ws/threads_and_messages/post_thread.rs
index 0be297cfb..92d9015fe 100644
--- a/api/src/routes/ws/threads_and_messages/post_thread.rs
+++ b/api/src/routes/ws/threads_and_messages/post_thread.rs
@@ -43,6 +43,9 @@ pub async fn post_thread(
                         ThreadEvent::GeneratingTitle => {
                             WsEvent::Threads(WSThreadEvent::GeneratingTitle)
                         }
+                        ThreadEvent::InitializeChat => {
+                            WsEvent::Threads(WSThreadEvent::InitializeChat)
+                        }
                     };
 
                     let response = WsResponseMessage::new_no_user(