fix: flash problem

app/api/chat/route.ts

@@ -75,7 +75,16 @@ ${lastMessage.content}
     tools: {
       // Client-side tool that will be executed on the client
       display_diagram: {
-        description: "Display a diagram on draw.io",
+        description: `Display a diagram on draw.io. You only need to pass the nodes inside the <root> tag (including the <root> tag itself) in the XML string.
+        For example:
+        <root>
+          <mxCell id="0"/>
+          <mxCell id="1" parent="0"/>
+          <mxGeometry x="20" y="20" width="100" height="100" as="geometry"/>
+          <mxCell id="2" value="Hello, World!" style="shape=rectangle" parent="1">
+            <mxGeometry x="20" y="20" width="100" height="100" as="geometry"/>
+          </mxCell>
+        </root>`,
         parameters: z.object({
           xml: z.string().describe("XML string to be displayed on draw.io")
         })

app/page.tsx

@@ -2,7 +2,6 @@
 import { DrawIoEmbed, DrawIoEmbedRef } from "react-drawio";
 
 import { useRef, useState } from "react";
-import { Button } from "@/components/ui/button";
 import { extractDiagramXML } from "./extract_xml";
 import ChatPanel from "@/components/chat-panel";
 
@@ -28,6 +27,7 @@ export default function Home() {
 
     const loadDiagram = (chart: string) => {
         if (drawioRef.current) {
+            console.log("xml before load", chart);
             drawioRef.current.load({
                 xml: chart,
             });
@@ -65,6 +65,7 @@ export default function Home() {
             </div>
             <div className="w-1/3 p-1 border-gray-300">
                 <ChatPanel
+                    chartXML={chartXML}
                     onDisplayChart={(xml) => loadDiagram(xml)}
                     onFetchChart={() => {
                         return new Promise<string>((resolve) => {

components/chat-message-display.tsx

@@ -6,9 +6,10 @@ import Image from "next/image";
 import { ScrollArea } from "@/components/ui/scroll-area";
 import ExamplePanel from "./chat-example-panel";
 import { Message } from "ai";
-import { convertToLegalXml } from "@/lib/utils";
+import { convertToLegalXml, replaceNodes } from "@/lib/utils";
 
 interface ChatMessageDisplayProps {
+    chartXML: string;
     messages: Message[];
     error?: Error | null;
     setInput: (input: string) => void;
@@ -17,6 +18,7 @@ interface ChatMessageDisplayProps {
 }
 
 export function ChatMessageDisplay({
+    chartXML,
     messages,
     error,
     setInput,
@@ -24,7 +26,7 @@ export function ChatMessageDisplay({
     onDisplayChart,
 }: ChatMessageDisplayProps) {
     const messagesEndRef = useRef<HTMLDivElement>(null);
-    const stepCounterRef = useRef<number>(0);
+    const previousXML = useRef<string>("");
     useEffect(() => {
         if (messagesEndRef.current) {
             messagesEndRef.current.scrollIntoView({ behavior: "smooth" });
@@ -42,19 +44,26 @@ export function ChatMessageDisplay({
                             const currentXml = toolInvocation.args?.xml || "";
 
                             // Increment the step counter
-                            stepCounterRef.current += 1;
 
                             // Determine whether to show details based on a simple threshold
-                            if (stepCounterRef.current % 20 === 0) {
-                                const convertedXml =
-                                    convertToLegalXml(currentXml);
+                            const convertedXml = convertToLegalXml(currentXml);
+                            if (convertedXml !== previousXML.current) {
+                                previousXML.current = convertedXml;
                                 // if "/root" in convertedXml
-                                if (convertedXml.includes("/root")) {
-                                    onDisplayChart(convertedXml);
+                                const replacedXML = replaceNodes(
+                                    chartXML,
+                                    convertedXml
+                                );
+                                console.log("currentXml", currentXml);
                                 console.log("converted xml", convertedXml);
-                                }
+                                console.log("replaced xml", replacedXML);
+                                onDisplayChart(replacedXML);
+
                                 // if convertedXml changed
                             }
+                            // if "/root" in convertedXml
+
+                            // if convertedXml changed
                         }
                         return (
                             <div

components/chat-panel.tsx

@@ -14,6 +14,7 @@ import { useChat } from "@ai-sdk/react";
 import { ChatInput } from "@/components/chat-input";
 import { ChatMessageDisplay } from "./chat-message-display";
 interface ChatPanelProps {
+    chartXML: string;
     onDisplayChart: (xml: string) => void;
     onFetchChart: () => Promise<string>;
     diagramHistory?: { svg: string; xml: string }[];
@@ -21,9 +22,9 @@ interface ChatPanelProps {
 }
 
 export default function ChatPanel({
+    chartXML,
     onDisplayChart,
     onFetchChart,
-    mergeXML,
     diagramHistory = [],
     onAddToHistory = () => {},
 }: ChatPanelProps) {
@@ -49,7 +50,8 @@ export default function ChatPanel({
         async onToolCall({ toolCall }) {
             if (toolCall.toolName === "display_diagram") {
                 const { xml } = toolCall.args as { xml: string };
-                onDisplayChart(xml);
+                // do nothing because we will handle this streamingly in the ChatMessageDisplay component
+                // onDisplayChart(xml);
                 return "Successfully displayed the flowchart.";
             }
         },
@@ -105,6 +107,7 @@ export default function ChatPanel({
             </CardHeader>
             <CardContent className="flex-grow overflow-hidden px-2">
                 <ChatMessageDisplay
+                    chartXML={chartXML}
                     messages={messages}
                     error={error}
                     setInput={setInput}

lib/utils.ts

@@ -6,55 +6,124 @@ export function cn(...inputs: ClassValue[]) {
 }
 
 /** 
- * Efficiently converts a potentially incomplete XML string to a legal XML string
- * by closing any open tags properly.
- * 
+ * Efficiently converts a potentially incomplete XML string to a legal XML string by closing any open tags properly.
+ * Additionally, if an <mxCell> tag does not have an mxGeometry child (e.g. <mxCell id="3">),
+ * it removes that tag from the output.
  * @param xmlString The potentially incomplete XML string
- * @returns A legal XML string with properly closed tags
+ * @returns A legal XML string with properly closed tags and removed incomplete mxCell elements.
  */
 export function convertToLegalXml(xmlString: string): string {
-  const stack: string[] = [];
-  let result = '';
-  let tagStart = -1;
+  // This regex will match either self-closing <mxCell .../> or a block element 
+  // <mxCell ...> ... </mxCell>. Unfinished ones are left out because they don't match.
+  const regex = /<mxCell\b[^>]*(?:\/>|>([\s\S]*?)<\/mxCell>)/g;
+  let match: RegExpExecArray | null;
+  let result = "<root>\n";
 
-  for (let i = 0; i < xmlString.length; i++) {
-    const char = xmlString[i];
-    result += char;
-
-    if (char === '<' && tagStart === -1) {
-      // Start of a new tag
-      tagStart = i;
-    } else if (char === '>' && tagStart !== -1) {
-      // End of a tag
-      const tagContent = xmlString.substring(tagStart + 1, i);
-
-      if (tagContent.startsWith('/')) {
-        // Closing tag
-        const tagName = tagContent.substring(1).trim().split(/\s+/)[0];
-        if (stack.length && stack[stack.length - 1] === tagName) {
-          stack.pop();
-        }
-      } else if (!tagContent.endsWith('/') && !tagContent.startsWith('?') && !tagContent.startsWith('!')) {
-        // Opening tag (not self-closing, processing instruction, or comment)
-        const tagName = tagContent.trim().split(/\s+/)[0];
-        stack.push(tagName);
-      }
-
-      tagStart = -1;
-    }
-  }
-
-  // If we have an incomplete tag at the end, don't include it in the result
-  if (tagStart !== -1) {
-    result = result.substring(0, tagStart);
-  }
-
-  // Close all remaining open tags
-  for (let j = stack.length - 1; j >= 0; j--) {
-    result += `</${stack[j]}>`;
+  while ((match = regex.exec(xmlString)) !== null) {
+    // match[0] contains the entire matched mxCell block
+    // Indent each line of the matched block for readability.
+    const formatted = match[0].split('\n').map(line => "    " + line.trim()).join('\n');
+    result += formatted + "\n";
   }
+  result += "</root>";
 
   return result;
 }
 
 
+/**
+ * Replace nodes in a Draw.io XML diagram
+ * @param currentXML - The original Draw.io XML string
+ * @param nodes - The XML string containing new nodes to replace in the diagram
+ * @returns The updated XML string with replaced nodes
+ */
+export function replaceNodes(currentXML: string, nodes: string): string {
+  // Check for valid inputs
+  if (!currentXML || !nodes) {
+    throw new Error("Both currentXML and nodes must be provided");
+  }
+
+  try {
+    // Parse the XML strings to create DOM objects
+    const parser = new DOMParser();
+    const currentDoc = parser.parseFromString(currentXML, "text/xml");
+
+    // Handle nodes input - if it doesn't contain <root>, wrap it
+    let nodesString = nodes;
+    if (!nodes.includes("<root>")) {
+      nodesString = `<root>${nodes}</root>`;
+    }
+
+    const nodesDoc = parser.parseFromString(nodesString, "text/xml");
+
+    // Find the root element in the current document
+    let currentRoot = currentDoc.querySelector("mxGraphModel > root");
+    if (!currentRoot) {
+      // If no root element is found, create the proper structure
+      const mxGraphModel = currentDoc.querySelector("mxGraphModel") ||
+        currentDoc.createElement("mxGraphModel");
+
+      if (!currentDoc.contains(mxGraphModel)) {
+        currentDoc.appendChild(mxGraphModel);
+      }
+
+      currentRoot = currentDoc.createElement("root");
+      mxGraphModel.appendChild(currentRoot);
+    }
+
+    // Find the root element in the nodes document
+    const nodesRoot = nodesDoc.querySelector("root");
+    if (!nodesRoot) {
+      throw new Error("Invalid nodes: Could not find or create <root> element");
+    }
+
+    // Clear all existing child elements from the current root
+    while (currentRoot.firstChild) {
+      currentRoot.removeChild(currentRoot.firstChild);
+    }
+
+    // Ensure the base cells exist
+    const hasCell0 = Array.from(nodesRoot.childNodes).some(
+      node => node.nodeName === "mxCell" &&
+        (node as Element).getAttribute("id") === "0"
+    );
+
+    const hasCell1 = Array.from(nodesRoot.childNodes).some(
+      node => node.nodeName === "mxCell" &&
+        (node as Element).getAttribute("id") === "1"
+    );
+
+    // Copy all child nodes from the nodes root to the current root
+    Array.from(nodesRoot.childNodes).forEach(node => {
+      const importedNode = currentDoc.importNode(node, true);
+      currentRoot.appendChild(importedNode);
+    });
+
+    // Add default cells if they don't exist
+    if (!hasCell0) {
+      const cell0 = currentDoc.createElement("mxCell");
+      cell0.setAttribute("id", "0");
+      currentRoot.insertBefore(cell0, currentRoot.firstChild);
+    }
+
+    if (!hasCell1) {
+      const cell1 = currentDoc.createElement("mxCell");
+      cell1.setAttribute("id", "1");
+      cell1.setAttribute("parent", "0");
+
+      // Insert after cell0 if possible
+      const cell0 = currentRoot.querySelector('mxCell[id="0"]');
+      if (cell0 && cell0.nextSibling) {
+        currentRoot.insertBefore(cell1, cell0.nextSibling);
+      } else {
+        currentRoot.appendChild(cell1);
+      }
+    }
+
+    // Convert the modified DOM back to a string
+    const serializer = new XMLSerializer();
+    return serializer.serializeToString(currentDoc);
+  } catch (error) {
+    throw new Error(`Error replacing nodes: ${error}`);
+  }
+}