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

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 }];