From f431f431911d60645f4056b8a3cd74f44ad6c05e Mon Sep 17 00:00:00 2001
From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com>
Date: Sat, 7 Mar 2026 04:29:52 +0530
Subject: [PATCH] feat: enhance message handling in NewChatPage by updating
 message IDs upon saving, ensuring consistent state management and improved
 user experience

---
 .../new-chat/[[...chat_id]]/page.tsx                 | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx
index 47a2b69e6..1df131cd3 100644
--- a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx
+++ b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx
@@ -557,7 +557,17 @@ export default function NewChatPage() {
 				role: "user",
 				content: persistContent,
 			})
-				.then(() => {
+				.then((savedMessage) => {
+					const newUserMsgId = `msg-${savedMessage.id}`;
+					setMessages((prev) =>
+						prev.map((m) => (m.id === userMsgId ? { ...m, id: newUserMsgId } : m))
+					);
+					setMessageDocumentsMap((prev) => {
+						const docs = prev[userMsgId];
+						if (!docs) return prev;
+						const { [userMsgId]: _, ...rest } = prev;
+						return { ...rest, [newUserMsgId]: docs };
+					});
 					if (isNewThread) {
 						queryClient.invalidateQueries({ queryKey: ["threads", String(searchSpaceId)] });
 					}