vzcode 1.38.0 → 1.40.0

package/src/server/aiChatHandler/{chatOperations.js → chatOperations.ts}

@@ -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,131 @@ 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: string,
+  shareDBDoc: any,
+) => {
+  const files = shareDBDoc.data.files;
+
+  // Search through all files to find matching name
+  for (const [fileId, file] of Object.entries(files)) {
+    if ((file as any).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 → index.ts}

@@ -3,20 +3,30 @@ 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';
+import { ShareDBDoc } from '../../types.js';
+import { VizContent } from '@vizhub/viz-types';
 
-const debug = false;
+const DEBUG = false;
 
 export const handleAIChatMessage =
-  (shareDBDoc, options = {}) =>
-  async (req, res) => {
+  ({
+    shareDBDoc,
+    createVizBotLocalPresence,
+    onCreditDeduction,
+  }: {
+    shareDBDoc: ShareDBDoc<VizContent>;
+    createVizBotLocalPresence: () => any;
+    onCreditDeduction?: any;
+  }) =>
+  async (req: any, res: any) => {
     const { content, chatId } = req.body;
 
-    if (debug) {
+    if (DEBUG) {
       console.log(
         '[handleAIChatMessage] content:',
         content,
@@ -31,9 +41,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,38 +49,35 @@ export const handleAIChatMessage =
       addUserMessage(shareDBDoc, chatId, content);
 
       // Create LLM function for streaming
-      const llmFunction = createLLMFunction(
+      const llmFunction = createLLMFunction({
         shareDBDoc,
+        createVizBotLocalPresence,
         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 &&
-        editResult.upstreamCostCents
+        onCreditDeduction &&
+        (editResult as any).upstreamCostCents
       ) {
         try {
-          await options.onCreditDeduction({
-            upstreamCostCents: editResult.upstreamCostCents,
-            provider: editResult.provider,
-            inputTokens: editResult.inputTokens,
-            outputTokens: editResult.outputTokens,
+          await onCreditDeduction({
+            upstreamCostCents: (editResult as any)
+              .upstreamCostCents,
+            provider: (editResult as any).provider,
+            inputTokens: (editResult as any).inputTokens,
+            outputTokens: (editResult as any).outputTokens,
           });
         } catch (creditError) {
           console.error(
@@ -84,7 +88,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.ts

@@ -0,0 +1,223 @@
+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;
+
+/**
+ * Creates and configures the LLM function for streaming
+ */
+export const createLLMFunction = ({
+  shareDBDoc,
+  createVizBotLocalPresence,
+  chatId,
+}) => {
+  return async (fullPrompt) => {
+    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 ||
+        'anthropic/claude-sonnet-4',
+      configuration: {
+        apiKey: process.env.VIZHUB_EDIT_WITH_AI_API_KEY,
+        baseURL: process.env.VIZHUB_EDIT_WITH_AI_BASE_URL,
+        defaultHeaders: {
+          'HTTP-Referer': 'https://vizhub.com',
+          'X-Title': 'VizHub',
+        },
+      },
+      streaming: true,
+    });
+
+    let fullContent = '';
+    let generationId = '';
+    let currentEditingFileId = null;
+    let currentEditingFileName = null;
+
+    // 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 &&
+          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,
+          'Editing ' + fileName,
+        );
+      },
+      onCodeLine: (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) => {
+        // 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);
+        chunks.push(chunkContent);
+        parser.processChunk(chunkContent);
+      }
+
+      if (!generationId && chunk.lc_kwargs?.id) {
+        generationId = chunk.lc_kwargs.id;
+      }
+    }
+    parser.flushRemaining();
+    reportFileEdited();
+    updateAIStatus(shareDBDoc, chatId, 'Done editing.');
+    updateAIScratchpad(shareDBDoc, chatId, fullContent);
+
+    // Clear VizBot presence when done
+    DEBUG &&
+      console.log(
+        'AI editing done, clearing VizBot presence',
+      );
+    localPresence.submit(null, (error) => {
+      DEBUG && console.log('VizBot presence cleared');
+      if (error) {
+        console.warn(
+          'VizBot presence cleanup error:',
+          error,
+        );
+      }
+    });
+
+    // 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 {
+      content: fullContent,
+      generationId: generationId,
+    };
+  };
+};

package/src/server/{generateAIResponse.js → generateAIResponse.ts}

@@ -28,7 +28,8 @@ const slowdown = false;
 
 // The options passed into the OpenAI client
 // new OpenAI(openAIOptions)
-const openAIOptions = {};
+const openAIOptions: { apiKey?: string; baseURL?: string } =
+  {};
 
 // Support specifying the API key via an environment variable
 // If VZCODE_AI_API_KEY is not set, note that the OpenAI client
@@ -71,7 +72,7 @@ const opComesFromAIAssist = (ops, source) =>
 
 // Keeps track of the currently ongoing AI streams.
 // There could be many streams at the same time.
-const streams = {};
+export const streams: Record<string, any> = {};
 
 export const generateAIResponse = async ({
   inputText,

package/src/server/{handleAIAssist.js → handleAIAssist.ts}

@@ -1,9 +1,12 @@
-import { generateAIResponse } from './generateAIResponse.js';
+import {
+  generateAIResponse,
+  streams,
+} from './generateAIResponse.js';
 
 const debug = false;
 
 export const handleAIAssist =
-  (shareDBDoc) => async (req, res) => {
+  (shareDBDoc: any) => async (req: any, res: any) => {
     const { inputText, insertionCursor, fileId } = req.body;
 
     if (debug) {
@@ -16,17 +19,21 @@ export const handleAIAssist =
     }
 
     try {
+      // Generate a unique streamId for this request
+      const streamId = `stream_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`;
+
       await generateAIResponse({
         inputText,
         insertionCursor,
         fileId,
+        streamId,
         shareDBDoc,
       });
 
       res
         .status(200)
         .send({ message: 'Operation successful!' });
-    } catch (error) {
+    } catch (error: any) {
       console.error('handleAIAssist error:', error);
       res.status(500).send({
         message: 'Internal Server Error',
@@ -35,7 +42,7 @@ export const handleAIAssist =
     }
   };
 
-export function haltGeneration(streamId) {
+export function haltGeneration(streamId: string) {
   const stream = streams[streamId];
 
   // Stream can be undefined here if the user

package/src/server/{index.js → index.ts}

@@ -4,7 +4,6 @@ import bodyParser from 'body-parser';
 import express from 'express';
 import fs from 'fs';
 import http from 'http';
-import ngrok from 'ngrok';
 import open from 'open';
 import path from 'path';
 import ShareDB from 'sharedb';
@@ -51,13 +50,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 +74,7 @@ wss.on('connection', (ws) => {
   });
 
   // Handle disconnections
-  ws.on('close', (code) => {
+  ws.on('close', () => {
     clientStream.end();
   });
 });
@@ -97,6 +89,22 @@ 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 createVizBotLocalPresence = () =>
+  docPresence.create(generateVizBotId());
+
 shareDBDoc.create(initialDocument, json1Presence.type.uri);
 
 // Handle AI Assist requests.
@@ -117,7 +125,11 @@ app.post(
 app.post(
   '/ai-chat-message',
   bodyParser.json(),
-  handleAIChatMessage(shareDBDoc),
+  handleAIChatMessage({
+    shareDBDoc,
+    createVizBotLocalPresence,
+    onCreditDeduction: undefined,
+  }),
 );
 
 // Livekit Token Generator
@@ -325,21 +337,13 @@ shareDBDoc.subscribe(() => {
 });
 
 server.listen(port, async () => {
-  if (process.env.NGROK_TOKEN) {
-    (async function () {
-      await ngrok.authtoken(process.env.NGROK_TOKEN);
-      const url = await ngrok.connect(port);
-      console.log(`Editor is live at ${url}`);
-      open(url);
-    })();
-  } else {
-    // Sets the port to the one specified in the environment
-    // variable (for development) or the default port.
-    let livePort = process.env.EDITOR_PORT || port;
-    console.log(`EDITOR_PORT: ${process.env.EDITOR_PORT}`);
-    console.log(
-      `Editor is live at http://localhost:${livePort}`,
-    );
-    open(`http://localhost:${livePort}`);
-  }
+  // Note: ngrok support can be added when the package is properly installed
+  // Sets the port to the one specified in the environment
+  // variable (for development) or the default port.
+  let livePort = process.env.EDITOR_PORT || port;
+  console.log(`EDITOR_PORT: ${process.env.EDITOR_PORT}`);
+  console.log(
+    `Editor is live at http://localhost:${livePort}`,
+  );
+  open(`http://localhost:${livePort}`);
 });

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);
+    }
+  };
+};

package/src/submitOperation.ts

@@ -0,0 +1,25 @@
+import { diff } from './client/diff';
+
+/**
+ * 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 = <T>(shareDBDoc: {
+  data: T;
+  submitOp: (op: any) => void;
+}): ((next: (data: T) => T) => void) => {
+  return (next) => {
+    const data: T = shareDBDoc.data;
+    const op = diff(data, next(data));
+    if (op && shareDBDoc) {
+      shareDBDoc.submitOp(op);
+    }
+  };
+};
+
+/**
+ * Type definition for the submitOperation function
+ */
+export type SubmitOperationFunction<T> = (
+  next: (data: T) => T,
+) => void;