Fix issues around rewinding (#28)

app/components/chat/Chat.client.tsx

@@ -213,7 +213,7 @@ export const ChatImpl = memo(
     }, [searchParams]);
 
     const { enhancingPrompt, promptEnhanced, enhancePrompt, resetEnhancer } = usePromptEnhancer();
-    const { parsedMessages, parseMessages } = useMessageParser();
+    const { parsedMessages, setParsedMessages, parseMessages } = useMessageParser();
 
     const TEXTAREA_MAX_HEIGHT = chatStarted ? 400 : 200;
 
@@ -437,6 +437,11 @@ export const ChatImpl = memo(
 
       const messageIndex = messages.findIndex((message) => message.id === messageId);
       if (messageIndex >= 0) {
+        const newParsedMessages = { ...parsedMessages };
+        for (let i = messageIndex + 1; i < messages.length; i++) {
+          delete newParsedMessages[i];
+        }
+        setParsedMessages(newParsedMessages);
         setMessages(messages.slice(0, messageIndex + 1));
       }
     };

app/components/chat/Messages.client.tsx

@@ -28,23 +28,33 @@ export function saveProjectContents(messageId: string, contents: ProjectContents
   gProjectContentsByMessageId.set(messageId, contents);
 }
 
+function hasFileModifications(content: string) {
+  return content.includes("__boltArtifact__");
+}
+
 export const Messages = React.forwardRef<HTMLDivElement, MessagesProps>((props: MessagesProps, ref) => {
   const { id, isStreaming = false, messages = [], onRewind } = props;
 
   const getLastMessageProjectContents = (index: number) => {
     // The message index is for the model response, and the project
     // contents will be associated with the last message present when
-    // the user prompt was sent to the model. So look back two messages
-    // for the previous contents.
-    if (index < 2) {
-      return undefined;
-    }
-    const previousMessage = messages[index - 2];
-    const contents = gProjectContentsByMessageId.get(previousMessage.id);
+    // the user prompt was sent to the model. This could be either two
+    // or three messages back, depending on whether the "fix bug"
+    // button was clicked.
+    const beforeUserMessage = messages[index - 2];
+    const contents = gProjectContentsByMessageId.get(beforeUserMessage?.id);
     if (!contents) {
-      return undefined;
+      const priorMessage = messages[index - 3];
+      const priorContents = gProjectContentsByMessageId.get(priorMessage?.id);
+      if (!priorContents) {
+        return undefined;
+      }
+
+      // We still rewind to just before the user message to retain any
+      // explanation from the Nut API.
+      return { messageId: beforeUserMessage.id, contents: priorContents };
     }
-    return { messageId: previousMessage.id, contents };
+    return { messageId: beforeUserMessage.id, contents };
   };
 
   return (
@@ -84,7 +94,11 @@ export const Messages = React.forwardRef<HTMLDivElement, MessagesProps>((props:
                       <AssistantMessage content={content} annotations={message.annotations} />
                     )}
                   </div>
-                  {!isUserMessage && messageId && onRewind && getLastMessageProjectContents(index) && (
+                  {!isUserMessage &&
+                   messageId &&
+                   onRewind &&
+                   getLastMessageProjectContents(index) &&
+                   hasFileModifications(content) && (
                     <div className="flex gap-2 flex-col lg:flex-row">
                       <WithTooltip tooltip="Undo changes in this message">
                         <button

app/lib/hooks/useMessageParser.ts

@@ -66,5 +66,5 @@ export function useMessageParser() {
     }
   }, []);
 
-  return { parsedMessages, parseMessages };
+  return { parsedMessages, setParsedMessages, parseMessages };
 }

app/lib/replay/SimulationPrompt.ts

@@ -48,7 +48,10 @@ class ChatManager {
       await this.client.initialize();
 
       const { chatId } = (await this.client.sendCommand({ method: "Nut.startChat", params: {} })) as { chatId: string };
-      return chatId;  
+
+      console.log("ChatStarted", new Date().toISOString(), chatId);
+
+      return chatId;
     })();
   }
 

app/lib/stores/workbench.ts

@@ -463,7 +463,7 @@ export class WorkbenchStore {
       const artifact = fileArtifacts.find((artifact) => artifact.path === filePath);
       const artifactContent = artifact?.content ?? "";
 
-      const actionId = `restore-contents-action-${messageId}-${filePath}`;
+      const actionId = `restore-contents-action-${messageId}-${filePath}-${Math.random().toString()}`;
       const data: ActionCallbackData = {
         actionId,
         messageId,