feat: improve XML handling and edit_diagram tool

- Add formatXML function to format single-line XML with proper indentation
- Format chartXml after fetching to ensure consistency
- Update replaceXMLParts to handle single-line XML with substring fallback
- Improve edit_diagram tool guidance with SEARCH/REPLACE best practices
- Add concrete examples to help AI use minimal, targeted edits

app/api/chat/route.ts

@@ -9,6 +9,7 @@ import { createGoogleGenerativeAI } from '@ai-sdk/google';
 import { createOpenAI } from '@ai-sdk/openai';
 
 import { z } from "zod/v3";
+import { replaceXMLParts } from "@/lib/utils";
 
 export const maxDuration = 60
 const openrouter = createOpenRouter({ apiKey: process.env.OPENROUTER_API_KEY });
@@ -16,20 +17,27 @@ const openrouter = createOpenRouter({ apiKey: process.env.OPENROUTER_API_KEY });
 export async function POST(req: Request) {
   const body = await req.json();
 
-  const { messages, data = {} } = body;
+  // Extract messages and xml directly from the body
+  const { messages, xml } = body;
   const guide = readFileSync(resolve('./app/api/chat/xml_guide.md'), 'utf8');
 
   // Read and escape the guide content
   const systemMessage = `
 You are an expert diagram creation assistant specializing in draw.io XML generation. Your primary function is crafting clear, well-organized visual diagrams through precise XML specifications.
 You can see the image that user uploaded.
-You utilize the following tool:
+You utilize the following tools:
 ---Tool1---
 tool name: display_diagram
 description: Display a diagram on draw.io
 parameters: {
   xml: string
 }
+---Tool2---
+tool name: edit_diagram
+description: Edit specific parts of the current diagram
+parameters: {
+  edits: Array<{search: string, replace: string}>
+}
 ---End of tools---
 
 Core capabilities:
@@ -47,6 +55,12 @@ Note that:
 - **Don't** write out the XML string. Just return the XML string in the tool call.
 - If user asks you to replicate a diagram based on an image, remember to match the diagram style and layout as closely as possible. Especially, pay attention to the lines and shapes, for example, if the lines are straight or curved, and if the shapes are rounded or square.
 
+When using edit_diagram tool:
+- Keep edits minimal - only include the specific line being changed plus 1-2 context lines
+- Example GOOD edit: {"search": "  <mxCell id=\"2\" value=\"Old Text\">", "replace": "  <mxCell id=\"2\" value=\"New Text\">"}
+- Example BAD edit: Including 10+ unchanged lines just to change one attribute
+- For multiple changes, use separate edits: [{"search": "line1", "replace": "new1"}, {"search": "line2", "replace": "new2"}]
+
 here is a guide for the XML format: ${guide}
 `;
 
@@ -58,7 +72,7 @@ here is a guide for the XML format: ${guide}
   const formattedContent = `
 Current diagram XML:
 """xml
-${data.xml || ''}
+${xml || ''}
 """
 User input:
 """md
@@ -67,10 +81,10 @@ ${lastMessageText}
 
   // Convert UIMessages to ModelMessages and add system message
   const modelMessages = convertToModelMessages(messages);
-  let enhancedMessages = [{ role: "system" as const, content: systemMessage }, ...modelMessages];
+  let enhancedMessages = [...modelMessages];
 
   // Update the last message with formatted content if it's a user message
-  if (enhancedMessages.length > 1) {
+  if (enhancedMessages.length >= 1) {
     const lastModelMessage = enhancedMessages[enhancedMessages.length - 1];
     if (lastModelMessage.role === 'user') {
       enhancedMessages = [
@@ -80,12 +94,13 @@ ${lastMessageText}
     }
   }
 
-  // console.log("Enhanced messages:", enhancedMessages);
+  console.log("Enhanced messages:", enhancedMessages);
 
   const result = streamText({
     // model: google("gemini-2.5-flash-preview-05-20"),
     // model: google("gemini-2.5-pro"),
     // model: bedrock('anthropic.claude-sonnet-4-20250514-v1:0'),
+    system: systemMessage,
     model: openai.chat('gpt-5'),
     // model: openrouter('moonshotai/kimi-k2:free'),
     // model: model,
@@ -119,6 +134,21 @@ ${lastMessageText}
           xml: z.string().describe("XML string to be displayed on draw.io")
         })
       },
+      edit_diagram: {
+        description: `Edit specific parts of the current diagram by replacing exact line matches. Use this tool to make targeted fixes without regenerating the entire XML. 
+        
+IMPORTANT: Keep edits concise:
+- Only include the lines that are changing, plus 1-2 surrounding lines for context if needed
+- Break large changes into multiple smaller edits
+- Each search must contain complete lines (never truncate mid-line)
+- First match only - be specific enough to target the right element`,
+        inputSchema: z.object({
+          edits: z.array(z.object({
+            search: z.string().describe("Exact lines to search for (including whitespace and indentation)"),
+            replace: z.string().describe("Replacement lines")
+          })).describe("Array of search/replace pairs to apply sequentially")
+        })
+      },
     },
     temperature: 0,
   });

components/chat-message-display.tsx

@@ -96,6 +96,7 @@ export function ChatMessageDisplay({
         const callId = part.toolCallId;
         const { state, input } = part;
         const isExpanded = expandedTools[callId] ?? true;
+        const toolName = part.type?.replace("tool-", "");
 
         const toggleExpanded = () => {
             setExpandedTools((prev) => ({
@@ -111,7 +112,7 @@ export function ChatMessageDisplay({
             >
                 <div className="flex flex-col gap-2">
                     <div className="flex items-center justify-between">
-                        <div className="text-xs">Tool: display_diagram</div>
+                        <div className="text-xs">Tool: {toolName}</div>
                         {input && Object.keys(input).length > 0 && (
                             <button
                                 onClick={toggleExpanded}
@@ -133,11 +134,19 @@ export function ChatMessageDisplay({
                             <div className="h-4 w-4 border-2 border-primary border-t-transparent rounded-full animate-spin" />
                         ) : state === "output-available" ? (
                             <div className="text-green-600">
-                                Diagram generated
+                                {toolName === "display_diagram"
+                                    ? "Diagram generated"
+                                    : toolName === "edit_diagram"
+                                    ? "Diagram edited"
+                                    : "Tool executed"}
                             </div>
                         ) : state === "output-error" ? (
                             <div className="text-red-600">
-                                Error generating diagram
+                                {toolName === "display_diagram"
+                                    ? "Error generating diagram"
+                                    : toolName === "edit_diagram"
+                                    ? "Error editing diagram"
+                                    : "Tool error"}
                             </div>
                         ) : null}
                     </div>

components/chat-panel.tsx

@@ -16,7 +16,7 @@ import { DefaultChatTransport } from "ai";
 import { ChatInput } from "@/components/chat-input";
 import { ChatMessageDisplay } from "./chat-message-display";
 import { useDiagram } from "@/contexts/diagram-context";
-import { replaceNodes } from "@/lib/utils";
+import { replaceNodes, formatXML } from "@/lib/utils";
 
 export default function ChatPanel() {
     const {
@@ -70,6 +70,34 @@ export default function ChatPanel() {
                         toolCallId: toolCall.toolCallId,
                         output: "Successfully displayed the flowchart.",
                     });
+                } else if (toolCall.toolName === "edit_diagram") {
+                    const { edits } = toolCall.input as {
+                        edits: Array<{ search: string; replace: string }>;
+                    };
+
+                    try {
+                        // Fetch current chart XML
+                        const currentXml = await onFetchChart();
+
+                        // Apply edits using the utility function
+                        const { replaceXMLParts } = await import("@/lib/utils");
+                        const editedXml = replaceXMLParts(currentXml, edits);
+
+                        // Load the edited diagram
+                        onDisplayChart(editedXml);
+
+                        addToolResult({
+                            tool: "edit_diagram",
+                            toolCallId: toolCall.toolCallId,
+                            output: `Successfully applied ${edits.length} edit(s) to the diagram.`,
+                        });
+                    } catch (error) {
+                        addToolResult({
+                            tool: "edit_diagram",
+                            toolCallId: toolCall.toolCallId,
+                            output: `Error editing diagram: ${error}`,
+                        });
+                    }
                 }
             },
             onError: (error) => {
@@ -91,7 +119,10 @@ export default function ChatPanel() {
         if (input.trim() && status !== "streaming") {
             try {
                 // Fetch chart data before sending message
-                const chartXml = await onFetchChart();
+                let chartXml = await onFetchChart();
+
+                // Format the XML to ensure consistency
+                chartXml = formatXML(chartXml);
 
                 // Create message parts
                 const parts: any[] = [{ type: "text", text: input }];

lib/utils.ts

@@ -6,6 +6,53 @@ export function cn(...inputs: ClassValue[]) {
   return twMerge(clsx(inputs))
 }
 
+/**
+ * Format XML string with proper indentation and line breaks
+ * @param xml - The XML string to format
+ * @param indent - The indentation string (default: '  ')
+ * @returns Formatted XML string
+ */
+export function formatXML(xml: string, indent: string = '  '): string {
+  let formatted = '';
+  let pad = 0;
+
+  // Remove existing whitespace between tags
+  xml = xml.replace(/>\s*</g, '><').trim();
+
+  // Split on tags
+  const tags = xml.split(/(?=<)|(?<=>)/g).filter(Boolean);
+
+  tags.forEach((node) => {
+    if (node.match(/^<\/\w/)) {
+      // Closing tag - decrease indent
+      pad = Math.max(0, pad - 1);
+      formatted += indent.repeat(pad) + node + '\n';
+    } else if (node.match(/^<\w[^>]*[^\/]>.*$/)) {
+      // Opening tag
+      formatted += indent.repeat(pad) + node;
+      // Only add newline if next item is a tag
+      const nextIndex = tags.indexOf(node) + 1;
+      if (nextIndex < tags.length && tags[nextIndex].startsWith('<')) {
+        formatted += '\n';
+        if (!node.match(/^<\w[^>]*\/>$/)) {
+          pad++;
+        }
+      }
+    } else if (node.match(/^<\w[^>]*\/>$/)) {
+      // Self-closing tag
+      formatted += indent.repeat(pad) + node + '\n';
+    } else if (node.startsWith('<')) {
+      // Other tags (like <?xml)
+      formatted += indent.repeat(pad) + node + '\n';
+    } else {
+      // Text content
+      formatted += node;
+    }
+  });
+
+  return formatted.trim();
+}
+
 /** 
  * 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">),
@@ -129,7 +176,135 @@ export function replaceNodes(currentXML: string, nodes: string): string {
   }
 }
 
+/**
+ * Replace specific parts of XML content using search and replace pairs
+ * @param xmlContent - The original XML string
+ * @param searchReplacePairs - Array of {search: string, replace: string} objects
+ * @returns The updated XML string with replacements applied
+ */
+export function replaceXMLParts(
+  xmlContent: string,
+  searchReplacePairs: Array<{ search: string; replace: string }>
+): string {
+  // Format the XML first to ensure consistent line breaks
+  let result = formatXML(xmlContent);
+  let lastProcessedIndex = 0;
 
+  for (const { search, replace } of searchReplacePairs) {
+    // Also format the search content for consistency
+    const formattedSearch = formatXML(search);
+    const searchLines = formattedSearch.split('\n');
+
+    // Split into lines for exact line matching
+    const resultLines = result.split('\n');
+
+    // Remove trailing empty line if exists (from the trailing \n in search content)
+    if (searchLines[searchLines.length - 1] === '') {
+      searchLines.pop();
+    }
+
+    // Find the line number where lastProcessedIndex falls
+    let startLineNum = 0;
+    let currentIndex = 0;
+    while (currentIndex < lastProcessedIndex && startLineNum < resultLines.length) {
+      currentIndex += resultLines[startLineNum].length + 1; // +1 for \n
+      startLineNum++;
+    }
+
+    // Try to find exact match starting from lastProcessedIndex
+    let matchFound = false;
+    let matchStartLine = -1;
+    let matchEndLine = -1;
+
+    // First try: exact match
+    for (let i = startLineNum; i <= resultLines.length - searchLines.length; i++) {
+      let matches = true;
+
+      for (let j = 0; j < searchLines.length; j++) {
+        if (resultLines[i + j] !== searchLines[j]) {
+          matches = false;
+          break;
+        }
+      }
+
+      if (matches) {
+        matchStartLine = i;
+        matchEndLine = i + searchLines.length;
+        matchFound = true;
+        break;
+      }
+    }
+
+    // Second try: line-trimmed match (fallback)
+    if (!matchFound) {
+      for (let i = startLineNum; i <= resultLines.length - searchLines.length; i++) {
+        let matches = true;
+
+        for (let j = 0; j < searchLines.length; j++) {
+          const originalTrimmed = resultLines[i + j].trim();
+          const searchTrimmed = searchLines[j].trim();
+
+          if (originalTrimmed !== searchTrimmed) {
+            matches = false;
+            break;
+          }
+        }
+
+        if (matches) {
+          matchStartLine = i;
+          matchEndLine = i + searchLines.length;
+          matchFound = true;
+          break;
+        }
+      }
+    }
+
+    // Third try: substring match as last resort (for single-line XML)
+    if (!matchFound) {
+      // Try to find as a substring in the entire content
+      const searchStr = search.trim();
+      const resultStr = result;
+      const index = resultStr.indexOf(searchStr);
+
+      if (index !== -1) {
+        // Found as substring - replace it
+        result = resultStr.substring(0, index) + replace.trim() + resultStr.substring(index + searchStr.length);
+        // Re-format after substring replacement
+        result = formatXML(result);
+        continue; // Skip the line-based replacement below
+      }
+    }
+
+    if (!matchFound) {
+      throw new Error(`Search block not found:\n${search}\n...does not match anything in the file.`);
+    }
+
+    // Replace the matched lines
+    const replaceLines = replace.split('\n');
+
+    // Remove trailing empty line if exists
+    if (replaceLines[replaceLines.length - 1] === '') {
+      replaceLines.pop();
+    }
+
+    // Perform the replacement
+    const newResultLines = [
+      ...resultLines.slice(0, matchStartLine),
+      ...replaceLines,
+      ...resultLines.slice(matchEndLine)
+    ];
+
+    result = newResultLines.join('\n');
+
+    // Update lastProcessedIndex to the position after the replacement
+    lastProcessedIndex = 0;
+    for (let i = 0; i < matchStartLine + replaceLines.length; i++) {
+      lastProcessedIndex += newResultLines[i].length + 1;
+    }
+  }
+
+  return result;
+}
 
 export function extractDiagramXML(xml_svg_string: string): string {
   try {