From b7b602016e550bdd9ee3a63ed7b91bc694178b68 Mon Sep 17 00:00:00 2001 From: Brian Hackett Date: Fri, 7 Feb 2025 11:52:19 -0800 Subject: [PATCH] Improve support for streaming simulation data to backend (#16) --- app/components/chat/Chat.client.tsx | 90 ++++++-- app/components/workbench/Preview.tsx | 4 +- app/lib/replay/Recording.ts | 119 ++++++----- app/lib/replay/ReplayProtocolClient.ts | 1 + app/lib/replay/SimulationData.ts | 3 +- app/lib/replay/SimulationPrompt.ts | 282 +++++++++++++++++-------- app/lib/runtime/action-runner.ts | 2 + 7 files changed, 335 insertions(+), 166 deletions(-) diff --git a/app/components/chat/Chat.client.tsx b/app/components/chat/Chat.client.tsx index 8997f924..4861094f 100644 --- a/app/components/chat/Chat.client.tsx +++ b/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()); - const { recordingId, recordingMessage } = await createRecording(simulationData, contentBase64); + gLockSimulationData = true; + try { + await flushSimulationData(); - if (numAbortsAtStart != gNumAborts) { - return; - } + const createRecordingPromise = createRecording(); + const enhancedPromptPromise = getEnhancedPrompt(_input); - console.log("RecordingMessage", recordingMessage); - setInjectedMessages([...injectedMessages, { message: recordingMessage, previousId: messages[messages.length - 1].id }]); - - if (recordingId) { - const info = await getEnhancedPrompt(recordingId, _input); + const { recordingId, recordingMessage } = await createRecordingPromise; if (numAbortsAtStart != gNumAborts) { return; } - - simulationEnhancedPrompt = info.enhancedPrompt; - console.log("EnhancedPromptMessage", info.enhancedPromptMessage); - setInjectedMessages([...injectedMessages, { message: info.enhancedPromptMessage, previousId: messages[messages.length - 1].id }]); + console.log("RecordingMessage", recordingMessage); + 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 }); }; diff --git a/app/components/workbench/Preview.tsx b/app/components/workbench/Preview.tsx index d1de04ba..24ac710c 100644 --- a/app/components/workbench/Preview.tsx +++ b/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); diff --git a/app/lib/replay/Recording.ts b/app/lib/replay/Recording.ts index 569c2077..8564cb4a 100644 --- a/app/lib/replay/Recording.ts +++ b/app/lib/replay/Recording.ts @@ -6,6 +6,7 @@ import type { LocalStorageAccess, NetworkResource, SimulationData, + SimulationPacket, UserInteraction, } from './SimulationData'; @@ -55,6 +56,8 @@ function sendIframeRequest( }); } +let gMessageCount = 0; + export async function getIFrameSimulationData(iframe: HTMLIFrameElement): Promise { 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() { - const resources: NetworkResource[] = []; - const interactions: UserInteraction[] = []; - const indexedDBAccesses: IndexedDBAccess[] = []; - const localStorageAccesses: LocalStorageAccess[] = []; +function addRecordingMessageHandler(messageHandlerId: string) { + const simulationData: SimulationData = []; + let numSimulationPacketsSent = 0; + + 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) { - const url = new URL(info.url, window.location.href).href; - resources.push({ + const url = (new URL(info.url, window.location.href)).href; + 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 { - const data: SimulationData = []; - - /* - * 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, - }); - } - + console.log("GetSimulationData", simulationData.length, numSimulationPacketsSent); + const data = simulationData.slice(numSimulationPacketsSent); + numSimulationPacketsSent = simulationData.length; 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), diff --git a/app/lib/replay/ReplayProtocolClient.ts b/app/lib/replay/ReplayProtocolClient.ts index df379f8f..7e0c4f66 100644 --- a/app/lib/replay/ReplayProtocolClient.ts +++ b/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); } } diff --git a/app/lib/replay/SimulationData.ts b/app/lib/replay/SimulationData.ts index 0bc2fc9d..bc58c457 100644 --- a/app/lib/replay/SimulationData.ts +++ b/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[]; diff --git a/app/lib/replay/SimulationPrompt.ts b/app/lib/replay/SimulationPrompt.ts index 589605ea..1505982a 100644 --- a/app/lib/replay/SimulationPrompt.ts +++ b/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( - interactionData: SimulationData, - repositoryContents: string -): Promise { - const client = new ProtocolClient(); - 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(); - } +function createRepositoryContentsPacket(contents: string) { + return { + kind: "repositoryContents", + contents, + time: new Date().toISOString(), + }; } 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; + + // Resolves when the recording has been created. + recordingIdPromise: Promise | 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 = ""; + 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 { + 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 { - const client = new ProtocolClient(); - await client.initialize(); - try { - const { chatId } = await client.sendCommand({ method: "Nut.startChat", params: {} }) as { chatId: string }; + assert(gChatManager, "Chat not started"); + assert(gChatManager.simulationFinished, "Simulation not finished"); - await client.sendCommand({ - method: "Nut.addRecording", - params: { chatId, recordingId }, - }); - - let system = SystemPrompt; - if (mouseData) { - system += `The user pointed to an element on the page `; - } - - 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 = ""; - await client.sendCommand({ - method: "Nut.sendChatMessage", - params: { chatId, responseId, messages }, - }); - - removeListener(); - return response; - } finally { - client.close(); + let system = SystemPrompt; + if (mouseData) { + system += `The user pointed to an element on the page `; } + + const messages: ProtocolMessage[] = [ + { + role: "system", + type: "text", + content: system, + }, + { + role: "user", + type: "text", + content: userMessage, + }, + ]; + + console.log("ChatSendMessage", messages); + + return gChatManager.sendChatMessage(messages); } diff --git a/app/lib/runtime/action-runner.ts b/app/lib/runtime/action-runner.ts index c7bd29b1..94c977cf 100644 --- a/app/lib/runtime/action-runner.ts +++ b/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);