diff --git a/app/api/chat/route.ts b/app/api/chat/route.ts index b2500bb..9d06327 100644 --- a/app/api/chat/route.ts +++ b/app/api/chat/route.ts @@ -3,10 +3,12 @@ import { convertToModelMessages, createUIMessageStream, createUIMessageStreamResponse, + InvalidToolInputError, LoadAPIKeyError, stepCountIs, streamText, } from "ai" +import { jsonrepair } from "jsonrepair" import { z } from "zod" import { getAIModel, supportsPromptCaching } from "@/lib/ai-providers" import { findCachedResponse } from "@/lib/cached-responses" @@ -320,6 +322,31 @@ ${userInputText} maxOutputTokens: parseInt(process.env.MAX_OUTPUT_TOKENS, 10), }), stopWhen: stepCountIs(5), + // Repair truncated tool calls when maxOutputTokens is reached mid-JSON + experimental_repairToolCall: async ({ toolCall, error }) => { + // Only attempt repair for invalid tool input (broken JSON from truncation) + if ( + error instanceof InvalidToolInputError || + error.name === "AI_InvalidToolInputError" + ) { + try { + // Use jsonrepair to fix truncated JSON + const repairedInput = jsonrepair(toolCall.input) + console.log( + `[repairToolCall] Repaired truncated JSON for tool: ${toolCall.toolName}`, + ) + return { ...toolCall, input: repairedInput } + } catch (repairError) { + console.warn( + `[repairToolCall] Failed to repair JSON for tool: ${toolCall.toolName}`, + repairError, + ) + return null + } + } + // Don't attempt to repair other errors (like NoSuchToolError) + return null + }, messages: allMessages, ...(providerOptions && { providerOptions }), // This now includes all reasoning configs ...(headers && { headers }), @@ -411,6 +438,26 @@ IMPORTANT: Keep edits concise: ), }), }, + append_diagram: { + description: `Continue generating diagram XML when previous display_diagram output was truncated due to length limits. + +WHEN TO USE: Only call this tool after display_diagram was truncated (you'll see an error message about truncation). + +CRITICAL INSTRUCTIONS: +1. Do NOT include or tags - they already exist in the partial +2. Continue from EXACTLY where your previous output stopped +3. Generate the remaining XML including closing tags +4. If still truncated, call append_diagram again with the next fragment + +Example: If previous output ended with '...' and complete the remaining elements.`, + inputSchema: z.object({ + xml: z + .string() + .describe( + "Continuation XML fragment to append (NO wrapper tags)", + ), + }), + }, }, ...(process.env.TEMPERATURE !== undefined && { temperature: parseFloat(process.env.TEMPERATURE), @@ -435,6 +482,7 @@ IMPORTANT: Keep edits concise: return { inputTokens: totalInputTokens, outputTokens: usage.outputTokens ?? 0, + finishReason: (part as any).finishReason, } } return undefined diff --git a/components/chat-panel.tsx b/components/chat-panel.tsx index 8d55ace..20eddee 100644 --- a/components/chat-panel.tsx +++ b/components/chat-panel.tsx @@ -67,7 +67,7 @@ const MAX_AUTO_RETRY_COUNT = 1 /** * Check if auto-resubmit should happen based on tool errors. - * Does NOT handle retry count or quota - those are handled by the caller. + * Only checks the LAST tool part (most recent tool call), not all tool parts. */ function hasToolErrors(messages: ChatMessage[]): boolean { const lastMessage = messages[messages.length - 1] @@ -84,7 +84,12 @@ function hasToolErrors(messages: ChatMessage[]): boolean { return false } - return toolParts.some((part) => part.state === TOOL_ERROR_STATE) + const lastToolPart = toolParts[toolParts.length - 1] + const hasError = lastToolPart?.state === TOOL_ERROR_STATE + console.log( + `[hasToolErrors] lastToolPart state: ${lastToolPart?.state}, hasError: ${hasError}`, + ) + return hasError } export default function ChatPanel({ @@ -192,6 +197,13 @@ export default function ChatPanel({ // Ref to track consecutive auto-retry count (reset on user action) const autoRetryCountRef = useRef(0) + // Ref to accumulate partial XML when output is truncated due to maxOutputTokens + const partialXmlRef = useRef("") + + // Ref to track if we're in continuation mode (truncation, not error) + // This allows unlimited retries for continuation vs limited for errors + const isContinuationModeRef = useRef(false) + // Persist processed tool call IDs so collapsing the chat doesn't replay old tool outputs const processedToolCallsRef = useRef>(new Set()) @@ -216,14 +228,83 @@ export default function ChatPanel({ if (toolCall.toolName === "display_diagram") { const { xml } = toolCall.input as { xml: string } - if (DEBUG) { + + // Always log truncation-related info for debugging + console.log(`[display_diagram] === TRUNCATION DEBUG ===`) + console.log( + `[display_diagram] Received XML length: ${xml.length}`, + ) + console.log( + `[display_diagram] XML starts with: ${xml.substring(0, 100)}...`, + ) + console.log( + `[display_diagram] XML ends with: ...${xml.substring(xml.length - 100)}`, + ) + console.log( + `[display_diagram] Has : ${xml.includes("")}`, + ) + console.log( + `[display_diagram] partialXmlRef.current length: ${partialXmlRef.current.length}`, + ) + + // Check if XML is truncated (missing indicates incomplete output) + // This happens when maxOutputTokens is reached mid-generation + const isTruncated = + !xml.includes("") && !xml.trim().endsWith("/>") + + // Check if this is a fresh start vs a continuation + // Fresh start indicators: , , or 0 + + console.log(`[display_diagram] isTruncated: ${isTruncated}`) + console.log(`[display_diagram] isFreshStart: ${isFreshStart}`) + console.log( + `[display_diagram] hadPreviousPartial: ${hadPreviousPartial}`, + ) + + if (isTruncated) { + // Store the partial XML for continuation via append_diagram + // Always reset to current xml since this is the first truncation + partialXmlRef.current = xml + isContinuationModeRef.current = true // Mark as continuation (not error) + console.log( - `[display_diagram] Received XML length: ${xml.length}`, + `[display_diagram] XML truncated (${xml.length} chars). Instructing LLM to use append_diagram.`, ) + + // Tell LLM to use append_diagram to continue + // Use "output-error" to trigger auto-retry, but isContinuationModeRef tracks it's not a real error + const partialEnding = partialXmlRef.current.slice(-500) + addToolOutput({ + tool: "display_diagram", + toolCallId: toolCall.toolCallId, + state: "output-error", + errorText: `Output was truncated due to length limits. Use the append_diagram tool to continue. + +Your output ended with: +\`\`\` +${partialEnding} +\`\`\` + +NEXT STEP: Call append_diagram with the continuation XML. +- Do NOT include or tags +- Start from EXACTLY where you stopped +- Continue until complete with `, + }) + return } + // Complete XML received - use it directly + // (continuation is now handled via append_diagram tool) + const finalXml = xml + partialXmlRef.current = "" // Reset any partial from previous truncation + // Wrap raw XML with full mxfile structure for draw.io - const fullXml = wrapWithMxFile(xml) + const fullXml = wrapWithMxFile(finalXml) // loadDiagram validates and returns error if invalid const validationError = onDisplayChart(fullXml) @@ -249,7 +330,7 @@ Please fix the XML issues and call display_diagram again with corrected XML. Your failed XML: \`\`\`xml -${xml} +${finalXml} \`\`\``, }) } else { @@ -353,6 +434,120 @@ ${currentXml || "No XML available"} Please retry with an adjusted search pattern or use display_diagram if retries are exhausted.`, }) } + } else if (toolCall.toolName === "append_diagram") { + const { xml } = toolCall.input as { xml: string } + + console.log(`[append_diagram] === APPEND DEBUG ===`) + console.log( + `[append_diagram] Received fragment length: ${xml.length}`, + ) + console.log( + `[append_diagram] Fragment starts with: ${xml.substring(0, 100)}...`, + ) + console.log( + `[append_diagram] Fragment ends with: ...${xml.substring(xml.length - 100)}`, + ) + console.log( + `[append_diagram] Current partialXmlRef length: ${partialXmlRef.current.length}`, + ) + + // Detect if LLM incorrectly started fresh instead of continuing + const isFreshStart = + xml.trim().startsWith(", , or . + +Continue from EXACTLY where the partial ended: +\`\`\` +${partialXmlRef.current.slice(-500)} +\`\`\` + +Start your continuation with the NEXT character after where it stopped.`, + }) + return + } + + // Append to accumulated XML + partialXmlRef.current += xml + console.log( + `[append_diagram] After append, total length: ${partialXmlRef.current.length}`, + ) + + // Check if XML is now complete + const isComplete = partialXmlRef.current.includes("") + console.log(`[append_diagram] isComplete: ${isComplete}`) + + if (isComplete) { + // Wrap and display the complete diagram + const finalXml = partialXmlRef.current + partialXmlRef.current = "" // Reset + isContinuationModeRef.current = false // Continuation complete + + console.log( + `[append_diagram] XML complete! Final length: ${finalXml.length}`, + ) + + const fullXml = wrapWithMxFile(finalXml) + const validationError = onDisplayChart(fullXml) + + if (validationError) { + console.warn( + `[append_diagram] Validation error after assembly:`, + validationError, + ) + addToolOutput({ + tool: "append_diagram", + toolCallId: toolCall.toolCallId, + state: "output-error", + errorText: `Validation error after assembly: ${validationError} + +Assembled XML: +\`\`\`xml +${finalXml.substring(0, 2000)}... +\`\`\` + +Please use display_diagram with corrected XML.`, + }) + } else { + console.log( + `[append_diagram] Success! Diagram displayed.`, + ) + addToolOutput({ + tool: "append_diagram", + toolCallId: toolCall.toolCallId, + output: "Diagram assembly complete and displayed successfully.", + }) + } + } else { + // Still incomplete - signal to continue (stay in continuation mode) + console.log( + `[append_diagram] Still incomplete, asking for more.`, + ) + // isContinuationModeRef.current stays true + addToolOutput({ + tool: "append_diagram", + toolCallId: toolCall.toolCallId, + state: "output-error", + errorText: `XML still incomplete (missing ). Call append_diagram again to continue. + +Current ending: +\`\`\` +${partialXmlRef.current.slice(-500)} +\`\`\` + +Continue from EXACTLY where you stopped.`, + }) + } } }, onError: (error) => { @@ -398,6 +593,12 @@ Please retry with an adjusted search pattern or use display_diagram if retries a const metadata = message?.metadata as | Record | undefined + + // Log finish reason for debugging truncation + console.log(`[onFinish] === FINISH DEBUG ===`) + console.log(`[onFinish] finishReason: ${metadata?.finishReason}`) + console.log(`[onFinish] outputTokens: ${metadata?.outputTokens}`) + if (metadata) { // Use Number.isFinite to guard against NaN (typeof NaN === 'number' is true) const inputTokens = Number.isFinite(metadata.inputTokens) @@ -414,65 +615,91 @@ Please retry with an adjusted search pattern or use display_diagram if retries a } }, sendAutomaticallyWhen: ({ messages }) => { + console.log(`[sendAutomaticallyWhen] === RETRY DEBUG ===`) + console.log( + `[sendAutomaticallyWhen] isContinuationMode: ${isContinuationModeRef.current}`, + ) + console.log( + `[sendAutomaticallyWhen] partialXmlRef.current length: ${partialXmlRef.current.length}`, + ) + console.log( + `[sendAutomaticallyWhen] autoRetryCountRef.current: ${autoRetryCountRef.current}`, + ) + const shouldRetry = hasToolErrors( messages as unknown as ChatMessage[], ) + console.log(`[sendAutomaticallyWhen] hasToolErrors: ${shouldRetry}`) if (!shouldRetry) { - // No error, reset retry count + // No error, reset retry count and clear state + console.log( + `[sendAutomaticallyWhen] No errors - resetting state`, + ) autoRetryCountRef.current = 0 - if (DEBUG) { - console.log("[sendAutomaticallyWhen] No errors, stopping") - } + partialXmlRef.current = "" + isContinuationModeRef.current = false 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.`, + // Continuation mode: unlimited retries (truncation continuation, not real errors) + // Server limits to 5 steps via stepCountIs(5) + if (isContinuationModeRef.current) { + console.log( + `[sendAutomaticallyWhen] Continuation mode - allowing retry without counting`, + ) + // Don't count against retry limit for continuation + // Quota checks still apply below + } else { + // Regular error: check retry count limit + if (autoRetryCountRef.current >= MAX_AUTO_RETRY_COUNT) { + console.log( + `[sendAutomaticallyWhen] Max error 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 + partialXmlRef.current = "" + isContinuationModeRef.current = false + return false + } + // Increment retry count for actual errors + autoRetryCountRef.current++ + console.log( + `[sendAutomaticallyWhen] Error retry ${autoRetryCountRef.current}/${MAX_AUTO_RETRY_COUNT}`, ) - 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", - ) - } + console.log( + "[sendAutomaticallyWhen] Token limit exceeded, stopping", + ) quotaManager.showTokenLimitToast(tokenLimitCheck.used) autoRetryCountRef.current = 0 + partialXmlRef.current = "" + isContinuationModeRef.current = false return false } const tpmCheck = quotaManager.checkTPMLimit() if (!tpmCheck.allowed) { - if (DEBUG) { - console.log( - "[sendAutomaticallyWhen] TPM limit exceeded, stopping", - ) - } + console.log( + "[sendAutomaticallyWhen] TPM limit exceeded, stopping", + ) quotaManager.showTPMLimitToast() autoRetryCountRef.current = 0 + partialXmlRef.current = "" + isContinuationModeRef.current = false return false } - // Increment retry count and allow retry - autoRetryCountRef.current++ - if (DEBUG) { - console.log( - `[sendAutomaticallyWhen] Retrying (${autoRetryCountRef.current}/${MAX_AUTO_RETRY_COUNT})`, - ) - } + // Allow retry + console.log( + `[sendAutomaticallyWhen] Allowing retry${isContinuationModeRef.current ? " [continuation mode]" : ""}`, + ) return true }, }) @@ -817,8 +1044,10 @@ 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 + // Reset all retry/continuation state on user-initiated message autoRetryCountRef.current = 0 + partialXmlRef.current = "" + isContinuationModeRef.current = false const config = getAIConfig() diff --git a/lib/system-prompts.ts b/lib/system-prompts.ts index 0f3571f..e99ac96 100644 --- a/lib/system-prompts.ts +++ b/lib/system-prompts.ts @@ -42,11 +42,18 @@ description: Edit specific parts of the EXISTING diagram. Use this when making s parameters: { edits: Array<{search: string, replace: string}> } +---Tool3--- +tool name: append_diagram +description: Continue generating diagram XML when display_diagram was truncated due to output length limits. Only use this after display_diagram truncation. +parameters: { + xml: string // Continuation fragment (NO wrapper tags like or ) +} ---End of tools--- IMPORTANT: Choose the right tool: - Use display_diagram for: Creating new diagrams, major restructuring, or when the current diagram XML is empty - Use edit_diagram for: Small modifications, adding/removing elements, changing text/colors, repositioning items +- Use append_diagram for: ONLY when display_diagram was truncated due to output length - continue generating from where you stopped Core capabilities: - Generate valid, well-formed XML strings for draw.io diagrams @@ -174,6 +181,18 @@ const EXTENDED_ADDITIONS = ` \`\`\` +### append_diagram Details + +**WHEN TO USE:** Only call this tool when display_diagram output was truncated (you'll see an error message about truncation). + +**CRITICAL RULES:** +1. Do NOT include , , or - they already exist in the partial +2. Continue from EXACTLY where your previous output stopped +3. Generate the remaining XML including closing tags +4. If still truncated, call append_diagram again with the next fragment + +**Example:** If previous output ended with \`...\` and complete the remaining elements. + ### edit_diagram Details **CRITICAL RULES:** diff --git a/package-lock.json b/package-lock.json index e29857e..545842b 100644 --- a/package-lock.json +++ b/package-lock.json @@ -40,6 +40,7 @@ "clsx": "^2.1.1", "js-tiktoken": "^1.0.21", "jsdom": "^26.0.0", + "jsonrepair": "^3.13.1", "lucide-react": "^0.483.0", "motion": "^12.23.25", "next": "^16.0.7", @@ -9199,6 +9200,15 @@ "node": ">=6" } }, + "node_modules/jsonrepair": { + "version": "3.13.1", + "resolved": "https://registry.npmjs.org/jsonrepair/-/jsonrepair-3.13.1.tgz", + "integrity": "sha512-WJeiE0jGfxYmtLwBTEk8+y/mYcaleyLXWaqp5bJu0/ZTSeG0KQq/wWQ8pmnkKenEdN6pdnn6QtcoSUkbqDHWNw==", + "license": "ISC", + "bin": { + "jsonrepair": "bin/cli.js" + } + }, "node_modules/jsx-ast-utils": { "version": "3.3.5", "resolved": "https://registry.npmjs.org/jsx-ast-utils/-/jsx-ast-utils-3.3.5.tgz", diff --git a/package.json b/package.json index 0e14032..3da31b7 100644 --- a/package.json +++ b/package.json @@ -44,6 +44,7 @@ "clsx": "^2.1.1", "js-tiktoken": "^1.0.21", "jsdom": "^26.0.0", + "jsonrepair": "^3.13.1", "lucide-react": "^0.483.0", "motion": "^12.23.25", "next": "^16.0.7",