From a0ae79a993cc612d3846f425feb4163cdd92387a Mon Sep 17 00:00:00 2001
From: eduardruzga <wonderwhy.er@gmail.com>
Date: Sat, 16 Nov 2024 21:22:00 +0200
Subject: [PATCH] Add ability to return to older chat message state

---
 app/components/chat/Messages.client.tsx | 21 ++++++++++++++++++++-
 app/lib/persistence/useChatHistory.ts   | 11 +++++++++--
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/app/components/chat/Messages.client.tsx b/app/components/chat/Messages.client.tsx
index 2f35f49d4..2809a5f96 100644
--- a/app/components/chat/Messages.client.tsx
+++ b/app/components/chat/Messages.client.tsx
@@ -3,6 +3,7 @@ import React from 'react';
 import { classNames } from '~/utils/classNames';
 import { AssistantMessage } from './AssistantMessage';
 import { UserMessage } from './UserMessage';
+import { useLocation, useNavigate } from '@remix-run/react';
 
 interface MessagesProps {
   id?: string;
@@ -13,12 +14,21 @@ interface MessagesProps {
 
 export const Messages = React.forwardRef<HTMLDivElement, MessagesProps>((props: MessagesProps, ref) => {
   const { id, isStreaming = false, messages = [] } = props;
+  const location = useLocation();
+  const navigate = useNavigate();
+
+  const handleRewind = (messageId: string) => {
+    const searchParams = new URLSearchParams(location.search);
+    searchParams.set('rewindId', messageId);
+    window.location.search = searchParams.toString();
+    //navigate(`${location.pathname}?${searchParams.toString()}`);
+  };
 
   return (
     <div id={id} ref={ref} className={props.className}>
       {messages.length > 0
         ? messages.map((message, index) => {
-            const { role, content } = message;
+            const { role, content, id: messageId } = message;
             const isUserMessage = role === 'user';
             const isFirst = index === 0;
             const isLast = index === messages.length - 1;
@@ -41,6 +51,15 @@ export const Messages = React.forwardRef<HTMLDivElement, MessagesProps>((props:
                 <div className="grid grid-col-1 w-full">
                   {isUserMessage ? <UserMessage content={content} /> : <AssistantMessage content={content} />}
                 </div>
+                {messageId && (
+                  <button
+                    onClick={() => handleRewind(messageId)}
+                    className="self-start p-2 text-bolt-elements-textSecondary hover:text-bolt-elements-textPrimary"
+                    title="Rewind to this message"
+                  >
+                    <div className="i-ph:arrow-counter-clockwise text-xl"></div>
+                  </button>
+                )}
               </div>
             );
           })
diff --git a/app/lib/persistence/useChatHistory.ts b/app/lib/persistence/useChatHistory.ts
index e56275327..4b4160315 100644
--- a/app/lib/persistence/useChatHistory.ts
+++ b/app/lib/persistence/useChatHistory.ts
@@ -1,4 +1,4 @@
-import { useLoaderData, useNavigate } from '@remix-run/react';
+import { useLoaderData, useNavigate, useSearchParams } from '@remix-run/react';
 import { useState, useEffect } from 'react';
 import { atom } from 'nanostores';
 import type { Message } from 'ai';
@@ -24,6 +24,7 @@ export const description = atom<string | undefined>(undefined);
 export function useChatHistory() {
   const navigate = useNavigate();
   const { id: mixedId } = useLoaderData<{ id?: string }>();
+  const [searchParams] = useSearchParams();
 
   const [initialMessages, setInitialMessages] = useState<Message[]>([]);
   const [ready, setReady] = useState<boolean>(false);
@@ -44,7 +45,13 @@ export function useChatHistory() {
       getMessages(db, mixedId)
         .then((storedMessages) => {
           if (storedMessages && storedMessages.messages.length > 0) {
-            setInitialMessages(storedMessages.messages);
+            const rewindId = searchParams.get('rewindId');
+            const filteredMessages = rewindId 
+              ? storedMessages.messages.slice(0, 
+                  storedMessages.messages.findIndex(m => m.id === rewindId) + 1)
+              : storedMessages.messages;
+            
+            setInitialMessages(filteredMessages);
             setUrlId(storedMessages.urlId);
             description.set(storedMessages.description);
             chatId.set(storedMessages.id);