Improve support for streaming simulation data to backend (#16)

app/components/chat/Chat.client.tsx

@@ -22,9 +22,8 @@ import { useSettings } from '~/lib/hooks/useSettings';
 import { useSearchParams } from '@remix-run/react';
 import { createSampler } from '~/utils/sampler';
 import { saveProjectContents } from './Messages.client';
-import { getSimulationRecording, getSimulationEnhancedPrompt } from '~/lib/replay/SimulationPrompt';
+import { getSimulationRecording, getSimulationEnhancedPrompt, simulationAddData, simulationRepositoryUpdated } from '~/lib/replay/SimulationPrompt';
 import { getIFrameSimulationData } from '~/lib/replay/Recording';
-import type { SimulationData } from '~/lib/replay/SimulationData';
 import { getCurrentIFrame } from '../workbench/Preview';
 import { getCurrentMouseData } from '../workbench/PointSelector';
 import { anthropicNumFreeUsesCookieName, anthropicApiKeyCookieName, MaxFreeUses } from '~/utils/freeUses';
@@ -38,6 +37,43 @@ const toastAnimation = cssTransition({
 
 const logger = createScopedLogger('Chat');
 
+// Debounce things after file writes to avoid creating a bunch of chats.
+let gResetChatFileWrittenTimeout: NodeJS.Timeout | undefined;
+
+export function resetChatFileWritten() {
+  clearTimeout(gResetChatFileWrittenTimeout);
+  gResetChatFileWrittenTimeout = setTimeout(async () => {
+    const { contentBase64 } = await workbenchStore.generateZipBase64();
+    await simulationRepositoryUpdated(contentBase64);
+  }, 500);
+}
+
+async function flushSimulationData() {
+  console.log("FlushSimulationData");
+
+  const iframe = getCurrentIFrame();
+  if (!iframe) {
+    return;
+  }
+  const simulationData = await getIFrameSimulationData(iframe);
+  if (!simulationData.length) {
+    return;
+  }
+
+  console.log("HaveSimulationData", simulationData.length);
+
+  // Add the simulation data to the chat.
+  await simulationAddData(simulationData);
+}
+
+let gLockSimulationData = false;
+
+setInterval(async () => {
+  if (!gLockSimulationData) {
+    flushSimulationData();
+  }
+}, 1000);
+
 export function Chat() {
   renderLogger.trace('Chat');
 
@@ -262,10 +298,10 @@ export const ChatImpl = memo(
       setChatStarted(true);
     };
 
-    const createRecording = async (simulationData: SimulationData, repositoryContents: string) => {
+    const createRecording = async () => {
       let recordingId, message;
       try {
-        recordingId = await getSimulationRecording(simulationData, repositoryContents);
+        recordingId = await getSimulationRecording();
         message = `[Recording of the bug](https://app.replay.io/recording/${recordingId})\n\n`;
       } catch (e) {
         console.error("Error creating recording", e);
@@ -281,11 +317,11 @@ export const ChatImpl = memo(
       return { recordingId, recordingMessage };
     };
 
-    const getEnhancedPrompt = async (recordingId: string, userMessage: string) => {
+    const getEnhancedPrompt = async (userMessage: string) => {
       let enhancedPrompt, message;
       try {
         const mouseData = getCurrentMouseData();
-        enhancedPrompt = await getSimulationEnhancedPrompt(recordingId, messages, userMessage, mouseData);
+        enhancedPrompt = await getSimulationEnhancedPrompt(messages, userMessage, mouseData);
         message = `Explanation of the bug:\n\n${enhancedPrompt}`;
       } catch (e) {
         console.error("Error enhancing prompt", e);
@@ -331,35 +367,42 @@ export const ChatImpl = memo(
        */
       await workbenchStore.saveAllFiles();
 
-      const { contentBase64 } = await workbenchStore.generateZipBase64();
-
       let simulationEnhancedPrompt: string | undefined;
 
       if (simulation) {
-        const simulationData = await getIFrameSimulationData(getCurrentIFrame());
+        gLockSimulationData = true;
-        const { recordingId, recordingMessage } = await createRecording(simulationData, contentBase64);
+        try {
+          await flushSimulationData();
 
-        if (numAbortsAtStart != gNumAborts) {
+          const createRecordingPromise = createRecording();
-          return;
+          const enhancedPromptPromise = getEnhancedPrompt(_input);
-        }
 
-        console.log("RecordingMessage", recordingMessage);
+          const { recordingId, recordingMessage } = await createRecordingPromise;
-        setInjectedMessages([...injectedMessages, { message: recordingMessage, previousId: messages[messages.length - 1].id }]);
-
-        if (recordingId) {
-          const info = await getEnhancedPrompt(recordingId, _input);
 
           if (numAbortsAtStart != gNumAborts) {
             return;
           }
-  
-          simulationEnhancedPrompt = info.enhancedPrompt;
 
-          console.log("EnhancedPromptMessage", info.enhancedPromptMessage);
+          console.log("RecordingMessage", recordingMessage);
-          setInjectedMessages([...injectedMessages, { message: info.enhancedPromptMessage, previousId: messages[messages.length - 1].id }]);
+          setInjectedMessages([...injectedMessages, { message: recordingMessage, previousId: messages[messages.length - 1].id }]);
+
+          if (recordingId) {
+            const info = await enhancedPromptPromise;
+
+            if (numAbortsAtStart != gNumAborts) {
+              return;
+            }
+
+            simulationEnhancedPrompt = info.enhancedPrompt;
+
+            console.log("EnhancedPromptMessage", info.enhancedPromptMessage);
+            setInjectedMessages([...injectedMessages, { message: info.enhancedPromptMessage, previousId: messages[messages.length - 1].id }]);
+          }
+        } finally {
+          gLockSimulationData = false;
         }
       }
-  
+
       const fileModifications = workbenchStore.getFileModifcations();
 
       chatStore.setKey('aborted', false);
@@ -404,6 +447,7 @@ export const ChatImpl = memo(
       // The project contents are associated with the last message present when
       // the user message is added.
       const lastMessage = messages[messages.length - 1];
+      const { contentBase64 } = await workbenchStore.generateZipBase64();
       saveProjectContents(lastMessage.id, { content: contentBase64 });
     };
 

app/components/workbench/Preview.tsx

@@ -2,16 +2,15 @@ import { useStore } from '@nanostores/react';
 import { memo, useCallback, useEffect, useRef, useState } from 'react';
 import { IconButton } from '~/components/ui/IconButton';
 import { workbenchStore } from '~/lib/stores/workbench';
+import { simulationReloaded } from '~/lib/replay/SimulationPrompt';
 import { PortDropdown } from './PortDropdown';
 import { PointSelector } from './PointSelector';
-import { assert } from '~/lib/replay/ReplayProtocolClient';
 
 type ResizeSide = 'left' | 'right' | null;
 
 let gCurrentIFrame: HTMLIFrameElement | undefined;
 
 export function getCurrentIFrame() {
-  assert(gCurrentIFrame);
   return gCurrentIFrame;
 }
 
@@ -125,6 +124,7 @@ export const Preview = memo(() => {
 
   const reloadPreview = () => {
     if (iframeRef.current) {
+      simulationReloaded();
       iframeRef.current.src = iframeRef.current.src;
     }
     setIsSelectionMode(false);

app/lib/replay/Recording.ts

@@ -6,6 +6,7 @@ import type {
   LocalStorageAccess,
   NetworkResource,
   SimulationData,
+  SimulationPacket,
   UserInteraction,
 } from './SimulationData';
 
@@ -55,6 +56,8 @@ function sendIframeRequest<K extends keyof RequestMap>(
   });
 }
 
+let gMessageCount = 0;
+
 export async function getIFrameSimulationData(iframe: HTMLIFrameElement): Promise<SimulationData> {
   const buffer = await sendIframeRequest(iframe, { request: 'recording-data' });
   const decoder = new TextDecoder();
@@ -76,14 +79,30 @@ export async function getMouseData(iframe: HTMLIFrameElement, position: { x: num
 }
 
 // Add handlers to the current iframe's window.
-function addRecordingMessageHandler() {
+function addRecordingMessageHandler(messageHandlerId: string) {
-  const resources: NetworkResource[] = [];
+  const simulationData: SimulationData = [];
-  const interactions: UserInteraction[] = [];
+  let numSimulationPacketsSent = 0;
-  const indexedDBAccesses: IndexedDBAccess[] = [];
+
-  const localStorageAccesses: LocalStorageAccess[] = [];
+  function pushSimulationData(packet: SimulationPacket) {
+    packet.time = new Date().toISOString();
+    simulationData.push(packet);
+  }
 
   const startTime = Date.now();
 
+  pushSimulationData({
+    kind: 'viewport',
+    size: { width: window.innerWidth, height: window.innerHeight },
+  });
+  pushSimulationData({
+    kind: "locationHref",
+    href: window.location.href,
+  });
+  pushSimulationData({
+    kind: "documentURL",
+    url: window.location.href,
+  });
+
   interface RequestInfo {
     url: string;
     requestBody: string;
@@ -93,9 +112,16 @@ function addRecordingMessageHandler() {
     return Math.min(Math.max(value, min), max);
   }
 
+  function addNetworkResource(resource: NetworkResource) {
+    pushSimulationData({
+      kind: "resource",
+      resource,
+    });
+  }
+
   function addTextResource(info: RequestInfo, text: string, responseHeaders: Record<string, string>) {
-    const url = new URL(info.url, window.location.href).href;
+    const url = (new URL(info.url, window.location.href)).href;
-    resources.push({
+    addNetworkResource({
       url,
       requestBodyBase64: stringToBase64(info.requestBody),
       responseBodyBase64: stringToBase64(text),
@@ -104,50 +130,31 @@ function addRecordingMessageHandler() {
     });
   }
 
+  function addInteraction(interaction: UserInteraction) {
+    pushSimulationData({
+      kind: "interaction",
+      interaction,
+    });
+  }
+
+  function addIndexedDBAccess(access: IndexedDBAccess) {
+    pushSimulationData({
+      kind: "indexedDB",
+      access,
+    });
+  }
+
+  function addLocalStorageAccess(access: LocalStorageAccess) {
+    pushSimulationData({
+      kind: "localStorage",
+      access,
+    });
+  }
+
   async function getSimulationData(): Promise<SimulationData> {
-    const data: SimulationData = [];
+    console.log("GetSimulationData", simulationData.length, numSimulationPacketsSent);
-
+    const data = simulationData.slice(numSimulationPacketsSent);
-    /*
+    numSimulationPacketsSent = simulationData.length;
-     * for now we only store the viewport size at the time of the simulation data request
-     * we don't deal with resizes during lifetime of the app
-     */
-    data.push({
-      kind: 'viewport',
-      size: { width: window.innerWidth, height: window.innerHeight },
-    });
-    data.push({
-      kind: 'locationHref',
-      href: window.location.href,
-    });
-    data.push({
-      kind: 'documentURL',
-      url: window.location.href,
-    });
-    for (const resource of resources) {
-      data.push({
-        kind: 'resource',
-        resource,
-      });
-    }
-    for (const interaction of interactions) {
-      data.push({
-        kind: 'interaction',
-        interaction,
-      });
-    }
-    for (const indexedDBAccess of indexedDBAccesses) {
-      data.push({
-        kind: 'indexedDB',
-        access: indexedDBAccess,
-      });
-    }
-    for (const localStorageAccess of localStorageAccesses) {
-      data.push({
-        kind: 'localStorage',
-        access: localStorageAccess,
-      });
-    }
-
     return data;
   }
 
@@ -265,7 +272,7 @@ function addRecordingMessageHandler() {
     'click',
     (event) => {
       if (event.target) {
-        interactions.push({
+        addInteraction({
           kind: 'click',
           time: Date.now() - startTime,
           ...getMouseEventTargetData(event),
@@ -281,7 +288,7 @@ function addRecordingMessageHandler() {
     'pointermove',
     (event) => {
       if (event.target) {
-        interactions.push({
+        addInteraction({
           kind: 'pointermove',
           time: Date.now() - startTime,
           ...getMouseEventTargetData(event),
@@ -295,7 +302,7 @@ function addRecordingMessageHandler() {
     'keydown',
     (event) => {
       if (event.key) {
-        interactions.push({
+        addInteraction({
           kind: 'keydown',
           time: Date.now() - startTime,
           ...getKeyboardEventTargetData(event),
@@ -348,7 +355,7 @@ function addRecordingMessageHandler() {
   };
 
   function pushIndexedDBAccess(request: IDBRequest, kind: IndexedDBAccess['kind'], key: any, item: any) {
-    indexedDBAccesses.push({
+    addIndexedDBAccess({
       kind,
       key,
       item,
@@ -393,7 +400,7 @@ function addRecordingMessageHandler() {
   };
 
   function pushLocalStorageAccess(kind: LocalStorageAccess['kind'], key: string, value?: string) {
-    localStorageAccesses.push({ kind, key, value });
+    addLocalStorageAccess({ kind, key, value });
   }
 
   const StorageMethods = {
@@ -506,7 +513,7 @@ function addRecordingMessageHandler() {
       responseToRequestInfo.set(rv, requestInfo);
       return createProxy(rv);
     } catch (error) {
-      resources.push({
+      addNetworkResource({
         url,
         requestBodyBase64: stringToBase64(requestBody),
         error: String(error),

app/lib/replay/ReplayProtocolClient.ts

@@ -2,6 +2,7 @@ const replayWsServer = "wss://dispatch.replay.io";
 
 export function assert(condition: any, message: string = "Assertion failed!"): asserts condition {
   if (!condition) {
+    debugger;
     throw new Error(message);
   }
 }

app/lib/replay/SimulationData.ts

@@ -166,7 +166,7 @@ interface SimulationPacketLocalStorage {
   access: LocalStorageAccess;
 }
 
-export type SimulationPacket =
+type SimulationPacketBase =
   | SimulationPacketServerURL
   | SimulationPacketRepositoryContents
   | SimulationPacketViewport
@@ -178,4 +178,5 @@ export type SimulationPacket =
   | SimulationPacketIndexedDB
   | SimulationPacketLocalStorage;
 
+export type SimulationPacket = SimulationPacketBase & { time?: string };
 export type SimulationData = SimulationPacket[];

app/lib/replay/SimulationPrompt.ts

@@ -2,48 +2,17 @@
 // the AI developer prompt.
 
 import type { Message } from 'ai';
-import type { SimulationData } from './SimulationData';
+import type { SimulationData, SimulationPacket } from './SimulationData';
 import { SimulationDataVersion } from './SimulationData';
 import { assert, ProtocolClient } from './ReplayProtocolClient';
 import type { MouseData } from './Recording';
 
-export async function getSimulationRecording(
+function createRepositoryContentsPacket(contents: string) {
-  interactionData: SimulationData,
+  return {
-  repositoryContents: string
+    kind: "repositoryContents",
-): Promise<string> {
+    contents,
-  const client = new ProtocolClient();
+    time: new Date().toISOString(),
-  await client.initialize();
+  };
-  try {
-    const { chatId } = await client.sendCommand({ method: "Nut.startChat", params: {} }) as { chatId: string };
-
-    const repositoryContentsPacket = {
-      kind: "repositoryContents",
-      contents: repositoryContents,
-    };
-
-    const simulationData = [repositoryContentsPacket, ...interactionData];
-
-    console.log("SimulationData", JSON.stringify(simulationData));
-
-    const { recordingId } = await client.sendCommand({
-      method: "Nut.addSimulation",
-      params: {
-        chatId,
-        version: SimulationDataVersion,
-        simulationData,
-        completeData: true,
-        saveRecording: true,
-      },
-    }) as { recordingId: string | undefined };
-
-    if (!recordingId) {
-      throw new Error("Expected recording ID in result");
-    }
-
-    return recordingId;
-  } finally {
-    client.close();
-  }
 }
 
 type ProtocolMessage = {
@@ -52,6 +21,175 @@ type ProtocolMessage = {
   content: string;
 };
 
+class ChatManager {
+  // Empty if this chat has been destroyed.
+  client: ProtocolClient | undefined;
+
+  // Resolves when the chat has started.
+  chatIdPromise: Promise<string>;
+
+  // Resolves when the recording has been created.
+  recordingIdPromise: Promise<string> | undefined;
+
+  // Whether all simulation data has been sent.
+  simulationFinished?: boolean;
+
+  // Any repository contents we sent up for this chat.
+  repositoryContents?: string;
+
+  // Simulation data for the page itself and any user interactions.
+  pageData: SimulationData = [];
+
+  constructor() {
+    this.client = new ProtocolClient();
+    this.chatIdPromise = (async () => {
+      assert(this.client, "Chat has been destroyed");
+
+      await this.client.initialize();
+      await this.client.sendCommand({
+        method: "Recording.globalExperimentalCommand",
+        params: { name: "enableOperatorPods" },
+      });
+  
+      const { chatId } = (await this.client.sendCommand({ method: "Nut.startChat", params: {} })) as { chatId: string };
+      return chatId;  
+    })();
+  }
+
+  destroy() {
+    this.client?.close();
+    this.client = undefined;
+  }
+
+  async setRepositoryContents(contents: string) {
+    assert(this.client, "Chat has been destroyed");
+    this.repositoryContents = contents;
+
+    const packet = createRepositoryContentsPacket(contents);
+
+    const chatId = await this.chatIdPromise;
+    await this.client.sendCommand({
+      method: "Nut.addSimulation",
+      params: {
+        chatId,
+        version: SimulationDataVersion,
+        simulationData: [packet],
+        completeData: false,
+        saveRecording: true,
+      },
+    });
+  }
+
+  async addPageData(data: SimulationData) {
+    assert(this.client, "Chat has been destroyed");
+    assert(this.repositoryContents, "Expected repository contents");
+
+    this.pageData.push(...data);
+
+    // If page data comes in while we are waiting for the chat to finish
+    // we remember it but don't update the existing chat.
+    if (this.simulationFinished) {
+      return;
+    }
+
+    const chatId = await this.chatIdPromise;
+    await this.client.sendCommand({
+      method: "Nut.addSimulationData",
+      params: { chatId, simulationData: data },
+    });
+  }
+
+  finishSimulationData() {
+    assert(this.client, "Chat has been destroyed");
+    assert(!this.simulationFinished, "Simulation has been finished");
+    assert(this.repositoryContents, "Expected repository contents");
+
+    this.recordingIdPromise = (async () => {
+      assert(this.client, "Chat has been destroyed");
+  
+      const chatId = await this.chatIdPromise;
+      const { recordingId } = await this.client.sendCommand({
+        method: "Nut.finishSimulationData",
+        params: { chatId },
+      }) as { recordingId: string | undefined };
+
+      assert(recordingId, "Recording ID not set");
+      return recordingId;
+    })();  
+
+    const allData = [createRepositoryContentsPacket(this.repositoryContents), ...this.pageData];
+    this.simulationFinished = true;
+    return allData;
+  }
+
+  async sendChatMessage(messages: ProtocolMessage[]) {
+    assert(this.client, "Chat has been destroyed");
+
+    let response: string = "";
+    this.client.listenForMessage("Nut.chatResponsePart", ({ message }: { message: ProtocolMessage }) => {
+      console.log("ChatResponsePart", message);
+      response += message.content;
+    });
+
+    const responseId = "<response-id>";
+    const chatId = await this.chatIdPromise;
+    await this.client.sendCommand({
+      method: "Nut.sendChatMessage",
+      params: { chatId, responseId, messages },
+    });
+
+    return response;
+  }
+}
+
+// There is only one chat active at a time.
+let gChatManager: ChatManager | undefined;
+
+function startChat(repositoryContents: string, pageData: SimulationData) {
+  if (gChatManager) {
+    gChatManager.destroy();
+  }
+  gChatManager = new ChatManager();
+
+  gChatManager.setRepositoryContents(repositoryContents);
+  if (pageData.length) {
+    gChatManager.addPageData(pageData);
+  }
+}
+
+// Called when the repository contents have changed. We'll start a new chat
+// with the same interaction data as any existing chat.
+export async function simulationRepositoryUpdated(repositoryContents: string) {
+  startChat(repositoryContents, gChatManager?.pageData ?? []);
+}
+
+// Called when the page gathering interaction data has been reloaded. We'll
+// start a new chat with the same repository contents as any existing chat.
+export async function simulationReloaded() {
+  assert(gChatManager, "Expected to have an active chat");
+
+  const repositoryContents = gChatManager.repositoryContents;
+  assert(repositoryContents, "Expected active chat to have repository contents");
+
+  startChat(repositoryContents, []);
+}
+
+export async function simulationAddData(data: SimulationData) {
+  assert(gChatManager, "Expected to have an active chat");
+  gChatManager.addPageData(data);
+}
+
+export async function getSimulationRecording(): Promise<string> {
+  assert(gChatManager, "Expected to have an active chat");
+
+  const simulationData = gChatManager.finishSimulationData();
+
+  console.log("SimulationData", new Date().toISOString(), JSON.stringify(simulationData));
+
+  assert(gChatManager.recordingIdPromise, "Expected recording promise");
+  return gChatManager.recordingIdPromise;
+}
+
 const SystemPrompt = `
 The following user message describes a bug or other problem on the page which needs to be fixed.
 You must respond with a useful explanation that will help the user understand the source of the problem.
@@ -59,56 +197,32 @@ Do not describe the specific fix needed.
 `;
 
 export async function getSimulationEnhancedPrompt(
-  recordingId: string,
   chatMessages: Message[],
   userMessage: string,
   mouseData: MouseData | undefined
 ): Promise<string> {
-  const client = new ProtocolClient();
+  assert(gChatManager, "Chat not started");
-  await client.initialize();
+  assert(gChatManager.simulationFinished, "Simulation not finished");
-  try {
-    const { chatId } = await client.sendCommand({ method: "Nut.startChat", params: {} }) as { chatId: string };
 
-    await client.sendCommand({
+  let system = SystemPrompt;
-      method: "Nut.addRecording",
+  if (mouseData) {
-      params: { chatId, recordingId },
+    system += `The user pointed to an element on the page <element selector=${JSON.stringify(mouseData.selector)} height=${mouseData.height} width=${mouseData.width} x=${mouseData.x} y=${mouseData.y} />`;
-    });
-
-    let system = SystemPrompt;
-    if (mouseData) {
-      system += `The user pointed to an element on the page <element selector=${JSON.stringify(mouseData.selector)} height=${mouseData.height} width=${mouseData.width} x=${mouseData.x} y=${mouseData.y} />`;
-    }
-
-    const messages = [
-      {
-        role: "system",
-        type: "text",
-        content: system,
-      },
-      {
-        role: "user",
-        type: "text",
-        content: userMessage,
-      },
-    ];
-
-    console.log("ChatSendMessage", messages);
-
-    let response: string = "";
-    const removeListener = client.listenForMessage("Nut.chatResponsePart", ({ message }: { message: ProtocolMessage }) => {
-      console.log("ChatResponsePart", message);
-      response += message.content;
-    });
-
-    const responseId = "<response-id>";
-    await client.sendCommand({
-      method: "Nut.sendChatMessage",
-      params: { chatId, responseId, messages },
-    });
-
-    removeListener();
-    return response;
-  } finally {
-    client.close();
   }
+
+  const messages: ProtocolMessage[] = [
+    {
+      role: "system",
+      type: "text",
+      content: system,
+    },
+    {
+      role: "user",
+      type: "text",
+      content: userMessage,
+    },
+  ];
+
+  console.log("ChatSendMessage", messages);
+
+  return gChatManager.sendChatMessage(messages);
 }

app/lib/runtime/action-runner.ts

@@ -6,6 +6,7 @@ import { createScopedLogger } from '~/utils/logger';
 import { unreachable } from '~/utils/unreachable';
 import type { ActionCallbackData } from './message-parser';
 import type { BoltShell } from '~/utils/shell';
+import { resetChatFileWritten } from '~/components/chat/Chat.client';
 
 const logger = createScopedLogger('ActionRunner');
 
@@ -294,6 +295,7 @@ export class ActionRunner {
 
     try {
       await webcontainer.fs.writeFile(relativePath, action.content);
+      resetChatFileWritten();
       logger.debug(`File written ${relativePath}`);
     } catch (error) {
       logger.error('Failed to write file\n\n', error);