vzcode 2.27.0 → 2.29.0

package/src/server/aiChatHandler/index.ts

@@ -4,6 +4,7 @@ import {
   ensureChatExists,
   addUserMessage,
   setAIStatus,
+  setChatAIMetadata,
 } from '../../llm-streaming-server/chatOperations.js';
 import { createLLMFunction } from '../../llm-streaming-server/llmStreaming.js';
 import { performAIEditing } from '../../llm-streaming-server/aiEditing.js';
@@ -23,15 +24,31 @@ export const handleAIChatMessage =
   ({
     shareDBDoc,
     onCreditDeduction,
+    onGenerationFinished,
     model,
     aiRequestOptions,
     enableReasoningTokens,
+    baseCommitId,
+    escalationLevel,
   }: {
     shareDBDoc: ShareDBDoc<VizContent>;
     onCreditDeduction?: any;
+    onGenerationFinished?: (result: {
+      success: boolean;
+      editResult?: any;
+      metrics?: any;
+      error?: any;
+    }) => Promise<void> | void;
     model?: string;
     aiRequestOptions?: any;
     enableReasoningTokens?: boolean;
+    // Phase 3: explicit escalation base. Persisted on the chat so that
+    // repeated "Try Harder" clicks are idempotent and do not depend on
+    // the fragile `parent(currentCommit)` heuristic.
+    baseCommitId?: string;
+    // Explicit escalation level, persisted on the chat so it survives
+    // reloads and is shared across clients.
+    escalationLevel?: number;
   }) =>
   async (req: any, res: any) => {
     const { content, chatId } = req.body;
@@ -60,6 +77,13 @@ export const handleAIChatMessage =
       // Add user message to chat
       addUserMessage(shareDBDoc, chatId, content);
 
+      // Persist explicit escalation metadata on the chat, so that
+      // retries are idempotent and the level survives reloads.
+      setChatAIMetadata(shareDBDoc, chatId, {
+        baseCommitId,
+        escalationLevel,
+      });
+
       // Return success immediately - AI generation continues in background
       res.status(200).json('success');
 
@@ -72,6 +96,7 @@ export const handleAIChatMessage =
         aiRequestOptions,
         enableReasoningTokens,
         onCreditDeduction,
+        onGenerationFinished,
       }).catch((error) => {
         console.error(
           'Background AI processing error:',
@@ -96,6 +121,7 @@ const processAIRequestAsync = async ({
   aiRequestOptions,
   enableReasoningTokens,
   onCreditDeduction,
+  onGenerationFinished,
 }: {
   shareDBDoc: ShareDBDoc<VizContent>;
   chatId: string;
@@ -104,6 +130,12 @@ const processAIRequestAsync = async ({
   aiRequestOptions?: any;
   enableReasoningTokens?: boolean;
   onCreditDeduction?: any;
+  onGenerationFinished?: (result: {
+    success: boolean;
+    editResult?: any;
+    metrics?: any;
+    error?: any;
+  }) => Promise<void> | void;
 }) => {
   try {
     // Create LLM function for streaming
@@ -128,23 +160,40 @@ const processAIRequestAsync = async ({
       runCode,
     });
 
-    // Handle credit deduction if callback is provided
+    // Billing is best-effort. It must NEVER prevent the edit from
+    // being committed. Any metadata failure is logged and swallowed.
+    let metrics: any = null;
     if (onCreditDeduction && editResult.generationId) {
       try {
-        await onCreditDeduction(
-          await getGenerationMetadata({
-            apiKey:
-              aiRequestOptions?.apiKey ||
-              process.env.VZCODE_EDIT_WITH_AI_API_KEY,
-            generationId: editResult.generationId,
-          }),
-        );
+        metrics = await getGenerationMetadata({
+          apiKey:
+            aiRequestOptions?.apiKey ||
+            process.env.VZCODE_EDIT_WITH_AI_API_KEY,
+          generationId: editResult.generationId,
+        });
+        await onCreditDeduction(metrics);
       } catch (creditError) {
         console.error(
-          'Credit deduction error:',
+          'Credit deduction error (edit will still be finalized):',
           creditError,
         );
-        // Don't fail the request if credit deduction fails
+      }
+    }
+
+    // Always notify settlement on success, regardless of billing
+    // outcome, so the caller can commit the edit and release locks.
+    if (onGenerationFinished) {
+      try {
+        await onGenerationFinished({
+          success: true,
+          editResult,
+          metrics,
+        });
+      } catch (settleError) {
+        console.error(
+          'onGenerationFinished (success) error:',
+          settleError,
+        );
       }
     }
 
@@ -153,6 +202,23 @@ const processAIRequestAsync = async ({
   } catch (error) {
     // Set error status and add error message to chat
     setAIStatus(shareDBDoc, chatId, 'error');
+
+    // Always notify settlement on failure so the caller can roll back
+    // the pre-restore snapshot and release locks.
+    if (onGenerationFinished) {
+      try {
+        await onGenerationFinished({
+          success: false,
+          error,
+        });
+      } catch (settleError) {
+        console.error(
+          'onGenerationFinished (failure) error:',
+          settleError,
+        );
+      }
+    }
+
     handleBackgroundError(shareDBDoc, chatId, error);
   }
 };

package/src/types.ts

@@ -262,6 +262,9 @@ export interface ExtendedVizChat extends VizChat {
   currentStatus?: string;
   isStreaming?: boolean;
   model?: string; // The LLM model used for this chat
+  // Phase 3: explicit escalation metadata persisted on the chat.
+  baseCommitId?: string; // Commit the current AI attempt is applied on top of
+  escalationLevel?: number; // How many times "Try Harder" has been pressed
 }
 
 // Extended VizChatMessage with progressive rendering support