From 1f940391b19e7187890623150136d24f0adb6322 Mon Sep 17 00:00:00 2001 From: Anirban Kar Date: Wed, 5 Mar 2025 10:59:48 +0530 Subject: [PATCH] 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 | 7 + app/lib/persistence/useChatHistory.ts | 218 +++++++++++++++++++++++++- 2 files changed, 217 insertions(+), 8 deletions(-) create mode 100644 app/lib/persistence/types.ts diff --git a/app/lib/persistence/types.ts b/app/lib/persistence/types.ts new file mode 100644 index 00000000..56dacd61 --- /dev/null +++ b/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; +} diff --git a/app/lib/persistence/useChatHistory.ts b/app/lib/persistence/useChatHistory.ts index 7baefa56..b8b5c833 100644 --- a/app/lib/persistence/useChatHistory.ts +++ b/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([]); const [initialMessages, setInitialMessages] = useState([]); const [ready, setReady] = useState(false); const [urlId, setUrlId] = useState(); @@ -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 + + ${Object.entries(snapshot?.files || {}) + .filter((x) => !x[0].endsWith('lock.json')) + .map(([key, value]) => { + if (value?.type === 'file') { + return ` + +${value.content} + + `; + } else { + return ``; + } + }) + .join('\n')} + + `, + 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)) {