feat: restoring project from snapshot on reload (#444)

* feat:(project-snapshot) restoring project from snapshot on reload

* minor bugfix

* updated message

* added snapshot reload with auto run dev commands

* added message context

* snapshot updated

app/lib/persistence/types.ts

@@ -0,0 +1,7 @@
+import type { FileMap } from '~/lib/stores/files';
+
+export interface Snapshot {
+  chatIndex: string;
+  files: FileMap;
+  summary?: string;
+}

app/lib/persistence/useChatHistory.ts

@@ -1,7 +1,7 @@
 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 { generateId, type JSONValue, type Message } from 'ai';
 import { toast } from 'react-toastify';
 import { workbenchStore } from '~/lib/stores/workbench';
 import { logStore } from '~/lib/stores/logs'; // Import logStore
@@ -15,6 +15,11 @@ import {
   createChatFromMessages,
   type IChatMetadata,
 } from './db';
+import type { FileMap } from '~/lib/stores/files';
+import type { Snapshot } from './types';
+import { webcontainer } from '~/lib/webcontainer';
+import { createCommandsMessage, detectProjectCommands } from '~/utils/projectCommands';
+import type { ContextAnnotation } from '~/types/context';
 
 export interface ChatHistoryItem {
   id: string;
@@ -37,6 +42,7 @@ export function useChatHistory() {
   const { id: mixedId } = useLoaderData<{ id?: string }>();
   const [searchParams] = useSearchParams();
 
+  const [archivedMessages, setArchivedMessages] = useState<Message[]>([]);
   const [initialMessages, setInitialMessages] = useState<Message[]>([]);
   const [ready, setReady] = useState<boolean>(false);
   const [urlId, setUrlId] = useState<string | undefined>();
@@ -56,14 +62,128 @@ export function useChatHistory() {
 
     if (mixedId) {
       getMessages(db, mixedId)
-        .then((storedMessages) => {
+        .then(async (storedMessages) => {
           if (storedMessages && storedMessages.messages.length > 0) {
+            const snapshotStr = localStorage.getItem(`snapshot:${mixedId}`);
+            const snapshot: Snapshot = snapshotStr ? JSON.parse(snapshotStr) : { chatIndex: 0, files: {} };
+            const summary = snapshot.summary;
+
             const rewindId = searchParams.get('rewindTo');
-            const filteredMessages = rewindId
-              ? storedMessages.messages.slice(0, storedMessages.messages.findIndex((m) => m.id === rewindId) + 1)
-              : storedMessages.messages;
+            let startingIdx = 0;
+            const endingIdx = rewindId
+              ? storedMessages.messages.findIndex((m) => m.id === rewindId) + 1
+              : storedMessages.messages.length;
+            const snapshotIndex = storedMessages.messages.findIndex((m) => m.id === snapshot.chatIndex);
+
+            if (snapshotIndex >= 0 && snapshotIndex < endingIdx) {
+              startingIdx = snapshotIndex;
+            }
+
+            if (snapshotIndex > 0 && storedMessages.messages[snapshotIndex].id == rewindId) {
+              startingIdx = 0;
+            }
+
+            let filteredMessages = storedMessages.messages.slice(startingIdx + 1, endingIdx);
+            let archivedMessages: Message[] = [];
+
+            if (startingIdx > 0) {
+              archivedMessages = storedMessages.messages.slice(0, startingIdx + 1);
+            }
+
+            setArchivedMessages(archivedMessages);
+
+            if (startingIdx > 0) {
+              const files = Object.entries(snapshot?.files || {})
+                .map(([key, value]) => {
+                  if (value?.type !== 'file') {
+                    return null;
+                  }
+
+                  return {
+                    content: value.content,
+                    path: key,
+                  };
+                })
+                .filter((x) => !!x);
+              const projectCommands = await detectProjectCommands(files);
+              const commands = createCommandsMessage(projectCommands);
+
+              filteredMessages = [
+                {
+                  id: generateId(),
+                  role: 'user',
+                  content: `Restore project from snapshot
+                  `,
+                  annotations: ['no-store', 'hidden'],
+                },
+                {
+                  id: storedMessages.messages[snapshotIndex].id,
+                  role: 'assistant',
+                  content: ` 📦 Chat Restored from snapshot, You can revert this message to load the full chat history
+                  <boltArtifact id="imported-files" title="Project Files Snapshot" type="bundled">
+                  ${Object.entries(snapshot?.files || {})
+                    .filter((x) => !x[0].endsWith('lock.json'))
+                    .map(([key, value]) => {
+                      if (value?.type === 'file') {
+                        return `
+                      <boltAction type="file" filePath="${key}">
+${value.content}
+                      </boltAction>
+                      `;
+                      } else {
+                        return ``;
+                      }
+                    })
+                    .join('\n')}
+                  </boltArtifact>
+                  `,
+                  annotations: [
+                    'no-store',
+                    ...(summary
+                      ? [
+                          {
+                            chatId: storedMessages.messages[snapshotIndex].id,
+                            type: 'chatSummary',
+                            summary,
+                          } satisfies ContextAnnotation,
+                        ]
+                      : []),
+                  ],
+                },
+                ...(commands !== null
+                  ? [
+                      {
+                        id: `${storedMessages.messages[snapshotIndex].id}-2`,
+                        role: 'user' as const,
+                        content: `setup project`,
+                        annotations: ['no-store', 'hidden'],
+                      },
+                      {
+                        ...commands,
+                        id: `${storedMessages.messages[snapshotIndex].id}-3`,
+                        annotations: [
+                          'no-store',
+                          ...(commands.annotations || []),
+                          ...(summary
+                            ? [
+                                {
+                                  chatId: `${storedMessages.messages[snapshotIndex].id}-3`,
+                                  type: 'chatSummary',
+                                  summary,
+                                } satisfies ContextAnnotation,
+                              ]
+                            : []),
+                        ],
+                      },
+                    ]
+                  : []),
+                ...filteredMessages,
+              ];
+              restoreSnapshot(mixedId);
+            }
 
             setInitialMessages(filteredMessages);
+
             setUrlId(storedMessages.urlId);
             description.set(storedMessages.description);
             chatId.set(storedMessages.id);
@@ -75,10 +195,64 @@ export function useChatHistory() {
           setReady(true);
         })
         .catch((error) => {
+          console.error(error);
+
           logStore.logError('Failed to load chat messages', error);
           toast.error(error.message);
         });
     }
+  }, [mixedId]);
+
+  const takeSnapshot = useCallback(
+    async (chatIdx: string, files: FileMap, _chatId?: string | undefined, chatSummary?: string) => {
+      const id = _chatId || chatId;
+
+      if (!id) {
+        return;
+      }
+
+      const snapshot: Snapshot = {
+        chatIndex: chatIdx,
+        files,
+        summary: chatSummary,
+      };
+      localStorage.setItem(`snapshot:${id}`, JSON.stringify(snapshot));
+    },
+    [chatId],
+  );
+
+  const restoreSnapshot = useCallback(async (id: string) => {
+    const snapshotStr = localStorage.getItem(`snapshot:${id}`);
+    const 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 {
@@ -105,14 +279,34 @@ export function useChatHistory() {
       }
 
       const { firstArtifact } = workbenchStore;
+      messages = messages.filter((m) => !m.annotations?.includes('no-store'));
+
+      let _urlId = urlId;
 
       if (!urlId && firstArtifact?.id) {
         const urlId = await getUrlId(db, firstArtifact.id);
-
+        _urlId = urlId;
         navigateChat(urlId);
         setUrlId(urlId);
       }
 
+      let chatSummary: string | undefined = undefined;
+      const lastMessage = messages[messages.length - 1];
+
+      if (lastMessage.role === 'assistant') {
+        const annotations = lastMessage.annotations as JSONValue[];
+        const filteredAnnotations = (annotations?.filter(
+          (annotation: JSONValue) =>
+            annotation && typeof annotation === 'object' && Object.keys(annotation).includes('type'),
+        ) || []) as { type: string; value: any } & { [key: string]: any }[];
+
+        if (filteredAnnotations.find((annotation) => annotation.type === 'chatSummary')) {
+          chatSummary = filteredAnnotations.find((annotation) => annotation.type === 'chatSummary')?.summary;
+        }
+      }
+
+      takeSnapshot(messages[messages.length - 1].id, workbenchStore.files.get(), _urlId, chatSummary);
+
       if (!description.get() && firstArtifact?.title) {
         description.set(firstArtifact?.title);
       }
@@ -127,7 +321,15 @@ export function useChatHistory() {
         }
       }
 
-      await setMessages(db, chatId.get() as string, messages, urlId, description.get(), undefined, chatMetadata.get());
+      await setMessages(
+        db,
+        chatId.get() as string,
+        [...archivedMessages, ...messages],
+        urlId,
+        description.get(),
+        undefined,
+        chatMetadata.get(),
+      );
     },
     duplicateCurrentChat: async (listItemId: string) => {
       if (!db || (!mixedId && !listItemId)) {