vzcode 2.17.0 → 2.21.0

package/src/client/VZSidebar/AIChat/DiffView.tsx

@@ -9,6 +9,8 @@ import {
   UnifiedFilesDiff,
   parseUnifiedDiffStats,
   combineUnifiedDiffs,
+  isFileDeletion,
+  getDeletedFileName,
 } from '../../../utils/fileDiff';
 import * as Diff2Html from 'diff2html';
 import 'diff2html/bundles/css/diff2html.min.css';
@@ -41,21 +43,33 @@ export const DiffView = forwardRef<
     useContext(VZCodeContext);
   const diffContainerRef = useRef<HTMLDivElement>(null);
 
-  const unifiedDiffs = Object.values(diffData).filter(
+  const allDiffs = Object.values(diffData).filter(
     (diff) => diff.length > 0,
   );
 
-  // Calculate statistics from all unified diffs
+  // Separate deleted files from regular diffs
+  const deletedFiles: string[] = [];
+  const regularDiffs: string[] = [];
+
+  for (const diff of allDiffs) {
+    if (isFileDeletion(diff)) {
+      deletedFiles.push(diff);
+    } else {
+      regularDiffs.push(diff);
+    }
+  }
+
+  // Calculate statistics from regular unified diffs only
   let totalAdditions = 0;
   let totalDeletions = 0;
 
-  for (const unifiedDiff of unifiedDiffs) {
+  for (const unifiedDiff of regularDiffs) {
     const stats = parseUnifiedDiffStats(unifiedDiff);
     totalAdditions += stats.additions;
     totalDeletions += stats.deletions;
   }
 
-  // Combine all unified diffs and convert to HTML using diff2html
+  // Combine regular unified diffs and convert to HTML using diff2html
   const combinedUnifiedDiff = combineUnifiedDiffs(diffData);
   const diffHtml = Diff2Html.html(combinedUnifiedDiff, {
     drawFileList: false,
@@ -155,7 +169,7 @@ export const DiffView = forwardRef<
     };
   }, [diffHtml, content, openTab, setIsAIChatOpen]);
 
-  if (unifiedDiffs.length === 0) {
+  if (allDiffs.length === 0) {
     return null;
   }
 
@@ -164,8 +178,8 @@ export const DiffView = forwardRef<
       <div className="diff-summary" id="diff-summary">
         <div className="diff-stats">
           <span className="files-changed">
-            {unifiedDiffs.length} file
-            {unifiedDiffs.length !== 1 ? 's' : ''} changed
+            {allDiffs.length} file
+            {allDiffs.length !== 1 ? 's' : ''} changed
           </span>
           {totalAdditions > 0 && (
             <span className="additions">
@@ -177,18 +191,43 @@ export const DiffView = forwardRef<
               -{totalDeletions}
             </span>
           )}
+          {deletedFiles.length > 0 && (
+            <span className="deletions">
+              {deletedFiles.length} deleted
+            </span>
+          )}
         </div>
       </div>
 
-      <div
-        className="diff-files"
-        ref={diffContainerRef}
-        tabIndex={-1}
-        role="region"
-        aria-label="Code diff content"
-        aria-describedby="diff-summary"
-        dangerouslySetInnerHTML={{ __html: diffHtml }}
-      />
+      {/* Render deleted files */}
+      {deletedFiles.map((deletionMarker, index) => {
+        const fileName = getDeletedFileName(deletionMarker);
+        return (
+          <div key={index} className="deleted-file">
+            <div className="deleted-file-header">
+              <span className="deleted-file-name">
+                {fileName}
+              </span>
+              <span className="deleted-file-status">
+                File deleted
+              </span>
+            </div>
+          </div>
+        );
+      })}
+
+      {/* Render regular diffs */}
+      {regularDiffs.length > 0 && (
+        <div
+          className="diff-files"
+          ref={diffContainerRef}
+          tabIndex={-1}
+          role="region"
+          aria-label="Code diff content"
+          aria-describedby="diff-summary"
+          dangerouslySetInnerHTML={{ __html: diffHtml }}
+        />
+      )}
     </div>
   );
 });

package/src/client/featureFlags.ts

@@ -9,3 +9,10 @@ export const enableAskMode = false;
 
 // Phase 0: Feature flag for minimal AI edit flow
 export const enableMinimalEditFlow = true;
+
+// If true, only include the minimal set of extensions,
+// namely only the JSON1 OT extension and that's it.
+// This is useful for testing and debugging the OT functionality
+// without any other extensions getting in the way,
+// just to rule them out as the source of a bug.
+export const MINIMAL_EXTENSIONS = false;

package/src/server/aiChatHandler/aiEditing.ts

@@ -7,9 +7,6 @@ import {
   createFilesSnapshot,
   generateFilesUnifiedDiff,
 } from '../../utils/fileDiff.js';
-import { formatFiles } from '../prettier.js';
-import { createSubmitOperation } from '../../submitOperation.js';
-import { VizContent } from '@vizhub/viz-types';
 
 // Dev flag for waiting 1 second before starting the LLM function.
 // Useful for debugging and testing purposes, e.g. checking the typing indicator.
@@ -60,7 +57,7 @@ export const performAIEditing = async ({
   llmFunction,
   runCode,
 }) => {
-  // 1. Capture the current state of files before editing
+  // Capture the current state of files before editing
   const beforeFiles = createFilesSnapshot(
     shareDBDoc.data.files,
   );
@@ -70,7 +67,7 @@ export const performAIEditing = async ({
   );
   const filesContext = formatMarkdownFiles(files);
 
-  // 2. Assemble the final prompt
+  // Assemble the final prompt
   const fullPrompt = assembleFullPrompt({
     filesContext,
     prompt,
@@ -83,55 +80,9 @@ export const performAIEditing = async ({
     );
   }
 
-  // Call the LLM function which will handle streaming and incremental file updates
+  // Call the LLM function which will handle streaming, incremental file updates, and Prettier formatting
   const result = await llmFunction(fullPrompt);
 
-  // 3. Run Prettier on changed files before running code
-  const afterLLMFiles = createFilesSnapshot(
-    shareDBDoc.data.files,
-  );
-
-  // Find files that were changed by the AI
-  const changedFileIds = Object.keys(afterLLMFiles).filter(
-    (fileId) =>
-      beforeFiles[fileId]?.text !==
-      afterLLMFiles[fileId]?.text,
-  );
-
-  if (changedFileIds.length > 0) {
-    // Format the changed files
-    const formattedFiles = await formatFiles(
-      shareDBDoc.data.files,
-      changedFileIds,
-    );
-
-    // Apply formatted versions to shareDBDoc if formatting succeeded
-    if (Object.keys(formattedFiles).length > 0) {
-      const submitOperation =
-        createSubmitOperation(shareDBDoc);
-      submitOperation((document: VizContent) => {
-        const updatedFiles = { ...document.files };
-
-        // Apply each formatted file
-        Object.entries(formattedFiles).forEach(
-          ([fileId, formattedText]) => {
-            if (updatedFiles[fileId]) {
-              updatedFiles[fileId] = {
-                ...updatedFiles[fileId],
-                text: formattedText,
-              };
-            }
-          },
-        );
-
-        return {
-          ...document,
-          files: updatedFiles,
-        };
-      });
-    }
-  }
-
   runCode();
 
   // 4. Capture the state of files after editing and formatting, then generate diff

package/src/server/aiChatHandler/chatOperations.ts

@@ -51,6 +51,7 @@ export const ensureChatExists = (
 
 /**
  * Adds a user message to the chat
+ * Clears old messages to reflect that each prompt is a self-contained code transformation
  */
 export const addUserMessage = (
   shareDBDoc: ShareDBDoc<VizContent>,
@@ -70,10 +71,7 @@ export const addUserMessage = (
       ...shareDBDoc.data.chats,
       [chatId]: {
         ...shareDBDoc.data.chats[chatId],
-        messages: [
-          ...shareDBDoc.data.chats[chatId].messages,
-          userMessage,
-        ],
+        messages: [userMessage], // Replace old messages with just the new user message
         updatedAt: dateToTimestamp(new Date()),
       },
     },
@@ -376,7 +374,10 @@ export const updateFiles = (
     ...shareDBDoc.data,
     files,
   });
+  DEBUG && console.log('updateFiles op:');
+  DEBUG && console.log(JSON.stringify(filesOp, null, 2));
   shareDBDoc.submitOp(filesOp);
+  return filesOp;
 };
 
 /**

package/src/server/aiChatHandler/llmStreaming.ts

@@ -4,8 +4,11 @@ import {
   StreamingMarkdownParser,
 } from 'llm-code-format';
 import { mergeFileChanges } from 'editcodewithai';
-import { VizChatId } from '@vizhub/viz-types';
-import { generateRunId } from '@vizhub/viz-utils';
+import {
+  FileCollection,
+  VizChatId,
+  VizFiles,
+} from '@vizhub/viz-types';
 import {
   updateFiles,
   updateAIScratchpad,
@@ -14,17 +17,24 @@ import {
   updateStreamingStatus,
   finalizeStreamingMessage,
 } from './chatOperations.js';
-import { diff } from '../../ot.js';
 import {
   ShareDBDoc,
   ExtendedVizContent,
 } from '../../types.js';
+import { formatFiles } from '../prettier.js';
 
+// Verbose logs
 const DEBUG = false;
 
 // Useful for testing/debugging the streaming behavior
 const slowMode = false;
 
+// If the `EMIT_FIXTURES` variable is true,
+// then an output file in the `test/fixtures` folder
+// with the before and after file states for testing purposes.
+// This feeds into tests in codemirror-ot.
+const EMIT_FIXTURES = false;
+
 /**
  * Creates and configures the LLM function for streaming with reasoning tokens
  */
@@ -107,7 +117,7 @@ export const createLLMFunction = ({
     const completeFileEditing = async (
       fileName: string,
     ) => {
-      if (fileName && currentFileContent) {
+      if (fileName) {
         DEBUG &&
           console.log(
             `LLMStreaming: Completing file editing for ${fileName}`,
@@ -183,6 +193,38 @@ export const createLLMFunction = ({
           }
         }
       },
+      onFileDelete: async (fileName: string) => {
+        DEBUG &&
+          console.log(
+            `LLMStreaming: File marked for deletion: ${fileName}`,
+          );
+
+        // Emit any accumulated text chunk first
+        await emitTextChunk();
+
+        // Complete previous file if any
+        if (currentEditingFileName) {
+          await completeFileEditing(currentEditingFileName);
+        }
+
+        // Reset current editing state
+        currentEditingFileName = null;
+        currentFileContent = '';
+
+        // Emit file delete event
+        await addStreamingEvent(shareDBDoc, chatId, {
+          type: 'file_delete',
+          fileName,
+          timestamp: Date.now(),
+        });
+
+        // Update status
+        updateStreamingStatus(
+          shareDBDoc,
+          chatId,
+          `Deleting ${fileName}...`,
+        );
+      },
     };
 
     const parser = new StreamingMarkdownParser(callbacks);
@@ -285,51 +327,73 @@ export const createLLMFunction = ({
       await completeFileEditing(currentEditingFileName);
     }
 
+    // // Capture the current state of files before applying changes
+    // const beforeFiles = createFilesSnapshot(
+    //   shareDBDoc.data.files,
+    // );
+
+    // Parse the full content to extract file changes
+    // export type FileCollection = Record<string, string>;
+    const newFilesUnformatted: FileCollection =
+      parseMarkdownFiles(fullContent, 'bold').files;
+
+    // Run Prettier on `newFiles` before applying them,
+    // preserving empty files as empty
+    // since that is the cue to delete a file.
+    const newFilesFormatted = await formatFiles(
+      newFilesUnformatted,
+    );
+
+    // Capture the current state of files before applying changes
+    let vizFilesBefore: VizFiles;
+    if (EMIT_FIXTURES) {
+      vizFilesBefore = JSON.parse(
+        JSON.stringify(shareDBDoc.data.files),
+      );
+    }
+
     // Apply all the edits at once
-    updateFiles(
-      shareDBDoc,
-      mergeFileChanges(
-        shareDBDoc.data.files,
-        parseMarkdownFiles(fullContent, 'bold').files,
-      ),
+    const vizFilesAfter: VizFiles = mergeFileChanges(
+      shareDBDoc.data.files,
+      newFilesFormatted,
     );
 
-    // Finalize streaming message
-    await finalizeStreamingMessage(shareDBDoc, chatId);
-
-    // Generate a new runId to trigger a run when AI finishes editing
-    // This will trigger a re-run without hot reloading
-    const newRunId = generateRunId();
-    const runIdOp = diff(shareDBDoc.data, {
-      ...shareDBDoc.data,
-      runId: newRunId,
-    });
-    shareDBDoc.submitOp(runIdOp, (error) => {
-      if (error) {
-        console.warn(
-          'Error setting runId after AI editing:',
-          error,
-        );
-      } else {
-        DEBUG &&
-          console.log(
-            'Set new runId after AI editing:',
-            newRunId,
-          );
+    const filesOp = updateFiles(shareDBDoc, vizFilesAfter);
+
+    if (EMIT_FIXTURES) {
+      const fs = await import('fs');
+      const path = await import('path');
+      const testCasesDir = path.resolve(
+        process.cwd(),
+        '../',
+        'fixtures',
+      );
+      if (!fs.existsSync(testCasesDir)) {
+        fs.mkdirSync(testCasesDir, { recursive: true });
       }
-    });
+      const timestamp = new Date()
+        .toISOString()
+        .replace(/[:.]/g, '-');
+      const testCasePath = path.join(
+        testCasesDir,
+        `ai-chat-${timestamp}.json`,
+      );
+      const testCaseData = {
+        vizFilesBefore,
+        vizFilesAfter,
+        filesOp,
+      };
+      fs.writeFileSync(
+        testCasePath,
+        JSON.stringify(testCaseData, null, 2),
+      );
+      console.log(
+        `AI chat test case written to ${testCasePath}`,
+      );
+    }
 
-    // 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}`,
-    //   );
-    // }
+    // Finalize streaming message
+    finalizeStreamingMessage(shareDBDoc, chatId);
 
     return {
       content: fullContent,

package/src/server/prettier.ts

@@ -5,7 +5,7 @@ import * as prettierPluginHtml from 'prettier/plugins/html';
 import * as prettierPluginMarkdown from 'prettier/plugins/markdown';
 import * as prettierPluginCSS from 'prettier/plugins/postcss';
 import * as prettierPluginTypescript from 'prettier/plugins/typescript';
-import { VizFiles, VizFileId } from '@vizhub/viz-types';
+import { FileCollection } from '@vizhub/viz-types';
 
 // Parser mappings - matches client-side implementation
 const parsers = {
@@ -82,30 +82,45 @@ export const formatFile = async (
   }
 };
 
+// Keys are file names, values are file text contents
+//export type FileCollection = Record<string, string>;
+
 /**
  * Formats multiple files using Prettier
  * Only formats files that have supported extensions
  * Returns a map of fileId -> formatted text for successfully formatted files
  */
 export const formatFiles = async (
-  files: VizFiles,
-  fileIds?: VizFileId[],
-): Promise<{ [fileId: VizFileId]: string }> => {
-  const results: { [fileId: VizFileId]: string } = {};
-  const targetFileIds = fileIds || Object.keys(files);
+  fileCollection: FileCollection,
+): Promise<FileCollection> => {
+  const results: FileCollection = {};
+  const targetFileNames = Object.keys(fileCollection);
 
   // Process files in parallel for better performance
-  const formatPromises = targetFileIds.map(
-    async (fileId) => {
-      const file = files[fileId];
-      if (!file) return;
+  const formatPromises = targetFileNames.map(
+    async (fileName: string) => {
+      const fileText = fileCollection[fileName];
+      results[fileName] = fileText;
+
+      // Preserve empty files as empty,
+      // since this is the cue to delete a file.
+      if (!fileText || fileText.trim() === '') {
+        return;
+      }
 
-      const formatted = await formatFile(
-        file.text,
-        file.name,
-      );
-      if (formatted !== null && formatted !== file.text) {
-        results[fileId] = formatted;
+      try {
+        const formatted = await formatFile(
+          fileText,
+          fileName,
+        );
+        if (formatted !== null && formatted !== fileText) {
+          results[fileName] = formatted;
+        }
+      } catch (error) {
+        console.error(
+          `Error formatting ${fileName}:`,
+          error,
+        );
       }
     },
   );

package/src/types.ts

@@ -245,6 +245,11 @@ export type StreamingEvent =
       afterContent: string;
       timestamp: number;
     }
+  | {
+      type: 'file_delete';
+      fileName: string;
+      timestamp: number;
+    }
   | {
       type: 'status_update';
       status: string;

package/src/utils/fileDiff.ts

@@ -2,7 +2,33 @@ import { createTwoFilesPatch } from 'diff';
 import { VizFiles, VizFileId } from '@vizhub/viz-types';
 
 export interface UnifiedFilesDiff {
-  [fileId: VizFileId]: string; // Unified diff string
+  [fileId: VizFileId]: string; // Unified diff string or deletion marker
+}
+
+// Special marker to indicate a file deletion
+export const FILE_DELETION_MARKER = '__FILE_DELETED__';
+
+/**
+ * Check if a diff string represents a file deletion
+ */
+export function isFileDeletion(
+  diffString: string,
+): boolean {
+  return diffString.startsWith(FILE_DELETION_MARKER);
+}
+
+/**
+ * Extract file name from a deletion marker
+ */
+export function getDeletedFileName(
+  diffString: string,
+): string {
+  if (!isFileDeletion(diffString)) {
+    return '';
+  }
+  return diffString.substring(
+    FILE_DELETION_MARKER.length + 1,
+  );
 }
 
 /**
@@ -19,6 +45,11 @@ export function generateFileUnifiedDiff(
     return '';
   }
 
+  // Handle file deletion case - when file existed before but is now empty/deleted
+  if (beforeContent && !afterContent) {
+    return `${FILE_DELETION_MARKER}:${fileName}`;
+  }
+
   // Use the diff library's createTwoFilesPatch function to generate unified diff
   const unifiedDiff = createTwoFilesPatch(
     fileName,
@@ -109,9 +140,12 @@ export function parseUnifiedDiffStats(
 
 /**
  * Combine multiple unified diffs into a single diff string
+ * Note: File deletion markers are excluded from combination
  */
 export function combineUnifiedDiffs(
   unifiedDiffs: UnifiedFilesDiff,
 ): string {
-  return Object.values(unifiedDiffs).join('\n');
+  return Object.values(unifiedDiffs)
+    .filter((diff) => !isFileDeletion(diff))
+    .join('\n');
 }