vzcode 1.37.0 → 1.39.0

package/src/runCode.ts

@@ -0,0 +1,58 @@
+import { SubmitOperation } from './types';
+import { VizContent } from '@vizhub/viz-types';
+import { createSubmitOperation } from './submitOperation';
+
+/**
+ * Creates a runCode function that triggers code execution by flashing `isInteracting` to `true`.
+ * This works for both client-side (with submitOperation) and server-side (with ShareDB document).
+ */
+export const createRunCodeFunction = (
+  submitOperationOrDoc:
+    | SubmitOperation
+    | { data: any; submitOp: (ops: any) => void },
+) => {
+  return () => {
+    let submitOperation: SubmitOperation;
+
+    // Check if this is a client-side submitOperation or server-side ShareDB document
+    if (typeof submitOperationOrDoc === 'function') {
+      // Client-side: already a submitOperation function
+      submitOperation =
+        submitOperationOrDoc as SubmitOperation;
+    } else {
+      // Server-side: create submitOperation from ShareDB document
+      submitOperation = createSubmitOperation(
+        submitOperationOrDoc,
+      );
+    }
+
+    // Use the unified submitOperation approach for both client and server
+    submitOperation((content: VizContent) => ({
+      ...content,
+      isInteracting: true,
+    }));
+
+    setTimeout(() => {
+      // This somewhat cryptic logic
+      // deletes the `isInteracting` property
+      // from the document.
+      submitOperation(
+        ({ isInteracting, ...newDocument }) => newDocument,
+      );
+    }, 0);
+  };
+};
+
+/**
+ * Creates a runCodeRef object compatible with React refs.
+ * This is useful for maintaining compatibility with existing code that expects a ref.
+ */
+export const createRunCodeRef = (
+  submitOperationOrDoc:
+    | SubmitOperation
+    | { data: any; submitOp: (ops: any) => void },
+) => {
+  return {
+    current: createRunCodeFunction(submitOperationOrDoc),
+  };
+};

package/src/server/aiChatHandler/aiEditing.js

@@ -1,42 +1,44 @@
-import { performAiEdit } from 'editcodewithai';
 import {
-  clearAIScratchpadAndStatus,
-  updateFiles,
-  setIsInteracting,
-} from './chatOperations.js';
+  assembleFullPrompt,
+  prepareFilesForPrompt,
+} from 'editcodewithai';
+import { formatMarkdownFiles } from 'llm-code-format';
 
 /**
- * Performs AI editing operations on the files
+ * Performs AI editing operations using streaming with incremental OT operations
  */
-export const performAIEditing = async (
+export const performAIEditing = async ({
+  prompt,
   shareDBDoc,
-  chatId,
-  content,
-  files,
   llmFunction,
-) => {
-  const editResult = await performAiEdit({
-    prompt: content,
-    files,
-    llmFunction,
-    apiKey: process.env.VIZHUB_EDIT_WITH_AI_API_KEY,
-  });
-
-  // Clear the scratchpad and update status
-  clearAIScratchpadAndStatus(
-    shareDBDoc,
-    chatId,
-    'Done editing with AI.',
+  runCode,
+}) => {
+  const preparedFiles = prepareFilesForPrompt(
+    shareDBDoc.data.files,
   );
+  const filesContext = formatMarkdownFiles(preparedFiles);
 
-  // Apply AI edits to files
-  updateFiles(shareDBDoc, editResult.changedFiles);
+  // 2. Assemble the final prompt
+  const fullPrompt = assembleFullPrompt({
+    filesContext,
+    prompt,
+    editFormat: 'whole',
+  });
+
+  // Call the LLM function which will handle streaming and incremental file updates
+  const result = await llmFunction(fullPrompt);
 
-  // Wait for propagation
-  await new Promise((resolve) => setTimeout(resolve, 100));
+  // Clear the scratchpad and update status
+  // clearAIScratchpadAndStatus(
+  //   shareDBDoc,
+  //   chatId,
+  //   'Done editing with AI.',
+  // );
 
-  // Unset isInteracting
-  setIsInteracting(shareDBDoc, false);
+  runCode();
 
-  return editResult;
+  return {
+    content: result.content,
+    generationId: result.generationId,
+  };
 };

package/src/server/aiChatHandler/chatOperations.js

@@ -1,5 +1,6 @@
 import { dateToTimestamp } from '@vizhub/viz-utils';
 import { diff } from '../../client/diff.js';
+import { randomId } from '../../randomId.js';
 
 /**
  * Ensures the chats object exists in the ShareDB document
@@ -108,7 +109,13 @@ export const updateAIScratchpad = (
       },
     },
   });
-  shareDBDoc.submitOp(op);
+
+  // op is `null` if there are no changes
+  // This can happen if the content is the same as before
+  // In that case, we don't need to submit an operation
+  if (op) {
+    shareDBDoc.submitOp(op);
+  }
 };
 
 /**
@@ -174,8 +181,7 @@ export const addAIMessage = (
 export const updateFiles = (shareDBDoc, files) => {
   const filesOp = diff(shareDBDoc.data, {
     ...shareDBDoc.data,
-    files: files,
-    isInteracting: true,
+    files,
   });
   shareDBDoc.submitOp(filesOp);
 };
@@ -187,9 +193,128 @@ export const setIsInteracting = (
   shareDBDoc,
   isInteracting,
 ) => {
-  const interactingOp = diff(
-    { isInteracting: !isInteracting },
-    { isInteracting: 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
+ */
+export const resolveFileId = (fileName, shareDBDoc) => {
+  const files = shareDBDoc.data.files;
+
+  // Search through all files to find matching name
+  for (const [fileId, file] of Object.entries(files)) {
+    if (file.name === fileName) {
+      return fileId;
+    }
+  }
+
+  // If file doesn't exist, return null
+  return null;
+};
+
+/**
+ * Creates a new file with a random ID
+ */
+export const createNewFile = (shareDBDoc, fileName) => {
+  // Generate a new random file ID
+  const newFileId = randomId();
+
+  const newState = {
+    ...shareDBDoc.data,
+    files: {
+      ...shareDBDoc.data.files,
+      [newFileId]: {
+        name: fileName,
+        text: '',
+      },
+    },
+  };
+
+  const op = diff(shareDBDoc.data, newState);
+  shareDBDoc.submitOp(op);
+  return newFileId;
+};
+
+/**
+ * Ensures a file exists, creating it if necessary
+ */
+export const ensureFileExists = (shareDBDoc, fileName) => {
+  let fileId = resolveFileId(fileName, shareDBDoc);
+
+  if (!fileId) {
+    // File doesn't exist, create it
+    fileId = createNewFile(shareDBDoc, fileName);
+  }
+
+  return fileId;
+};
+
+/**
+ * Clears the content of a file
+ */
+export const clearFileContent = (shareDBDoc, fileId) => {
+  const currentFile = shareDBDoc.data.files[fileId];
+
+  if (currentFile && currentFile.text) {
+    // Clear the file content
+    const newState = {
+      ...shareDBDoc.data,
+      files: {
+        ...shareDBDoc.data.files,
+        [fileId]: {
+          ...currentFile,
+          text: '',
+        },
+      },
+    };
+
+    const op = diff(shareDBDoc.data, newState);
+    shareDBDoc.submitOp(op);
+  }
+};
+
+/**
+ * Appends a line to a file using OT operations
+ */
+export const appendLineToFile = (
+  shareDBDoc,
+  fileId,
+  line,
+) => {
+  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,
+    },
+  };
+
+  // Generate OT operation using the diff utility
+  const op = diff(shareDBDoc.data, newDocState);
+  shareDBDoc.submitOp(op);
+};

package/src/server/aiChatHandler/index.js

@@ -3,20 +3,20 @@ import {
   ensureChatsExist,
   ensureChatExists,
   addUserMessage,
-  addAIMessage,
 } from './chatOperations.js';
 import { createLLMFunction } from './llmStreaming.js';
 import { performAIEditing } from './aiEditing.js';
 import { handleError } from './errorHandling.js';
+import { createRunCodeFunction } from '../../runCode.js';
 
-const debug = false;
+const DEBUG = false;
 
 export const handleAIChatMessage =
-  (shareDBDoc, options = {}) =>
+  ({ shareDBDoc, localPresence, onCreditDeduction }) =>
   async (req, res) => {
     const { content, chatId } = req.body;
 
-    if (debug) {
+    if (DEBUG) {
       console.log(
         '[handleAIChatMessage] content:',
         content,
@@ -31,9 +31,6 @@ export const handleAIChatMessage =
     }
 
     try {
-      // Get existing files from ShareDB doc
-      const files = shareDBDoc.data.files;
-
       // Ensure chats structure exists
       ensureChatsExist(shareDBDoc);
       ensureChatExists(shareDBDoc, chatId);
@@ -42,34 +39,31 @@ export const handleAIChatMessage =
       addUserMessage(shareDBDoc, chatId, content);
 
       // Create LLM function for streaming
-      const llmFunction = createLLMFunction(
+      const llmFunction = createLLMFunction({
         shareDBDoc,
+        localPresence,
         chatId,
-      );
+      });
+
+      // Create server-side runCode function using shared module
+      const runCode = createRunCodeFunction(shareDBDoc);
 
       // Perform AI editing
-      const editResult = await performAIEditing(
+      const editResult = await performAIEditing({
+        prompt: content,
         shareDBDoc,
         chatId,
-        content,
-        files,
         llmFunction,
-      );
-
-      // Add AI response message
-      const aiResponse = addAIMessage(
-        shareDBDoc,
-        chatId,
-        editResult.content,
-      );
+        runCode,
+      });
 
       // Handle credit deduction if callback is provided
       if (
-        options.onCreditDeduction &&
+        onCreditDeduction &&
         editResult.upstreamCostCents
       ) {
         try {
-          await options.onCreditDeduction({
+          await onCreditDeduction({
             upstreamCostCents: editResult.upstreamCostCents,
             provider: editResult.provider,
             inputTokens: editResult.inputTokens,
@@ -84,7 +78,7 @@ export const handleAIChatMessage =
         }
       }
 
-      res.status(200).json(aiResponse);
+      res.status(200).json('success');
     } catch (error) {
       handleError(shareDBDoc, chatId, error, res);
     }

package/src/server/aiChatHandler/llmStreaming.js

@@ -1,17 +1,40 @@
 import { StreamingMarkdownParser } from 'llm-code-format';
 import { ChatOpenAI } from '@langchain/openai';
+import fs from 'fs';
 import {
   updateAIStatus,
   updateAIScratchpad,
+  ensureFileExists,
+  clearFileContent,
+  appendLineToFile,
 } from './chatOperations.js';
 
-const debug = false;
+const DEBUG = false;
 
 /**
  * Creates and configures the LLM function for streaming
  */
-export const createLLMFunction = (shareDBDoc, chatId) => {
+export const createLLMFunction = ({
+  shareDBDoc,
+  localPresence,
+  chatId,
+}) => {
   return async (fullPrompt) => {
+    // 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 ||
@@ -29,22 +52,58 @@ export const createLLMFunction = (shareDBDoc, chatId) => {
 
     let fullContent = '';
     let generationId = '';
+    let currentEditingFileId = null;
+    let currentEditingFileName = null;
 
-    // Track expected state to avoid race conditions
-    let expectedAiScratchpad =
-      shareDBDoc.data.chats[chatId]?.aiScratchpad || '';
-
-    // Throttle updates to avoid "op too long" errors
-    let lastUpdateTime = 0;
-    const UPDATE_THROTTLE_MS = 100;
+    // 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.
+    const reportFileEdited = () => {
+      if (currentEditingFileName) {
+        fullContent += ` * Edited ${currentEditingFileName}\n`;
+        updateAIScratchpad(shareDBDoc, chatId, fullContent);
+        currentEditingFileName = null;
+      }
+    };
 
     // Define callbacks for streaming parser
     const callbacks = {
       onFileNameChange: (fileName, format) => {
-        debug &&
+        DEBUG &&
           console.log(
             `File changed to: ${fileName} (${format})`,
           );
+
+        // Find existing file or create new one
+        currentEditingFileId = ensureFileExists(
+          shareDBDoc,
+          fileName,
+        );
+
+        reportFileEdited();
+        currentEditingFileName = fileName;
+
+        // Clear the file content to start fresh
+        // (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,
           chatId,
@@ -52,37 +111,70 @@ export const createLLMFunction = (shareDBDoc, chatId) => {
         );
       },
       onCodeLine: (line) => {
-        debug && console.log(`Code line: ${line}`);
+        DEBUG && console.log(`Code line: ${line}`);
+
+        if (currentEditingFileId) {
+          // Apply OT operation for this line immediately
+          appendLineToFile(
+            shareDBDoc,
+            currentEditingFileId,
+            line,
+          );
+
+          // Update VizBot presence to show cursor at the end of the file
+          const currentFile =
+            shareDBDoc.data.files[currentEditingFileId];
+          if (currentFile && currentFile.text) {
+            const textLength = currentFile.text.length;
+            const filePresence = {
+              username: 'VizBot',
+              start: [
+                'files',
+                currentEditingFileId,
+                'text',
+                textLength,
+              ],
+              end: [
+                'files',
+                currentEditingFileId,
+                'text',
+                textLength,
+              ],
+            };
+
+            localPresence.submit(filePresence, (error) => {
+              if (error) {
+                console.warn(
+                  'VizBot line presence submission error:',
+                  error,
+                );
+              }
+            });
+          }
+        }
       },
       onNonCodeLine: (line) => {
-        debug && console.log(`Comment/text: ${line}`);
+        // We want to report a file edited only if the line is not empty,
+        // because sometimes the LLMs leave a newline between the file name
+        // declaration and th
+        if (line.trim() !== '') {
+          reportFileEdited();
+        }
+        fullContent += line + '\n';
+        updateAIScratchpad(shareDBDoc, chatId, fullContent);
       },
     };
 
     const parser = new StreamingMarkdownParser(callbacks);
 
+    const chunks = [];
+
     // Stream the response
     const stream = await chatModel.stream(fullPrompt);
     for await (const chunk of stream) {
       if (chunk.content) {
         const chunkContent = String(chunk.content);
-        fullContent += chunkContent;
-
-        // Throttle updates
-        const now = Date.now();
-        const shouldUpdate =
-          now - lastUpdateTime >= UPDATE_THROTTLE_MS;
-
-        if (shouldUpdate) {
-          updateAIScratchpad(
-            shareDBDoc,
-            chatId,
-            fullContent,
-          );
-          expectedAiScratchpad = fullContent;
-          lastUpdateTime = now;
-        }
-
+        chunks.push(chunkContent);
         parser.processChunk(chunkContent);
       }
 
@@ -91,10 +183,30 @@ export const createLLMFunction = (shareDBDoc, chatId) => {
       }
     }
     parser.flushRemaining();
+    reportFileEdited();
+    updateAIStatus(shareDBDoc, chatId, 'Done editing.');
+    updateAIScratchpad(shareDBDoc, chatId, fullContent);
+
+    // Clear VizBot presence when done
+    localPresence.destroy((error) => {
+      if (error) {
+        console.warn(
+          'VizBot presence cleanup error:',
+          error,
+        );
+      }
+    });
 
-    // Submit final update if there's pending content
-    if (expectedAiScratchpad !== fullContent) {
-      updateAIScratchpad(shareDBDoc, chatId, fullContent);
+    // Write chunks file for debugging
+    if (DEBUG) {
+      const chunksFileJSONpath = `./ai-chunks-${chatId}.json`;
+      fs.writeFileSync(
+        chunksFileJSONpath,
+        JSON.stringify(chunks, null, 2),
+      );
+      console.log(
+        `AI chunks written to ${chunksFileJSONpath}`,
+      );
     }
 
     return {

package/src/server/index.js

@@ -51,13 +51,6 @@ ShareDB.types.register(json1Presence.type);
 
 const app = express();
 
-// TODO make this configurable
-// See https://github.com/vizhub-core/vzcode/issues/95
-app.post('/saveTime', (req, res) => {
-  //autoSaveDebounceTimeMS = req.body.autoSaveDebounceTimeMS;
-  console.log('autoSaveDebounceTimeMS', req.body);
-});
-
 // Use ShareDB over WebSocket
 const shareDBBackend = new ShareDB({
   // Enable presence
@@ -82,7 +75,7 @@ wss.on('connection', (ws) => {
   });
 
   // Handle disconnections
-  ws.on('close', (code) => {
+  ws.on('close', () => {
     clientStream.end();
   });
 });
@@ -97,6 +90,23 @@ app.use(express.static(dir));
 // which is a representation of files on disk.
 const shareDBConnection = shareDBBackend.connect();
 const shareDBDoc = shareDBConnection.get('documents', '1');
+
+// Set up presence for VizBot following the same pattern as useShareDB.ts
+const docPresence = shareDBConnection.getDocPresence(
+  'documents',
+  '1',
+);
+
+// Create local presence for VizBot with a unique ID
+const generateVizBotId = () => {
+  const timestamp = Date.now().toString(36);
+  return `vizbot-${timestamp}`;
+};
+
+const localPresence = docPresence.create(
+  generateVizBotId(),
+);
+
 shareDBDoc.create(initialDocument, json1Presence.type.uri);
 
 // Handle AI Assist requests.
@@ -117,7 +127,7 @@ app.post(
 app.post(
   '/ai-chat-message',
   bodyParser.json(),
-  handleAIChatMessage(shareDBDoc),
+  handleAIChatMessage({ shareDBDoc, localPresence }),
 );
 
 // Livekit Token Generator

package/src/submitOperation.js

@@ -0,0 +1,16 @@
+import { diff } from './client/diff.js';
+// TODO migrate the server-side code to use TypeScript,
+// and delete this file once the migration is complete (use src/submitOperation.ts instead).
+/**
+ * Creates a submitOperation function that can be used to submit diff-based operations to ShareDB.
+ * This is the core logic extracted from useSubmitOperation for reuse across client and server.
+ */
+export const createSubmitOperation = (shareDBDoc) => {
+  return (next) => {
+    const data = shareDBDoc.data;
+    const op = diff(data, next(data));
+    if (op && shareDBDoc) {
+      shareDBDoc.submitOp(op);
+    }
+  };
+};