vzcode 1.45.0 → 1.46.0

package/src/client/VZSidebar/aiCopyPaste.ts

@@ -7,8 +7,12 @@ import {
   mergeFileChanges,
 } from 'editcodewithai';
 import { VizFiles } from '@vizhub/viz-types';
-import { vizFilesToFileCollection } from '@vizhub/viz-utils';
+import {
+  generateRunId,
+  vizFilesToFileCollection,
+} from '@vizhub/viz-utils';
 import JSZip from 'jszip';
+import { generate } from '@langchain/core/dist/utils/fast-json-patch';
 
 export const createAICopyPasteHandlers = (
   files: VizFiles,
@@ -106,6 +110,7 @@ export const createAICopyPasteHandlers = (
         submitOperation((document) => ({
           ...document,
           files: mergedFiles,
+          runId: generateRunId(),
         }));
 
         const fileCount = Object.keys(parsed.files).length;

package/src/client/VZSidebar/index.tsx

@@ -110,7 +110,7 @@ export const VZSidebar = ({
   ),
   aiChatToolTipText = (
     <div>
-      <strong>AI Chat (beta)</strong>
+      <strong>VizBot - AI Code Editor</strong>
     </div>
   ),
 }: {
@@ -184,7 +184,7 @@ export const VZSidebar = ({
     [files, submitOperation],
   );
 
-  const { sidebarWidth } = useContext(
+  const { sidebarWidth, setSidebarView } = useContext(
     SplitPaneResizeContext,
   );
 
@@ -282,6 +282,11 @@ export const VZSidebar = ({
     previousPendingRef.current = pending;
   }, [pending, connected, isConnecting]);
 
+  // Automatically adjust sidebar width when view changes
+  useEffect(() => {
+    setSidebarView(isAIChatOpen);
+  }, [isAIChatOpen, setSidebarView]);
+
   return (
     <div
       className="vz-sidebar"
@@ -307,6 +312,7 @@ export const VZSidebar = ({
               onClick={() => {
                 setIsSearchOpen(false);
                 setIsAIChatOpen(false);
+                setSidebarView(false); // Switch to files view
               }}
             >
               <FolderSVG />
@@ -327,6 +333,7 @@ export const VZSidebar = ({
               onClick={() => {
                 setIsSearchOpen(true);
                 setIsAIChatOpen(false);
+                setSidebarView(false); // Switch to files view (search uses same width as files)
               }}
             >
               <SearchSVG />
@@ -348,6 +355,7 @@ export const VZSidebar = ({
                 onClick={() => {
                   setIsAIChatOpen(true);
                   setIsSearchOpen(false);
+                  setSidebarView(true); // Switch to AI chat view
                 }}
               >
                 <SparklesSVG />

package/src/client/bootstrap.ts

@@ -3,6 +3,8 @@
 // Otherwise we get `Error [ERR_UNSUPPORTED_DIR_IMPORT]: Directory import`
 import Form from 'react-bootstrap/cjs/Form.js';
 import Button from 'react-bootstrap/cjs/Button.js';
+import ButtonGroup from 'react-bootstrap/cjs/ButtonGroup.js';
+import ToggleButton from 'react-bootstrap/cjs/ToggleButton.js';
 import Modal from 'react-bootstrap/cjs/Modal.js';
 import OverlayTrigger from 'react-bootstrap/cjs/OverlayTrigger.js';
 import Tooltip from 'react-bootstrap/cjs/Tooltip.js';
@@ -10,4 +12,12 @@ import Tooltip from 'react-bootstrap/cjs/Tooltip.js';
 // Pull in custom CSS from vizhub-ui.
 import 'vizhub-ui/dist/vizhub-ui.css';
 
-export { Form, Button, Modal, OverlayTrigger, Tooltip };
+export {
+  Form,
+  Button,
+  ButtonGroup,
+  ToggleButton,
+  Modal,
+  OverlayTrigger,
+  Tooltip,
+};

package/src/client/useActions.ts

@@ -168,6 +168,17 @@ export const useActions = (
     });
   }, [dispatch]);
 
+  // Set the AI chat mode (ask or edit)
+  const setAIChatMode = useCallback(
+    (mode: 'ask' | 'edit') => {
+      dispatch({
+        type: 'set_ai_chat_mode',
+        mode,
+      });
+    },
+    [dispatch],
+  );
+
   // True to show the settings modal.
   const setIsSettingsOpen = useCallback(
     (value: boolean) => {
@@ -251,6 +262,7 @@ export const useActions = (
     toggleSearchFocused,
     setIsAIChatOpen,
     toggleAIChatFocused,
+    setAIChatMode,
     setIsSettingsOpen,
     setIsDocOpen,
     closeSettings,

package/src/client/usePrettier/index.ts

@@ -1,7 +1,5 @@
-import { shouldTriggerRun } from '../shouldTriggerRun';
-import { JSONOp } from '../../ot';
 import { VizFileId, VizContent } from '@vizhub/viz-types';
-import { ShareDBDoc } from '../../types';
+import { JSONOp, ShareDBDoc } from '../../types';
 import { useEffect, useRef, useState } from 'react';
 
 // The time in milliseconds by which auto-saving is debounced.

package/src/client/vzReducer/aiChatReducer.ts

@@ -29,3 +29,16 @@ export const toggleAIChatFocusedReducer = (
   }
   return state;
 };
+
+export const setAIChatModeReducer = (
+  state: VZState,
+  action: VZAction,
+): VZState => {
+  if (action.type === 'set_ai_chat_mode') {
+    return {
+      ...state,
+      aiChatMode: action.mode,
+    };
+  }
+  return state;
+};

package/src/client/vzReducer/createInitialState.ts

@@ -27,6 +27,7 @@ export const createInitialState = ({
   isSearchOpen: false,
   isAIChatOpen: false,
   aiChatFocused: false,
+  aiChatMode: 'edit',
   isSettingsOpen: false,
   isDocOpen: false,
   editorWantsFocus: false,

package/src/client/vzReducer/index.ts

@@ -33,6 +33,7 @@ import {
 import {
   setIsAIChatOpenReducer,
   toggleAIChatFocusedReducer,
+  setAIChatModeReducer,
 } from './aiChatReducer';
 import { toggleAutoFollowReducer } from './toggleAutoFollowReducer';
 import { updatePresenceIndicatorReducer } from './updatePresenceIndicatorReducer';
@@ -62,6 +63,9 @@ export type VZState = {
   // True if the AI chat input should focus on the next render.
   aiChatFocused: boolean;
 
+  // The current AI chat mode: 'ask' or 'edit'
+  aiChatMode: 'ask' | 'edit';
+
   // True to show the settings modal.
   isSettingsOpen: boolean;
 
@@ -133,6 +137,10 @@ export type VZAction =
   // * Toggles focused variable to trigger AI chat input focus
   | { type: 'toggle_ai_chat_focused' }
 
+  // `set_ai_chat_mode`
+  // * Sets the AI chat mode (ask or edit)
+  | { type: 'set_ai_chat_mode'; mode: 'ask' | 'edit' }
+
   // `set_search`
   //  * Sets the current search pattern
   | { type: 'set_search'; value: string }
@@ -212,6 +220,7 @@ const reducers = [
   setIsSearchOpenReducer,
   setIsAIChatOpenReducer,
   toggleAIChatFocusedReducer,
+  setAIChatModeReducer,
   setIsSettingsOpenReducer,
   setIsDocOpenReducer,
   editorNoLongerWantsFocusReducer,

package/src/client/vzReducer/searchReducer.test.ts

@@ -2,10 +2,52 @@ import { describe, expect, it } from 'vitest';
 import { setSearchResultsReducer } from './searchReducer';
 import { VZAction, VZState, createInitialState } from '.';
 import { defaultTheme } from '../themes';
+import { ShareDBDoc } from '../../types';
+import { VizContent } from '@vizhub/viz-types';
 
 // Mock ShareDBDoc for testing
-const createMockShareDBDoc = (files: any) => ({
-  data: { files },
+const createMockShareDBDoc = (
+  files: any,
+): ShareDBDoc<VizContent> => ({
+  data: { files, id: 'mock-id' },
+  ingestSnapshot: (snapshot: any, callback: any) => {
+    // No-op for testing
+    if (callback) callback();
+  },
+  subscribe: (callback: any) => {
+    // No-op for testing
+    if (callback) callback();
+  },
+  on: (
+    event: string,
+    callback: (op: any, source: boolean) => void,
+  ) => {
+    // No-op for testing
+  },
+  off: (
+    event: string,
+    callback: (op: any, source: boolean) => void,
+  ) => {
+    // No-op for testing
+  },
+  removeListener: (
+    event: string,
+    callback: (op: any, source: boolean) => void,
+  ) => {
+    // No-op for testing
+  },
+  submitOp: (
+    op: any,
+    options?: any,
+    callback?: () => void,
+  ) => {
+    // No-op for testing
+    if (callback) callback();
+  },
+  whenNothingPending: (callback: () => void) => {
+    // No-op for testing
+    if (callback) callback();
+  },
 });
 
 const initialState: VZState = createInitialState({

package/src/runCode.ts

@@ -1,30 +1,18 @@
+import { generateRunId } from '@vizhub/viz-utils';
 import { SubmitOperation } from './types';
 import { VizContent } from '@vizhub/viz-types';
 
 /**
- * Creates a runCode function that triggers code execution by flashing `isInteracting` to `true`.
+ * Creates a runCode function that triggers code execution by changing the `runId` in the content.
  * This works for both client-side (with submitOperation) and server-side (with ShareDB document).
  */
-export const createRunCodeFunction = (
-  submitOperation: SubmitOperation,
-) => {
-  return () => {
-    // Use the unified submitOperation approach for both client and server
+export const createRunCodeFunction =
+  (submitOperation: SubmitOperation) => () => {
     submitOperation((content: VizContent) => ({
       ...content,
-      isInteracting: true,
+      runId: generateRunId(),
     }));
-
-    setTimeout(() => {
-      // This somewhat cryptic logic
-      // deletes the `isInteracting` property
-      // from the document.
-      submitOperation(
-        ({ isInteracting, ...newDocument }) => newDocument,
-      );
-    }, 100);
   };
-};
 
 /**
  * Creates a runCodeRef object compatible with React refs.

package/src/server/aiChatHandler/aiEditing.ts

@@ -12,6 +12,42 @@ import {
 // Useful for debugging and testing purposes, e.g. checking the typing indicator.
 const delayStart = false;
 
+/**
+ * Performs AI chat without editing - just generates a response
+ */
+export const performAIChat = async ({
+  prompt,
+  shareDBDoc,
+  llmFunction,
+}) => {
+  const preparedFiles = prepareFilesForPrompt(
+    shareDBDoc.data.files,
+  );
+  const filesContext = formatMarkdownFiles(preparedFiles);
+
+  // 2. Assemble the final prompt for Q&A mode
+  const fullPrompt = assembleFullPrompt({
+    filesContext,
+    prompt,
+    editFormat: 'whole',
+  });
+
+  if (delayStart) {
+    await new Promise((resolve) =>
+      setTimeout(resolve, 1000),
+    );
+  }
+
+  // Call the LLM function which will handle streaming but won't edit files
+  const result = await llmFunction(fullPrompt);
+
+  return {
+    content: result.content,
+    generationId: result.generationId,
+    diffData: {}, // No file changes in ask mode
+  };
+};
+
 /**
  * Performs AI editing operations using streaming with incremental OT operations
  */
@@ -62,5 +98,6 @@ export const performAIEditing = async ({
     content: result.content,
     generationId: result.generationId,
     diffData,
+    beforeFiles, // Store the snapshot for undo functionality
   };
 };

package/src/server/aiChatHandler/chatOperations.ts

@@ -1,4 +1,7 @@
-import { dateToTimestamp } from '@vizhub/viz-utils';
+import {
+  dateToTimestamp,
+  generateRunId,
+} from '@vizhub/viz-utils';
 import { randomId } from '../../randomId.js';
 import { ShareDBDoc } from '../../types.js';
 import { diff } from '../../ot.js';
@@ -253,6 +256,7 @@ export const addDiffToAIMessage = (
   shareDBDoc: ShareDBDoc<VizContent>,
   chatId: VizChatId,
   diffData: any,
+  beforeFiles?: VizFiles, // Optional snapshot for undo functionality
 ) => {
   const chat = shareDBDoc.data.chats[chatId];
   const messages = [...chat.messages];
@@ -266,6 +270,7 @@ export const addDiffToAIMessage = (
     const newMessage = {
       ...messages[lastAIMessageIndex],
       diffData,
+      ...(beforeFiles && { beforeFiles }), // Add beforeFiles only if provided
     };
     // Use type assertion to extend the message with diffData
     (messages[lastAIMessageIndex] as any) = newMessage;
@@ -447,3 +452,35 @@ export const appendLineToFile = (
 
   shareDBDoc.submitOp(diff(shareDBDoc.data, newDocState));
 };
+
+/**
+ * Undoes the last AI edit by restoring files and removing the AI message
+ */
+export const undoAIEdit = (
+  shareDBDoc: ShareDBDoc<VizContent>,
+  chatId: VizChatId,
+  messageId: string,
+  beforeFiles: VizFiles,
+) => {
+  const chat = shareDBDoc.data.chats[chatId];
+  const messages = chat.messages.filter(
+    (msg) => msg.id !== messageId,
+  );
+
+  const op = diff(shareDBDoc.data, {
+    ...shareDBDoc.data,
+    files: beforeFiles,
+    chats: {
+      ...shareDBDoc.data.chats,
+      [chatId]: {
+        ...chat,
+        messages,
+        updatedAt: dateToTimestamp(new Date()),
+      },
+    },
+    // Trigger a re-run
+    runId: generateRunId(),
+  });
+
+  shareDBDoc.submitOp(op);
+};

package/src/server/aiChatHandler/index.ts

@@ -6,7 +6,10 @@ import {
   addDiffToAIMessage,
 } from './chatOperations.js';
 import { createLLMFunction } from './llmStreaming.js';
-import { performAIEditing } from './aiEditing.js';
+import {
+  performAIEditing,
+  performAIChat,
+} from './aiEditing.js';
 import { handleError } from './errorHandling.js';
 import { createRunCodeFunction } from '../../runCode.js';
 import { ShareDBDoc } from '../../types.js';
@@ -26,7 +29,7 @@ export const handleAIChatMessage =
     onCreditDeduction?: any;
   }) =>
   async (req: any, res: any) => {
-    const { content, chatId } = req.body;
+    const { content, chatId, mode = 'edit' } = req.body;
 
     if (DEBUG) {
       console.log(
@@ -65,13 +68,20 @@ export const handleAIChatMessage =
       const runCode =
         createRunCodeFunction(submitOperation);
 
-      // Perform AI editing
-      const editResult = await performAIEditing({
-        prompt: content,
-        shareDBDoc,
-        llmFunction,
-        runCode,
-      });
+      // 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 (
@@ -82,6 +92,7 @@ export const handleAIChatMessage =
           shareDBDoc,
           chatId,
           editResult.diffData,
+          (editResult as any).beforeFiles, // Pass the beforeFiles snapshot for undo (only available for edit mode)
         );
       }
 

package/src/server/aiChatHandler/llmStreaming.ts

@@ -3,7 +3,9 @@ import {
   StreamingMarkdownParser,
 } from 'llm-code-format';
 import { ChatOpenAI } from '@langchain/openai';
+import OpenAI from 'openai';
 import fs from 'fs';
+import { generateRunId } from '@vizhub/viz-utils';
 import {
   updateAIStatus,
   createAIMessage,
@@ -13,8 +15,12 @@ import {
   clearFileContent,
   appendLineToFile,
   updateFiles,
+  updateAIScratchpad,
 } from './chatOperations.js';
 import { mergeFileChanges } from 'editcodewithai';
+import { diff } from '../../ot.js';
+import { VizChatId, VizContent } from '@vizhub/viz-types';
+import { ShareDBDoc } from '../../types.js';
 
 const DEBUG = false;
 
@@ -34,28 +40,30 @@ const DEBUG = false;
 const enableStreamingEditing = false;
 
 /**
- * Creates and configures the LLM function for streaming
+ * Creates and configures the LLM function for streaming with reasoning tokens
  */
 export const createLLMFunction = ({
   shareDBDoc,
   createVizBotLocalPresence,
   chatId,
+}: {
+  shareDBDoc: ShareDBDoc<VizContent>;
+  createVizBotLocalPresence: () => any;
+  chatId: VizChatId;
 }) => {
   return async (fullPrompt: string) => {
     const localPresence = createVizBotLocalPresence();
-    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',
-        },
+
+    // Create OpenRouter client for reasoning token support
+    const openRouterClient = new OpenAI({
+      apiKey: process.env.VIZHUB_EDIT_WITH_AI_API_KEY,
+      baseURL:
+        process.env.VIZHUB_EDIT_WITH_AI_BASE_URL ||
+        'https://openrouter.ai/api/v1',
+      defaultHeaders: {
+        'HTTP-Referer': 'https://vizhub.com',
+        'X-Title': 'VizHub',
       },
-      streaming: true,
     });
 
     let fullContent = '';
@@ -177,12 +185,58 @@ export const createLLMFunction = ({
     const parser = new StreamingMarkdownParser(callbacks);
 
     const chunks = [];
+    let reasoningContent = '';
+
+    // Stream the response with reasoning tokens
+    const modelName =
+      process.env.VIZHUB_EDIT_WITH_AI_MODEL_NAME ||
+      'anthropic/claude-3.5-sonnet';
+    const stream = await (
+      openRouterClient.chat.completions.create as any
+    )({
+      model: modelName,
+      messages: [{ role: 'user', content: fullPrompt }],
+      max_tokens: 8192,
+      reasoning: {
+        effort: 'medium',
+        exclude: false,
+      },
+      usage: { include: true },
+      stream: true,
+    });
+
+    let reasoningStarted = false;
+    let contentStarted = false;
 
-    // Stream the response
-    const stream = await chatModel.stream(fullPrompt);
     for await (const chunk of stream) {
-      if (chunk && chunk.content) {
-        const chunkContent = String(chunk.content);
+      const delta = chunk.choices[0]?.delta as any; // Type assertion for OpenRouter-specific reasoning fields
+
+      if (delta?.reasoning) {
+        // Handle reasoning tokens (thinking)
+        if (!reasoningStarted) {
+          reasoningStarted = true;
+          updateAIStatus(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
+          contentStarted = true;
+          updateAIScratchpad(shareDBDoc, chatId, '');
+          updateAIStatus(
+            shareDBDoc,
+            chatId,
+            'Generating response...',
+          );
+        }
+
+        const chunkContent = delta.content;
         chunks.push(chunkContent);
 
         if (enableStreamingEditing) {
@@ -196,14 +250,20 @@ export const createLLMFunction = ({
             fullContent,
           );
         }
+      } else if (chunk.usage) {
+        // Handle usage information
+        DEBUG && console.log('Usage:', chunk.usage);
       }
 
-      if (!generationId && chunk.lc_kwargs?.id) {
-        generationId = chunk.lc_kwargs.id;
+      if (!generationId && chunk.id) {
+        generationId = chunk.id;
       }
     }
     await parser.flushRemaining();
     reportFileEdited();
+
+    // Final cleanup - clear scratchpad and set final status
+    updateAIScratchpad(shareDBDoc, chatId, '');
     updateAIStatus(shareDBDoc, chatId, 'Done editing.');
     updateAIMessageContent(
       shareDBDoc,
@@ -242,6 +302,28 @@ export const createLLMFunction = ({
       );
     }
 
+    // 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,
+          );
+      }
+    });
+
     // Write chunks file for debugging
     if (DEBUG) {
       const chunksFileJSONpath = `./ai-chunks-${chatId}.json`;

package/src/server/generateAIResponse.ts

@@ -150,6 +150,10 @@ export const generateAIResponse = async ({
     // model: 'gpt-3.5-turbo',
     model: 'gpt-4o',
     messages,
+    reasoning: {
+      effort: 'medium',
+      exclude: false,
+    } as any, // Type assertion for OpenRouter-specific reasoning parameter
     stream: true,
   });