Update server/wsHandler.js

server/wsHandler.js

@@ -48,6 +48,7 @@ import { pendingTurnstileTokens } from './index.js';
 
 const activeStreams = new Map();
 const VERSION_META_FIELDS = ['toolCalls', 'responseEdits', 'responseSegments', 'error'];
+const ROOT_JSON_KEY = '__historyRootJson';
 const CONTINUE_ASSISTANT_PROMPT =
   'Continue your previous response exactly where it left off. Do not restart, summarize, or repeat the opening. Preserve the same formatting and only add the missing continuation.';
 const FREE_WEB_SEARCH_LIMIT = 15;
@@ -454,7 +455,13 @@ const handlers = {
           await sessionStore.updateUserSession(client.userId, client.accessToken, sessionId, { history: newHistory, name: newName });
         else sessionStore.updateTempSession(client.tempId, sessionId, { history: newHistory, name: newName });
 
-        safeSend(ws, { type: aborted ? 'chat:aborted' : 'chat:done', sessionId, name: newName, history: extractFlatHistory(newRootMessage) });
+        safeSend(ws, {
+          type: aborted ? 'chat:aborted' : 'chat:done',
+          sessionId,
+          name: newName,
+          history: extractFlatHistory(newRootMessage),
+          [ROOT_JSON_KEY]: JSON.stringify(newRootMessage),
+        });
       },
       onError(err) {
         activeStreams.delete(ws);
@@ -525,7 +532,15 @@ const handlers = {
       return safeSend(ws, { type: 'error', message: 'Failed to apply edit - message lost' });
     }
     
-    safeSend(ws, { type: 'chat:messageEdited', sessionId, messageId: targetMsg.id, messageIndex, message: updatedTargetMsg, history: updatedFlatHistory });
+    safeSend(ws, {
+      type: 'chat:messageEdited',
+      sessionId,
+      messageId: targetMsg.id,
+      messageIndex,
+      message: updatedTargetMsg,
+      history: updatedFlatHistory,
+      [ROOT_JSON_KEY]: JSON.stringify(newRoot),
+    });
   },
 
   'chat:selectVersion': async (ws, msg, client) => {
@@ -560,7 +575,14 @@ const handlers = {
     }
     
     // Send back with messageId for clarity
-    safeSend(ws, { type: 'chat:versionSelected', sessionId, messageId: targetMsg.id, messageIndex, history: extractFlatHistory(newRoot) });
+    safeSend(ws, {
+      type: 'chat:versionSelected',
+      sessionId,
+      messageId: targetMsg.id,
+      messageIndex,
+      history: extractFlatHistory(newRoot),
+      [ROOT_JSON_KEY]: JSON.stringify(newRoot),
+    });
   },
 
   'chat:assistantAction': async (ws, msg, client) => {
@@ -709,7 +731,13 @@ const handlers = {
           sessionStore.updateTempSession(client.tempId, sessionId, { history: newHistory, name: newName });
         }
 
-        safeSend(ws, { type: 'chat:done', sessionId, name: newName, history: extractFlatHistory(newRoot) });
+        safeSend(ws, {
+          type: 'chat:done',
+          sessionId,
+          name: newName,
+          history: extractFlatHistory(newRoot),
+          [ROOT_JSON_KEY]: JSON.stringify(newRoot),
+        });
       },
       onError(err) {
         activeStreams.delete(ws);
@@ -853,7 +881,22 @@ const handlers = {
   },
 };
 
-function ser(s) { return { id: s.id, name: s.name, created: s.created, history: s.history || [], model: s.model }; }
+export function serializeSessionForClient(session) {
+  const rootMessage = Array.isArray(session?.history) && session.history[0]
+    ? cloneAndRepairTree(session.history[0])
+    : null;
+  return {
+    id: session?.id,
+    name: session?.name,
+    created: session?.created,
+    history: rootMessage ? extractFlatHistory(rootMessage) : [],
+    model: session?.model || null,
+    updatedAt: session?.updatedAt || null,
+    [ROOT_JSON_KEY]: rootMessage ? JSON.stringify(rootMessage) : null,
+  };
+}
+
+function ser(s) { return serializeSessionForClient(s); }
 
 function getClientOwner(client) {
   return client.userId
@@ -1093,6 +1136,10 @@ function cloneVersionMetaValue(value) {
   return JSON.parse(JSON.stringify(value));
 }
 
+function cloneJson(value) {
+  return JSON.parse(JSON.stringify(value));
+}
+
 function syncMessageFromActiveVersion(message) {
   if (!message) return message;
   const currentVersion = getActiveVersion(message);
@@ -1143,16 +1190,23 @@ function appendConversationTurn(rootMessage, userEntry, assistantEntry, mediaEnt
 
 function extractFlatHistory(rootMessage) {
   if (!rootMessage) return [];
-  
-  // Helper to ensure message has valid content
-  const ensureValidContent = (msg) => {
-    if (msg.content === undefined || msg.content === null) {
-      msg.content = '';
+
+  const toFlatEntry = (message) => {
+    const cloned = cloneJson(message);
+    if (cloned.content === undefined || cloned.content === null) {
+      cloned.content = '';
+    }
+    syncMessageFromActiveVersion(cloned);
+    if (Array.isArray(cloned.versions)) {
+      cloned.versions = cloned.versions.map((version) => ({
+        ...version,
+        tail: [],
+      }));
     }
-    return syncMessageFromActiveVersion(msg);
+    return cloned;
   };
-  
-  const history = [ensureValidContent(rootMessage)];
+
+  const history = [toFlatEntry(rootMessage)];
   const currentVerIdx = rootMessage.currentVersionIdx ?? 0;
   
   if (!Array.isArray(rootMessage.versions)) {
@@ -1171,7 +1225,11 @@ function extractFlatHistory(rootMessage) {
     const walkTail = (tail) => {
       for (let i = 0; i < tail.length; i++) {
         const msg = tail[i];
-        history.push(ensureValidContent(msg));
+        if (msg?.content === undefined || msg?.content === null) {
+          msg.content = '';
+        }
+        syncMessageFromActiveVersion(msg);
+        history.push(toFlatEntry(msg));
         const ver = msg.versions?.[msg.currentVersionIdx ?? 0];
         if (ver?.tail && Array.isArray(ver.tail)) {
           walkTail(ver.tail);