UniqAI/bolt.diy
1
0
Fork 0
mirror of https://github.com/stackblitz-labs/bolt.diy synced 2025-06-26 18:26:38 +00:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Brian Hackett 2025-02-07 11:52:19 -08:00 committed by GitHub
parent d143863285
commit b7b602016e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 335 additions and 166 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

86
app/components/chat/Chat.client.tsx
View File

@ -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,32 +367,39 @@ 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("RecordingMessage", recordingMessage);
setInjectedMessages([...injectedMessages, { message: recordingMessage, previousId: messages[messages.length - 1].id }]);
console.log("EnhancedPromptMessage", info.enhancedPromptMessage);
setInjectedMessages([...injectedMessages, { message: info.enhancedPromptMessage, 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;
}
}
@ -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 });
};

4
app/components/workbench/Preview.tsx
View File

@ -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);

119
app/lib/replay/Recording.ts
View File

@ -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() {
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<string, string>) {
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<SimulationData> {
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),

1
app/lib/replay/ReplayProtocolClient.ts
View File

@ -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);
}
}

3
app/lib/replay/SimulationData.ts
View File

@ -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[];

282
app/lib/replay/SimulationPrompt.ts
View File

@ -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<string> {
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<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();
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 <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();
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: ProtocolMessage[] = [
{
role: "system",
type: "text",
content: system,
},
{
role: "user",
type: "text",
content: userMessage,
},
];
console.log("ChatSendMessage", messages);
return gChatManager.sendChatMessage(messages);
}

2
app/lib/runtime/action-runner.ts
View File

@ -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);