i/next-ai-draw-io
1
0
Fork 0
mirror of https://github.com/DayuanJiang/next-ai-draw-io.git synced 2026-01-02 14:22:28 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
- 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
This commit is contained in:
Dayuan Jiang
2025-12-07 14:38:15 +09:00
committed by GitHub
parent a54068fec2
commit 0baf21fadb
5 changed files with 109 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

63
app/api/chat/route.ts
View File

@@ -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)

3
components/chat-message-display.tsx
View File

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

40
components/chat-panel.tsx
View File

@@ -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

3
components/history-dialog.tsx
View File

@@ -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()
}
}

24
contexts/diagram-context.tsx
View File

@@ -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)
setChartXML(emptyDiagram)
// Skip validation for trusted internal template (loadDiagram also sets chartXML)
loadDiagram(emptyDiagram, true)
setLatestSvg("")
setDiagramHistory([])
}