mirror of
https://github.com/stackblitz-labs/bolt.diy
synced 2025-06-26 18:26:38 +00:00
278 lines
7.6 KiB
TypeScript
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);
|
|
}
|