vzcode 2.2.0 → 2.3.0

package/src/server/aiChatHandler/index.ts

@@ -1,200 +0,0 @@
-import { validateRequest } from './validation.js';
-import {
-  ensureChatsExist,
-  ensureChatExists,
-  addUserMessage,
-  addDiffToAIMessage,
-  setAIStatus,
-} from './chatOperations.js';
-import { createLLMFunction } from './llmStreaming.js';
-import {
-  performAIEditing,
-  performAIChat,
-} from './aiEditing.js';
-import {
-  handleError,
-  handleBackgroundError,
-} from './errorHandling.js';
-import { createRunCodeFunction } from '../../runCode.js';
-import { ShareDBDoc } from '../../types.js';
-import { VizContent } from '@vizhub/viz-types';
-import { createSubmitOperation } from '../../submitOperation.js';
-
-const DEBUG = false;
-
-export const handleAIChatMessage =
-  ({
-    shareDBDoc,
-    createAIEditLocalPresence,
-    onCreditDeduction,
-    getCurrentCommitId,
-    model,
-    aiRequestOptions,
-  }: {
-    shareDBDoc: ShareDBDoc<VizContent>;
-    createAIEditLocalPresence: () => any;
-    onCreditDeduction?: any;
-    getCurrentCommitId?: () => string | null;
-    model?: string;
-    aiRequestOptions?: any;
-  }) =>
-  async (req: any, res: any) => {
-    const { content, chatId, mode = 'edit' } = req.body;
-
-    if (DEBUG) {
-      console.log(
-        '[handleAIChatMessage] content:',
-        content,
-        'chatId:',
-        chatId,
-        'shareDBDoc:',
-        shareDBDoc,
-      );
-    }
-
-    // Validate request
-    if (!validateRequest(req, res)) {
-      return;
-    }
-
-    try {
-      // Ensure chats structure exists
-      ensureChatsExist(shareDBDoc);
-      ensureChatExists(shareDBDoc, chatId);
-
-      // Add user message to chat
-      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,
-        createAIEditLocalPresence,
-        getCurrentCommitId,
-        model,
-        aiRequestOptions,
-        onCreditDeduction,
-      }).catch((error) => {
-        console.error(
-          'Background AI processing error:',
-          error,
-        );
-        // Handle error without HTTP response
-        handleBackgroundError(shareDBDoc, chatId, error);
-      });
-    } catch (error) {
-      handleError(shareDBDoc, chatId, error, res);
-    }
-  };
-
-/**
- * Processes the AI request asynchronously in the background
- */
-const processAIRequestAsync = async ({
-  shareDBDoc,
-  chatId,
-  content,
-  mode,
-  createAIEditLocalPresence,
-  getCurrentCommitId,
-  model,
-  aiRequestOptions,
-  onCreditDeduction,
-}: {
-  shareDBDoc: ShareDBDoc<VizContent>;
-  chatId: string;
-  content: string;
-  mode: string;
-  createAIEditLocalPresence: () => any;
-  getCurrentCommitId?: () => string | null;
-  model?: string;
-  aiRequestOptions?: any;
-  onCreditDeduction?: any;
-}) => {
-  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,
-      createAIEditLocalPresence,
-      chatId,
-      model,
-      aiRequestOptions,
-    });
-
-    // Create server-side runCode function using shareDBDoc
-    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,
-        (editResult as any).beforeFiles, // Pass the beforeFiles snapshot for undo (legacy)
-        beforeCommitId, // Pass the commit ID before AI changes (for VizHub integration)
-      );
-    }
-
-    // Handle credit deduction if callback is provided
-    if (
-      onCreditDeduction &&
-      (editResult as any).upstreamCostCents
-    ) {
-      try {
-        await onCreditDeduction({
-          upstreamCostCents: (editResult as any)
-            .upstreamCostCents,
-          provider: (editResult as any).provider,
-          inputTokens: (editResult as any).inputTokens,
-          outputTokens: (editResult as any).outputTokens,
-        });
-      } catch (creditError) {
-        console.error(
-          'Credit deduction error:',
-          creditError,
-        );
-        // Don't fail the request if credit deduction fails
-      }
-    }
-
-    // Clear the AI status to indicate completion
-    setAIStatus(shareDBDoc, chatId, undefined);
-  } catch (error) {
-    // Set error status and add error message to chat
-    setAIStatus(shareDBDoc, chatId, 'error');
-    handleBackgroundError(shareDBDoc, chatId, error);
-  }
-};

package/src/server/aiChatHandler/llmStreaming.ts

@@ -1,419 +0,0 @@
-import {
-  parseMarkdownFiles,
-  StreamingMarkdownParser,
-} from 'llm-code-format';
-import OpenAI from 'openai';
-import fs from 'fs';
-import { generateRunId } from '@vizhub/viz-utils';
-import {
-  updateAIStatus,
-  createAIMessage,
-  updateAIMessageContent,
-  finalizeAIMessage,
-  ensureFileExists,
-  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;
-
-// 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 = 100;
-
-// 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
- */
-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,
-  aiRequestOptions,
-}: {
-  shareDBDoc: ShareDBDoc<VizContent>;
-  createAIEditLocalPresence: () => any;
-  chatId: VizChatId;
-  enableReasoningTokens?: boolean;
-  model?: string;
-  aiRequestOptions?: any;
-}) => {
-  return async (fullPrompt: string) => {
-    const localPresence = enableStreamingEditing
-      ? createAIEditLocalPresence()
-      : null;
-
-    // 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',
-      },
-    });
-
-    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: NodeJS.Timeout | null = null;
-
-      function invoke() {
-        updateAIMessageContent(
-          shareDBDoc,
-          chatId,
-          aiMessageId,
-          latestContent,
-        );
-        lastCall = Date.now();
-      }
-
-      const fn = (content: string) => {
-        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),
-          );
-        }
-      };
-
-      // 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;
-      }
-    };
-
-    // Define callbacks for streaming parser
-    const callbacks = {
-      onFileNameChange: async (
-        fileName: string,
-        format: string,
-      ) => {
-        DEBUG &&
-          console.log(
-            `File changed to: ${fileName} (${format})`,
-          );
-
-        // Find existing file or create new one
-        currentEditingFileId = ensureFileExists(
-          shareDBDoc,
-          fileName,
-        );
-
-        reportFileEdited();
-        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,
-        );
-      },
-      onCodeLine: async (line: string) => {
-        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,
-          );
-
-          // 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: 'AI Editor',
-              start: [
-                'files',
-                currentEditingFileId,
-                'text',
-                textLength,
-              ],
-              end: [
-                'files',
-                currentEditingFileId,
-                'text',
-                textLength,
-              ],
-            };
-
-            if (localPresence) {
-              localPresence.submit(
-                filePresence,
-                (error) => {
-                  if (error) {
-                    console.warn(
-                      'AI Editor line presence submission error:',
-                      error,
-                    );
-                  }
-                },
-              );
-            }
-          }
-        }
-      },
-      onNonCodeLine: async (line: string) => {
-        // 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
-        if (line.trim() !== '') {
-          reportFileEdited();
-        }
-        fullContent += line + '\n';
-        throttledUpdateAIMessageContent(fullContent);
-      },
-    };
-
-    const parser = new StreamingMarkdownParser(callbacks);
-
-    const chunks = [];
-    let reasoningContent = '';
-
-    // Stream the response with reasoning tokens
-    const modelName =
-      model ||
-      process.env.VIZHUB_EDIT_WITH_AI_MODEL_NAME ||
-      'anthropic/claude-3.5-sonnet';
-
-    // Configure reasoning tokens based on enableReasoningTokens flag
-    const requestConfig: any = {
-      model: modelName,
-      messages: [{ role: 'user', content: fullPrompt }],
-      usage: { include: true },
-      stream: true,
-      ...aiRequestOptions,
-    };
-
-    // 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;
-
-    for await (const chunk of stream) {
-      if (slowMode) {
-        await new Promise((resolve) =>
-          setTimeout(resolve, 500),
-        );
-      }
-      const delta = chunk.choices[0]?.delta as any; // Type assertion for OpenRouter-specific reasoning fields
-
-      if (delta?.reasoning && enableReasoningTokens) {
-        // Handle reasoning tokens (thinking) - only if enabled
-        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) {
-          await parser.processChunk(chunkContent);
-        } else {
-          fullContent += chunkContent;
-          throttledUpdateAIMessageContent(fullContent);
-        }
-      } else if (chunk.usage) {
-        // Handle usage information
-        DEBUG && console.log('Usage:', chunk.usage);
-      }
-
-      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.');
-    throttledUpdateAIMessageContent(fullContent);
-
-    // Flush to ensure the final content is written immediately
-    throttledUpdateAIMessageContent.flush();
-
-    // Finalize the AI message by clearing temporary fields
-    finalizeAIMessage(shareDBDoc, chatId);
-
-    // Clear AI Editor presence when done
-    DEBUG &&
-      console.log(
-        'AI editing done, clearing AI Editor presence',
-      );
-    if (localPresence) {
-      localPresence.submit(null, (error) => {
-        DEBUG && console.log('AI Editor presence cleared');
-        if (error) {
-          console.warn(
-            'AI Editor presence cleanup error:',
-            error,
-          );
-        }
-      });
-    }
-
-    // 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,
-        ),
-      );
-    }
-
-    // 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`;
-      fs.writeFileSync(
-        chunksFileJSONpath,
-        JSON.stringify(chunks, null, 2),
-      );
-      console.log(
-        `AI chunks written to ${chunksFileJSONpath}`,
-      );
-    }
-
-    return {
-      content: fullContent,
-      generationId: generationId,
-    };
-  };
-};

package/src/server/aiChatHandler/undoHandler.ts

@@ -1,97 +0,0 @@
-import { undoAIEdit } from './chatOperations.js';
-import { handleError } from './errorHandling.js';
-import { ShareDBDoc } from '../../types.js';
-import { VizContent } from '@vizhub/viz-types';
-
-const DEBUG = false;
-
-/**
- * Validates incoming request data for AI chat undo
- */
-const validateUndoRequest = (req, res) => {
-  const { chatId, messageId } = req.body;
-
-  if (!chatId || typeof chatId !== 'string') {
-    res.status(400).json({
-      error:
-        'Invalid request: chatId is required and must be a string',
-    });
-    return false;
-  }
-
-  if (!messageId || typeof messageId !== 'string') {
-    res.status(400).json({
-      error:
-        'Invalid request: messageId is required and must be a string',
-    });
-    return false;
-  }
-
-  return true;
-};
-
-export const handleAIChatUndo =
-  ({
-    shareDBDoc,
-  }: {
-    shareDBDoc: ShareDBDoc<VizContent>;
-  }) =>
-  async (req: any, res: any) => {
-    const { chatId, messageId } = req.body;
-
-    if (DEBUG) {
-      console.log(
-        '[handleAIChatUndo] chatId:',
-        chatId,
-        'messageId:',
-        messageId,
-        'shareDBDoc:',
-        shareDBDoc,
-      );
-    }
-
-    // Validate request
-    if (!validateUndoRequest(req, res)) {
-      return;
-    }
-
-    try {
-      const chat = shareDBDoc.data.chats?.[chatId];
-      if (!chat) {
-        res.status(404).json({
-          error: 'Chat not found',
-        });
-        return;
-      }
-
-      const message = chat.messages.find(
-        (msg) => msg.id === messageId,
-      );
-
-      if (
-        !message ||
-        message.role !== 'assistant' ||
-        !(message as any).beforeFiles
-      ) {
-        res.status(400).json({
-          error: 'Invalid message for undo',
-        });
-        return;
-      }
-
-      // Perform the undo operation
-      undoAIEdit(
-        shareDBDoc,
-        chatId,
-        messageId,
-        (message as any).beforeFiles,
-      );
-
-      res.status(200).json({ success: true });
-    } catch (error) {
-      if (DEBUG) {
-        console.error('[handleAIChatUndo] Error:', error);
-      }
-      handleError(shareDBDoc, chatId, error, res);
-    }
-  };