feat: add append_diagram tool and improve truncation handling (#252)

* feat: add append_diagram tool for truncation continuation

When LLM output hits maxOutputTokens mid-generation, instead of
failing with an error loop, the system now:

1. Detects truncation (missing </root> in XML)
2. Stores partial XML and tells LLM to use new append_diagram tool
3. LLM continues generating from where it stopped
4. Fragments are accumulated until XML is complete
5. Server limits to 5 steps via stepCountIs(5)

Key changes:
- Add append_diagram tool definition in route.ts
- Add append_diagram handler in chat-panel.tsx
- Track continuation mode separately from error mode
- Continuation mode has unlimited retries (not counted against limit)
- Error mode still limited to MAX_AUTO_RETRY_COUNT (1)
- Update system prompts to document append_diagram tool

* fix: show friendly message and yellow badge for truncated output

- Add yellow 'Truncated' badge in UI instead of red 'Error' when XML is incomplete
- Show friendly error message for toolUse.input is invalid errors
- Built on top of append_diagram continuation feature

* refactor: remove debug logs and simplify truncation state

- Remove all debug console.log statements
- Remove isContinuationModeRef, derive from partialXmlRef.current.length > 0

* docs: fix append_diagram instructions for consistency

- Change 'Do NOT include' to 'Do NOT start with' (clearer intent)
- Add <mxCell id="0"> to prohibited start patterns
- Change 'closing tags </root></mxGraphModel>' to just '</root>' (wrapWithMxFile handles the rest)

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 start with <mxGraphModel>, <root>, or <mxCell id="0"> - they already exist in the partial
+2. Continue from EXACTLY where your previous output stopped
+3. End with the closing </root> tag to complete the diagram
+4. If still truncated, call append_diagram again with the next fragment
+
+Example: If previous output ended with '<mxCell id="x" style="rounded=1', continue with ';" vertex="1">...' 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

components/chat-message-display.tsx

@@ -451,11 +451,24 @@ export function ChatMessageDisplay({
                                 Complete
                             </span>
                         )}
-                        {state === "output-error" && (
+                        {state === "output-error" &&
+                            (() => {
+                                // Check if this is a truncation (incomplete XML) vs real error
+                                const isTruncated =
+                                    (toolName === "display_diagram" ||
+                                        toolName === "append_diagram") &&
+                                    (!input?.xml ||
+                                        !input.xml.includes("</root>"))
+                                return isTruncated ? (
+                                    <span className="text-xs font-medium text-yellow-600 bg-yellow-50 px-2 py-0.5 rounded-full">
+                                        Truncated
+                                    </span>
+                                ) : (
                                     <span className="text-xs font-medium text-red-600 bg-red-50 px-2 py-0.5 rounded-full">
                                         Error
                                     </span>
-                        )}
+                                )
+                            })()}
                         {input && Object.keys(input).length > 0 && (
                             <button
                                 type="button"
@@ -488,11 +501,23 @@ export function ChatMessageDisplay({
                         ) : null}
                     </div>
                 )}
-                {output && state === "output-error" && (
+                {output &&
-                    <div className="px-4 py-3 border-t border-border/40 text-sm text-red-600">
+                    state === "output-error" &&
-                        {output}
+                    (() => {
+                        const isTruncated =
+                            (toolName === "display_diagram" ||
+                                toolName === "append_diagram") &&
+                            (!input?.xml || !input.xml.includes("</root>"))
+                        return (
+                            <div
+                                className={`px-4 py-3 border-t border-border/40 text-sm ${isTruncated ? "text-yellow-600" : "text-red-600"}`}
+                            >
+                                {isTruncated
+                                    ? "Output truncated due to length limits. Try a simpler request or increase the maxOutputLength."
+                                    : output}
                             </div>
-                )}
+                        )
+                    })()}
             </div>
         )
     }

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,8 @@ function hasToolErrors(messages: ChatMessage[]): boolean {
         return false
     }
 
-    return toolParts.some((part) => part.state === TOOL_ERROR_STATE)
+    const lastToolPart = toolParts[toolParts.length - 1]
+    return lastToolPart?.state === TOOL_ERROR_STATE
 }
 
 export default function ChatPanel({
@@ -192,6 +193,10 @@ 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
+    // When partialXmlRef.current.length > 0, we're in continuation mode
+    const partialXmlRef = useRef<string>("")
+
     // Persist processed tool call IDs so collapsing the chat doesn't replay old tool outputs
     const processedToolCallsRef = useRef<Set<string>>(new Set())
 
@@ -216,14 +221,43 @@ export default function ChatPanel({
 
             if (toolCall.toolName === "display_diagram") {
                 const { xml } = toolCall.input as { xml: string }
-                if (DEBUG) {
+
-                    console.log(
+                // Check if XML is truncated (missing </root> indicates incomplete output)
-                        `[display_diagram] Received XML length: ${xml.length}`,
+                const isTruncated =
-                    )
+                    !xml.includes("</root>") && !xml.trim().endsWith("/>")
+
+                if (isTruncated) {
+                    // Store the partial XML for continuation via append_diagram
+                    partialXmlRef.current = xml
+
+                    // Tell LLM to use append_diagram to continue
+                    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 start with <mxGraphModel>, <root>, or <mxCell id="0"> (they already exist)
+- Start from EXACTLY where you stopped
+- End with the closing </root> tag to complete the diagram`,
+                    })
+                    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 +283,7 @@ Please fix the XML issues and call display_diagram again with corrected XML.
 
 Your failed XML:
 \`\`\`xml
-${xml}
+${finalXml}
 \`\`\``,
                     })
                 } else {
@@ -277,27 +311,13 @@ ${xml}
 
                 let currentXml = ""
                 try {
-                    console.log("[edit_diagram] Starting...")
                     // Use chartXML from ref directly - more reliable than export
-                    // especially on Vercel where DrawIO iframe may have latency issues
-                    // Using ref to avoid stale closure in callback
                     const cachedXML = chartXMLRef.current
                     if (cachedXML) {
                         currentXml = cachedXML
-                        console.log(
-                            "[edit_diagram] Using cached chartXML, length:",
-                            currentXml.length,
-                        )
                     } else {
                         // Fallback to export only if no cached XML
-                        console.log(
-                            "[edit_diagram] No cached XML, fetching from DrawIO...",
-                        )
                         currentXml = await onFetchChart(false)
-                        console.log(
-                            "[edit_diagram] Got XML from export, length:",
-                            currentXml.length,
-                        )
                     }
 
                     const { replaceXMLParts } = await import("@/lib/utils")
@@ -331,7 +351,6 @@ Please fix the edit to avoid structural issues (e.g., duplicate IDs, invalid ref
                         toolCallId: toolCall.toolCallId,
                         output: `Successfully applied ${edits.length} edit(s) to the diagram.`,
                     })
-                    console.log("[edit_diagram] Success")
                 } catch (error) {
                     console.error("[edit_diagram] Failed:", error)
 
@@ -353,6 +372,83 @@ ${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 }
+
+                // Detect if LLM incorrectly started fresh instead of continuing
+                const isFreshStart =
+                    xml.trim().startsWith("<mxGraphModel") ||
+                    xml.trim().startsWith("<root") ||
+                    xml.trim().startsWith('<mxCell id="0"')
+
+                if (isFreshStart) {
+                    addToolOutput({
+                        tool: "append_diagram",
+                        toolCallId: toolCall.toolCallId,
+                        state: "output-error",
+                        errorText: `ERROR: You started fresh with wrapper tags. Do NOT include <mxGraphModel>, <root>, or <mxCell id="0">.
+
+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
+
+                // Check if XML is now complete
+                const isComplete = partialXmlRef.current.includes("</root>")
+
+                if (isComplete) {
+                    // Wrap and display the complete diagram
+                    const finalXml = partialXmlRef.current
+                    partialXmlRef.current = "" // Reset
+
+                    const fullXml = wrapWithMxFile(finalXml)
+                    const validationError = onDisplayChart(fullXml)
+
+                    if (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 {
+                        addToolOutput({
+                            tool: "append_diagram",
+                            toolCallId: toolCall.toolCallId,
+                            output: "Diagram assembly complete and displayed successfully.",
+                        })
+                    }
+                } else {
+                    // Still incomplete - signal to continue
+                    addToolOutput({
+                        tool: "append_diagram",
+                        toolCallId: toolCall.toolCallId,
+                        state: "output-error",
+                        errorText: `XML still incomplete (missing </root>). Call append_diagram again to continue.
+
+Current ending:
+\`\`\`
+${partialXmlRef.current.slice(-500)}
+\`\`\`
+
+Continue from EXACTLY where you stopped.`,
+                    })
+                }
             }
         },
         onError: (error) => {
@@ -371,6 +467,12 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
                 friendlyMessage = "Network error. Please check your connection."
             }
 
+            // Truncated tool input error (model output limit too low)
+            if (friendlyMessage.includes("toolUse.input is invalid")) {
+                friendlyMessage =
+                    "Output was truncated before the diagram could be generated. Try a simpler request or increase the maxOutputLength."
+            }
+
             // Translate image not supported error
             if (friendlyMessage.includes("image content block")) {
                 friendlyMessage = "This model doesn't support image input."
@@ -398,6 +500,7 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
             const metadata = message?.metadata as
                 | Record<string, unknown>
                 | undefined
+
             if (metadata) {
                 // Use Number.isFinite to guard against NaN (typeof NaN === 'number' is true)
                 const inputTokens = Number.isFinite(metadata.inputTokens)
@@ -414,65 +517,55 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
             }
         },
         sendAutomaticallyWhen: ({ messages }) => {
+            const isInContinuationMode = partialXmlRef.current.length > 0
+
             const shouldRetry = hasToolErrors(
                 messages as unknown as ChatMessage[],
             )
 
             if (!shouldRetry) {
-                // No error, reset retry count
+                // No error, reset retry count and clear state
                 autoRetryCountRef.current = 0
-                if (DEBUG) {
+                partialXmlRef.current = ""
-                    console.log("[sendAutomaticallyWhen] No errors, stopping")
-                }
                 return false
             }
 
-            // Check retry count limit
+            // Continuation mode: unlimited retries (truncation continuation, not real errors)
+            // Server limits to 5 steps via stepCountIs(5)
+            if (isInContinuationMode) {
+                // 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) {
-                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
+                    partialXmlRef.current = ""
                     return false
                 }
+                // Increment retry count for actual errors
+                autoRetryCountRef.current++
+            }
 
             // 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
+                partialXmlRef.current = ""
                 return false
             }
 
             const tpmCheck = quotaManager.checkTPMLimit()
             if (!tpmCheck.allowed) {
-                if (DEBUG) {
-                    console.log(
-                        "[sendAutomaticallyWhen] TPM limit exceeded, stopping",
-                    )
-                }
                 quotaManager.showTPMLimitToast()
                 autoRetryCountRef.current = 0
+                partialXmlRef.current = ""
                 return false
             }
 
-            // Increment retry count and allow retry
-            autoRetryCountRef.current++
-            if (DEBUG) {
-                console.log(
-                    `[sendAutomaticallyWhen] Retrying (${autoRetryCountRef.current}/${MAX_AUTO_RETRY_COUNT})`,
-                )
-            }
             return true
         },
     })
@@ -817,8 +910,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
+        // Reset all retry/continuation state on user-initiated message
         autoRetryCountRef.current = 0
+        partialXmlRef.current = ""
 
         const config = getAIConfig()
 

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 <mxGraphModel> or <root>)
+}
 ---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 = `
 </root>
 \`\`\`
 
+### 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 start with <mxGraphModel>, <root>, or <mxCell id="0"> - they already exist in the partial
+2. Continue from EXACTLY where your previous output stopped
+3. End with the closing </root> tag to complete the diagram
+4. If still truncated, call append_diagram again with the next fragment
+
+**Example:** If previous output ended with \`<mxCell id="x" style="rounded=1\`, continue with \`;" vertex="1">...\` and complete the remaining elements.
+
 ### edit_diagram Details
 
 **CRITICAL RULES:**

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",

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",