fix: validate XML before displaying diagram to catch duplicate IDs (#147)

- Add validation to loadDiagram in diagram-context, returns error or null
- display_diagram and edit_diagram tools now check validation result
- Return error to AI agent with state: output-error so it can retry
- Skip validation for trusted sources (localStorage, history, internal templates)
- Add debug logging for tool call inputs to diagnose Bedrock API issues

app/api/chat/route.ts

@@ -67,17 +67,42 @@ function isMinimalDiagram(xml: string): boolean {
 // Helper function to fix tool call inputs for Bedrock API
 // Bedrock requires toolUse.input to be a JSON object, not a string
 function fixToolCallInputs(messages: any[]): any[] {
-    return messages.map((msg) => {
+    return messages.map((msg, msgIndex) => {
         if (msg.role !== "assistant" || !Array.isArray(msg.content)) {
             return msg
         }
-        const fixedContent = msg.content.map((part: any) => {
+        const fixedContent = msg.content.map((part: any, partIndex: number) => {
-            if (part.type === "tool-call" && typeof part.input === "string") {
+            if (part.type === "tool-call") {
-                try {
+                console.log(
-                    return { ...part, input: JSON.parse(part.input) }
+                    `[fixToolCallInputs] msg[${msgIndex}].content[${partIndex}] tool-call:`,
-                } catch {
+                    {
-                    // If parsing fails, wrap the string in an object
+                        toolName: part.toolName,
-                    return { ...part, input: { rawInput: part.input } }
+                        inputType: typeof part.input,
+                        input: part.input,
+                    },
+                )
+                if (typeof part.input === "string") {
+                    try {
+                        const parsed = JSON.parse(part.input)
+                        console.log(
+                            `[fixToolCallInputs] Parsed string input to JSON:`,
+                            parsed,
+                        )
+                        return { ...part, input: parsed }
+                    } catch {
+                        // If parsing fails, wrap the string in an object
+                        console.log(
+                            `[fixToolCallInputs] Failed to parse, wrapping in object`,
+                        )
+                        return { ...part, input: { rawInput: part.input } }
+                    }
+                }
+                // Input is already an object, but verify it's not null/undefined
+                if (part.input === null || part.input === undefined) {
+                    console.log(
+                        `[fixToolCallInputs] Input is null/undefined, using empty object`,
+                    )
+                    return { ...part, input: {} }
                 }
             }
             return part
@@ -212,6 +237,28 @@ ${lastMessageText}
     // Convert UIMessages to ModelMessages and add system message
     const modelMessages = convertToModelMessages(messages)
 
+    // Debug: log raw messages to see what's coming in
+    console.log(
+        "[DEBUG] Raw UI messages:",
+        JSON.stringify(
+            messages.map((m: any, i: number) => ({
+                index: i,
+                role: m.role,
+                partsCount: m.parts?.length,
+                parts: m.parts?.map((p: any) => ({
+                    type: p.type,
+                    toolName: p.toolName,
+                    toolCallId: p.toolCallId,
+                    state: p.state,
+                    inputType: p.input ? typeof p.input : undefined,
+                    input: p.input,
+                })),
+            })),
+            null,
+            2,
+        ),
+    )
+
     // Fix tool call inputs for Bedrock API (requires JSON objects, not strings)
     const fixedMessages = fixToolCallInputs(modelMessages)
 

components/chat-message-display.tsx

@@ -185,7 +185,8 @@ export function ChatMessageDisplay({
                 const validationError = validateMxCellStructure(replacedXML)
                 if (!validationError) {
                     previousXML.current = convertedXml
-                    onDisplayChart(replacedXML)
+                    // Skip validation in loadDiagram since we already validated above
+                    onDisplayChart(replacedXML, true)
                 } else {
                     console.log(
                         "[ChatMessageDisplay] XML validation failed:",

components/chat-panel.tsx

@@ -32,7 +32,7 @@ const STORAGE_SESSION_ID_KEY = "next-ai-draw-io-session-id"
 const STORAGE_DIAGRAM_XML_KEY = "next-ai-draw-io-diagram-xml"
 
 import { useDiagram } from "@/contexts/diagram-context"
-import { formatXML, validateMxCellStructure } from "@/lib/utils"
+import { formatXML } from "@/lib/utils"
 import { ChatMessageDisplay } from "./chat-message-display"
 
 interface ChatPanelProps {
@@ -142,8 +142,8 @@ export default function ChatPanel({
             if (toolCall.toolName === "display_diagram") {
                 const { xml } = toolCall.input as { xml: string }
 
-                // Validate the final XML result
+                // loadDiagram validates and returns error if invalid
-                const validationError = validateMxCellStructure(xml)
+                const validationError = onDisplayChart(xml)
 
                 if (validationError) {
                     console.warn(
@@ -206,7 +206,28 @@ ${xml}
                     const { replaceXMLParts } = await import("@/lib/utils")
                     const editedXml = replaceXMLParts(currentXml, edits)
 
-                    onDisplayChart(editedXml)
+                    // loadDiagram validates and returns error if invalid
+                    const validationError = onDisplayChart(editedXml)
+                    if (validationError) {
+                        console.warn(
+                            "[edit_diagram] Validation error:",
+                            validationError,
+                        )
+                        addToolOutput({
+                            tool: "edit_diagram",
+                            toolCallId: toolCall.toolCallId,
+                            state: "output-error",
+                            errorText: `Edit produced invalid XML: ${validationError}
+
+Current diagram XML:
+\`\`\`xml
+${currentXml}
+\`\`\`
+
+Please fix the edit to avoid structural issues (e.g., duplicate IDs, invalid references).`,
+                        })
+                        return
+                    }
 
                     addToolOutput({
                         tool: "edit_diagram",
@@ -330,7 +351,8 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
                     "[ChatPanel] Loading saved diagram XML, length:",
                     savedDiagramXml.length,
                 )
-                onDisplayChart(savedDiagramXml)
+                // Skip validation for trusted saved diagrams
+                onDisplayChart(savedDiagramXml, true)
                 chartXMLRef.current = savedDiagramXml
             }
         } catch (error) {
@@ -525,8 +547,8 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
             return
         }
 
-        // Restore the diagram to the saved state
+        // Restore the diagram to the saved state (skip validation for trusted snapshots)
-        onDisplayChart(savedXml)
+        onDisplayChart(savedXml, true)
 
         // Update ref directly to ensure edit_diagram has the correct XML
         chartXMLRef.current = savedXml
@@ -579,8 +601,8 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
             return
         }
 
-        // Restore the diagram to the saved state
+        // Restore the diagram to the saved state (skip validation for trusted snapshots)
-        onDisplayChart(savedXml)
+        onDisplayChart(savedXml, true)
 
         // Update ref directly to ensure edit_diagram has the correct XML
         chartXMLRef.current = savedXml

components/history-dialog.tsx

@@ -32,7 +32,8 @@ export function HistoryDialog({
 
     const handleConfirmRestore = () => {
         if (selectedIndex !== null) {
-            onDisplayChart(diagramHistory[selectedIndex].xml)
+            // Skip validation for trusted history snapshots
+            onDisplayChart(diagramHistory[selectedIndex].xml, true)
             handleClose()
         }
     }

contexts/diagram-context.tsx

@@ -4,13 +4,13 @@ import type React from "react"
 import { createContext, useContext, useRef, useState } from "react"
 import type { DrawIoEmbedRef } from "react-drawio"
 import type { ExportFormat } from "@/components/save-dialog"
-import { extractDiagramXML } from "../lib/utils"
+import { extractDiagramXML, validateMxCellStructure } from "../lib/utils"
 
 interface DiagramContextType {
     chartXML: string
     latestSvg: string
     diagramHistory: { svg: string; xml: string }[]
-    loadDiagram: (chart: string) => void
+    loadDiagram: (chart: string, skipValidation?: boolean) => string | null
     handleExport: () => void
     handleExportWithoutHistory: () => void
     resolverRef: React.Ref<((value: string) => void) | null>
@@ -73,7 +73,19 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
         }
     }
 
-    const loadDiagram = (chart: string) => {
+    const loadDiagram = (
+        chart: string,
+        skipValidation?: boolean,
+    ): string | null => {
+        // Validate XML structure before loading (unless skipped for internal use)
+        if (!skipValidation) {
+            const validationError = validateMxCellStructure(chart)
+            if (validationError) {
+                console.warn("[loadDiagram] Validation error:", validationError)
+                return validationError
+            }
+        }
+
         // Keep chartXML in sync even when diagrams are injected (e.g., display_diagram tool)
         setChartXML(chart)
 
@@ -82,6 +94,8 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
                 xml: chart,
             })
         }
+
+        return null
     }
 
     const handleDiagramExport = (data: any) => {
@@ -121,8 +135,8 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
 
     const clearDiagram = () => {
         const emptyDiagram = `<mxfile><diagram name="Page-1" id="page-1"><mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/></root></mxGraphModel></diagram></mxfile>`
-        loadDiagram(emptyDiagram)
+        // Skip validation for trusted internal template (loadDiagram also sets chartXML)
-        setChartXML(emptyDiagram)
+        loadDiagram(emptyDiagram, true)
         setLatestSvg("")
         setDiagramHistory([])
     }