From cd76fa615e928e14fc510b94cb7db94f304cbe57 Mon Sep 17 00:00:00 2001 From: Dayuan Jiang <34411969+DayuanJiang@users.noreply.github.com> Date: Mon, 15 Dec 2025 21:28:31 +0900 Subject: [PATCH] fix: edit_diagram streaming and JSON repair improvements (#271) - 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 --- app/api/chat/route.ts | 169 +++++++++++++++++++++++--- components/chat-message-display.tsx | 177 +++++++++++++++++++++++----- components/chat-panel.tsx | 74 ++++++++++-- contexts/diagram-context.tsx | 7 -- lib/utils.ts | 28 ----- 5 files changed, 364 insertions(+), 91 deletions(-) diff --git a/app/api/chat/route.ts b/app/api/chat/route.ts index 5cb2ee5..cc84465 100644 --- a/app/api/chat/route.ts +++ b/app/api/chat/route.ts @@ -70,29 +70,41 @@ function isMinimalDiagram(xml: string): boolean { // Helper function to replace historical tool call XML with placeholders // This reduces token usage and forces LLM to rely on the current diagram XML (source of truth) +// Also fixes invalid/undefined inputs from interrupted streaming function replaceHistoricalToolInputs(messages: any[]): any[] { return messages.map((msg) => { if (msg.role !== "assistant" || !Array.isArray(msg.content)) { return msg } - const replacedContent = msg.content.map((part: any) => { - if (part.type === "tool-call") { - const toolName = part.toolName - if ( - toolName === "display_diagram" || - toolName === "edit_diagram" - ) { - return { - ...part, - input: { - placeholder: - "[XML content replaced - see current diagram XML in system context]", - }, + const replacedContent = msg.content + .map((part: any) => { + if (part.type === "tool-call") { + const toolName = part.toolName + // Fix invalid/undefined inputs from interrupted streaming + if ( + !part.input || + typeof part.input !== "object" || + Object.keys(part.input).length === 0 + ) { + // Skip tool calls with invalid inputs entirely + return null + } + if ( + toolName === "display_diagram" || + toolName === "edit_diagram" + ) { + return { + ...part, + input: { + placeholder: + "[XML content replaced - see current diagram XML in system context]", + }, + } } } - } - return part - }) + return part + }) + .filter(Boolean) // Remove null entries (invalid tool calls) return { ...msg, content: replacedContent } }) } @@ -231,6 +243,36 @@ ${userInputText} // Convert UIMessages to ModelMessages and add system message const modelMessages = convertToModelMessages(messages) + // DEBUG: Log incoming messages structure + console.log("[route.ts] Incoming messages count:", messages.length) + messages.forEach((msg: any, idx: number) => { + console.log( + `[route.ts] Message ${idx} role:`, + msg.role, + "parts count:", + msg.parts?.length, + ) + if (msg.parts) { + msg.parts.forEach((part: any, partIdx: number) => { + if ( + part.type === "tool-invocation" || + part.type === "tool-result" + ) { + console.log(`[route.ts] Part ${partIdx}:`, { + 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, + }) + } + }) + } + }) + // Replace historical tool call XML with placeholders to reduce tokens // Disabled by default - some models (e.g. minimax) copy placeholders instead of generating XML const enableHistoryReplace = @@ -246,6 +288,63 @@ ${userInputText} msg.content && Array.isArray(msg.content) && msg.content.length > 0, ) + // Filter out tool-calls with invalid inputs (from failed repair or interrupted streaming) + // Bedrock API rejects messages where toolUse.input is not a valid JSON object + enhancedMessages = enhancedMessages + .map((msg: any) => { + if (msg.role !== "assistant" || !Array.isArray(msg.content)) { + return msg + } + const filteredContent = msg.content.filter((part: any) => { + if (part.type === "tool-call") { + // Check if input is a valid object (not null, undefined, or empty) + if ( + !part.input || + typeof part.input !== "object" || + Object.keys(part.input).length === 0 + ) { + console.warn( + `[route.ts] Filtering out tool-call with invalid input:`, + { toolName: part.toolName, input: part.input }, + ) + return false + } + } + return true + }) + return { ...msg, content: filteredContent } + }) + .filter((msg: any) => msg.content && msg.content.length > 0) + + // DEBUG: Log modelMessages structure (what's being sent to AI) + console.log("[route.ts] Model messages count:", enhancedMessages.length) + enhancedMessages.forEach((msg: any, idx: number) => { + console.log( + `[route.ts] ModelMsg ${idx} role:`, + msg.role, + "content count:", + msg.content?.length, + ) + if (msg.content) { + msg.content.forEach((part: any, partIdx: number) => { + if (part.type === "tool-call" || part.type === "tool-result") { + console.log(`[route.ts] Content ${partIdx}:`, { + type: part.type, + toolName: part.toolName, + hasInput: !!part.input, + inputType: typeof part.input, + inputValue: + part.input === undefined + ? "undefined" + : part.input === null + ? "null" + : "object", + }) + } + }) + } + }) + // Update the last message with user input only (XML moved to separate cached system message) if (enhancedMessages.length >= 1) { const lastModelMessage = enhancedMessages[enhancedMessages.length - 1] @@ -327,14 +426,30 @@ ${userInputText} stopWhen: stepCountIs(5), // Repair truncated tool calls when maxOutputTokens is reached mid-JSON experimental_repairToolCall: async ({ toolCall, error }) => { + // DEBUG: Log what we're trying to repair + console.log(`[repairToolCall] Tool: ${toolCall.toolName}`) + console.log( + `[repairToolCall] Error: ${error.name} - ${error.message}`, + ) + console.log(`[repairToolCall] Input type: ${typeof toolCall.input}`) + console.log(`[repairToolCall] Input value:`, toolCall.input) + // Only attempt repair for invalid tool input (broken JSON from truncation) if ( error instanceof InvalidToolInputError || error.name === "AI_InvalidToolInputError" ) { try { + // Pre-process to fix common LLM JSON errors that jsonrepair can't handle + let inputToRepair = toolCall.input + if (typeof inputToRepair === "string") { + // Fix `:=` instead of `: ` (LLM sometimes generates this) + inputToRepair = inputToRepair.replace(/:=/g, ": ") + // Fix `= "` instead of `: "` + inputToRepair = inputToRepair.replace(/=\s*"/g, ': "') + } // Use jsonrepair to fix truncated JSON - const repairedInput = jsonrepair(toolCall.input) + const repairedInput = jsonrepair(inputToRepair) console.log( `[repairToolCall] Repaired truncated JSON for tool: ${toolCall.toolName}`, ) @@ -344,6 +459,26 @@ ${userInputText} `[repairToolCall] Failed to repair JSON for tool: ${toolCall.toolName}`, repairError, ) + // Return a placeholder input to avoid API errors in multi-step + // The tool will fail gracefully on client side + if (toolCall.toolName === "edit_diagram") { + return { + ...toolCall, + input: { + operations: [], + _error: "JSON repair failed - no operations to apply", + }, + } + } + if (toolCall.toolName === "display_diagram") { + return { + ...toolCall, + input: { + xml: "", + _error: "JSON repair failed - empty diagram", + }, + } + } return null } } diff --git a/components/chat-message-display.tsx b/components/chat-message-display.tsx index 3be4322..b2dafad 100644 --- a/components/chat-message-display.tsx +++ b/components/chat-message-display.tsx @@ -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> + editDiagramOriginalXmlRef: MutableRefObject> 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 | null>( + null, + ) const [expandedTools, setExpandedTools] = useState>( {}, ) @@ -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 || `` - 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 diff --git a/components/chat-panel.tsx b/components/chat-panel.tsx index 9a4fac6..b7408b0 100644 --- a/components/chat-panel.tsx +++ b/components/chat-panel.tsx @@ -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>(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>(new Map()) + // Debounce timeout for localStorage writes (prevents blocking during streaming) const localStorageDebounceRef = useRef { + 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} diff --git a/contexts/diagram-context.tsx b/contexts/diagram-context.tsx index 2f630cf..0f68863 100644 --- a/contexts/diagram-context.tsx +++ b/contexts/diagram-context.tsx @@ -86,20 +86,16 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) { chart: string, skipValidation?: boolean, ): string | null => { - console.time("perf:loadDiagram") let xmlToLoad = chart // Validate XML structure before loading (unless skipped for internal use) if (!skipValidation) { - console.time("perf:loadDiagram-validation") const validation = validateAndFixXml(chart) - console.timeEnd("perf:loadDiagram-validation") if (!validation.valid) { console.warn( "[loadDiagram] Validation error:", validation.error, ) - console.timeEnd("perf:loadDiagram") return validation.error } // Use fixed XML if auto-fix was applied @@ -116,14 +112,11 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) { setChartXML(xmlToLoad) if (drawioRef.current) { - console.time("perf:drawio-iframe-load") drawioRef.current.load({ xml: xmlToLoad, }) - console.timeEnd("perf:drawio-iframe-load") } - console.timeEnd("perf:loadDiagram") return null } diff --git a/lib/utils.ts b/lib/utils.ts index 12cbd1f..bff1cdc 100644 --- a/lib/utils.ts +++ b/lib/utils.ts @@ -743,8 +743,6 @@ function checkNestedMxCells(xml: string): string | null { * @returns null if valid, error message string if invalid */ export function validateMxCellStructure(xml: string): string | null { - console.time("perf:validateMxCellStructure") - console.log(`perf:validateMxCellStructure XML size: ${xml.length} bytes`) // Size check for performance if (xml.length > MAX_XML_SIZE) { console.warn( @@ -754,18 +752,10 @@ export function validateMxCellStructure(xml: string): string | null { // 0. First use DOM parser to catch syntax errors (most accurate) try { - console.time("perf:validate-DOMParser") const parser = new DOMParser() const doc = parser.parseFromString(xml, "text/xml") - console.timeEnd("perf:validate-DOMParser") const parseError = doc.querySelector("parsererror") if (parseError) { - const actualError = parseError.textContent || "Unknown parse error" - console.log( - "[validateMxCellStructure] DOMParser error:", - actualError, - ) - console.timeEnd("perf:validateMxCellStructure") return `Invalid XML: The XML contains syntax errors (likely unescaped special characters like <, >, & in attribute values). Please escape special characters: use < for <, > for >, & for &, " for ". Regenerate the diagram with properly escaped values.` } @@ -774,7 +764,6 @@ export function validateMxCellStructure(xml: string): string | null { for (const cell of allCells) { if (cell.parentElement?.tagName === "mxCell") { const id = cell.getAttribute("id") || "unknown" - console.timeEnd("perf:validateMxCellStructure") return `Invalid XML: Found nested mxCell (id="${id}"). Cells should be siblings, not nested inside other mxCell elements.` } } @@ -788,16 +777,12 @@ export function validateMxCellStructure(xml: string): string | null { // 1. Check for CDATA wrapper (invalid at document root) if (/^\s* from end" } // 2. Check for duplicate structural attributes - console.time("perf:checkDuplicateAttributes") const dupAttrError = checkDuplicateAttributes(xml) - console.timeEnd("perf:checkDuplicateAttributes") if (dupAttrError) { - console.timeEnd("perf:validateMxCellStructure") return dupAttrError } @@ -807,33 +792,25 @@ export function validateMxCellStructure(xml: string): string | null { while ((attrValMatch = attrValuePattern.exec(xml)) !== null) { const value = attrValMatch[1] if (/]*\sid\s*=\s*["']\s*["'][^>]*>/g.test(xml)) { - console.timeEnd("perf:validateMxCellStructure") return "Invalid XML: Found mxCell element(s) with empty id attribute" } // 10. Check for nested mxCell tags const nestedCellError = checkNestedMxCells(xml) if (nestedCellError) { - console.timeEnd("perf:validateMxCellStructure") return nestedCellError } - console.timeEnd("perf:validateMxCellStructure") return null }