vzcode 1.57.0 → 1.59.0

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

@@ -1,42 +1,28 @@
-import {
-  useContext,
-  useState,
-  useCallback,
-  useMemo,
-  useEffect,
-} from 'react';
+import { useContext, useMemo, useEffect } from 'react';
 import { VZCodeContext } from '../../VZCodeContext';
-import { v4 as uuidv4 } from 'uuid';
 import { MessageList } from './MessageList';
 import { ChatInput } from './ChatInput';
 import './styles.scss';
 
-const defaultAIChatEndpoint = '/api/ai-chat/';
-
 const DEBUG = false;
 
-const showSuggestedRequests = false;
+const showSuggestedRequests = true;
 
 export const AIChat = () => {
-  const { aiChatMessage, setAIChatMessage } =
-    useContext(VZCodeContext);
-
-  const [isLoading, setIsLoading] = useState(false);
-  const [currentChatId] = useState(() => uuidv4());
-  const [errorMessage, setErrorMessage] = useState<
-    string | null
-  >(null);
-
   const {
     aiChatFocused,
     content,
-    aiChatEndpoint = defaultAIChatEndpoint,
-    aiChatOptions = {},
     aiChatMode,
+    aiChatMessage,
+    isLoading,
+    currentChatId,
+    aiErrorMessage,
     setAIChatMode,
-    autoForkAndRetryAI,
     clearStoredAIPrompt,
     getStoredAIPrompt,
+    setAIChatMessage,
+    handleSendMessage,
+    setAIErrorMessage,
   } = useContext(VZCodeContext);
 
   // Get current chat data from content
@@ -58,118 +44,6 @@ export const AIChat = () => {
   // Check if this is the first time opening the chat (no messages)
   const isEmptyState = rawMessages.length === 0;
 
-  const handleSendMessage = useCallback(
-    async (messageToSend?: string) => {
-      DEBUG &&
-        console.log(
-          'AIChat: handleSendMessage called with:',
-          messageToSend,
-          'aiChatMessage:',
-          aiChatMessage,
-        );
-      const messageContent = messageToSend || aiChatMessage;
-
-      if (
-        !messageContent ||
-        typeof messageContent !== 'string' ||
-        !messageContent.trim() ||
-        isLoading
-      ) {
-        return;
-      }
-
-      const currentPrompt = aiChatMessage.trim();
-      setAIChatMessage('');
-      setIsLoading(true);
-      setErrorMessage(null); // Clear any previous errors
-
-      // Call backend endpoint for AI response
-      // The server will handle all ShareDB operations including adding the user message
-      try {
-        const response = await fetch(aiChatEndpoint, {
-          method: 'POST',
-          headers: {
-            'Content-Type': 'application/json',
-          },
-          body: JSON.stringify({
-            ...aiChatOptions,
-            vizId: aiChatOptions.vizId,
-            content: messageContent.trim(),
-            chatId: currentChatId,
-            mode: aiChatMode,
-          }),
-        });
-
-        if (!response.ok) {
-          throw new Error(
-            `HTTP error! status: ${response.status}`,
-          );
-        }
-
-        // Parse the response to check for errors
-        const responseData = await response.json();
-
-        // Check if the response contains a VizHub error
-        if (
-          responseData.outcome === 'failure' &&
-          responseData.error
-        ) {
-          const errorMessage = responseData.error.message;
-
-          // Check if this is the specific permission error that should trigger auto-fork
-          if (
-            errorMessage ===
-            'You do not have permission to use AI chat on this visualization. Only users with edit access can use this feature. Fork the viz to edit it.'
-          ) {
-            // Trigger auto-fork instead of showing error
-            try {
-              await autoForkAndRetryAI?.(
-                currentPrompt,
-                aiChatMode,
-              );
-              // If we reach here, the fork was successful and redirect should happen
-              return;
-            } catch (forkError) {
-              console.error('Auto-fork failed:', forkError);
-              setErrorMessage(
-                'Failed to fork visualization. Please try forking manually.',
-              );
-              return;
-            }
-          }
-
-          // For other errors, show the error message
-          setErrorMessage(errorMessage);
-          return;
-        }
-
-        // The backend handles all ShareDB operations for successful responses
-      } catch (error) {
-        console.error('Error getting AI response:', error);
-        // Fail silently in case of timeout,
-        // which happens frequently with the longer running LLM calls.
-        // TODO refactor this so that the network request to the server
-        // always returns immediately, and we track the `isLoading` state
-        // within the ShareDB document itself.
-        //
-        // setErrorMessage(
-        //   'Failed to send message. Please try again.',
-        // );
-      } finally {
-        setIsLoading(false);
-      }
-    },
-    [
-      aiChatMessage,
-      isLoading,
-      aiChatEndpoint,
-      aiChatOptions,
-      currentChatId,
-      aiChatMode,
-      autoForkAndRetryAI,
-    ],
-  );
-
   // Check for stored AI prompt on component mount (post-fork restoration)
   useEffect(() => {
     DEBUG &&
@@ -350,18 +224,18 @@ export const AIChat = () => {
               aiScratchpad={aiScratchpad}
             />
           )}
-          {errorMessage && (
+          {aiErrorMessage && (
             <div className="ai-chat-error">
               <div className="ai-chat-error-content">
                 <span className="ai-chat-error-icon">
                   ⚠️
                 </span>
                 <span className="ai-chat-error-message">
-                  {errorMessage}
+                  {aiErrorMessage}
                 </span>
                 <button
                   className="ai-chat-error-dismiss"
-                  onClick={() => setErrorMessage(null)}
+                  onClick={() => setAIErrorMessage(null)}
                   aria-label="Dismiss error"
                 >
                   ×

package/src/client/VZSidebar/AIChat/styles.scss

@@ -182,6 +182,7 @@
   .undo-button-container {
     display: flex;
     justify-content: flex-end;
+    gap: 8px;
     margin-top: 16px;
   }
 
@@ -560,6 +561,21 @@
     max-height: 200px;
     overflow-y: auto;
     background: var(--vh-color-neutral-02);
+    scroll-behavior: smooth;
+
+    &::-webkit-scrollbar {
+      width: 6px;
+    }
+
+    &::-webkit-scrollbar-track {
+      background: var(--vh-color-neutral-02);
+      border-radius: 3px;
+    }
+
+    &::-webkit-scrollbar-thumb {
+      background: var(--vh-color-neutral-04);
+      border-radius: 3px;
+    }
 
     // Style markdown content in thinking scratchpad
     p {

package/src/server/aiChatHandler/chatOperations.ts

@@ -256,7 +256,8 @@ export const addDiffToAIMessage = (
   shareDBDoc: ShareDBDoc<VizContent>,
   chatId: VizChatId,
   diffData: any,
-  beforeFiles?: VizFiles, // Optional snapshot for undo functionality
+  beforeFiles?: VizFiles, // Optional snapshot for undo functionality (legacy)
+  beforeCommitId?: string, // Commit ID before AI changes for VizHub integration
 ) => {
   const chat = shareDBDoc.data.chats[chatId];
   const messages = [...chat.messages];
@@ -270,7 +271,8 @@ export const addDiffToAIMessage = (
     const newMessage = {
       ...messages[lastAIMessageIndex],
       diffData,
-      ...(beforeFiles && { beforeFiles }), // Add beforeFiles only if provided
+      ...(beforeFiles && { beforeFiles }), // Add beforeFiles only if provided (legacy)
+      ...(beforeCommitId && { beforeCommitId }), // Add beforeCommitId for VizHub integration
     };
     // Use type assertion to extend the message with diffData
     (messages[lastAIMessageIndex] as any) = newMessage;

package/src/server/aiChatHandler/index.ts

@@ -23,10 +23,14 @@ export const handleAIChatMessage =
     shareDBDoc,
     createAIEditLocalPresence,
     onCreditDeduction,
+    getCurrentCommitId,
+    model,
   }: {
     shareDBDoc: ShareDBDoc<VizContent>;
     createAIEditLocalPresence: () => any;
     onCreditDeduction?: any;
+    getCurrentCommitId?: () => string | null;
+    model?: string;
   }) =>
   async (req: any, res: any) => {
     const { content, chatId, mode = 'edit' } = req.body;
@@ -52,6 +56,11 @@ export const handleAIChatMessage =
       ensureChatsExist(shareDBDoc);
       ensureChatExists(shareDBDoc, chatId);
 
+      // Capture the current commit ID before making changes (for VizHub integration)
+      const beforeCommitId = getCurrentCommitId
+        ? getCurrentCommitId()
+        : null;
+
       // Add user message to chat
       addUserMessage(shareDBDoc, chatId, content);
 
@@ -60,6 +69,7 @@ export const handleAIChatMessage =
         shareDBDoc,
         createAIEditLocalPresence,
         chatId,
+        model,
       });
 
       // Create server-side runCode function using shareDBDoc
@@ -92,7 +102,8 @@ export const handleAIChatMessage =
           shareDBDoc,
           chatId,
           editResult.diffData,
-          (editResult as any).beforeFiles, // Pass the beforeFiles snapshot for undo (only available for edit mode)
+          (editResult as any).beforeFiles, // Pass the beforeFiles snapshot for undo (legacy)
+          beforeCommitId, // Pass the commit ID before AI changes (for VizHub integration)
         );
       }
 

package/src/server/aiChatHandler/llmStreaming.ts

@@ -45,10 +45,17 @@ export const createLLMFunction = ({
   shareDBDoc,
   createAIEditLocalPresence,
   chatId,
+  // Feature flag to enable/disable reasoning tokens.
+  // When false, reasoning tokens are not requested from the API
+  // and reasoning content is not processed in the streaming response.
+  enableReasoningTokens = false,
+  model,
 }: {
   shareDBDoc: ShareDBDoc<VizContent>;
   createAIEditLocalPresence: () => any;
   chatId: VizChatId;
+  enableReasoningTokens?: boolean;
+  model?: string;
 }) => {
   return async (fullPrompt: string) => {
     const localPresence = enableStreamingEditing
@@ -195,23 +202,32 @@ export const createLLMFunction = ({
 
     // Stream the response with reasoning tokens
     const modelName =
+      model ||
       process.env.VIZHUB_EDIT_WITH_AI_MODEL_NAME ||
       'anthropic/claude-3.5-sonnet';
-    const stream = await (
-      openRouterClient.chat.completions.create as any
-    )({
+
+    // Configure reasoning tokens based on enableReasoningTokens flag
+    const requestConfig: any = {
       model: modelName,
       messages: [{ role: 'user', content: fullPrompt }],
-      reasoning: {
-        effort: 'low',
-        exclude: false,
-      },
       provider: {
         sort: 'price',
       },
       usage: { include: true },
       stream: true,
-    });
+    };
+
+    // Only include reasoning configuration if reasoning tokens are enabled
+    if (enableReasoningTokens) {
+      requestConfig.reasoning = {
+        effort: 'low',
+        exclude: false,
+      };
+    }
+
+    const stream = await (
+      openRouterClient.chat.completions.create as any
+    )(requestConfig);
 
     let reasoningStarted = false;
     let contentStarted = false;
@@ -219,8 +235,8 @@ export const createLLMFunction = ({
     for await (const chunk of stream) {
       const delta = chunk.choices[0]?.delta as any; // Type assertion for OpenRouter-specific reasoning fields
 
-      if (delta?.reasoning) {
-        // Handle reasoning tokens (thinking)
+      if (delta?.reasoning && enableReasoningTokens) {
+        // Handle reasoning tokens (thinking) - only if enabled
         if (!reasoningStarted) {
           reasoningStarted = true;
           updateAIStatus(shareDBDoc, chatId, 'Thinking...');