vzcode 1.41.0 → 1.43.0

package/src/server/aiChatHandler/aiEditing.ts

@@ -4,6 +4,10 @@ import {
 } from 'editcodewithai';
 import { formatMarkdownFiles } from 'llm-code-format';
 
+// Dev flag for waiting 1 second before starting the LLM function.
+// Useful for debugging and testing purposes, e.g. checking the typing indicator.
+const delayStart = false;
+
 /**
  * Performs AI editing operations using streaming with incremental OT operations
  */
@@ -25,16 +29,15 @@ export const performAIEditing = async ({
     editFormat: 'whole',
   });
 
+  if (delayStart) {
+    await new Promise((resolve) =>
+      setTimeout(resolve, 1000),
+    );
+  }
+
   // Call the LLM function which will handle streaming and incremental file updates
   const result = await llmFunction(fullPrompt);
 
-  // Clear the scratchpad and update status
-  // clearAIScratchpadAndStatus(
-  //   shareDBDoc,
-  //   chatId,
-  //   'Done editing with AI.',
-  // );
-
   runCode();
 
   return {

package/src/server/aiChatHandler/chatOperations.ts

@@ -1,6 +1,13 @@
 import { dateToTimestamp } from '@vizhub/viz-utils';
-import { diff } from '../../client/diff.js';
 import { randomId } from '../../randomId.js';
+import { ShareDBDoc } from '../../types.js';
+import { diff } from '../../ot.js';
+import {
+  VizChatId,
+  VizContent,
+  VizFileId,
+  VizFiles,
+} from '@vizhub/viz-types';
 
 /**
  * Ensures the chats object exists in the ShareDB document
@@ -18,7 +25,10 @@ export const ensureChatsExist = (shareDBDoc) => {
 /**
  * Ensures a specific chat exists in the ShareDB document
  */
-export const ensureChatExists = (shareDBDoc, chatId) => {
+export const ensureChatExists = (
+  shareDBDoc: ShareDBDoc<VizContent>,
+  chatId: VizChatId,
+) => {
   if (!shareDBDoc.data.chats[chatId]) {
     const op = diff(shareDBDoc.data, {
       ...shareDBDoc.data,
@@ -40,9 +50,9 @@ export const ensureChatExists = (shareDBDoc, chatId) => {
  * Adds a user message to the chat
  */
 export const addUserMessage = (
-  shareDBDoc,
-  chatId,
-  content,
+  shareDBDoc: ShareDBDoc<VizContent>,
+  chatId: VizChatId,
+  content: string,
 ) => {
   const userMessage = {
     id: `user-${Date.now()}`,
@@ -74,9 +84,9 @@ export const addUserMessage = (
  * Updates AI status in the chat
  */
 export const updateAIStatus = (
-  shareDBDoc,
-  chatId,
-  status,
+  shareDBDoc: ShareDBDoc<VizContent>,
+  chatId: VizChatId,
+  status: string,
 ) => {
   const op = diff(shareDBDoc.data, {
     ...shareDBDoc.data,
@@ -95,9 +105,9 @@ export const updateAIStatus = (
  * Updates AI scratchpad content
  */
 export const updateAIScratchpad = (
-  shareDBDoc,
-  chatId,
-  content,
+  shareDBDoc: ShareDBDoc<VizContent>,
+  chatId: VizChatId,
+  content: string,
 ) => {
   const op = diff(shareDBDoc.data, {
     ...shareDBDoc.data,
@@ -122,9 +132,9 @@ export const updateAIScratchpad = (
  * Clears AI scratchpad and updates status
  */
 export const clearAIScratchpadAndStatus = (
-  shareDBDoc,
-  chatId,
-  status,
+  shareDBDoc: ShareDBDoc<VizContent>,
+  chatId: VizChatId,
+  status: string,
 ) => {
   const op = diff(shareDBDoc.data, {
     ...shareDBDoc.data,
@@ -141,12 +151,108 @@ export const clearAIScratchpadAndStatus = (
 };
 
 /**
- * Adds an AI response message to the chat
+ * Creates an initial empty AI message for streaming
+ */
+export const createAIMessage = (
+  shareDBDoc: ShareDBDoc<VizContent>,
+  chatId: VizChatId,
+) => {
+  const aiMessage = {
+    id: `assistant-${Date.now()}`,
+    role: 'assistant',
+    content: '',
+    timestamp: dateToTimestamp(new Date()),
+  };
+
+  const messageOp = diff(shareDBDoc.data, {
+    ...shareDBDoc.data,
+    chats: {
+      ...shareDBDoc.data.chats,
+      [chatId]: {
+        ...shareDBDoc.data.chats[chatId],
+        messages: [
+          ...shareDBDoc.data.chats[chatId].messages,
+          aiMessage,
+        ],
+        updatedAt: dateToTimestamp(new Date()),
+      },
+    },
+  });
+  shareDBDoc.submitOp(messageOp);
+
+  return aiMessage.id;
+};
+
+/**
+ * Updates the content of an AI message during streaming
+ */
+export const updateAIMessageContent = (
+  shareDBDoc: ShareDBDoc<VizContent>,
+  chatId: VizChatId,
+  messageId: string,
+  content: string,
+) => {
+  const chat = shareDBDoc.data.chats[chatId];
+  const messageIndex = chat.messages.findIndex(
+    (msg) => msg.id === messageId,
+  );
+
+  if (messageIndex === -1) {
+    console.warn(
+      `AI message with id ${messageId} not found`,
+    );
+    return;
+  }
+
+  const updatedMessages = [...chat.messages];
+  updatedMessages[messageIndex] = {
+    ...updatedMessages[messageIndex],
+    content,
+  };
+
+  const messageOp = diff(shareDBDoc.data, {
+    ...shareDBDoc.data,
+    chats: {
+      ...shareDBDoc.data.chats,
+      [chatId]: {
+        ...chat,
+        messages: updatedMessages,
+        updatedAt: dateToTimestamp(new Date()),
+      },
+    },
+  });
+  shareDBDoc.submitOp(messageOp);
+};
+
+/**
+ * Finalizes an AI message by clearing temporary fields
+ */
+export const finalizeAIMessage = (
+  shareDBDoc: ShareDBDoc<VizContent>,
+  chatId: VizChatId,
+) => {
+  const op = diff(shareDBDoc.data, {
+    ...shareDBDoc.data,
+    chats: {
+      ...shareDBDoc.data.chats,
+      [chatId]: {
+        ...shareDBDoc.data.chats[chatId],
+        aiScratchpad: undefined,
+        aiStatus: undefined,
+        updatedAt: dateToTimestamp(new Date()),
+      },
+    },
+  });
+  shareDBDoc.submitOp(op);
+};
+
+/**
+ * Adds an AI response message to the chat (legacy function, kept for compatibility)
  */
 export const addAIMessage = (
-  shareDBDoc,
-  chatId,
-  content,
+  shareDBDoc: ShareDBDoc<VizContent>,
+  chatId: VizChatId,
+  content?: string,
 ) => {
   const aiResponse = {
     id: Date.now() + 1,
@@ -178,7 +284,10 @@ export const addAIMessage = (
 /**
  * Updates files in the ShareDB document
  */
-export const updateFiles = (shareDBDoc, files) => {
+export const updateFiles = (
+  shareDBDoc: ShareDBDoc<VizContent>,
+  files: VizFiles,
+) => {
   const filesOp = diff(shareDBDoc.data, {
     ...shareDBDoc.data,
     files,
@@ -186,30 +295,6 @@ export const updateFiles = (shareDBDoc, files) => {
   shareDBDoc.submitOp(filesOp);
 };
 
-/**
- * Sets isInteracting flag
- */
-export const setIsInteracting = (
-  shareDBDoc,
-  isInteracting,
-) => {
-  // Only generate an operation if the value is actually changing
-  const currentIsInteracting =
-    shareDBDoc.data.isInteracting;
-
-  if (currentIsInteracting === isInteracting) {
-    return;
-  }
-
-  const newState = {
-    ...shareDBDoc.data,
-    isInteracting: isInteracting,
-  };
-
-  const interactingOp = diff(shareDBDoc.data, newState);
-  shareDBDoc.submitOp(interactingOp);
-};
-
 /**
  * Finds a file ID by searching for a matching file name
  */
@@ -233,7 +318,10 @@ export const resolveFileId = (
 /**
  * Creates a new file with a random ID
  */
-export const createNewFile = (shareDBDoc, fileName) => {
+export const createNewFile = (
+  shareDBDoc: ShareDBDoc<VizContent>,
+  fileName: string,
+) => {
   // Generate a new random file ID
   const newFileId = randomId();
 
@@ -270,7 +358,10 @@ export const ensureFileExists = (shareDBDoc, fileName) => {
 /**
  * Clears the content of a file
  */
-export const clearFileContent = (shareDBDoc, fileId) => {
+export const clearFileContent = (
+  shareDBDoc: ShareDBDoc<VizContent>,
+  fileId: VizFileId,
+) => {
   const currentFile = shareDBDoc.data.files[fileId];
 
   if (currentFile && currentFile.text) {
@@ -295,29 +386,24 @@ export const clearFileContent = (shareDBDoc, fileId) => {
  * Appends a line to a file using OT operations
  */
 export const appendLineToFile = (
-  shareDBDoc,
-  fileId,
-  line,
+  shareDBDoc: ShareDBDoc<VizContent>,
+  fileId: VizFileId,
+  line: string,
 ) => {
   const currentFile = shareDBDoc.data.files[fileId];
   const currentContent = currentFile?.text || '';
   const newContent = currentContent + line + '\n';
 
-  // Create the new file state
-  const newFileState = {
-    ...currentFile,
-    text: newContent,
-  };
-
   const newDocState = {
     ...shareDBDoc.data,
     files: {
       ...shareDBDoc.data.files,
-      [fileId]: newFileState,
+      [fileId]: {
+        ...currentFile,
+        text: newContent,
+      },
     },
   };
 
-  // Generate OT operation using the diff utility
-  const op = diff(shareDBDoc.data, newDocState);
-  shareDBDoc.submitOp(op);
+  shareDBDoc.submitOp(diff(shareDBDoc.data, newDocState));
 };

package/src/server/aiChatHandler/errorHandling.ts

@@ -1,5 +1,5 @@
 import { dateToTimestamp } from '@vizhub/viz-utils';
-import { diff } from '../../client/diff.js';
+import { diff } from '../../ot.js';
 
 /**
  * Handles errors by adding an error message to the chat and clearing AI state

package/src/server/aiChatHandler/index.ts

@@ -10,6 +10,7 @@ import { handleError } from './errorHandling.js';
 import { createRunCodeFunction } from '../../runCode.js';
 import { ShareDBDoc } from '../../types.js';
 import { VizContent } from '@vizhub/viz-types';
+import { createSubmitOperation } from '../../submitOperation.js';
 
 const DEBUG = false;
 
@@ -32,6 +33,8 @@ export const handleAIChatMessage =
         content,
         'chatId:',
         chatId,
+        'shareDBDoc:',
+        shareDBDoc,
       );
     }
 
@@ -55,8 +58,11 @@ export const handleAIChatMessage =
         chatId,
       });
 
-      // Create server-side runCode function using shared module
-      const runCode = createRunCodeFunction(shareDBDoc);
+      // Create server-side runCode function using shareDBDoc
+      const submitOperation =
+        createSubmitOperation(shareDBDoc);
+      const runCode =
+        createRunCodeFunction(submitOperation);
 
       // Perform AI editing
       const editResult = await performAIEditing({

package/src/server/aiChatHandler/llmStreaming.ts

@@ -1,16 +1,38 @@
-import { StreamingMarkdownParser } from 'llm-code-format';
+import {
+  parseMarkdownFiles,
+  StreamingMarkdownParser,
+} from 'llm-code-format';
 import { ChatOpenAI } from '@langchain/openai';
 import fs from 'fs';
 import {
   updateAIStatus,
-  updateAIScratchpad,
+  createAIMessage,
+  updateAIMessageContent,
+  finalizeAIMessage,
   ensureFileExists,
   clearFileContent,
   appendLineToFile,
+  updateFiles,
 } from './chatOperations.js';
+import { mergeFileChanges } from 'editcodewithai';
 
 const DEBUG = false;
 
+// Feature flag to enable/disable streaming editing.
+// * If `true`, the AI streaming response will be used to
+//   edit files in real-time by submitting ShareDB ops.
+// * If `false`, the updates to code files will be applied
+// only after the AI has finished generating the entire response.
+//
+// Current status: there's a tricky bug with the streaming
+// where the AI edits sometimes don't apply correctly in CodeMirror.
+// It seems that sometimes the op that clears the file content
+// is not applied correctly in the front end, leading to
+// a situation where the AI streaming edits are concatenated into the middle
+// of the file instead of replacing it.
+// See https://github.com/codemirror/codemirror.next/issues/1234
+const enableStreamingEditing = false;
+
 /**
  * Creates and configures the LLM function for streaming
  */
@@ -19,23 +41,8 @@ export const createLLMFunction = ({
   createVizBotLocalPresence,
   chatId,
 }) => {
-  return async (fullPrompt) => {
+  return async (fullPrompt: string) => {
     const localPresence = createVizBotLocalPresence();
-    // Submit initial presence for VizBot
-    const vizBotPresence = {
-      username: 'VizBot',
-      start: ['files'], // Indicate VizBot is working on files
-      end: ['files'],
-    };
-
-    localPresence.submit(vizBotPresence, (error) => {
-      if (error) {
-        console.warn(
-          'VizBot presence submission error:',
-          error,
-        );
-      }
-    });
     const chatModel = new ChatOpenAI({
       modelName:
         process.env.VIZHUB_EDIT_WITH_AI_MODEL_NAME ||
@@ -56,13 +63,21 @@ export const createLLMFunction = ({
     let currentEditingFileId = null;
     let currentEditingFileName = null;
 
+    // Create initial AI message for streaming
+    const aiMessageId = createAIMessage(shareDBDoc, chatId);
+
     // Function to report file edited
     // This is called when the AI has finished editing a file
-    // and we want to update the scratchpad with the file name.
+    // and we want to update the message content with the file name.
     const reportFileEdited = () => {
       if (currentEditingFileName) {
         fullContent += ` * Edited ${currentEditingFileName}\n`;
-        updateAIScratchpad(shareDBDoc, chatId, fullContent);
+        updateAIMessageContent(
+          shareDBDoc,
+          chatId,
+          aiMessageId,
+          fullContent,
+        );
         currentEditingFileName = null;
       }
     };
@@ -91,22 +106,6 @@ export const createLLMFunction = ({
         // (AI will regenerate the entire file content)
         clearFileContent(shareDBDoc, currentEditingFileId);
 
-        // Update VizBot presence to show it's editing this specific file
-        const filePresence = {
-          username: 'VizBot',
-          start: ['files', currentEditingFileId, 'text', 0],
-          end: ['files', currentEditingFileId, 'text', 0],
-        };
-
-        localPresence.submit(filePresence, (error) => {
-          if (error) {
-            console.warn(
-              'VizBot file presence submission error:',
-              error,
-            );
-          }
-        });
-
         // Update AI status
         updateAIStatus(
           shareDBDoc,
@@ -117,6 +116,7 @@ export const createLLMFunction = ({
       onCodeLine: async (line: string) => {
         DEBUG && console.log(`Code line: ${line}`);
 
+        // If streaming is enabled, we apply the line immediately
         if (currentEditingFileId) {
           // Apply OT operation for this line immediately
           appendLineToFile(
@@ -165,7 +165,12 @@ export const createLLMFunction = ({
           reportFileEdited();
         }
         fullContent += line + '\n';
-        updateAIScratchpad(shareDBDoc, chatId, fullContent);
+        updateAIMessageContent(
+          shareDBDoc,
+          chatId,
+          aiMessageId,
+          fullContent,
+        );
       },
     };
 
@@ -176,10 +181,21 @@ export const createLLMFunction = ({
     // Stream the response
     const stream = await chatModel.stream(fullPrompt);
     for await (const chunk of stream) {
-      if (chunk.content) {
+      if (chunk && chunk.content) {
         const chunkContent = String(chunk.content);
         chunks.push(chunkContent);
-        await parser.processChunk(chunkContent);
+
+        if (enableStreamingEditing) {
+          await parser.processChunk(chunkContent);
+        } else {
+          fullContent += chunkContent;
+          updateAIMessageContent(
+            shareDBDoc,
+            chatId,
+            aiMessageId,
+            fullContent,
+          );
+        }
       }
 
       if (!generationId && chunk.lc_kwargs?.id) {
@@ -189,7 +205,15 @@ export const createLLMFunction = ({
     await parser.flushRemaining();
     reportFileEdited();
     updateAIStatus(shareDBDoc, chatId, 'Done editing.');
-    updateAIScratchpad(shareDBDoc, chatId, fullContent);
+    updateAIMessageContent(
+      shareDBDoc,
+      chatId,
+      aiMessageId,
+      fullContent,
+    );
+
+    // Finalize the AI message by clearing temporary fields
+    finalizeAIMessage(shareDBDoc, chatId);
 
     // Clear VizBot presence when done
     DEBUG &&
@@ -206,6 +230,18 @@ export const createLLMFunction = ({
       }
     });
 
+    // If streaming editing is not enabled, we need to
+    // apply all the edits at once
+    if (!enableStreamingEditing) {
+      updateFiles(
+        shareDBDoc,
+        mergeFileChanges(
+          shareDBDoc.data.files,
+          parseMarkdownFiles(fullContent, 'bold').files,
+        ),
+      );
+    }
+
     // Write chunks file for debugging
     if (DEBUG) {
       const chunksFileJSONpath = `./ai-chunks-${chatId}.json`;

package/src/submitOperation.ts

@@ -1,4 +1,4 @@
-import { diff } from './client/diff';
+import { diff } from './ot.js';
 
 /**
  * Creates a submitOperation function that can be used to submit diff-based operations to ShareDB.