Versions

server/index.js

@@ -95,7 +95,10 @@ wss.on('connection', (ws, req) => {
 
   ws.on('message', async raw => {
     try { await handleWsMessage(ws, JSON.parse(raw.toString()), wsClients); }
-    catch (ex) { safeSend(ws, { type: 'error', message: 'Invalid message' }); console.log("Invalid message: " + ex); }
+    catch (ex) { 
+      console.error("Invalid message error:", ex.message, "\nStack:", ex.stack);
+      safeSend(ws, { type: 'error', message: 'Invalid message: ' + ex.message }); 
+    }
   });
   ws.on('close', () => {
     const c = wsClients.get(ws);

server/wsHandler.js

@@ -13,25 +13,31 @@ import { streamChat, extractSessionName } from './chatStream.js';
 import crypto from 'crypto';
 
 /**
- * ─ Message Structure Documentation ─
+ * Message Structure: Tree-based with versioned tails
+ * 
+ * Each message has versions, and each version has a complete tail of subsequent messages.
+ * Messages only exist within parent tails (no separate flat array).
  * 
- * Each message in session.history contains:
  * {
- *   id: "msg-123...",           // Unique identifier for this message
+ *   id: "msg-123",
  *   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
+ *       tail: [  // Full message objects
+ *         { id, role, content, timestamp, versions: [...], currentVersionIdx, ... },
+ *         ...
+ *       ],
  *       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 */
+ *   currentVersionIdx: 0,
+ *   toolCalls?: [...]
+ * }
+ */
 
 const activeStreams = new Map();
 
@@ -56,7 +62,6 @@ function bcast(wsClients, userId, data, excludeWs) {
 const handlers = {
   'ping': (ws) => { safeSend(ws, { type: 'pong' }); },
 
-  // Verify turnstile token sent over websocket
   'turnstile:verify': async (ws, msg, client) => {
     try {
       const token = msg?.token;
@@ -81,7 +86,6 @@ const handlers = {
     client.deviceToken = deviceSessionStore.create(user.id, client.ip, client.userAgent);
     sessionStore.markOnline(user.id, ws);
 
-    // If the client provided a persistent tempId, use and save it on the ws client
     if (clientTempId) client.tempId = clientTempId;
     const tId = client.tempId;
     await sessionStore.transferTempToUser(tId, user.id, accessToken);
@@ -106,7 +110,6 @@ const handlers = {
   },
 
   'auth:guest': (ws, msg, client) => {
-    // Honor a client-supplied persistent tempId and persist it on the ws client
     const t = msg.tempId || client.tempId;
     client.tempId = t;
     sessionStore.initTemp(t);
@@ -191,9 +194,17 @@ const handlers = {
     let fullText = '';
     const assetsCollected = [], toolCallsCollected = [];
 
+    // Extract flat history from tree structure
+    const rootMessage = session.history?.[0];
+    const flatHistory = rootMessage ? extractFlatHistory(rootMessage) : [];
+
     await streamChat(ws, {
-      history: normalizeHistoryStructure(session.history || []), userMessage: content, tools: tools || {},
-      accessToken: client.accessToken, clientId: msg.clientId, abortSignal: abort.signal,
+      history: flatHistory, 
+      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) {
         safeSend(ws, { type: 'chat:toolCall', call, sessionId });
@@ -211,34 +222,39 @@ const handlers = {
           return { ...c, state: resolved.state || 'resolved', result: resolved.result };
         });
         const asstEntry = buildEntry('assistant', finalText, mergedCalls);
-        let newHistory = normalizeHistoryStructure([...(session.history || []), userEntry, asstEntry]);
-
-        // 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) {
-              if (!version.nextMessageId) {
-                version.nextMessageId = nextMsg.id;
-              }
-            }
-          }
+
+        // Rebuild tree structure with new messages appended
+        let newRootMessage = rootMessage ? JSON.parse(JSON.stringify(rootMessage)) : null;
+        
+        if (!newRootMessage) {
+          // First message in session
+          newRootMessage = userEntry;
+          const asstWrap = { ...asstEntry };
+          newRootMessage.versions[0].tail = [asstWrap];
+        } else {
+          // Append to current tail
+          const currentVerIdx = newRootMessage.currentVersionIdx ?? 0;
+          let currentTail = newRootMessage.versions[currentVerIdx].tail || [];
+          currentTail = JSON.parse(JSON.stringify(currentTail));
+          currentTail.push(userEntry);
+          currentTail.push(asstEntry);
+          newRootMessage.versions[currentVerIdx].tail = currentTail;
         }
 
-        // Use session name from XML tag; fall back to existing name or default
+        const newHistory = [newRootMessage];
+
         let newName = session.name;
         if (sessionNameFromTag) {
           newName = sessionNameFromTag;
         } else if (!session.history?.length || session.name === 'New Chat') {
-          // No tag and first message — keep as "New Chat" (no extra API call)
           newName = session.name;
         }
 
         if (client.userId)
           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: newHistory });
+        
+        safeSend(ws, { type: aborted ? 'chat:aborted' : 'chat:done', sessionId, name: newName, history: extractFlatHistory(newRootMessage) });
       },
       onError(err) { activeStreams.delete(ws); safeSend(ws, { type: 'chat:error', error: String(err), sessionId }); },
     });
@@ -252,25 +268,41 @@ const handlers = {
       ? sessionStore.getUserSession(client.userId, sessionId)
       : sessionStore.getTempSession(client.tempId, sessionId);
     if (!session) return;
-    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);
+    
+    const rootMessage = session.history?.[0];
+    if (!rootMessage) return;
+    
+    const flatHistory = extractFlatHistory(rootMessage);
+    const targetMsg = flatHistory[messageIndex];
+    if (!targetMsg) return;
+    
+    // Find the target message in the tree and add new version
+    const newRoot = JSON.parse(JSON.stringify(rootMessage));
+    const found = findAndUpdateMessage(newRoot, targetMsg.id, (msgInTree) => {
+      // Get current tail
+      const currentVerIdx = msgInTree.currentVersionIdx ?? 0;
+      const currentTail = msgInTree.versions[currentVerIdx]?.tail || [];
+      
+      // Add new version with same tail as current
+      msgInTree.versions.push({
+        content: newContent,
+        tail: JSON.parse(JSON.stringify(currentTail)),
+        timestamp: Date.now()
+      });
+      msgInTree.currentVersionIdx = msgInTree.versions.length - 1;
+      msgInTree.content = newContent;
+    });
+    
+    if (!found) return;
+    
+    const newHistory = [newRoot];
     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 });
+    
+    safeSend(ws, { type: 'chat:messageEdited', sessionId, messageIndex, message: targetMsg, history: extractFlatHistory(newRoot) });
   },
 
   'chat:selectVersion': async (ws, msg, client) => {
@@ -279,18 +311,32 @@ const handlers = {
       ? sessionStore.getUserSession(client.userId, sessionId)
       : sessionStore.getTempSession(client.tempId, sessionId);
     if (!session) return;
-    const history = normalizeHistoryStructure(session.history || []);
-    const m = history[messageIndex];
-    if (!m?.versions || versionIdx < 0 || versionIdx >= m.versions.length) return;
-    m.currentVersionIdx = versionIdx;
-    m.content = m.versions[versionIdx].content;
-    const newHistory = reconstructHistoryFromVersion(history, messageIndex, versionIdx);
+    
+    const rootMessage = session.history?.[0];
+    if (!rootMessage) return;
+    
+    const flatHistory = extractFlatHistory(rootMessage);
+    const targetMsg = flatHistory[messageIndex];
+    if (!targetMsg || !targetMsg.versions || versionIdx >= targetMsg.versions.length) return;
+    
+    // Find and update the message in tree, switching to specified version
+    const newRoot = JSON.parse(JSON.stringify(rootMessage));
+    const found = findAndUpdateMessage(newRoot, targetMsg.id, (msgInTree) => {
+      msgInTree.currentVersionIdx = versionIdx;
+      msgInTree.content = msgInTree.versions[versionIdx].content;
+      // Tail is automatically correct since each version has its own tail
+    });
+    
+    if (!found) return;
+    
+    const newHistory = [newRoot];
     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 });
+    
+    safeSend(ws, { type: 'chat:versionSelected', sessionId, history: extractFlatHistory(newRoot) });
   },
 
   'settings:get': async (ws, msg, client) => {
@@ -341,45 +387,53 @@ function buildEntry(role, content, toolCalls = []) {
     role, 
     content, 
     timestamp: Date.now(),
-    versions: [{ content, nextMessageId: null, timestamp: Date.now() }], 
+    versions: [{ content, tail: [], 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;
+function extractFlatHistory(rootMessage) {
+  if (!rootMessage) return [];
+  const history = [rootMessage];
+  const currentVerIdx = rootMessage.currentVersionIdx ?? 0;
+  const currentTail = rootMessage.versions[currentVerIdx]?.tail;
+  
+  if (currentTail && Array.isArray(currentTail)) {
+    const walkTail = (tail) => {
+      for (const msg of tail) {
+        history.push(msg);
+        const ver = msg.versions?.[msg.currentVersionIdx ?? 0];
+        if (ver?.tail && Array.isArray(ver.tail)) {
+          walkTail(ver.tail);
         }
-        if (!v.nextMessageId) v.nextMessageId = null;
-        return v;
-      });
-    }
-    return msg;
-  });
-  return normalized;
+      }
+    };
+    walkTail(currentTail);
+  }
+  return history;
 }
 
-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;
+function findAndUpdateMessage(rootMessage, targetId, updateFn) {
+  if (rootMessage.id === targetId) {
+    updateFn(rootMessage);
+    return true;
   }
-  return reconstructed;
-}
+  
+  const search = (msg) => {
+    const verIdx = msg.currentVersionIdx ?? 0;
+    const tail = msg.versions?.[verIdx]?.tail;
+    if (!tail || !Array.isArray(tail)) return false;
+    
+    for (const child of tail) {
+      if (child.id === targetId) {
+        updateFn(child);
+        return true;
+      }
+      if (search(child)) return true;
+    }
+    return false;
+  };
+  
+  return search(rootMessage);
+}