Send simulation prompt commands from client (#8)

app/components/chat/Chat.client.tsx

@@ -3,10 +3,10 @@
  * Preventing TS checks with files presented in the video for a better presentation.
  */
 import { useStore } from '@nanostores/react';
-import type { Message } from 'ai';
+import type { CreateMessage, Message } from 'ai';
 import { useChat } from 'ai/react';
 import { useAnimate } from 'framer-motion';
-import { memo, useCallback, useEffect, useRef, useState } from 'react';
+import { memo, useCallback, useEffect, useMemo, useRef, useState } from 'react';
 import { cssTransition, toast, ToastContainer } from 'react-toastify';
 import { useMessageParser, usePromptEnhancer, useShortcuts, useSnapScroll } from '~/lib/hooks';
 import { description, useChatHistory } from '~/lib/persistence';
@@ -22,10 +22,11 @@ import { useSettings } from '~/lib/hooks/useSettings';
 import { useSearchParams } from '@remix-run/react';
 import { createSampler } from '~/utils/sampler';
 import { saveProjectContents } from './Messages.client';
-import type { SimulationPromptClientData } from '~/lib/replay/SimulationPrompt';
+import { getSimulationRecording, getSimulationEnhancedPrompt } from '~/lib/replay/SimulationPrompt';
-import { getIFrameSimulationData } from '~/lib/replay/Recording';
+import { getIFrameSimulationData, type SimulationData } from '~/lib/replay/Recording';
 import { getCurrentIFrame } from '../workbench/Preview';
 import { getCurrentMouseData } from '../workbench/PointSelector';
+import { assert } from '~/lib/replay/ReplayProtocolClient';
 
 const toastAnimation = cssTransition({
   enter: 'animated fadeInRight',
@@ -108,6 +109,26 @@ interface ChatProps {
   description?: string;
 }
 
+let gNumAborts = 0;
+
+interface InjectedMessage {
+  message: Message;
+  previousId: string;
+}
+
+function handleInjectMessages(baseMessages: Message[], injectedMessages: InjectedMessage[]) {
+  const messages = [];
+  for (const message of baseMessages) {
+    messages.push(message);
+    for (const injectedMessage of injectedMessages) {
+      if (injectedMessage.previousId === message.id) {
+        messages.push(injectedMessage.message);
+      }
+    }
+  }
+  return messages;
+}
+
 export const ChatImpl = memo(
   ({ description, initialMessages, storeMessageHistory, importChat, exportChat }: ChatProps) => {
     useShortcuts();
@@ -117,6 +138,8 @@ export const ChatImpl = memo(
     const [uploadedFiles, setUploadedFiles] = useState<File[]>([]); // Move here
     const [imageDataList, setImageDataList] = useState<string[]>([]); // Move here
     const [searchParams, setSearchParams] = useSearchParams();
+    const [injectedMessages, setInjectedMessages] = useState<InjectedMessage[]>([]);
+    const [simulationLoading, setSimulationLoading] = useState(false);
     const files = useStore(workbenchStore.files);
     const { promptId } = useSettings();
 
@@ -124,7 +147,7 @@ export const ChatImpl = memo(
 
     const [animationScope, animate] = useAnimate();
 
-    const { messages, isLoading, input, handleInputChange, setInput, stop, append } = useChat({
+    const { messages: baseMessages, isLoading, input, handleInputChange, setInput, stop, append } = useChat({
       api: '/api/chat',
       body: {
         files,
@@ -137,20 +160,14 @@ export const ChatImpl = memo(
           'There was an error processing your request: ' + (error.message ? error.message : 'No details were returned'),
         );
       },
-      onFinish: (message, response) => {
-        const usage = response.usage;
-
-        if (usage) {
-          console.log('Token usage:', usage);
-
-          // You can now use the usage data as needed
-        }
-
-        logger.debug('Finished streaming');
-      },
       initialMessages,
       initialInput: Cookies.get(PROMPT_COOKIE_KEY) || '',
     });
+
+    const messages = useMemo(() => {
+      return handleInjectMessages(baseMessages, injectedMessages);
+    }, [baseMessages, injectedMessages]);
+
     useEffect(() => {
       const prompt = searchParams.get('prompt');
 
@@ -198,8 +215,10 @@ export const ChatImpl = memo(
 
     const abort = () => {
       stop();
+      gNumAborts++;
       chatStore.setKey('aborted', true);
       workbenchStore.abortAllActions();
+      setSimulationLoading(false);
     };
 
     useEffect(() => {
@@ -230,13 +249,54 @@ export const ChatImpl = memo(
       setChatStarted(true);
     };
 
+    const createRecording = async (simulationData: SimulationData, repositoryContents: string) => {
+      let recordingId, message;
+      try {
+        recordingId = await getSimulationRecording(simulationData, repositoryContents);
+        message = `[Recording of the bug](https://app.replay.io/recording/${recordingId})\n\n`;
+      } catch (e) {
+        console.error("Error creating recording", e);
+        message = "Error creating recording.";
+      }
+
+      const recordingMessage: Message = {
+        id: `create-recording-${messages.length}`,
+        role: 'assistant',
+        content: message,
+      };
+
+      return { recordingId, recordingMessage };
+    };
+
+    const getEnhancedPrompt = async (recordingId: string, repositoryContents: string) => {
+      let enhancedPrompt, message;
+      try {
+        enhancedPrompt = await getSimulationEnhancedPrompt(recordingId, repositoryContents);
+        message = `Explanation of the bug: ${enhancedPrompt}`;
+      } catch (e) {
+        console.error("Error enhancing prompt", e);
+        message = "Error enhancing prompt.";
+      }
+
+      const enhancedPromptMessage: Message = {
+        id: `enhanced-prompt-${messages.length}`,
+        role: 'assistant',
+        content: message,
+      };
+
+      return { enhancedPrompt, enhancedPromptMessage };
+    }
+
     const sendMessage = async (_event: React.UIEvent, messageInput?: string, simulation?: boolean) => {
       const _input = messageInput || input;
+      const numAbortsAtStart = gNumAborts;
 
       if (_input.length === 0 || isLoading) {
         return;
       }
 
+      setSimulationLoading(true);
+
       /**
        * @note (delm) Usually saving files shouldn't take long but it may take longer if there
        * many unsaved files. In that case we need to block user input and show an indicator
@@ -248,16 +308,31 @@ export const ChatImpl = memo(
 
       const { contentBase64 } = await workbenchStore.generateZipBase64();
 
-      let simulationClientData: SimulationPromptClientData | undefined;
+      let simulationEnhancedPrompt: string | undefined;
+
       if (simulation) {
         const simulationData = await getIFrameSimulationData(getCurrentIFrame());
-        const mouseData = getCurrentMouseData();
+        const { recordingId, recordingMessage } = await createRecording(simulationData, contentBase64);
 
-        simulationClientData = {
+        if (numAbortsAtStart != gNumAborts) {
-          simulationData: simulationData,
+          return;
-          repositoryContents: contentBase64,
+        }
-          mouseData: mouseData,
+
-        };
+        console.log("RecordingMessage", recordingMessage);
+        setInjectedMessages([...injectedMessages, { message: recordingMessage, previousId: messages[messages.length - 1].id }]);
+
+        if (recordingId) {
+          const info = await getEnhancedPrompt(recordingId, contentBase64);
+
+          if (numAbortsAtStart != gNumAborts) {
+            return;
+          }
+  
+          simulationEnhancedPrompt = info.enhancedPrompt;
+
+          console.log("EnhancedPromptMessage", info.enhancedPromptMessage);
+          setInjectedMessages([...injectedMessages, { message: info.enhancedPromptMessage, previousId: messages[messages.length - 1].id }]);
+        }
       }
   
       const fileModifications = workbenchStore.getFileModifcations();
@@ -266,14 +341,8 @@ export const ChatImpl = memo(
 
       runAnimation();
 
-      if (fileModifications !== undefined) {
+      setSimulationLoading(false);
-        /**
+
-         * If we have file modifications we append a new user message manually since we have to prefix
-         * the user input with the file modifications and we don't want the new user input to appear
-         * in the prompt. Using `append` is almost the same as `handleSubmit` except that we have to
-         * manually reset the input and we'd have to manually pass in file attachments. However, those
-         * aren't relevant here.
-         */
       append({
         role: 'user',
         content: [
@@ -286,27 +355,14 @@ export const ChatImpl = memo(
             image: imageData,
           })),
         ] as any, // Type assertion to bypass compiler check
-        });
+      }, { body: { simulationEnhancedPrompt } });
 
+      if (fileModifications !== undefined) {
         /**
          * After sending a new message we reset all modifications since the model
          * should now be aware of all the changes.
          */
         workbenchStore.resetAllFileModifications();
-      } else {
-        append({
-          role: 'user',
-          content: [
-            {
-              type: 'text',
-              text: _input,
-            },
-            ...imageDataList.map((imageData) => ({
-              type: 'image',
-              image: imageData,
-            })),
-          ] as any, // Type assertion to bypass compiler check
-        }, { body: { simulationClientData } });
       }
 
       setInput('');
@@ -355,7 +411,7 @@ export const ChatImpl = memo(
         input={input}
         showChat={showChat}
         chatStarted={chatStarted}
-        isStreaming={isLoading}
+        isStreaming={isLoading || simulationLoading}
         enhancingPrompt={enhancingPrompt}
         promptEnhanced={promptEnhanced}
         sendMessage={sendMessage}

app/lib/replay/SimulationPrompt.ts

@@ -1,19 +1,10 @@
-// Core logic for using simulation data from remote recording to enhance
+// Core logic for using simulation data from a remote recording to enhance
 // the AI developer prompt.
 
-// Currently the simulation prompt is sent from the server.
-
 import { type SimulationData, type MouseData } from './Recording';
 import { assert, ProtocolClient, sendCommandDedicatedClient } from './ReplayProtocolClient';
 import JSZip from 'jszip';
 
-// Data supplied by the client for a simulation prompt, separate from the chat input.
-export interface SimulationPromptClientData {
-  simulationData: SimulationData;
-  repositoryContents: string; // base64 encoded zip file
-  mouseData?: MouseData;
-}
-
 interface RerecordGenerateParams {
   rerecordData: SimulationData;
   repositoryContents: string;

app/routes/about.tsx

@@ -11,7 +11,11 @@ function AboutPage() {
         <BackgroundRays />
         <Header />
         <ClientOnly>{() => <Menu />}</ClientOnly>
-        <div>Hello World! About Page</div>
+        <div>
+          Nut is an open source fork of Bolt.new designed to help you more easily fix bugs
+          and make improvements to your app which AI developers struggle with. We want to be better
+          at cracking tough nuts, so to speak.
+        </div>
       </div>
     </TooltipProvider>
   );

app/routes/api.chat.ts

@@ -1,5 +1,4 @@
 import { type ActionFunctionArgs } from '@remix-run/cloudflare';
-import { type SimulationPromptClientData, getSimulationEnhancedPrompt, getSimulationRecording } from '~/lib/replay/SimulationPrompt';
 import { ChatStreamController } from '~/utils/chatStreamController';
 import { assert } from '~/lib/replay/ReplayProtocolClient';
 import { getStreamTextArguments, type Messages } from '~/lib/.server/llm/stream-text';
@@ -16,17 +15,17 @@ Focus specifically on fixing this bug. Do not guess about other problems.
 `;
 
 async function chatAction({ context, request }: ActionFunctionArgs) {
-  const { messages, files, promptId, simulationClientData } = await request.json<{
+  const { messages, files, promptId, simulationEnhancedPrompt } = await request.json<{
     messages: Messages;
     files: any;
     promptId?: string;
-    simulationClientData?: SimulationPromptClientData;
+    simulationEnhancedPrompt?: string;
   }>();
 
   let finished: (v?: any) => void;
   context.cloudflare.ctx.waitUntil(new Promise((resolve) => finished = resolve));
 
-  console.log("SimulationClientData", simulationClientData);
+  console.log("SimulationEnhancedPrompt", simulationEnhancedPrompt);
 
   try {
     const { system, messages: coreMessages } = await getStreamTextArguments({
@@ -47,37 +46,13 @@ async function chatAction({ context, request }: ActionFunctionArgs) {
       async start(controller) {
         const chatController = new ChatStreamController(controller);
 
-        let recordingId: string | undefined;
+        if (simulationEnhancedPrompt) {
-        if (simulationClientData) {
-          try {
-            const { simulationData, repositoryContents } = simulationClientData;
-            recordingId = await getSimulationRecording(simulationData, repositoryContents);
-            chatController.writeText(`[Recording of the bug](https://app.replay.io/recording/${recordingId})\n\n`);
-          } catch (e) {
-            console.error("Error creating recording", e);
-            chatController.writeText("Error creating recording.");
-          }
-        }
-
-        let enhancedPrompt: string | undefined;
-        if (recordingId) {
-          try {
-            assert(simulationClientData, "SimulationClientData is required");
-            enhancedPrompt = await getSimulationEnhancedPrompt(recordingId, simulationClientData.repositoryContents);
-            chatController.writeText(`Enhanced prompt: ${enhancedPrompt}\n\n`);
-          } catch (e) {
-            console.error("Error enhancing prompt", e);
-            chatController.writeText("Error enhancing prompt.");
-          }
-        }
-
-        if (enhancedPrompt) {
           const lastMessage = coreMessages[coreMessages.length - 1];
           assert(lastMessage.role == "user", "Last message must be a user message");
           assert(lastMessage.content.length > 0, "Last message must have content");
           const lastContent = lastMessage.content[0];
           assert(typeof lastContent == "object" && lastContent.type == "text", "Last message content must be text");
-          lastContent.text += `\n\n${EnhancedPromptPrefix}\n\n${enhancedPrompt}`;
+          lastContent.text += `\n\n${EnhancedPromptPrefix}\n\n${simulationEnhancedPrompt}`;
         }
 
         try {