From 276d629a140c084b56317117b9f6b715d40314fd Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Thu, 19 Sep 2024 03:49:35 +0200
Subject: [PATCH] fix: message delete

---
 src/lib/components/chat/Chat.svelte     |  1 +
 src/lib/components/chat/Messages.svelte | 53 ++++++++++---------------
 2 files changed, 23 insertions(+), 31 deletions(-)

diff --git a/src/lib/components/chat/Chat.svelte b/src/lib/components/chat/Chat.svelte
index 0a9483d3b..5bd700712 100644
--- a/src/lib/components/chat/Chat.svelte
+++ b/src/lib/components/chat/Chat.svelte
@@ -1787,6 +1787,7 @@
 						{regenerateResponse}
 						{mergeResponses}
 						{chatActionHandler}
+						{showMessage}
 					/>
 				</div>
 			</div>
diff --git a/src/lib/components/chat/Messages.svelte b/src/lib/components/chat/Messages.svelte
index 3080bcac9..791723d96 100644
--- a/src/lib/components/chat/Messages.svelte
+++ b/src/lib/components/chat/Messages.svelte
@@ -21,6 +21,7 @@
 	export let regenerateResponse: Function;
 	export let mergeResponses: Function;
 	export let chatActionHandler: Function;
+	export let showMessage: Function = () => {};
 
 	export let user = $_user;
 	export let prompt;
@@ -244,49 +245,39 @@
 
 	const deleteMessageHandler = async (messageId) => {
 		const messageToDelete = history.messages[messageId];
-
 		const parentMessageId = messageToDelete.parentId;
 		const childMessageIds = messageToDelete.childrenIds ?? [];
 
-		const hasDescendantMessages = childMessageIds.some(
-			(childId) => history.messages[childId]?.childrenIds?.length > 0
+		// Collect all grandchildren
+		const grandchildrenIds = childMessageIds.flatMap(
+			(childId) => history.messages[childId]?.childrenIds ?? []
 		);
 
-		history.currentId = parentMessageId;
-		await tick();
+		// Update parent's children
+		if (parentMessageId && history.messages[parentMessageId]) {
+			history.messages[parentMessageId].childrenIds = [
+				...history.messages[parentMessageId].childrenIds.filter((id) => id !== messageId),
+				...grandchildrenIds
+			];
+		}
 
-		// Remove the message itself from the parent message's children array
-		history.messages[parentMessageId].childrenIds = history.messages[
-			parentMessageId
-		].childrenIds.filter((id) => id !== messageId);
-
-		await tick();
-
-		childMessageIds.forEach((childId) => {
-			const childMessage = history.messages[childId];
-
-			if (childMessage && childMessage.childrenIds) {
-				if (childMessage.childrenIds.length === 0 && !hasDescendantMessages) {
-					// If there are no other responses/prompts
-					history.messages[parentMessageId].childrenIds = [];
-				} else {
-					childMessage.childrenIds.forEach((grandChildId) => {
-						if (history.messages[grandChildId]) {
-							history.messages[grandChildId].parentId = parentMessageId;
-							history.messages[parentMessageId].childrenIds.push(grandChildId);
-						}
-					});
-				}
+		// Update grandchildren's parent
+		grandchildrenIds.forEach((grandchildId) => {
+			if (history.messages[grandchildId]) {
+				history.messages[grandchildId].parentId = parentMessageId;
 			}
+		});
 
-			// Remove child message id from the parent message's children array
-			history.messages[parentMessageId].childrenIds = history.messages[
-				parentMessageId
-			].childrenIds.filter((id) => id !== childId);
+		// Delete the message and its children
+		[messageId, ...childMessageIds].forEach((id) => {
+			delete history.messages[id];
 		});
 
 		await tick();
 
+		showMessage({ id: parentMessageId });
+
+		// Update the chat
 		await updateChatById(localStorage.token, chatId, {
 			messages: messages,
 			history: history