minor bugfix

app/lib/persistence/useChatHistory.ts

@@ -1,5 +1,5 @@
 import { useLoaderData, useNavigate, useSearchParams } from '@remix-run/react';
-import { useState, useEffect } from 'react';
+import { useState, useEffect, useCallback } from 'react';
 import { atom } from 'nanostores';
 import type { Message } from 'ai';
 import { toast } from 'react-toastify';
@@ -13,7 +13,9 @@ import {
   duplicateChat,
   createChatFromMessages,
 } from './db';
-import { useSnapshot } from './useSnapshot';
+import type { FileMap } from '../stores/files';
+import type { Snapshot } from './types';
+import { webcontainer } from '../webcontainer';
 
 export interface ChatHistoryItem {
   id: string;
@@ -33,9 +35,6 @@ export const description = atom<string | undefined>(undefined);
 export function useChatHistory() {
   const navigate = useNavigate();
   const { id: mixedId } = useLoaderData<{ id?: string }>();
-  console.log({mixedId});
-
-  const {snapshot, takeSnapshot,restoreSnapshot}=useSnapshot(mixedId);
   const [searchParams] = useSearchParams();
 
   const [archivedMessages, setArchivedMessages] = useState<Message[]>([]);
@@ -58,29 +57,37 @@ export function useChatHistory() {
       getMessages(db, mixedId)
         .then((storedMessages) => {
           if (storedMessages && storedMessages.messages.length > 0) {
+            let snapshotStr = localStorage.getItem(`snapshot:${mixedId}`);
+            const snapshot: Snapshot = snapshotStr ? JSON.parse(snapshotStr) : { chatIndex: 0, files: {} };
+
             const rewindId = searchParams.get('rewindTo');
             let startingIdx=0;
             let endingIdx = rewindId ? storedMessages.messages.findIndex((m) => m.id === rewindId) + 1 : storedMessages.messages.length;
-            if (snapshot?.chatIndex && snapshot.chatIndex <=endingIdx){
+            if (snapshot?.chatIndex && snapshot.chatIndex-1 <= endingIdx){
               startingIdx=snapshot.chatIndex;
             }
+            if(storedMessages.messages[snapshot.chatIndex-1].id==rewindId){
+              startingIdx=0;
+            }
             let filteredMessages = storedMessages.messages.slice(startingIdx, endingIdx);
 
+
+
             startingIdx>0? setArchivedMessages(storedMessages.messages.slice(0,startingIdx)):setArchivedMessages([]);
             if (startingIdx>0){
               filteredMessages=[
                 {
-                  "id":`${Date.now()}`,
+                  "id": storedMessages.messages[snapshot.chatIndex-1].id,
                   "role":"assistant",
                   "content":` 📦 Restoring chat from snapshot`, 
                   annotations:['no-store']
                 },
                 {
-                  "id":`${Date.now()}`,
+                  "id": `${Date.now()}`,
                   "role":"assistant",
                   "content":` Below are the files and content of the files restored:
                   ### Files ###
-                  ${Object.entries(snapshot?.files||{}).map(([key,value])=>{
+                  ${Object.entries(snapshot?.files||{}).filter(x=>!x[0].endsWith('lock.json')).map(([key,value])=>{  
                     if(value?.type==="file"){
                       return `
                       #### ${key}
@@ -94,9 +101,10 @@ export function useChatHistory() {
                 },
                 ...filteredMessages
               ]
-              restoreSnapshot();
+              restoreSnapshot(mixedId);
 
             }
+
             setInitialMessages(filteredMessages);
 
             setUrlId(storedMessages.urlId);
@@ -112,7 +120,48 @@ export function useChatHistory() {
           toast.error(error.message);
         });
     }
-  }, [snapshot]);
+  }, [mixedId]);
+
+
+
+  const takeSnapshot = useCallback(async (chatIdx: number, files: FileMap, _chatId?: string | undefined) => {
+    let id = _chatId || chatId;
+    if (!id) return;
+    const snapshot: Snapshot = {
+      chatIndex: chatIdx,
+      files
+    }
+    localStorage.setItem(`snapshot:${id}`, JSON.stringify(snapshot));
+  }, [chatId])
+
+  const restoreSnapshot = useCallback(async (id:string) => {
+    let snapshotStr = localStorage.getItem(`snapshot:${mixedId}`);
+    let container =await webcontainer;
+    // if (snapshotStr)setSnapshot(JSON.parse(snapshotStr)); 
+    const snapshot:Snapshot = snapshotStr ? JSON.parse(snapshotStr) : { chatIndex: 0, files: {} };
+    if (!(snapshot?.files)) return;
+    Object.entries(snapshot.files).forEach(async ([key, value]) => {
+      if (key.startsWith(container.workdir)) {
+        key = key.replace(container.workdir, '');
+      }
+      if (value?.type === "folder") {
+        await container.fs.mkdir(key,{recursive:true});
+      }
+    })
+    Object.entries(snapshot.files).forEach(async ([key, value]) => {
+      if (value?.type === "file") {
+        if (key.startsWith(container.workdir)){
+          key = key.replace(container.workdir,'');
+        }
+        await container.fs.writeFile(key, value.content,{encoding:value.isBinary?undefined:'utf8'});
+      }
+      else{
+
+      }
+    })
+    // workbenchStore.files.setKey(snapshot?.files)
+
+  }, [ ])
 
   return {
     ready: !mixedId || ready,
@@ -148,7 +197,7 @@ export function useChatHistory() {
         }
       }
 
-      await setMessages(db, chatId.get() as string, messages, urlId, description.get());
+      await setMessages(db, chatId.get() as string, [...archivedMessages,...messages], urlId, description.get());
     },
     duplicateCurrentChat: async (listItemId: string) => {
       if (!db || (!mixedId && !listItemId)) {