mirror of
https://github.com/DayuanJiang/next-ai-draw-io.git
synced 2026-01-02 22:32:27 +08:00
fix: limit auto-retry to 3 attempts and enforce quota checks
- 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.
This commit is contained in:
@@ -61,31 +61,15 @@ interface ChatPanelProps {
|
|||||||
// Constants for tool states
|
// Constants for tool states
|
||||||
const TOOL_ERROR_STATE = "output-error" as const
|
const TOOL_ERROR_STATE = "output-error" as const
|
||||||
const DEBUG = process.env.NODE_ENV === "development"
|
const DEBUG = process.env.NODE_ENV === "development"
|
||||||
|
const MAX_AUTO_RETRY_COUNT = 3
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* Custom auto-resubmit logic for the AI chat.
|
* Check if auto-resubmit should happen based on tool errors.
|
||||||
*
|
* Does NOT handle retry count or quota - those are handled by the caller.
|
||||||
* 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
|
|
||||||
*/
|
*/
|
||||||
function shouldAutoResubmit(messages: ChatMessage[]): boolean {
|
function hasToolErrors(messages: ChatMessage[]): boolean {
|
||||||
const lastMessage = messages[messages.length - 1]
|
const lastMessage = messages[messages.length - 1]
|
||||||
if (!lastMessage || lastMessage.role !== "assistant") {
|
if (!lastMessage || lastMessage.role !== "assistant") {
|
||||||
if (DEBUG) {
|
|
||||||
console.log(
|
|
||||||
"[sendAutomaticallyWhen] No assistant message, returning false",
|
|
||||||
)
|
|
||||||
}
|
|
||||||
return false
|
return false
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -95,31 +79,10 @@ function shouldAutoResubmit(messages: ChatMessage[]): boolean {
|
|||||||
) || []
|
) || []
|
||||||
|
|
||||||
if (toolParts.length === 0) {
|
if (toolParts.length === 0) {
|
||||||
if (DEBUG) {
|
|
||||||
console.log(
|
|
||||||
"[sendAutomaticallyWhen] No tool parts, returning false",
|
|
||||||
)
|
|
||||||
}
|
|
||||||
return false
|
return false
|
||||||
}
|
}
|
||||||
|
|
||||||
// Only auto-resubmit if ANY tool has an error
|
return toolParts.some((part) => part.state === TOOL_ERROR_STATE)
|
||||||
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
|
|
||||||
}
|
}
|
||||||
|
|
||||||
export default function ChatPanel({
|
export default function ChatPanel({
|
||||||
@@ -223,6 +186,9 @@ export default function ChatPanel({
|
|||||||
// Ref to hold stop function for use in onToolCall (avoids stale closure)
|
// Ref to hold stop function for use in onToolCall (avoids stale closure)
|
||||||
const stopRef = useRef<(() => void) | null>(null)
|
const stopRef = useRef<(() => void) | null>(null)
|
||||||
|
|
||||||
|
// Ref to track consecutive auto-retry count (reset on user action)
|
||||||
|
const autoRetryCountRef = useRef(0)
|
||||||
|
|
||||||
const {
|
const {
|
||||||
messages,
|
messages,
|
||||||
sendMessage,
|
sendMessage,
|
||||||
@@ -441,8 +407,68 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
sendAutomaticallyWhen: ({ messages }) =>
|
sendAutomaticallyWhen: ({ messages }) => {
|
||||||
shouldAutoResubmit(messages as unknown as ChatMessage[]),
|
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
|
// 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,
|
previousXml: string,
|
||||||
sessionId: string,
|
sessionId: string,
|
||||||
) => {
|
) => {
|
||||||
|
// Reset auto-retry count on user-initiated message
|
||||||
|
autoRetryCountRef.current = 0
|
||||||
|
|
||||||
const config = getAIConfig()
|
const config = getAIConfig()
|
||||||
|
|
||||||
sendMessage(
|
sendMessage(
|
||||||
|
|||||||
Reference in New Issue
Block a user