From 0baf21fadbc0bef7d7221429eb014648f11b8eef Mon Sep 17 00:00:00 2001 From: Dayuan Jiang <34411969+DayuanJiang@users.noreply.github.com> Date: Sun, 7 Dec 2025 14:38:15 +0900 Subject: [PATCH] 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 | 63 +++++++++++++++++++++++++---- components/chat-message-display.tsx | 3 +- components/chat-panel.tsx | 40 +++++++++++++----- components/history-dialog.tsx | 3 +- contexts/diagram-context.tsx | 24 ++++++++--- 5 files changed, 109 insertions(+), 24 deletions(-) diff --git a/app/api/chat/route.ts b/app/api/chat/route.ts index 42be311..99474b2 100644 --- a/app/api/chat/route.ts +++ b/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) => { - if (part.type === "tool-call" && typeof part.input === "string") { - try { - return { ...part, input: JSON.parse(part.input) } - } catch { - // If parsing fails, wrap the string in an object - return { ...part, input: { rawInput: part.input } } + const fixedContent = msg.content.map((part: any, partIndex: number) => { + if (part.type === "tool-call") { + console.log( + `[fixToolCallInputs] msg[${msgIndex}].content[${partIndex}] tool-call:`, + { + toolName: part.toolName, + 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) diff --git a/components/chat-message-display.tsx b/components/chat-message-display.tsx index d3b9243..aaa0e68 100644 --- a/components/chat-message-display.tsx +++ b/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:", diff --git a/components/chat-panel.tsx b/components/chat-panel.tsx index 6aa66ae..8425f05 100644 --- a/components/chat-panel.tsx +++ b/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 - const validationError = validateMxCellStructure(xml) + // loadDiagram validates and returns error if invalid + 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 - onDisplayChart(savedXml) + // Restore the diagram to the saved state (skip validation for trusted snapshots) + 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 - onDisplayChart(savedXml) + // Restore the diagram to the saved state (skip validation for trusted snapshots) + onDisplayChart(savedXml, true) // Update ref directly to ensure edit_diagram has the correct XML chartXMLRef.current = savedXml diff --git a/components/history-dialog.tsx b/components/history-dialog.tsx index 32500fd..6db5f59 100644 --- a/components/history-dialog.tsx +++ b/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() } } diff --git a/contexts/diagram-context.tsx b/contexts/diagram-context.tsx index 560ea43..ee6c92e 100644 --- a/contexts/diagram-context.tsx +++ b/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 = `` - loadDiagram(emptyDiagram) - setChartXML(emptyDiagram) + // Skip validation for trusted internal template (loadDiagram also sets chartXML) + loadDiagram(emptyDiagram, true) setLatestSvg("") setDiagramHistory([]) }