bolt.diy/app/lib/replay/SimulationPrompt.ts
2025-03-21 12:50:23 -07:00

278 lines
7.6 KiB
TypeScript

/*
* Core logic for using simulation data from a remote recording to enhance
* the AI developer prompt.
*/
import type { SimulationData, SimulationPacket } from './SimulationData';
import { simulationDataVersion } from './SimulationData';
import { assert, generateRandomId, ProtocolClient } from './ReplayProtocolClient';
import { updateDevelopmentServer } from './DevelopmentServer';
import type { Message } from '~/lib/persistence/message';
function createRepositoryIdPacket(repositoryId: string): SimulationPacket {
return {
kind: 'repositoryId',
repositoryId,
time: new Date().toISOString(),
};
}
interface ChatReferenceElement {
kind: 'element';
selector: string;
width: number;
height: number;
x: number;
y: number;
}
export type ChatReference = ChatReferenceElement;
type ChatResponsePartCallback = (message: Message) => void;
class ChatManager {
// Empty if this chat has been destroyed.
client: ProtocolClient | undefined;
// Resolves when the chat has started.
chatIdPromise: Promise<string>;
// Whether all simulation data has been sent.
simulationFinished?: boolean;
// Any repository ID we specified for this chat.
repositoryId?: string;
// Simulation data for the page itself and any user interactions.
pageData: SimulationData = [];
// State to ensure that the chat manager is not destroyed until all messages finish.
private pendingMessages = 0;
private mustDestroyAfterChatFinishes = false;
constructor() {
this.client = new ProtocolClient();
this.chatIdPromise = (async () => {
assert(this.client, 'Chat has been destroyed');
await this.client.initialize();
const { chatId } = (await this.client.sendCommand({ method: 'Nut.startChat', params: {} })) as { chatId: string };
console.log('ChatStarted', new Date().toISOString(), chatId);
return chatId;
})();
}
isValid() {
return !!this.client;
}
private destroy() {
this.client?.close();
this.client = undefined;
}
destroyAfterChatFinishes() {
if (this.pendingMessages == 0) {
this.destroy();
} else {
this.mustDestroyAfterChatFinishes = true;
}
}
async setRepositoryId(repositoryId: string) {
assert(this.client, 'Chat has been destroyed');
this.repositoryId = repositoryId;
const packet = createRepositoryIdPacket(repositoryId);
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.repositoryId, 'Expected repository ID');
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;
console.log('ChatAddPageData', new Date().toISOString(), chatId, data.length);
await this.client.sendCommand({
method: 'Nut.addSimulationData',
params: { chatId, simulationData: data },
});
}
async finishSimulationData() {
assert(this.client, 'Chat has been destroyed');
assert(!this.simulationFinished, 'Simulation has been finished');
this.simulationFinished = true;
const chatId = await this.chatIdPromise;
await this.client.sendCommand({
method: 'Nut.finishSimulationData',
params: { chatId },
});
}
async sendChatMessage(messages: Message[], references: ChatReference[], onResponsePart: ChatResponsePartCallback) {
assert(this.client, 'Chat has been destroyed');
this.pendingMessages++;
const responseId = `response-${generateRandomId()}`;
const removeResponseListener = this.client.listenForMessage(
'Nut.chatResponsePart',
({ responseId: eventResponseId, message }: { responseId: string; message: Message }) => {
if (responseId == eventResponseId) {
console.log('ChatResponse', chatId, message);
onResponsePart(message);
}
},
);
const chatId = await this.chatIdPromise;
console.log('ChatSendMessage', new Date().toISOString(), chatId, JSON.stringify({ messages, references }));
await this.client.sendCommand({
method: 'Nut.sendChatMessage',
params: { chatId, responseId, messages, references },
});
console.log('ChatMessageFinished', new Date().toISOString(), chatId);
removeResponseListener();
if (--this.pendingMessages == 0 && this.mustDestroyAfterChatFinishes) {
this.destroy();
}
}
}
// There is only one chat active at a time.
let gChatManager: ChatManager | undefined;
function startChat(repositoryId: string | undefined, pageData: SimulationData) {
// Any existing chat manager won't be used anymore for new messages, but it will
// not close until its messages actually finish and any future repository updates
// occur.
if (gChatManager) {
gChatManager.destroyAfterChatFinishes();
}
gChatManager = new ChatManager();
if (repositoryId) {
gChatManager.setRepositoryId(repositoryId);
}
if (pageData.length) {
gChatManager.addPageData(pageData);
}
}
/*
* Called when the repository has changed. We'll start a new chat
* and update the remote development server.
*/
export function simulationRepositoryUpdated(repositoryId: string) {
startChat(repositoryId, []);
updateDevelopmentServer(repositoryId);
}
/*
* 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 function simulationReloaded() {
assert(gChatManager, 'Expected to have an active chat');
const repositoryId = gChatManager.repositoryId;
assert(repositoryId, 'Expected active chat to have repository ID');
startChat(repositoryId, []);
}
/*
* Called when the current message has finished with no repository change.
* We'll start a new chat with the same simulation data as the previous chat.
*/
export function simulationReset() {
assert(gChatManager, 'Expected to have an active chat');
startChat(gChatManager.repositoryId, gChatManager.pageData);
}
export function simulationAddData(data: SimulationData) {
assert(gChatManager, 'Expected to have an active chat');
gChatManager.addPageData(data);
}
let gLastUserSimulationData: SimulationData | undefined;
export function simulationFinishData() {
if (gChatManager) {
gChatManager.finishSimulationData();
gLastUserSimulationData = [...gChatManager.pageData];
}
}
export function getLastUserSimulationData(): SimulationData | undefined {
return gLastUserSimulationData;
}
export function isSimulatingOrHasFinished(): boolean {
return gChatManager?.isValid() ?? false;
}
let gLastSimulationChatMessages: Message[] | undefined;
export function getLastSimulationChatMessages(): Message[] | undefined {
return gLastSimulationChatMessages;
}
let gLastSimulationChatReferences: ChatReference[] | undefined;
export function getLastSimulationChatReferences(): ChatReference[] | undefined {
return gLastSimulationChatReferences;
}
export async function sendChatMessage(
messages: Message[],
references: ChatReference[],
onResponsePart: ChatResponsePartCallback,
) {
if (!gChatManager) {
gChatManager = new ChatManager();
}
gLastSimulationChatMessages = messages;
gLastSimulationChatReferences = references;
await gChatManager.sendChatMessage(messages, references, onResponsePart);
}