npm - vzcode - Versions diffs - 2.15.0 → 2.17.0 - Mend

vzcode 2.15.0 → 2.17.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Files changed (51) hide show

package/dist/assets/{bindings_wasm_bg-9w8E-TmY.wasm → bindings_wasm_bg-D9vNLG9F.wasm} +0 -0
package/dist/assets/buildWorker-DL4wXpRr.js +695 -0
package/dist/assets/index-BiPmUreW.js +474 -0
package/dist/assets/{index-DLm5FQ0E.css → index-BvGPtSrr.css} +1 -1
package/dist/assets/worker-RJ-cjbld.js +327 -0
package/dist/index.html +2 -2
package/dist/server/aiChatHandler/chatOperations.js +168 -35
package/dist/server/aiChatHandler/index.js +11 -30
package/dist/server/aiChatHandler/llmStreaming.js +105 -118
package/package.json +11 -11
package/src/client/AIAssist/AIAssistWidget/index.tsx +12 -1
package/src/client/App/useShareDB.ts +1 -1
package/src/client/CodeEditor/index.tsx +9 -10
package/src/client/VZCodeContext/types.ts +3 -1
package/src/client/VZCodeContext/useVZCodeState.ts +7 -5
package/src/client/VZRight.tsx +10 -2
package/src/client/VZSidebar/AIChat/ChatInput.tsx +1 -7
package/src/client/VZSidebar/AIChat/DiffView.scss +1 -0
package/src/client/VZSidebar/AIChat/DiffView.tsx +72 -29
package/src/client/VZSidebar/AIChat/FileEditingIndicator.tsx +89 -0
package/src/client/VZSidebar/AIChat/IndividualFileDiff.tsx +86 -0
package/src/client/VZSidebar/AIChat/JumpToLatestButton.tsx +64 -0
package/src/client/VZSidebar/AIChat/Message.tsx +124 -56
package/src/client/VZSidebar/AIChat/MessageList.tsx +155 -114
package/src/client/VZSidebar/AIChat/ThinkingScratchpad.tsx +4 -118
package/src/client/VZSidebar/AIChat/index.tsx +59 -19
package/src/client/VZSidebar/AIChat/styles.scss +190 -43
package/src/client/VZSidebar/Item.tsx +2 -2
package/src/client/VZSidebar/Search.tsx +13 -6
package/src/client/VZSidebar/VisualEditor/VisualEditor.tsx +39 -23
package/src/client/VZSidebar/VisualEditor/utils.ts +1 -1
package/src/client/VZSidebar/index.tsx +12 -10
package/src/client/VZSidebar/useDragAndDrop.tsx +225 -214
package/src/client/featureFlags.ts +3 -0
package/src/client/hooks/useAutoScroll.ts +328 -0
package/src/client/tabsSearchParameters.ts +1 -17
package/src/client/useFileCRUD.ts +5 -2
package/src/client/useKeyboardShortcuts.ts +7 -0
package/src/client/useOpenDirectories.ts +2 -1
package/src/client/usePrettier/index.ts +1 -1
package/src/client/useURLSync.ts +1 -0
package/src/client/utils/scrollUtils.ts +170 -0
package/src/client/vzReducer/searchReducer.ts +6 -3
package/src/server/aiChatHandler/chatOperations.ts +224 -43
package/src/server/aiChatHandler/index.ts +8 -48
package/src/server/aiChatHandler/llmStreaming.ts +149 -170
package/src/types.ts +81 -1
package/dist/assets/buildWorker-Bi6wsfQk.js +0 -695
package/dist/assets/index-BpUpNto4.js +0 -461
package/dist/assets/worker-BSzbM0Uo.js +0 -330
package/src/client/VZSidebar/AIChat/StreamingMessage.tsx +0 -35

package/dist/index.html CHANGED Viewed

@@ -20,8 +20,8 @@
       href="https://fonts.googleapis.com/css2?family=Poppins:wght@300;400;500;600;700&display=swap"
       rel="stylesheet"
     />
-    <script type="module" crossorigin src="/assets/index-BpUpNto4.js"></script>
-    <link rel="stylesheet" crossorigin href="/assets/index-DLm5FQ0E.css">
+    <script type="module" crossorigin src="/assets/index-BiPmUreW.js"></script>
+    <link rel="stylesheet" crossorigin href="/assets/index-BvGPtSrr.css">
   </head>
   <body>
     <div id="root"></div>

package/dist/server/aiChatHandler/chatOperations.js CHANGED Viewed

@@ -60,10 +60,13 @@ export const addUserMessage = (shareDBDoc, chatId, content) => {
     shareDBDoc.submitOp(userMessageOp);
     return userMessage;
 };
+const DEBUG = false;
 /**
  * Updates AI status in the chat
  */
 export const updateAIStatus = (shareDBDoc, chatId, status) => {
+    DEBUG &&
+        console.log(`ChatOperations: updateAIStatus called with status: "${status}" for chatId: ${chatId}`);
     const op = diff(shareDBDoc.data, {
         ...shareDBDoc.data,
         chats: {
@@ -74,7 +77,11 @@ export const updateAIStatus = (shareDBDoc, chatId, status) => {
             },
         },
     });
+    DEBUG &&
+        console.log(`ChatOperations: Submitting operation for status update:`, op);
     shareDBDoc.submitOp(op);
+    DEBUG &&
+        console.log(`ChatOperations: Status update operation submitted successfully`);
 };
 /**
  * Updates AI scratchpad content
@@ -307,54 +314,180 @@ export const createNewFile = (shareDBDoc, fileName) => {
     shareDBDoc.submitOp(op);
     return newFileId;
 };
+// ============================================================================
+// Streaming Chat Operations
+// ============================================================================
 /**
- * Ensures a file exists, creating it if necessary
+ * Creates a streaming AI message with events array
  */
-export const ensureFileExists = (shareDBDoc, fileName) => {
-    let fileId = resolveFileId(fileName, shareDBDoc);
-    if (!fileId) {
-        // File doesn't exist, create it
-        fileId = createNewFile(shareDBDoc, fileName);
-    }
-    return fileId;
+export const createStreamingAIMessage = (shareDBDoc, chatId) => {
+    const aiMessage = {
+        id: `assistant-${Date.now()}`,
+        role: 'assistant',
+        content: '',
+        timestamp: dateToTimestamp(new Date()),
+        streamingEvents: [],
+        isProgressive: true,
+    };
+    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()),
+                isStreaming: true,
+            },
+        },
+    });
+    shareDBDoc.submitOp(messageOp);
+    return aiMessage.id;
 };
 /**
- * Clears the content of a file
+ * Adds a streaming event to the most recent AI message
  */
-export const clearFileContent = (shareDBDoc, fileId) => {
-    const currentFile = shareDBDoc.data.files[fileId];
-    if (currentFile && currentFile.text) {
-        // Clear the file content
-        const newState = {
+export const addStreamingEvent = (shareDBDoc, chatId, event) => {
+    DEBUG &&
+        console.log(`ChatOperations: Adding streaming event:`, event);
+    const chat = shareDBDoc.data.chats[chatId];
+    const messages = [...chat.messages];
+    const lastMessageIndex = messages.length - 1;
+    if (lastMessageIndex >= 0 &&
+        messages[lastMessageIndex].role === 'assistant') {
+        const lastMessage = messages[lastMessageIndex];
+        const updatedEvents = [
+            ...(lastMessage.streamingEvents || []),
+            event,
+        ];
+        messages[lastMessageIndex] = {
+            ...lastMessage,
+            streamingEvents: updatedEvents,
+        };
+        const messageOp = diff(shareDBDoc.data, {
             ...shareDBDoc.data,
-            files: {
-                ...shareDBDoc.data.files,
-                [fileId]: {
-                    ...currentFile,
-                    text: '',
+            chats: {
+                ...shareDBDoc.data.chats,
+                [chatId]: {
+                    ...chat,
+                    messages,
+                    updatedAt: dateToTimestamp(new Date()),
                 },
             },
-        };
-        const op = diff(shareDBDoc.data, newState);
-        shareDBDoc.submitOp(op);
+        });
+        shareDBDoc.submitOp(messageOp);
     }
 };
 /**
- * Appends a line to a file using OT operations
+ * Updates streaming status for a chat
  */
-export const appendLineToFile = (shareDBDoc, fileId, line) => {
-    const currentFile = shareDBDoc.data.files[fileId];
-    const currentContent = currentFile?.text || '';
-    const newContent = currentContent + line + '\n';
-    const newDocState = {
+export const updateStreamingStatus = (shareDBDoc, chatId, status, isStreaming = true) => {
+    DEBUG &&
+        console.log(`ChatOperations: Updating streaming status: "${status}"`);
+    const op = diff(shareDBDoc.data, {
         ...shareDBDoc.data,
-        files: {
-            ...shareDBDoc.data.files,
-            [fileId]: {
-                ...currentFile,
-                text: newContent,
+        chats: {
+            ...shareDBDoc.data.chats,
+            [chatId]: {
+                ...shareDBDoc.data.chats[chatId],
+                currentStatus: status,
+                isStreaming,
+                updatedAt: dateToTimestamp(new Date()),
             },
         },
-    };
-    shareDBDoc.submitOp(diff(shareDBDoc.data, newDocState));
+    });
+    shareDBDoc.submitOp(op);
+};
+/**
+ * Finalizes streaming message and clears streaming state
+ */
+export const finalizeStreamingMessage = (shareDBDoc, chatId) => {
+    DEBUG &&
+        console.log(`ChatOperations: Finalizing streaming message`);
+    const chat = shareDBDoc.data.chats[chatId];
+    const messages = [...chat.messages];
+    const lastMessageIndex = messages.length - 1;
+    if (lastMessageIndex >= 0 &&
+        messages[lastMessageIndex].role === 'assistant') {
+        const lastMessage = messages[lastMessageIndex];
+        messages[lastMessageIndex] = {
+            ...lastMessage,
+            isComplete: true,
+        };
+        const messageOp = diff(shareDBDoc.data, {
+            ...shareDBDoc.data,
+            chats: {
+                ...shareDBDoc.data.chats,
+                [chatId]: {
+                    ...chat,
+                    messages,
+                    currentStatus: 'Done',
+                    isStreaming: false,
+                    updatedAt: dateToTimestamp(new Date()),
+                },
+            },
+        });
+        shareDBDoc.submitOp(messageOp);
+    }
 };
+// /**
+//  * 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: ShareDBDoc<VizContent>,
+//   fileId: VizFileId,
+// ) => {
+//   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: ShareDBDoc<VizContent>,
+//   fileId: VizFileId,
+//   line: string,
+// ) => {
+//   const currentFile = shareDBDoc.data.files[fileId];
+//   const currentContent = currentFile?.text || '';
+//   const newContent = currentContent + line + '\n';
+//   const newDocState = {
+//     ...shareDBDoc.data,
+//     files: {
+//       ...shareDBDoc.data.files,
+//       [fileId]: {
+//         ...currentFile,
+//         text: newContent,
+//       },
+//     },
+//   };
+//   shareDBDoc.submitOp(diff(shareDBDoc.data, newDocState));
+// };

package/dist/server/aiChatHandler/index.js CHANGED Viewed

@@ -1,14 +1,14 @@
 import { validateRequest } from './validation.js';
-import { ensureChatsExist, ensureChatExists, addUserMessage, addDiffToAIMessage, setAIStatus, } from './chatOperations.js';
+import { ensureChatsExist, ensureChatExists, addUserMessage, setAIStatus, } from './chatOperations.js';
 import { createLLMFunction } from './llmStreaming.js';
-import { performAIEditing, performAIChat, } from './aiEditing.js';
+import { performAIEditing } from './aiEditing.js';
 import { handleError, handleBackgroundError, } from './errorHandling.js';
 import { createRunCodeFunction } from '../../runCode.js';
 import { createSubmitOperation } from '../../submitOperation.js';
 import { getGenerationMetadata } from 'editcodewithai';
 const DEBUG = false;
-export const handleAIChatMessage = ({ shareDBDoc, onCreditDeduction, getCurrentCommitId, model, aiRequestOptions, }) => async (req, res) => {
-    const { content, chatId, mode = 'edit' } = req.body;
+export const handleAIChatMessage = ({ shareDBDoc, onCreditDeduction, model, aiRequestOptions, }) => async (req, res) => {
+    const { content, chatId } = req.body;
     if (DEBUG) {
         console.log('[handleAIChatMessage] content:', content, 'chatId:', chatId, 'shareDBDoc:', shareDBDoc);
     }
@@ -24,15 +24,11 @@ export const handleAIChatMessage = ({ shareDBDoc, onCreditDeduction, getCurrentC
         addUserMessage(shareDBDoc, chatId, content);
         // Return success immediately - AI generation continues in background
         res.status(200).json('success');
-        // Set AI status to indicate generation is starting
-        setAIStatus(shareDBDoc, chatId, 'generating');
         // Continue AI processing in background (don't await)
         processAIRequestAsync({
             shareDBDoc,
             chatId,
             content,
-            mode,
-            getCurrentCommitId,
             model,
             aiRequestOptions,
             onCreditDeduction,
@@ -49,12 +45,8 @@ export const handleAIChatMessage = ({ shareDBDoc, onCreditDeduction, getCurrentC
 /**
  * Processes the AI request asynchronously in the background
  */
-const processAIRequestAsync = async ({ shareDBDoc, chatId, content, mode, getCurrentCommitId, model, aiRequestOptions, onCreditDeduction, }) => {
+const processAIRequestAsync = async ({ shareDBDoc, chatId, content, model, aiRequestOptions, onCreditDeduction, }) => {
     try {
-        // Capture the current commit ID before making changes (for VizHub integration)
-        const beforeCommitId = getCurrentCommitId
-            ? getCurrentCommitId()
-            : null;
         // Create LLM function for streaming
         const llmFunction = createLLMFunction({
             shareDBDoc,
@@ -66,23 +58,12 @@ const processAIRequestAsync = async ({ shareDBDoc, chatId, content, mode, getCur
         const submitOperation = createSubmitOperation(shareDBDoc);
         const runCode = createRunCodeFunction(submitOperation);
         // Perform AI editing or chat based on mode
-        const editResult = mode === 'ask'
-            ? await performAIChat({
-                prompt: content,
-                shareDBDoc,
-                llmFunction,
-            })
-            : await performAIEditing({
-                prompt: content,
-                shareDBDoc,
-                llmFunction,
-                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, beforeCommitId);
-        }
+        const editResult = await performAIEditing({
+            prompt: content,
+            shareDBDoc,
+            llmFunction,
+            runCode,
+        });
         // Handle credit deduction if callback is provided
         if (onCreditDeduction && editResult.generationId) {
             try {

package/dist/server/aiChatHandler/llmStreaming.js CHANGED Viewed

@@ -1,34 +1,12 @@
-import fs from 'fs';
 import OpenAI from 'openai';
 import { parseMarkdownFiles, StreamingMarkdownParser, } from 'llm-code-format';
 import { mergeFileChanges } from 'editcodewithai';
 import { generateRunId } from '@vizhub/viz-utils';
-import { updateAIStatus, createAIMessage, updateAIMessageContent, finalizeAIMessage, ensureFileExists, clearFileContent, appendLineToFile, updateFiles, updateAIScratchpad, } from './chatOperations.js';
+import { updateFiles, updateAIScratchpad, createStreamingAIMessage, addStreamingEvent, updateStreamingStatus, finalizeStreamingMessage, } from './chatOperations.js';
 import { diff } from '../../ot.js';
 const DEBUG = false;
 // Useful for testing/debugging the streaming behavior
 const slowMode = false;
-// Throttle the streaming updates, so that we don't
-// overwhelm the ShareDB server with too many updates.
-// It happened actually, before adding this.
-// MongoDB VizHub server got in fact overloaded with
-// too many updates from the AI streaming response, with
-// warning: "Replication Oplog Window has gone below 1 hour"
-const THROTTLE_INTERVAL_MS = 500;
-// 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 with reasoning tokens
  */
@@ -50,87 +28,92 @@ enableReasoningTokens = false, model, aiRequestOptions, }) => {
         });
         let fullContent = '';
         let generationId = '';
-        let currentEditingFileId = null;
         let currentEditingFileName = null;
-        // Create initial AI message for streaming
-        const aiMessageId = createAIMessage(shareDBDoc, chatId);
-        // --- throttle wrapper ---
-        function makeThrottledUpdater() {
-            let lastCall = 0;
-            let latestContent = '';
-            let timer = null;
-            function invoke() {
-                updateAIMessageContent(shareDBDoc, chatId, aiMessageId, latestContent);
-                lastCall = Date.now();
-            }
-            const fn = (content) => {
-                latestContent = content;
-                const now = Date.now();
-                if (now - lastCall >= THROTTLE_INTERVAL_MS) {
-                    // safe to call immediately
-                    invoke();
-                }
-                else if (!timer) {
-                    // schedule for later
-                    timer = setTimeout(() => {
-                        timer = null;
-                        invoke();
-                    }, THROTTLE_INTERVAL_MS - (now - lastCall));
+        let accumulatedTextChunk = '';
+        let currentFileContent = '';
+        // Create streaming AI message
+        createStreamingAIMessage(shareDBDoc, chatId);
+        // Set initial content generation status
+        updateStreamingStatus(shareDBDoc, chatId, 'Formulating a plan...');
+        // Helper to get original file content
+        const getOriginalFileContent = (fileName) => {
+            const files = shareDBDoc.data.files;
+            for (const file of Object.values(files)) {
+                if (file.name === fileName) {
+                    return file.text || '';
                 }
-            };
-            // expose a flush() helper to force an immediate write
-            fn.flush = () => {
-                if (timer) {
-                    clearTimeout(timer);
-                    timer = null;
-                }
-                invoke();
-            };
-            return fn;
-        }
-        const throttledUpdateAIMessageContent = makeThrottledUpdater();
-        // Function to report file edited
-        // This is called when the AI has finished editing a file
-        // and we want to update the message content with the file name.
-        const reportFileEdited = () => {
-            if (currentEditingFileName) {
-                fullContent += ` * Edited ${currentEditingFileName}\n`;
-                throttledUpdateAIMessageContent(fullContent);
-                currentEditingFileName = null;
+            }
+            return '';
+        };
+        // Helper to emit text chunk when accumulated
+        const emitTextChunk = async () => {
+            if (accumulatedTextChunk.trim()) {
+                DEBUG &&
+                    console.log('LLMStreaming: Emitting text chunk:', accumulatedTextChunk.substring(0, 100) + '...');
+                await addStreamingEvent(shareDBDoc, chatId, {
+                    type: 'text_chunk',
+                    content: accumulatedTextChunk,
+                    timestamp: Date.now(),
+                });
+                accumulatedTextChunk = '';
+            }
+        };
+        // Helper to complete file editing
+        const completeFileEditing = async (fileName) => {
+            if (fileName && currentFileContent) {
+                DEBUG &&
+                    console.log(`LLMStreaming: Completing file editing for ${fileName}`);
+                await addStreamingEvent(shareDBDoc, chatId, {
+                    type: 'file_complete',
+                    fileName,
+                    beforeContent: getOriginalFileContent(fileName),
+                    afterContent: currentFileContent,
+                    timestamp: Date.now(),
+                });
+                currentFileContent = '';
             }
         };
         // Define callbacks for streaming parser
         const callbacks = {
             onFileNameChange: async (fileName, format) => {
                 DEBUG &&
-                    console.log(`File changed to: ${fileName} (${format})`);
-                // Find existing file or create new one
-                currentEditingFileId = ensureFileExists(shareDBDoc, fileName);
-                reportFileEdited();
+                    console.log(`LLMStreaming: File changed to: ${fileName} (${format})`);
+                // Emit any accumulated text chunk first
+                await emitTextChunk();
+                // Complete previous file if any
+                if (currentEditingFileName) {
+                    await completeFileEditing(currentEditingFileName);
+                }
+                // Start new file
                 currentEditingFileName = fileName;
-                // Clear the file content to start fresh
-                // (AI will regenerate the entire file content)
-                clearFileContent(shareDBDoc, currentEditingFileId);
-                // Update AI status
-                updateAIStatus(shareDBDoc, chatId, 'Editing ' + fileName);
+                currentFileContent = '';
+                // Emit file start event
+                await addStreamingEvent(shareDBDoc, chatId, {
+                    type: 'file_start',
+                    fileName,
+                    timestamp: Date.now(),
+                });
+                // Update status
+                updateStreamingStatus(shareDBDoc, chatId, `Editing ${fileName}...`);
             },
             onCodeLine: async (line) => {
                 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(shareDBDoc, currentEditingFileId, line);
-                }
+                // Accumulate code content for the current file
+                currentFileContent += line + '\n';
             },
             onNonCodeLine: async (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
+                DEBUG && console.log(`Non-code line: ${line}`);
+                // Accumulate non-code content as text chunk
                 if (line.trim() !== '') {
-                    reportFileEdited();
+                    accumulatedTextChunk += line + '\n';
+                    // Update status for subsequent non-code chunks
+                    if (firstNonCodeChunkProcessed) {
+                        updateStreamingStatus(shareDBDoc, chatId, 'Describing changes...');
+                    }
+                    else {
+                        firstNonCodeChunkProcessed = true;
+                    }
                 }
-                fullContent += line + '\n';
-                throttledUpdateAIMessageContent(fullContent);
             },
         };
         const parser = new StreamingMarkdownParser(callbacks);
@@ -158,6 +141,7 @@ enableReasoningTokens = false, model, aiRequestOptions, }) => {
         const stream = await openRouterClient.chat.completions.create(requestConfig);
         let reasoningStarted = false;
         let contentStarted = false;
+        let firstNonCodeChunkProcessed = false;
         for await (const chunk of stream) {
             if (slowMode) {
                 await new Promise((resolve) => setTimeout(resolve, 500));
@@ -167,28 +151,30 @@ enableReasoningTokens = false, model, aiRequestOptions, }) => {
                 // Handle reasoning tokens (thinking) - only if enabled
                 if (!reasoningStarted) {
                     reasoningStarted = true;
-                    updateAIStatus(shareDBDoc, chatId, 'Thinking...');
+                    updateStreamingStatus(shareDBDoc, chatId, 'Thinking...');
                 }
                 reasoningContent += delta.reasoning;
                 updateAIScratchpad(shareDBDoc, chatId, reasoningContent);
             }
             else if (delta?.content) {
                 // Handle regular content tokens
-                if (reasoningStarted && !contentStarted) {
-                    // Clear reasoning when content starts
+                if (!contentStarted) {
                     contentStarted = true;
-                    updateAIScratchpad(shareDBDoc, chatId, '');
-                    updateAIStatus(shareDBDoc, chatId, 'Generating response...');
+                    if (reasoningStarted) {
+                        // Clear reasoning when content starts
+                        updateAIScratchpad(shareDBDoc, chatId, '');
+                    }
+                    // // Set initial content generation status
+                    // updateStreamingStatus(
+                    //   shareDBDoc,
+                    //   chatId,
+                    //   'Formulating a plan...',
+                    // );
                 }
                 const chunkContent = delta.content;
                 chunks.push(chunkContent);
-                if (enableStreamingEditing) {
-                    await parser.processChunk(chunkContent);
-                }
-                else {
-                    fullContent += chunkContent;
-                    throttledUpdateAIMessageContent(fullContent);
-                }
+                await parser.processChunk(chunkContent);
+                fullContent += chunkContent;
             }
             else if (chunk.usage) {
                 // Handle usage information
@@ -199,20 +185,16 @@ enableReasoningTokens = false, model, aiRequestOptions, }) => {
             }
         }
         await parser.flushRemaining();
-        reportFileEdited();
-        // Final cleanup - clear scratchpad and set final status
-        updateAIScratchpad(shareDBDoc, chatId, '');
-        updateAIStatus(shareDBDoc, chatId, 'Done editing.');
-        throttledUpdateAIMessageContent(fullContent);
-        // Flush to ensure the final content is written immediately
-        throttledUpdateAIMessageContent.flush();
-        // Finalize the AI message by clearing temporary fields
-        finalizeAIMessage(shareDBDoc, chatId);
-        // 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));
+        // Emit any remaining text chunk
+        await emitTextChunk();
+        // Complete final file if any
+        if (currentEditingFileName) {
+            await completeFileEditing(currentEditingFileName);
         }
+        // Apply all the edits at once
+        updateFiles(shareDBDoc, mergeFileChanges(shareDBDoc.data.files, parseMarkdownFiles(fullContent, 'bold').files));
+        // 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();
@@ -230,11 +212,16 @@ enableReasoningTokens = false, model, aiRequestOptions, }) => {
             }
         });
         // 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}`);
-        }
+        // 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,