bolt.diy/app/lib/replay/SimulationPrompt.ts

/*
 * 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);
}