From b9da24dd6d401293f552da36cd173c2f53b24599 Mon Sep 17 00:00:00 2001
From: Dayuan Jiang <34411969+DayuanJiang@users.noreply.github.com>
Date: Thu, 11 Dec 2025 17:56:40 +0900
Subject: [PATCH] fix: limit auto-retry to 3 attempts and enforce quota checks
 (#219)

- Add MAX_AUTO_RETRY_COUNT (3) to prevent infinite retry loops
- Check token and TPM limits before each auto-retry
- Reset retry counter on user-initiated messages
- Show toast notification when limits are reached

Fixes issue where models returning invalid tool inputs caused 45+ API
requests due to sendAutomaticallyWhen having no retry limit or quota check.
---
 components/chat-panel.tsx | 117 ++++++++++++++++++++++++--------------
 1 file changed, 73 insertions(+), 44 deletions(-)

diff --git a/components/chat-panel.tsx b/components/chat-panel.tsx
index 7424b56..c72c4bb 100644
--- a/components/chat-panel.tsx
+++ b/components/chat-panel.tsx
@@ -61,31 +61,15 @@ interface ChatPanelProps {
 // Constants for tool states
 const TOOL_ERROR_STATE = "output-error" as const
 const DEBUG = process.env.NODE_ENV === "development"
+const MAX_AUTO_RETRY_COUNT = 3
 
 /**
- * Custom auto-resubmit logic for the AI chat.
- *
- * Strategy:
- * - When tools return errors (e.g., invalid XML), automatically resubmit
- *   the conversation to let the AI retry with corrections
- * - When tools succeed (e.g., diagram displayed), stop without AI acknowledgment
- *   to prevent unnecessary regeneration cycles
- *
- * This fixes the issue where successful diagrams were being regenerated
- * multiple times because the previous logic (lastAssistantMessageIsCompleteWithToolCalls)
- * auto-resubmitted on BOTH success and error.
- *
- * @param messages - Current conversation messages from AI SDK
- * @returns true to auto-resubmit (for error recovery), false to stop
+ * Check if auto-resubmit should happen based on tool errors.
+ * Does NOT handle retry count or quota - those are handled by the caller.
  */
-function shouldAutoResubmit(messages: ChatMessage[]): boolean {
+function hasToolErrors(messages: ChatMessage[]): boolean {
     const lastMessage = messages[messages.length - 1]
     if (!lastMessage || lastMessage.role !== "assistant") {
-        if (DEBUG) {
-            console.log(
-                "[sendAutomaticallyWhen] No assistant message, returning false",
-            )
-        }
         return false
     }
 
@@ -95,31 +79,10 @@ function shouldAutoResubmit(messages: ChatMessage[]): boolean {
         ) || []
 
     if (toolParts.length === 0) {
-        if (DEBUG) {
-            console.log(
-                "[sendAutomaticallyWhen] No tool parts, returning false",
-            )
-        }
         return false
     }
 
-    // Only auto-resubmit if ANY tool has an error
-    const hasError = toolParts.some((part) => part.state === TOOL_ERROR_STATE)
-
-    if (DEBUG) {
-        if (hasError) {
-            console.log(
-                "[sendAutomaticallyWhen] Retrying due to errors in tools:",
-                toolParts
-                    .filter((p) => p.state === TOOL_ERROR_STATE)
-                    .map((p) => p.toolName),
-            )
-        } else {
-            console.log("[sendAutomaticallyWhen] No errors, stopping")
-        }
-    }
-
-    return hasError
+    return toolParts.some((part) => part.state === TOOL_ERROR_STATE)
 }
 
 export default function ChatPanel({
@@ -223,6 +186,9 @@ export default function ChatPanel({
     // Ref to hold stop function for use in onToolCall (avoids stale closure)
     const stopRef = useRef<(() => void) | null>(null)
 
+    // Ref to track consecutive auto-retry count (reset on user action)
+    const autoRetryCountRef = useRef(0)
+
     const {
         messages,
         sendMessage,
@@ -441,8 +407,68 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
                 }
             }
         },
-        sendAutomaticallyWhen: ({ messages }) =>
-            shouldAutoResubmit(messages as unknown as ChatMessage[]),
+        sendAutomaticallyWhen: ({ messages }) => {
+            const shouldRetry = hasToolErrors(
+                messages as unknown as ChatMessage[],
+            )
+
+            if (!shouldRetry) {
+                // No error, reset retry count
+                autoRetryCountRef.current = 0
+                if (DEBUG) {
+                    console.log("[sendAutomaticallyWhen] No errors, stopping")
+                }
+                return false
+            }
+
+            // Check retry count limit
+            if (autoRetryCountRef.current >= MAX_AUTO_RETRY_COUNT) {
+                if (DEBUG) {
+                    console.log(
+                        `[sendAutomaticallyWhen] Max retry count (${MAX_AUTO_RETRY_COUNT}) reached, stopping`,
+                    )
+                }
+                toast.error(
+                    `Auto-retry limit reached (${MAX_AUTO_RETRY_COUNT}). Please try again manually.`,
+                )
+                autoRetryCountRef.current = 0
+                return false
+            }
+
+            // Check quota limits before auto-retry
+            const tokenLimitCheck = quotaManager.checkTokenLimit()
+            if (!tokenLimitCheck.allowed) {
+                if (DEBUG) {
+                    console.log(
+                        "[sendAutomaticallyWhen] Token limit exceeded, stopping",
+                    )
+                }
+                quotaManager.showTokenLimitToast(tokenLimitCheck.used)
+                autoRetryCountRef.current = 0
+                return false
+            }
+
+            const tpmCheck = quotaManager.checkTPMLimit()
+            if (!tpmCheck.allowed) {
+                if (DEBUG) {
+                    console.log(
+                        "[sendAutomaticallyWhen] TPM limit exceeded, stopping",
+                    )
+                }
+                quotaManager.showTPMLimitToast()
+                autoRetryCountRef.current = 0
+                return false
+            }
+
+            // Increment retry count and allow retry
+            autoRetryCountRef.current++
+            if (DEBUG) {
+                console.log(
+                    `[sendAutomaticallyWhen] Retrying (${autoRetryCountRef.current}/${MAX_AUTO_RETRY_COUNT})`,
+                )
+            }
+            return true
+        },
     })
 
     // Update stopRef so onToolCall can access it
@@ -759,6 +785,9 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
         previousXml: string,
         sessionId: string,
     ) => {
+        // Reset auto-retry count on user-initiated message
+        autoRetryCountRef.current = 0
+
         const config = getAIConfig()
 
         sendMessage(