vzcode 1.50.0 → 1.52.0

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

@@ -3,6 +3,7 @@ import {
   useState,
   useCallback,
   useMemo,
+  useEffect,
 } from 'react';
 import { VZCodeContext } from '../../VZCodeContext';
 import { v4 as uuidv4 } from 'uuid';
@@ -12,6 +13,8 @@ import './styles.scss';
 
 const defaultAIChatEndpoint = '/api/ai-chat/';
 
+const DEBUG = false;
+
 export const AIChat = () => {
   const { aiChatMessage, setAIChatMessage } =
     useContext(VZCodeContext);
@@ -29,6 +32,9 @@ export const AIChat = () => {
     aiChatOptions = {},
     aiChatMode,
     setAIChatMode,
+    autoForkAndRetryAI,
+    clearStoredAIPrompt,
+    getStoredAIPrompt,
   } = useContext(VZCodeContext);
 
   // Get current chat data from content
@@ -50,63 +56,154 @@ export const AIChat = () => {
   // Check if this is the first time opening the chat (no messages)
   const isEmptyState = rawMessages.length === 0;
 
-  const handleSendMessage = useCallback(async () => {
-    if (!aiChatMessage.trim() || isLoading) return;
+  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;
 
-    setAIChatMessage('');
-    setIsLoading(true);
-    setErrorMessage(null); // Clear any previous errors
+      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: aiChatMessage.trim(),
-          chatId: currentChatId,
-          mode: aiChatMode,
-        }),
-      });
+      // 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}`,
-        );
-      }
+        if (!response.ok) {
+          throw new Error(
+            `HTTP error! status: ${response.status}`,
+          );
+        }
 
-      // Parse the response to check for errors
-      const responseData = await response.json();
+        // 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
-      ) {
-        setErrorMessage(responseData.error.message);
-        return;
+        // 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);
+        setErrorMessage(
+          'Failed to send message. Please try again.',
+        );
+      } finally {
+        setIsLoading(false);
       }
+    },
+    [
+      aiChatMessage,
+      isLoading,
+      aiChatEndpoint,
+      aiChatOptions,
+      currentChatId,
+      aiChatMode,
+      autoForkAndRetryAI,
+    ],
+  );
 
-      // The backend handles all ShareDB operations for successful responses
-    } catch (error) {
-      console.error('Error getting AI response:', error);
-      setErrorMessage(
-        'Failed to send message. Please try again.',
+  // Check for stored AI prompt on component mount (post-fork restoration)
+  useEffect(() => {
+    DEBUG &&
+      console.log('AIChat: Checking for stored AI prompt');
+    const storedPrompt = getStoredAIPrompt();
+    DEBUG &&
+      console.log(
+        'AIChat: Stored prompt result:',
+        storedPrompt,
       );
-    } finally {
-      setIsLoading(false);
+    if (storedPrompt) {
+      // Restore the prompt and mode
+      DEBUG &&
+        console.log('AIChat: Restoring prompt and mode');
+      setAIChatMessage(storedPrompt.prompt);
+      setAIChatMode(
+        storedPrompt.modelName === 'ask' ? 'ask' : 'edit',
+      );
+
+      // Clear the stored prompt
+      DEBUG &&
+        console.log('AIChat: Clearing stored prompt');
+      clearStoredAIPrompt();
+
+      // Auto-submit the restored prompt after a short delay
+      DEBUG &&
+        console.log('AIChat: Scheduling auto-submit');
+      setTimeout(() => {
+        DEBUG &&
+          console.log(
+            'AIChat: Auto-submitting restored prompt',
+          );
+        handleSendMessage(storedPrompt.prompt);
+      }, 100);
+    } else {
+      DEBUG &&
+        console.log('AIChat: No stored prompt found');
     }
   }, [
-    aiChatMessage,
-    isLoading,
-    aiChatEndpoint,
-    aiChatOptions,
-    currentChatId,
+    getStoredAIPrompt,
+    clearStoredAIPrompt,
+    setAIChatMessage,
+    setAIChatMode,
+    handleSendMessage,
   ]);
 
   return (
@@ -123,7 +220,7 @@ export const AIChat = () => {
           <div className="ai-chat-empty">
             <div className="ai-chat-empty-icon">✨</div>
             <h3 className="ai-chat-empty-title">
-              Hi, I'm VizBot!
+              Hi! I'm here to help you edit with AI.
             </h3>
             <div className="ai-chat-empty-text">
               How can I help you?
@@ -136,7 +233,7 @@ export const AIChat = () => {
                     <button
                       className="ai-chat-suggested-prompt"
                       onClick={() =>
-                        setAIChatMessage(
+                        handleSendMessage(
                           'Explain how this works',
                         )
                       }
@@ -146,7 +243,7 @@ export const AIChat = () => {
                     <button
                       className="ai-chat-suggested-prompt"
                       onClick={() =>
-                        setAIChatMessage(
+                        handleSendMessage(
                           'How could I change it so that the circles are bigger?',
                         )
                       }
@@ -157,7 +254,7 @@ export const AIChat = () => {
                     <button
                       className="ai-chat-suggested-prompt"
                       onClick={() =>
-                        setAIChatMessage(
+                        handleSendMessage(
                           'What does this function do?',
                         )
                       }
@@ -167,7 +264,7 @@ export const AIChat = () => {
                     <button
                       className="ai-chat-suggested-prompt"
                       onClick={() =>
-                        setAIChatMessage(
+                        handleSendMessage(
                           'How can I make this more accessible?',
                         )
                       }
@@ -183,7 +280,7 @@ export const AIChat = () => {
                     <button
                       className="ai-chat-suggested-prompt"
                       onClick={() =>
-                        setAIChatMessage(
+                        handleSendMessage(
                           'Change the circles to squares',
                         )
                       }
@@ -193,7 +290,7 @@ export const AIChat = () => {
                     <button
                       className="ai-chat-suggested-prompt"
                       onClick={() =>
-                        setAIChatMessage(
+                        handleSendMessage(
                           'Add a button that toggles the animation',
                         )
                       }
@@ -204,7 +301,7 @@ export const AIChat = () => {
                     <button
                       className="ai-chat-suggested-prompt"
                       onClick={() =>
-                        setAIChatMessage(
+                        handleSendMessage(
                           'Fix the CSS so the layout is responsive',
                         )
                       }
@@ -215,7 +312,7 @@ export const AIChat = () => {
                     <button
                       className="ai-chat-suggested-prompt"
                       onClick={() =>
-                        setAIChatMessage(
+                        handleSendMessage(
                           'Refactor this function to use async/await',
                         )
                       }

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

@@ -150,6 +150,33 @@
   line-height: 1.4;
   word-wrap: break-word;
 
+  // Undo button styles (moved from DiffView)
+  .undo-button-container {
+    display: flex;
+    justify-content: flex-end;
+    margin-top: 16px;
+  }
+
+  .undo-button {
+    background-color: #dc3545;
+    color: white;
+    border: none;
+    border-radius: 4px;
+    padding: 4px 8px;
+    font-size: 12px;
+    cursor: pointer;
+    transition: background-color 0.2s;
+
+    &:hover:not(:disabled) {
+      background-color: #c82333;
+    }
+
+    &:disabled {
+      background-color: #6c757d;
+      cursor: not-allowed;
+    }
+  }
+
   // Markdown styling for dark backgrounds
   pre {
     background-color: rgba(0, 0, 0, 0.3);

package/src/client/VZSidebar/index.tsx

@@ -110,7 +110,7 @@ export const VZSidebar = ({
   ),
   aiChatToolTipText = (
     <div>
-      <strong>VizBot - AI Code Editor</strong>
+      <strong>Edit with AI</strong>
     </div>
   ),
 }: {

package/src/server/aiChatHandler/index.ts

@@ -21,11 +21,11 @@ const DEBUG = false;
 export const handleAIChatMessage =
   ({
     shareDBDoc,
-    createVizBotLocalPresence,
+    createAIEditLocalPresence,
     onCreditDeduction,
   }: {
     shareDBDoc: ShareDBDoc<VizContent>;
-    createVizBotLocalPresence: () => any;
+    createAIEditLocalPresence: () => any;
     onCreditDeduction?: any;
   }) =>
   async (req: any, res: any) => {
@@ -58,7 +58,7 @@ export const handleAIChatMessage =
       // Create LLM function for streaming
       const llmFunction = createLLMFunction({
         shareDBDoc,
-        createVizBotLocalPresence,
+        createAIEditLocalPresence,
         chatId,
       });
 

package/src/server/aiChatHandler/llmStreaming.ts

@@ -44,15 +44,15 @@ const enableStreamingEditing = false;
  */
 export const createLLMFunction = ({
   shareDBDoc,
-  createVizBotLocalPresence,
+  createAIEditLocalPresence,
   chatId,
 }: {
   shareDBDoc: ShareDBDoc<VizContent>;
-  createVizBotLocalPresence: () => any;
+  createAIEditLocalPresence: () => any;
   chatId: VizChatId;
 }) => {
   return async (fullPrompt: string) => {
-    const localPresence = createVizBotLocalPresence();
+    const localPresence = createAIEditLocalPresence();
 
     // Create OpenRouter client for reasoning token support
     const openRouterClient = new OpenAI({
@@ -133,13 +133,13 @@ export const createLLMFunction = ({
             line,
           );
 
-          // Update VizBot presence to show cursor at the end of the file
+          // Update AI presence to show cursor at the end of the file
           const currentFile =
             shareDBDoc.data.files[currentEditingFileId];
           if (currentFile && currentFile.text) {
             const textLength = currentFile.text.length;
             const filePresence = {
-              username: 'VizBot',
+              username: 'AI Editor',
               start: [
                 'files',
                 currentEditingFileId,
@@ -157,7 +157,7 @@ export const createLLMFunction = ({
             localPresence.submit(filePresence, (error) => {
               if (error) {
                 console.warn(
-                  'VizBot line presence submission error:',
+                  'AI Editor line presence submission error:',
                   error,
                 );
               }
@@ -201,6 +201,9 @@ export const createLLMFunction = ({
         effort: 'medium',
         exclude: false,
       },
+      provider: {
+        sort: 'throughput',
+      }, // New parameter for OpenRouter routing
       usage: { include: true },
       stream: true,
     });
@@ -275,16 +278,16 @@ export const createLLMFunction = ({
     // Finalize the AI message by clearing temporary fields
     finalizeAIMessage(shareDBDoc, chatId);
 
-    // Clear VizBot presence when done
+    // Clear AI Editor presence when done
     DEBUG &&
       console.log(
-        'AI editing done, clearing VizBot presence',
+        'AI editing done, clearing AI Editor presence',
       );
     localPresence.submit(null, (error) => {
-      DEBUG && console.log('VizBot presence cleared');
+      DEBUG && console.log('AI Editor presence cleared');
       if (error) {
         console.warn(
-          'VizBot presence cleanup error:',
+          'AI Editor presence cleanup error:',
           error,
         );
       }

package/src/server/generateAIResponse.ts

@@ -154,6 +154,9 @@ export const generateAIResponse = async ({
       effort: 'medium',
       exclude: false,
     } as any, // Type assertion for OpenRouter-specific reasoning parameter
+    provider: {
+      sort: 'throughput',
+    } as any, // Type assertion for OpenRouter-specific provider parameter
     stream: true,
   });
 

package/src/server/handleAICopilot.ts

@@ -40,6 +40,9 @@ export const handleAICopilot = () => {
         apiKey: VZCODE_AI_COPILOT_API_KEY,
         baseURL: VZCODE_AI_COPILOT_BASE_URL,
       },
+      additionalParameters: {
+        provider: { sort: 'throughput' },
+      },
       streaming: false,
     };
     debug && console.log('chatModel options:', options);

package/src/server/index.ts

@@ -100,20 +100,20 @@ app.use(express.static(dir));
 const shareDBConnection = shareDBBackend.connect();
 const shareDBDoc = shareDBConnection.get('documents', '1');
 
-// Set up presence for VizBot following the same pattern as useShareDB.ts
+// Set up presence for AI editing following the same pattern as useShareDB.ts
 const docPresence = shareDBConnection.getDocPresence(
   'documents',
   '1',
 );
 
-// Create local presence for VizBot with a unique ID
-const generateVizBotId = () => {
+// Create local presence for AI editing with a unique ID
+const generateAIEditId = () => {
   const timestamp = Date.now().toString(36);
-  return `vizbot-${timestamp}`;
+  return `ai-edit-${timestamp}`;
 };
 
-const createVizBotLocalPresence = () =>
-  docPresence.create(generateVizBotId());
+const createAIEditLocalPresence = () =>
+  docPresence.create(generateAIEditId());
 
 shareDBDoc.create(initialDocument, json1Presence.type.uri);
 
@@ -137,7 +137,7 @@ app.post(
   bodyParser.json(),
   handleAIChatMessage({
     shareDBDoc,
-    createVizBotLocalPresence,
+    createAIEditLocalPresence,
     onCreditDeduction: undefined,
   }),
 );