Add new session storage format

server/sessionStore.js

@@ -45,13 +45,129 @@ function safeFileId(value, fallback = 'unknown') {
   return normalized.replace(/[^a-zA-Z0-9_.-]+/g, '_').slice(0, 160) || fallback;
 }
 
+function cloneJson(value) {
+  return JSON.parse(JSON.stringify(value));
+}
+
+function ensureMessageVersioningShape(message) {
+  if (!message || typeof message !== 'object') return message;
+  if (!Array.isArray(message.versions) || message.versions.length === 0) {
+    message.versions = [{
+      content: message.content ?? '',
+      tail: [],
+      timestamp: Number.isFinite(message.timestamp) ? message.timestamp : Date.now(),
+    }];
+    message.currentVersionIdx = 0;
+  }
+
+  const idx = Number.isInteger(message.currentVersionIdx)
+    ? Math.max(0, Math.min(message.currentVersionIdx, message.versions.length - 1))
+    : 0;
+  message.currentVersionIdx = idx;
+
+  const active = message.versions[idx] || {};
+  if (!Array.isArray(active.tail)) active.tail = [];
+  if (active.content === undefined || active.content === null) {
+    active.content = message.content ?? '';
+  }
+  if (!Number.isFinite(active.timestamp)) {
+    active.timestamp = Number.isFinite(message.timestamp) ? message.timestamp : Date.now();
+  }
+  message.versions[idx] = active;
+  message.content = active.content;
+  return message;
+}
+
+function getActiveVersionNode(message) {
+  if (!message || typeof message !== 'object') return null;
+  ensureMessageVersioningShape(message);
+  return message.versions[message.currentVersionIdx];
+}
+
+function cloneAndRepairTree(rootMessage) {
+  if (!rootMessage || typeof rootMessage !== 'object') return null;
+  const cloned = cloneJson(rootMessage);
+  const walk = (node) => {
+    if (!node || typeof node !== 'object') return;
+    ensureMessageVersioningShape(node);
+    const active = getActiveVersionNode(node);
+    for (const child of active?.tail || []) walk(child);
+  };
+  walk(cloned);
+  return cloned;
+}
+
+function getActiveLeafMessage(rootMessage) {
+  let current = rootMessage;
+  while (current) {
+    const active = getActiveVersionNode(current);
+    const tail = Array.isArray(active?.tail) ? active.tail : [];
+    if (!tail.length) return current;
+    current = tail[tail.length - 1];
+  }
+  return rootMessage;
+}
+
+function appendEntriesToActiveLeaf(rootMessage, entries = []) {
+  if (!rootMessage || !entries.length) return rootMessage;
+  const leaf = getActiveLeafMessage(rootMessage);
+  const active = getActiveVersionNode(leaf);
+  active.tail = [...(active.tail || []), ...entries];
+  return rootMessage;
+}
+
+function normalizeHistoryToTree(history) {
+  const rawHistory = Array.isArray(history) ? history : [];
+  if (!rawHistory.length) return [];
+
+  // Already tree-shaped ([rootMessage]).
+  if (
+    rawHistory.length === 1 &&
+    rawHistory[0] &&
+    typeof rawHistory[0] === 'object' &&
+    Array.isArray(rawHistory[0].versions)
+  ) {
+    const root = cloneAndRepairTree(rawHistory[0]);
+    return root ? [root] : [];
+  }
+
+  // Legacy flat history; rebuild a linear active-branch tree.
+  const nodes = rawHistory
+    .filter((entry) => entry && typeof entry === 'object')
+    .map((entry) => ensureMessageVersioningShape(cloneJson(entry)));
+  if (!nodes.length) return [];
+  const root = nodes[0];
+  for (let i = 1; i < nodes.length; i++) {
+    appendEntriesToActiveLeaf(root, [nodes[i]]);
+  }
+  return [root];
+}
+
 function ensureSessionShape(raw, fallbackId = null) {
   const created = Number.isFinite(raw?.created) ? raw.created : Date.now();
+  let history = [];
+
+  // Prefer desktop's legacy serialized tree when present.
+  const legacyRootRaw = typeof raw?.__historyRootJson === 'string'
+    ? raw.__historyRootJson.trim()
+    : '';
+  if (legacyRootRaw) {
+    try {
+      const parsedRoot = cloneAndRepairTree(JSON.parse(legacyRootRaw));
+      if (parsedRoot) history = [parsedRoot];
+    } catch {
+      history = [];
+    }
+  }
+  if (!history.length) {
+    history = normalizeHistoryToTree(raw?.history);
+  }
+
   return {
     id: raw?.id || fallbackId || crypto.randomUUID(),
     name: String(raw?.name || 'New Chat'),
     created,
-    history: Array.isArray(raw?.history) ? raw.history : [],
+    history,
     model: raw?.model || null,
   };
 }
@@ -197,8 +313,12 @@ async function loadTempStore() {
   const data = await loadEncryptedJson(TEMP_STORE_FILE, 'chat:temp:index');
   if (!data) return;
   for (const [id, d] of Object.entries(data)) {
+    const restoredSessions = {};
+    for (const [sessionId, sessionValue] of Object.entries(d.sessions || {})) {
+      restoredSessions[sessionId] = ensureSessionShape(sessionValue, sessionId);
+    }
     tempStore.set(id, {
-      sessions: new Map(Object.entries(d.sessions || {})),
+      sessions: new Map(Object.entries(restoredSessions)),
       msgCount: d.msgCount || 0,
       created: d.created || Date.now(),
       lastActive: d.lastActive || Date.now(),
@@ -263,7 +383,13 @@ export const sessionStore = {
   },
   createTempSession(t) {
     const d = this.initTemp(t);
-    const s = { id: crypto.randomUUID(), name: 'New Chat', created: Date.now(), history: [] };
+    const s = ensureSessionShape({
+      id: crypto.randomUUID(),
+      name: 'New Chat',
+      created: Date.now(),
+      history: [],
+      model: null,
+    });
     d.sessions.set(s.id, s);
     d.lastActive = Date.now();
     saveTempStore().catch((err) => console.error('Failed to save temp store:', err));
@@ -275,13 +401,15 @@ export const sessionStore = {
     const s = d.sessions.get(id);
     if (!s) return null;
     Object.assign(s, patch);
+    const shaped = ensureSessionShape(s, id);
+    d.sessions.set(id, shaped);
     d.lastActive = Date.now();
     saveTempStore().catch((err) => console.error('Failed to save temp store:', err));
-    return s;
+    return shaped;
   },
   restoreTempSession(t, session) {
     const d = this.initTemp(t);
-    const restored = JSON.parse(JSON.stringify(session));
+    const restored = ensureSessionShape(JSON.parse(JSON.stringify(session)));
     d.sessions.set(restored.id, restored);
     d.lastActive = Date.now();
     saveTempStore().catch((err) => console.error('Failed to save temp store:', err));

server/wsHandler.js

@@ -454,7 +454,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: newHistory,
+          flatHistory: extractFlatHistory(newRootMessage),
+        });
       },
       onError(err) {
         activeStreams.delete(ws);
@@ -525,7 +531,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: [newRoot],
+      flatHistory: updatedFlatHistory,
+    });
   },
 
   'chat:selectVersion': async (ws, msg, client) => {
@@ -560,7 +574,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: [newRoot],
+      flatHistory: extractFlatHistory(newRoot),
+    });
   },
 
   'chat:assistantAction': async (ws, msg, client) => {
@@ -709,7 +730,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: newHistory,
+          flatHistory: extractFlatHistory(newRoot),
+        });
       },
       onError(err) {
         activeStreams.delete(ws);