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: edit_diagram streaming and JSON repair improvements (#271)

Browse Source 
- Add shared editDiagramOriginalXmlRef between streaming preview and tool handler
  to avoid conflicts when applying operations (fixes "cell already exists" errors)
- Add JSON repair preprocessing to fix LLM-generated malformed JSON like `:=`
- Filter out tool calls with invalid/undefined inputs from interrupted streaming
- Remove perf console logs
This commit is contained in:
Dayuan Jiang
2025-12-15 21:28:31 +09:00
committed by GitHub
parent c527ce1520
commit cd76fa615e
5 changed files with 364 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -28,6 +28,7 @@ import {
} from "@/components/ai-elements/reasoning"
import { ScrollArea } from "@/components/ui/scroll-area"
import {
applyDiagramOperations,
convertToLegalXml,
isMxCellXmlComplete,
replaceNodes,
@@ -42,6 +43,23 @@ interface DiagramOperation {
new_xml?: string
}
// Helper to extract complete operations from streaming input
function getCompleteOperations(
operations: DiagramOperation[] | undefined,
): DiagramOperation[] {
if (!operations || !Array.isArray(operations)) return []
return operations.filter(
(op) =>
op &&
typeof op.type === "string" &&
["update", "add", "delete"].includes(op.type) &&
typeof op.cell_id === "string" &&
op.cell_id.length > 0 &&
// delete doesn't need new_xml, update/add do
(op.type === "delete" || typeof op.new_xml === "string"),
)
}
// Tool part interface for type safety
interface ToolPartLike {
type: string
@@ -167,6 +185,7 @@ interface ChatMessageDisplayProps {
setInput: (input: string) => void
setFiles: (files: File[]) => void
processedToolCallsRef: MutableRefObject<Set<string>>
editDiagramOriginalXmlRef: MutableRefObject<Map<string, string>>
sessionId?: string
onRegenerate?: (messageIndex: number) => void
onEditMessage?: (messageIndex: number, newText: string) => void
@@ -178,6 +197,7 @@ export function ChatMessageDisplay({
setInput,
setFiles,
processedToolCallsRef,
editDiagramOriginalXmlRef,
sessionId,
onRegenerate,
onEditMessage,
@@ -195,6 +215,14 @@ export function ChatMessageDisplay({
null,
)
const STREAMING_DEBOUNCE_MS = 150 // Only update diagram every 150ms during streaming
// Refs for edit_diagram streaming
const pendingEditRef = useRef<{
operations: DiagramOperation[]
toolCallId: string
} | null>(null)
const editDebounceTimeoutRef = useRef<ReturnType<typeof setTimeout> | null>(
null,
)
const [expandedTools, setExpandedTools] = useState<Record<string, boolean>>(
{},
)
@@ -286,15 +314,12 @@ export function ChatMessageDisplay({
const handleDisplayChart = useCallback(
(xml: string, showToast = false) => {
console.time("perf:handleDisplayChart")
const currentXml = xml || ""
const convertedXml = convertToLegalXml(currentXml)
if (convertedXml !== previousXML.current) {
// Parse and validate XML BEFORE calling replaceNodes
console.time("perf:DOMParser")
const parser = new DOMParser()
const testDoc = parser.parseFromString(convertedXml, "text/xml")
console.timeEnd("perf:DOMParser")
const parseError = testDoc.querySelector("parsererror")
if (parseError) {
@@ -310,7 +335,6 @@ export function ChatMessageDisplay({
"AI generated invalid diagram XML. Please try regenerating.",
)
}
console.timeEnd("perf:handleDisplayChart")
return // Skip this update
}
@@ -320,14 +344,10 @@ export function ChatMessageDisplay({
const baseXML =
chartXML ||
`<mxfile><diagram name="Page-1" id="page-1"><mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/></root></mxGraphModel></diagram></mxfile>`
console.time("perf:replaceNodes")
const replacedXML = replaceNodes(baseXML, convertedXml)
console.timeEnd("perf:replaceNodes")
// Validate and auto-fix the XML
console.time("perf:validateAndFixXml")
const validation = validateAndFixXml(replacedXML)
console.timeEnd("perf:validateAndFixXml")
if (validation.valid) {
previousXML.current = convertedXml
// Use fixed XML if available, otherwise use original
@@ -364,9 +384,6 @@ export function ChatMessageDisplay({
)
}
}
console.timeEnd("perf:handleDisplayChart")
} else {
console.timeEnd("perf:handleDisplayChart")
}
},
[chartXML, onDisplayChart],
@@ -385,11 +402,6 @@ export function ChatMessageDisplay({
}, [editingMessageId])
useEffect(() => {
console.time("perf:message-display-useEffect")
let processedCount = 0
let skippedCount = 0
let debouncedCount = 0
// Only process the last message for streaming performance
// Previous messages are already processed and won't change
const messagesToProcess =
@@ -419,7 +431,6 @@ export function ChatMessageDisplay({
const lastXml =
lastProcessedXmlRef.current.get(toolCallId)
if (lastXml === xml) {
skippedCount++
return // Skip redundant processing
}
@@ -439,9 +450,6 @@ export function ChatMessageDisplay({
debounceTimeoutRef.current = null
pendingXmlRef.current = null
if (pendingXml) {
console.log(
"perf:debounced-handleDisplayChart executing",
)
handleDisplayChart(
pendingXml,
false,
@@ -455,7 +463,6 @@ export function ChatMessageDisplay({
STREAMING_DEBOUNCE_MS,
)
}
debouncedCount++
} else if (
state === "output-available" &&
!processedToolCalls.current.has(toolCallId)
@@ -471,17 +478,129 @@ export function ChatMessageDisplay({
processedToolCalls.current.add(toolCallId)
// Clean up the ref entry - tool is complete, no longer needed
lastProcessedXmlRef.current.delete(toolCallId)
processedCount++
}
}
// Handle edit_diagram streaming - apply operations incrementally for preview
// Uses shared editDiagramOriginalXmlRef to coordinate with tool handler
if (
part.type === "tool-edit_diagram" &&
input?.operations
) {
const completeOps = getCompleteOperations(
input.operations as DiagramOperation[],
)
if (completeOps.length === 0) return
// Capture original XML when streaming starts (store in shared ref)
if (
!editDiagramOriginalXmlRef.current.has(
toolCallId,
)
) {
if (!chartXML) {
console.warn(
"[edit_diagram streaming] No chart XML available",
)
return
}
editDiagramOriginalXmlRef.current.set(
toolCallId,
chartXML,
)
}
const originalXml =
editDiagramOriginalXmlRef.current.get(
toolCallId,
)
if (!originalXml) return
// Skip if no change from last processed state
const lastCount = lastProcessedXmlRef.current.get(
toolCallId + "-opCount",
)
if (lastCount === String(completeOps.length)) return
if (
state === "input-streaming" ||
state === "input-available"
) {
// Queue the operations for debounced processing
pendingEditRef.current = {
operations: completeOps,
toolCallId,
}
if (!editDebounceTimeoutRef.current) {
editDebounceTimeoutRef.current = setTimeout(
() => {
const pending =
pendingEditRef.current
editDebounceTimeoutRef.current =
null
pendingEditRef.current = null
if (pending) {
const origXml =
editDiagramOriginalXmlRef.current.get(
pending.toolCallId,
)
if (!origXml) return
try {
const {
result: editedXml,
} = applyDiagramOperations(
origXml,
pending.operations,
)
handleDisplayChart(
editedXml,
false,
)
lastProcessedXmlRef.current.set(
pending.toolCallId +
"-opCount",
String(
pending.operations
.length,
),
)
} catch (e) {
console.warn(
`[edit_diagram streaming] Operation failed:`,
e instanceof Error
? e.message
: e,
)
}
}
},
STREAMING_DEBOUNCE_MS,
)
}
} else if (
state === "output-available" &&
!processedToolCalls.current.has(toolCallId)
) {
// Final state - cleanup streaming refs (tool handler does final application)
if (editDebounceTimeoutRef.current) {
clearTimeout(editDebounceTimeoutRef.current)
editDebounceTimeoutRef.current = null
}
lastProcessedXmlRef.current.delete(
toolCallId + "-opCount",
)
processedToolCalls.current.add(toolCallId)
// Note: Don't delete editDiagramOriginalXmlRef here - tool handler needs it
}
}
}
})
}
})
console.log(
`perf:message-display-useEffect processed=${processedCount} skipped=${skippedCount} debounced=${debouncedCount}`,
)
console.timeEnd("perf:message-display-useEffect")
// Cleanup: clear any pending debounce timeout on unmount
return () => {
@@ -489,8 +608,12 @@ export function ChatMessageDisplay({
clearTimeout(debounceTimeoutRef.current)
debounceTimeoutRef.current = null
}
if (editDebounceTimeoutRef.current) {
clearTimeout(editDebounceTimeoutRef.current)
editDebounceTimeoutRef.current = null
}
}
}, [messages, handleDisplayChart])
}, [messages, handleDisplayChart, chartXML])
const renderToolPart = (part: ToolPartLike) => {
const callId = part.toolCallId

74
components/chat-panel.tsx
View File

@@ -202,6 +202,10 @@ export default function ChatPanel({
// Persist processed tool call IDs so collapsing the chat doesn't replay old tool outputs
const processedToolCallsRef = useRef<Set<string>>(new Set())
// Store original XML for edit_diagram streaming - shared between streaming preview and tool handler
// Key: toolCallId, Value: original XML before any operations applied
const editDiagramOriginalXmlRef = useRef<Map<string, string>>(new Map())
// Debounce timeout for localStorage writes (prevents blocking during streaming)
const localStorageDebounceRef = useRef<ReturnType<
typeof setTimeout
@@ -333,13 +337,22 @@ ${finalXml}
let currentXml = ""
try {
// Use chartXML from ref directly - more reliable than export
const cachedXML = chartXMLRef.current
if (cachedXML) {
currentXml = cachedXML
// Use the original XML captured during streaming (shared with chat-message-display)
// This ensures we apply operations to the same base XML that streaming used
const originalXml = editDiagramOriginalXmlRef.current.get(
toolCall.toolCallId,
)
if (originalXml) {
currentXml = originalXml
} else {
// Fallback to export only if no cached XML
currentXml = await onFetchChart(false)
// Fallback: use chartXML from ref if streaming didn't capture original
const cachedXML = chartXMLRef.current
if (cachedXML) {
currentXml = cachedXML
} else {
// Last resort: export from iframe
currentXml = await onFetchChart(false)
}
}
const { applyDiagramOperations } = await import(
@@ -370,6 +383,10 @@ ${currentXml}
Please check the cell IDs and retry.`,
})
// Clean up the shared original XML ref
editDiagramOriginalXmlRef.current.delete(
toolCall.toolCallId,
)
return
}
@@ -393,6 +410,10 @@ ${currentXml}
Please fix the operations to avoid structural issues.`,
})
// Clean up the shared original XML ref
editDiagramOriginalXmlRef.current.delete(
toolCall.toolCallId,
)
return
}
onExport()
@@ -401,6 +422,10 @@ Please fix the operations to avoid structural issues.`,
toolCallId: toolCall.toolCallId,
output: `Successfully applied ${operations.length} operation(s) to the diagram.`,
})
// Clean up the shared original XML ref
editDiagramOriginalXmlRef.current.delete(
toolCall.toolCallId,
)
} catch (error) {
console.error("[edit_diagram] Failed:", error)
@@ -420,6 +445,10 @@ ${currentXml || "No XML available"}
Please check cell IDs and retry, or use display_diagram to regenerate.`,
})
// Clean up the shared original XML ref even on error
editDiagramOriginalXmlRef.current.delete(
toolCall.toolCallId,
)
}
} else if (toolCall.toolName === "append_diagram") {
const { xml } = toolCall.input as { xml: string }
@@ -508,6 +537,32 @@ Continue from EXACTLY where you stopped.`,
// Silence access code error in console since it's handled by UI
if (!error.message.includes("Invalid or missing access code")) {
console.error("Chat error:", error)
// Debug: Log messages structure when error occurs
console.log("[onError] messages count:", messages.length)
messages.forEach((msg, idx) => {
console.log(`[onError] Message ${idx}:`, {
role: msg.role,
partsCount: msg.parts?.length,
})
if (msg.parts) {
msg.parts.forEach((part: any, partIdx: number) => {
console.log(
`[onError] Part ${partIdx}:`,
JSON.stringify({
type: part.type,
toolName: part.toolName,
hasInput: !!part.input,
inputType: typeof part.input,
inputKeys:
part.input &&
typeof part.input === "object"
? Object.keys(part.input)
: null,
}),
)
})
}
})
}
// Translate technical errors into user-friendly messages
@@ -723,12 +778,10 @@ Continue from EXACTLY where you stopped.`,
// Debounce: save after 1 second of no changes
localStorageDebounceRef.current = setTimeout(() => {
try {
console.time("perf:localStorage-messages")
localStorage.setItem(
STORAGE_MESSAGES_KEY,
JSON.stringify(messages),
)
console.timeEnd("perf:localStorage-messages")
} catch (error) {
console.error("Failed to save messages to localStorage:", error)
}
@@ -754,9 +807,7 @@ Continue from EXACTLY where you stopped.`,
// Debounce: save after 1 second of no changes
xmlStorageDebounceRef.current = setTimeout(() => {
console.time("perf:localStorage-xml")
localStorage.setItem(STORAGE_DIAGRAM_XML_KEY, chartXML)
console.timeEnd("perf:localStorage-xml")
}, LOCAL_STORAGE_DEBOUNCE_MS)
return () => {
@@ -769,13 +820,11 @@ Continue from EXACTLY where you stopped.`,
// Save XML snapshots to localStorage whenever they change
const saveXmlSnapshots = useCallback(() => {
try {
console.time("perf:localStorage-snapshots")
const snapshotsArray = Array.from(xmlSnapshotsRef.current.entries())
localStorage.setItem(
STORAGE_XML_SNAPSHOTS_KEY,
JSON.stringify(snapshotsArray),
)
console.timeEnd("perf:localStorage-snapshots")
} catch (error) {
console.error(
"Failed to save XML snapshots to localStorage:",
@@ -1326,6 +1375,7 @@ Continue from EXACTLY where you stopped.`,
setInput={setInput}
setFiles={handleFileChange}
processedToolCallsRef={processedToolCallsRef}
editDiagramOriginalXmlRef={editDiagramOriginalXmlRef}
sessionId={sessionId}
onRegenerate={handleRegenerate}
status={status}