refactor: improve diagram handling and error messaging in chat components

app/api/chat/route.ts

@@ -20,22 +20,27 @@ export async function POST(req: Request) {
 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.
-When you need to generate diagram about aws architecture, use AWS 2025 icons.
+Note that when you need to generate diagram about aws architecture, use **AWS 2025 icons**.
+
 You utilize the following tools:
 ---Tool1---
 tool name: display_diagram
-description: Display a diagram on draw.io
+description: Display a NEW diagram on draw.io. Use this when creating a diagram from scratch or when major structural changes are needed.
 parameters: {
   xml: string
 }
 ---Tool2---
 tool name: edit_diagram
-description: Edit specific parts of the current diagram
+description: Edit specific parts of the EXISTING diagram. Use this when making small targeted changes like adding/removing elements, changing labels, or adjusting properties. This is more efficient than regenerating the entire diagram.
 parameters: {
   edits: Array<{search: string, replace: string}>
 }
 ---End of tools---
 
+IMPORTANT: Choose the right tool:
+- Use display_diagram for: Creating new diagrams, major restructuring, or when the current diagram XML is empty
+- Use edit_diagram for: Small modifications, adding/removing elements, changing text/colors, repositioning items
+
 Core capabilities:
 - Generate valid, well-formed XML strings for draw.io diagrams
 - Create professional flowcharts, mind maps, entity diagrams, and technical illustrations
@@ -45,17 +50,29 @@ Core capabilities:
 - Optimize element positioning to prevent overlapping and maintain readability
 - Structure complex systems into clear, organized visual components
 
+Layout constraints:
+- CRITICAL: Keep all diagram elements within a single page viewport to avoid page breaks
+- Position all elements with x coordinates between 0-800 and y coordinates between 0-600
+- Maximum width for containers (like AWS cloud boxes): 700 pixels
+- Maximum height for containers: 550 pixels
+- Use compact, efficient layouts that fit the entire diagram in one view
+- Start positioning from reasonable margins (e.g., x=40, y=40) and keep elements grouped closely
+- For large diagrams with many elements, use vertical stacking or grid layouts that stay within bounds
+- Avoid spreading elements too far apart horizontally - users should see the complete diagram without a page break line
+
 Note that:
 - Focus on producing clean, professional diagrams that effectively communicate the intended information through thoughtful layout and design choices.
 - When artistic drawings are requested, creatively compose them using standard diagram shapes and connectors while maintaining visual clarity.
 - Return XML only via tool calls, never in text responses.
 - 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.
+- Note that when you need to generate diagram about aws architecture, use **AWS 2025 icons**.
 
 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"}]
+- CRITICAL: If edit_diagram fails because the search pattern cannot be found, fall back to using display_diagram to regenerate the entire diagram with your changes. Do NOT keep trying edit_diagram with different search patterns.
 `;
 
     const lastMessage = messages[messages.length - 1];
@@ -140,7 +157,9 @@ ${lastMessageText}
             <mxCell id="2" value="Hello, World!" style="shape=rectangle" parent="1">
               <mxGeometry x="20" y="20" width="100" height="100" as="geometry"/>
             </mxCell>
-          </root>`,
+          </root>
+          - Note that when you need to generate diagram about aws architecture, use **AWS 2025 icons**.
+          `,
           inputSchema: z.object({
             xml: z.string().describe("XML string to be displayed on draw.io")
           })

components/chat-message-display.tsx

@@ -94,7 +94,7 @@ export function ChatMessageDisplay({
 
     const renderToolPart = (part: any) => {
         const callId = part.toolCallId;
-        const { state, input } = part;
+        const { state, input, output } = part;
         const isExpanded = expandedTools[callId] ?? true;
         const toolName = part.type?.replace("tool-", "");
 
@@ -134,19 +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">
-                                {toolName === "display_diagram"
+                                {output || (toolName === "display_diagram"
                                     ? "Diagram generated"
                                     : toolName === "edit_diagram"
                                     ? "Diagram edited"
-                                    : "Tool executed"}
+                                    : "Tool executed")}
                             </div>
                         ) : state === "output-error" ? (
                             <div className="text-red-600">
-                                {toolName === "display_diagram"
+                                {output || (toolName === "display_diagram"
                                     ? "Error generating diagram"
                                     : toolName === "edit_diagram"
                                     ? "Error editing diagram"
-                                    : "Tool error"}
+                                    : "Tool error")}
                             </div>
                         ) : null}
                     </div>

components/chat-panel.tsx

@@ -65,24 +65,21 @@ export default function ChatPanel() {
             }),
             async onToolCall({ toolCall }) {
                 if (toolCall.toolName === "display_diagram") {
-                    const { xml } = toolCall.input as { xml: string };
+                    // Diagram is handled streamingly in the ChatMessageDisplay component
-                    // do nothing because we will handle this streamingly in the ChatMessageDisplay component
-                    // onDisplayChart(replaceNodes(chartXML, xml));
-
-                    // Use addToolResult instead of returning a value
                     addToolResult({
                         tool: "display_diagram",
                         toolCallId: toolCall.toolCallId,
-                        output: "Successfully displayed the flowchart.",
+                        output: "Successfully displayed the diagram.",
                     });
                 } else if (toolCall.toolName === "edit_diagram") {
                     const { edits } = toolCall.input as {
                         edits: Array<{ search: string; replace: string }>;
                     };
 
+                    let currentXml = '';
                     try {
                         // Fetch current chart XML
-                        const currentXml = await onFetchChart();
+                        currentXml = await onFetchChart();
 
                         // Apply edits using the utility function
                         const { replaceXMLParts } = await import("@/lib/utils");
@@ -97,10 +94,14 @@ export default function ChatPanel() {
                             output: `Successfully applied ${edits.length} edit(s) to the diagram.`,
                         });
                     } catch (error) {
+                        console.error("Edit diagram failed:", error);
+
+                        const errorMessage = error instanceof Error ? error.message : String(error);
+
                         addToolResult({
                             tool: "edit_diagram",
                             toolCallId: toolCall.toolCallId,
-                            output: `Error editing diagram: ${error}`,
+                            output: `Failed to edit diagram: ${errorMessage}`,
                         });
                     }
                 }

lib/utils.ts

@@ -276,7 +276,7 @@ export function replaceXMLParts(
     }
 
     if (!matchFound) {
-      throw new Error(`Search block not found:\n${search}\n...does not match anything in the file.`);
+      throw new Error(`Search pattern not found in the diagram. The pattern may not exist in the current structure.`);
     }
 
     // Replace the matched lines