vzcode 1.44.0 → 1.45.0

package/src/server/aiChatHandler/aiEditing.ts

@@ -3,6 +3,10 @@ import {
   prepareFilesForPrompt,
 } from 'editcodewithai';
 import { formatMarkdownFiles } from 'llm-code-format';
+import {
+  createFilesSnapshot,
+  generateFilesUnifiedDiff,
+} from '../../utils/fileDiff.js';
 
 // Dev flag for waiting 1 second before starting the LLM function.
 // Useful for debugging and testing purposes, e.g. checking the typing indicator.
@@ -17,6 +21,11 @@ export const performAIEditing = async ({
   llmFunction,
   runCode,
 }) => {
+  // 1. Capture the current state of files before editing
+  const beforeFiles = createFilesSnapshot(
+    shareDBDoc.data.files,
+  );
+
   const preparedFiles = prepareFilesForPrompt(
     shareDBDoc.data.files,
   );
@@ -40,8 +49,18 @@ export const performAIEditing = async ({
 
   runCode();
 
+  // 3. Capture the state of files after editing and generate diff
+  const afterFiles = createFilesSnapshot(
+    shareDBDoc.data.files,
+  );
+  const diffData = generateFilesUnifiedDiff(
+    beforeFiles,
+    afterFiles,
+  );
+
   return {
     content: result.content,
     generationId: result.generationId,
+    diffData,
   };
 };

package/src/server/aiChatHandler/chatOperations.ts

@@ -246,6 +246,46 @@ export const finalizeAIMessage = (
   shareDBDoc.submitOp(op);
 };
 
+/**
+ * Adds diff data to the most recent AI message
+ */
+export const addDiffToAIMessage = (
+  shareDBDoc: ShareDBDoc<VizContent>,
+  chatId: VizChatId,
+  diffData: any,
+) => {
+  const chat = shareDBDoc.data.chats[chatId];
+  const messages = [...chat.messages];
+
+  // Find the most recent AI message
+  const lastAIMessageIndex = messages.length - 1;
+  if (
+    lastAIMessageIndex >= 0 &&
+    messages[lastAIMessageIndex].role === 'assistant'
+  ) {
+    const newMessage = {
+      ...messages[lastAIMessageIndex],
+      diffData,
+    };
+    // Use type assertion to extend the message with diffData
+    (messages[lastAIMessageIndex] as any) = newMessage;
+
+    const messageOp = diff(shareDBDoc.data, {
+      ...shareDBDoc.data,
+      chats: {
+        ...shareDBDoc.data.chats,
+        [chatId]: {
+          ...chat,
+          messages,
+          updatedAt: dateToTimestamp(new Date()),
+        },
+      },
+    });
+
+    shareDBDoc.submitOp(messageOp);
+  }
+};
+
 /**
  * Adds an AI response message to the chat (legacy function, kept for compatibility)
  */

package/src/server/aiChatHandler/index.ts

@@ -3,6 +3,7 @@ import {
   ensureChatsExist,
   ensureChatExists,
   addUserMessage,
+  addDiffToAIMessage,
 } from './chatOperations.js';
 import { createLLMFunction } from './llmStreaming.js';
 import { performAIEditing } from './aiEditing.js';
@@ -72,6 +73,18 @@ export const handleAIChatMessage =
         runCode,
       });
 
+      // Add diff data to the AI message if there are changes
+      if (
+        editResult.diffData &&
+        Object.keys(editResult.diffData).length > 0
+      ) {
+        addDiffToAIMessage(
+          shareDBDoc,
+          chatId,
+          editResult.diffData,
+        );
+      }
+
       // Handle credit deduction if callback is provided
       if (
         onCreditDeduction &&

package/src/server/computeInitialDocument.ts

@@ -10,6 +10,15 @@ import createIgnore from 'ignore';
 import { ignoreFilePattern, baseIgnore } from './config.js';
 import { isDirectory } from './isDirectory.js';
 
+// Import the image file utility
+const isImageFile = (fileName) => {
+  return (
+    fileName.match(
+      /\.(png|jpg|jpeg|gif|bmp|svg|webp)$/i,
+    ) !== null
+  );
+};
+
 /**
  * @param {string} fullPath - absolut path of the workspace root
  * @param {string} currentDirectoryPath - path where the ignore file is found, relative to fullPath
@@ -180,13 +189,23 @@ export const computeInitialDocument = ({ fullPath }) => {
 
   files.forEach((file) => {
     const id = randomId();
+    let text = null;
+
+    if (!isDirectory(file)) {
+      const filePath = path.join(fullPath, file);
+
+      if (isImageFile(file)) {
+        // Read image files as binary and convert to base64
+        const buffer = fs.readFileSync(filePath);
+        text = buffer.toString('base64');
+      } else {
+        // Read non-image files as UTF-8 text
+        text = fs.readFileSync(filePath, 'utf-8');
+      }
+    }
+
     initialDocument.files[id] = {
-      text: isDirectory(file)
-        ? null
-        : fs.readFileSync(
-            path.join(fullPath, file),
-            'utf-8',
-          ),
+      text,
       name: file,
     };
   });

package/src/server/index.ts

@@ -18,6 +18,15 @@ import { isDirectory } from './isDirectory.js';
 import { createToken } from './livekit.js';
 import './setupEnv.js';
 
+// Import the image file utility
+const isImageFile = (fileName) => {
+  return (
+    fileName.match(
+      /\.(png|jpg|jpeg|gif|bmp|svg|webp)$/i,
+    ) !== null
+  );
+};
+
 // The time in milliseconds by which auto-saving is debounced.
 const autoSaveDebounceTimeMS = 800;
 
@@ -161,7 +170,17 @@ const save = () => {
     if (previous && current) {
       // Handle changing of text content.
       if (previous.text !== current.text) {
-        fs.writeFileSync(current.name, current.text);
+        if (isImageFile(current.name)) {
+          // Write image files as binary from base64
+          const buffer = Buffer.from(
+            current.text,
+            'base64',
+          );
+          fs.writeFileSync(current.name, buffer);
+        } else {
+          // Write non-image files as text
+          fs.writeFileSync(current.name, current.text);
+        }
       }
 
       // // Handle renaming files.
@@ -256,7 +275,17 @@ const save = () => {
     if (!previous && current) {
       //File Creation
       if (!isDirectory(current.name)) {
-        fs.writeFileSync(current.name, current.text);
+        if (isImageFile(current.name)) {
+          // Write image files as binary from base64
+          const buffer = Buffer.from(
+            current.text,
+            'base64',
+          );
+          fs.writeFileSync(current.name, buffer);
+        } else {
+          // Write non-image files as text
+          fs.writeFileSync(current.name, current.text);
+        }
       } else {
         fs.mkdirSync(current.name, { recursive: true });
       }

package/src/utils/fileDiff.ts

@@ -0,0 +1,114 @@
+import { createTwoFilesPatch } from 'diff';
+import { VizFiles, VizFileId } from '@vizhub/viz-types';
+
+export interface UnifiedFilesDiff {
+  [fileId: VizFileId]: string; // Unified diff string
+}
+
+/**
+ * Generate a unified diff for a single file using the diff library
+ */
+export function generateFileUnifiedDiff(
+  fileId: VizFileId,
+  fileName: string,
+  beforeContent: string,
+  afterContent: string,
+): string {
+  // Only generate diff if there are actual changes
+  if (beforeContent === afterContent) {
+    return '';
+  }
+
+  // Use the diff library's createTwoFilesPatch function to generate unified diff
+  const unifiedDiff = createTwoFilesPatch(
+    fileName,
+    fileName,
+    beforeContent,
+    afterContent,
+  );
+
+  return unifiedDiff;
+}
+
+/**
+ * Generate unified diffs for multiple files
+ */
+export function generateFilesUnifiedDiff(
+  beforeFiles: VizFiles,
+  afterFiles: VizFiles,
+): UnifiedFilesDiff {
+  const result: UnifiedFilesDiff = {};
+  const allFileIds = new Set([
+    ...Object.keys(beforeFiles),
+    ...Object.keys(afterFiles),
+  ]);
+
+  for (const fileId of allFileIds) {
+    const beforeFile = beforeFiles[fileId];
+    const afterFile = afterFiles[fileId];
+
+    const beforeContent = beforeFile?.text || '';
+    const afterContent = afterFile?.text || '';
+    const fileName =
+      afterFile?.name || beforeFile?.name || fileId;
+
+    const unifiedDiff = generateFileUnifiedDiff(
+      fileId,
+      fileName,
+      beforeContent,
+      afterContent,
+    );
+
+    // Only include files that have changes
+    if (unifiedDiff) {
+      result[fileId] = unifiedDiff;
+    }
+  }
+
+  return result;
+}
+
+/**
+ * Create a snapshot of current files for diff comparison
+ */
+export function createFilesSnapshot(
+  files: VizFiles,
+): VizFiles {
+  return JSON.parse(JSON.stringify(files));
+}
+
+/**
+ * Parse unified diff to extract basic statistics
+ */
+export function parseUnifiedDiffStats(
+  unifiedDiff: string,
+): {
+  additions: number;
+  deletions: number;
+} {
+  const lines = unifiedDiff.split('\n');
+  let additions = 0;
+  let deletions = 0;
+
+  for (const line of lines) {
+    if (line.startsWith('+') && !line.startsWith('+++')) {
+      additions++;
+    } else if (
+      line.startsWith('-') &&
+      !line.startsWith('---')
+    ) {
+      deletions++;
+    }
+  }
+
+  return { additions, deletions };
+}
+
+/**
+ * Combine multiple unified diffs into a single diff string
+ */
+export function combineUnifiedDiffs(
+  unifiedDiffs: UnifiedFilesDiff,
+): string {
+  return Object.values(unifiedDiffs).join('\n');
+}