Update wsHandler.js

server/wsHandler.js

@@ -12,6 +12,27 @@ import {
 import { streamChat, extractSessionName } from './chatStream.js';
 import crypto from 'crypto';
 
+/**
+ * ─ Message Structure Documentation ─
+ * 
+ * Each message in session.history contains:
+ * {
+ *   id: "msg-123...",           // Unique identifier for this message
+ *   role: "user" | "assistant",
+ *   content: string | array,
+ *   timestamp: number,
+ *   versions: [
+ *     { 
+ *       content: string | array,
+ *       nextMessageId: "msg-456..." | null,  // ID of next message when this version is active
+ *       timestamp: number
+ *     },
+ *     ...
+ *   ],
+ *   currentVersionIdx: 0,        // Which version is currently active
+ *   toolCalls?: [...],           // For assistant messages with tool calls
+ * }\n *\n * VERSION HISTORY: Changed from storing full message arrays in \"tail\" field\n * to storing only nextMessageId references. This reduces payload size significantly\n * while maintaining backward compatibility through normalizeHistoryStructure().\n *\n * HISTORY RECONSTRUCTION: When switching versions, the message chain is rebuilt\n * by following nextMessageId pointers from the selected version.\n */
+
 const activeStreams = new Map();
 
 initGuestRequestLimiter().catch(err => console.error('Failed to initialize guest request limiter:', err));
@@ -171,7 +192,7 @@ const handlers = {
     const assetsCollected = [], toolCallsCollected = [];
 
     await streamChat(ws, {
-      history: session.history || [], userMessage: content, tools: tools || {},
+      history: normalizeHistoryStructure(session.history || []), userMessage: content, tools: tools || {},
       accessToken: client.accessToken, clientId: msg.clientId, abortSignal: abort.signal,
       onToken(t) { fullText += t; safeSend(ws, { type: 'chat:token', token: t, sessionId }); },
       onToolCall(call) {
@@ -190,14 +211,17 @@ const handlers = {
           return { ...c, state: resolved.state || 'resolved', result: resolved.result };
         });
         const asstEntry = buildEntry('assistant', finalText, mergedCalls);
-        const newHistory = [...(session.history || []), userEntry, asstEntry];
+        let newHistory = normalizeHistoryStructure([...(session.history || []), userEntry, asstEntry]);
 
-        // Update tails for all versioned messages to include new messages
-        for (let i = 0; i < newHistory.length; i++) {
+        // Update nextMessageId pointers for all versioned messages to link to the next message
+        for (let i = 0; i < newHistory.length - 1; i++) {
           const msg = newHistory[i];
+          const nextMsg = newHistory[i + 1];
           if (msg.versions) {
             for (const version of msg.versions) {
-              version.tail = newHistory.slice(i + 1);
+              if (!version.nextMessageId) {
+                version.nextMessageId = nextMsg.id;
+              }
             }
           }
         }
@@ -228,15 +252,24 @@ const handlers = {
       ? sessionStore.getUserSession(client.userId, sessionId)
       : sessionStore.getTempSession(client.tempId, sessionId);
     if (!session) return;
-    const history = session.history || [];
-    const m = history[messageIndex]; if (!m) return;
-    if (!m.versions) m.versions = [{ content: m.content, tail: history.slice(messageIndex + 1), timestamp: m.timestamp || Date.now() }];
-    m.versions.push({ content: newContent, tail: history.slice(messageIndex + 1), timestamp: Date.now() });
+    const history = normalizeHistoryStructure(session.history || []);
+    const m = history[messageIndex]; 
+    if (!m) return;
+    const currentVersionIdx = m.currentVersionIdx ?? 0;
+    const currentVersion = m.versions?.[currentVersionIdx];
+    const nextMessageId = currentVersion?.nextMessageId || (history[messageIndex + 1]?.id || null);
+    if (!m.versions) {
+      m.versions = [{ content: m.content, nextMessageId: nextMessageId, timestamp: m.timestamp || Date.now() }];
+    }
+    m.versions.push({ content: newContent, nextMessageId: nextMessageId, timestamp: Date.now() });
     m.currentVersionIdx = m.versions.length - 1;
     m.content = newContent;
     const newHistory = history.slice(0, messageIndex + 1);
-    if (client.userId) await sessionStore.updateUserSession(client.userId, client.accessToken, sessionId, { history: newHistory });
-    else sessionStore.updateTempSession(client.tempId, sessionId, { history: newHistory });
+    if (client.userId) {
+      await sessionStore.updateUserSession(client.userId, client.accessToken, sessionId, { history: newHistory });
+    } else {
+      sessionStore.updateTempSession(client.tempId, sessionId, { history: newHistory });
+    }
     safeSend(ws, { type: 'chat:messageEdited', sessionId, messageIndex, message: m, history: newHistory });
   },
 
@@ -246,14 +279,17 @@ const handlers = {
       ? sessionStore.getUserSession(client.userId, sessionId)
       : sessionStore.getTempSession(client.tempId, sessionId);
     if (!session) return;
-    const history = session.history || [];
+    const history = normalizeHistoryStructure(session.history || []);
     const m = history[messageIndex];
     if (!m?.versions || versionIdx < 0 || versionIdx >= m.versions.length) return;
-    const v = m.versions[versionIdx];
-    m.currentVersionIdx = versionIdx; m.content = v.content;
-    const newHistory = [...history.slice(0, messageIndex + 1), ...(v.tail || [])];
-    if (client.userId) await sessionStore.updateUserSession(client.userId, client.accessToken, sessionId, { history: newHistory });
-    else sessionStore.updateTempSession(client.tempId, sessionId, { history: newHistory });
+    m.currentVersionIdx = versionIdx;
+    m.content = m.versions[versionIdx].content;
+    const newHistory = reconstructHistoryFromVersion(history, messageIndex, versionIdx);
+    if (client.userId) {
+      await sessionStore.updateUserSession(client.userId, client.accessToken, sessionId, { history: newHistory });
+    } else {
+      sessionStore.updateTempSession(client.tempId, sessionId, { history: newHistory });
+    }
     safeSend(ws, { type: 'chat:versionSelected', sessionId, history: newHistory });
   },
 
@@ -288,6 +324,10 @@ const handlers = {
 
 function ser(s) { return { id: s.id, name: s.name, created: s.created, history: s.history || [], model: s.model }; }
 
+function generateMessageId() {
+  return `msg-${Date.now()}-${Math.random().toString(36).substr(2, 9)}`;
+}
+
 function buildEntry(role, content, toolCalls = []) {
   const normalizedCalls = toolCalls.map(c => ({
     id:     c.id,
@@ -296,7 +336,50 @@ function buildEntry(role, content, toolCalls = []) {
     state:  c.state || 'resolved',
     result: c.result,
   }));
-  return { role, content, timestamp: Date.now(),
-    versions: [{ content, tail: [], timestamp: Date.now() }], currentVersionIdx: 0,
-    ...(normalizedCalls.length ? { toolCalls: normalizedCalls } : {}) };
+  return { 
+    id: generateMessageId(),
+    role, 
+    content, 
+    timestamp: Date.now(),
+    versions: [{ content, nextMessageId: null, timestamp: Date.now() }], 
+    currentVersionIdx: 0,
+    ...(normalizedCalls.length ? { toolCalls: normalizedCalls } : {}) 
+  };
+}
+
+function normalizeHistoryStructure(history) {
+  if (!history || !Array.isArray(history)) return [];
+  const normalized = history.map((msg, idx) => {
+    if (!msg.id) msg.id = `msg-${msg.timestamp || Date.now()}-${idx}`;
+    if (msg.versions && Array.isArray(msg.versions)) {
+      msg.versions = msg.versions.map(v => {
+        if (v.tail !== undefined && !v.nextMessageId) {
+          v.nextMessageId = v.tail && v.tail.length > 0 ? v.tail[0].id : null;
+        }
+        if (!v.nextMessageId) v.nextMessageId = null;
+        return v;
+      });
+    }
+    return msg;
+  });
+  return normalized;
+}
+
+function reconstructHistoryFromVersion(history, messageIndex, versionIdx) {
+  const baseHistory = history.slice(0, messageIndex + 1);
+  const branchMessage = baseHistory[messageIndex];
+  if (!branchMessage?.versions || versionIdx >= branchMessage.versions.length) {
+    return baseHistory;
+  }
+  const version = branchMessage.versions[versionIdx];
+  if (!version.nextMessageId) return baseHistory;
+  const reconstructed = [...baseHistory];
+  let currentId = version.nextMessageId;
+  while (currentId) {
+    const nextMsg = history.find(m => m.id === currentId);
+    if (!nextMsg) break;
+    reconstructed.push(nextMsg);
+    currentId = nextMsg.versions?.[nextMsg.currentVersionIdx]?.nextMessageId || null;
+  }
+  return reconstructed;
 }