fix: address code review issues

- Remove verbose debug console.logs (kept only token usage log)
- Add basic input validation for messages array
- Fix XML comparison using formatXML for consistency
- Extract CACHED_TOOL_PREFIX constant to avoid magic string
- Fix addToolResult API to use correct tool/output params

app/api/chat/route.ts

@@ -1,10 +1,14 @@
 import { streamText, convertToModelMessages, createUIMessageStream, createUIMessageStreamResponse } from 'ai';
 import { getAIModel } from '@/lib/ai-providers';
 import { findCachedResponse } from '@/lib/cached-responses';
+import { formatXML } from '@/lib/utils';
 import { z } from "zod";
 
 export const maxDuration = 300;
 
+// Prefix for cached tool call IDs (used by client to detect cached responses)
+export const CACHED_TOOL_PREFIX = 'cached-';
+
 // Helper function to check if diagram is minimal/empty
 function isMinimalDiagram(xml: string): boolean {
   const stripped = xml.replace(/\s/g, '');
@@ -13,14 +17,18 @@ function isMinimalDiagram(xml: string): boolean {
 
 // Helper function to create cached stream response
 function createCachedStreamResponse(xml: string): Response {
-  const toolCallId = `cached-${Date.now()}`;
+  const toolCallId = `${CACHED_TOOL_PREFIX}${Date.now()}`;
 
   const stream = createUIMessageStream({
     execute: async ({ writer }) => {
       writer.write({ type: 'start' });
       writer.write({ type: 'tool-input-start', toolCallId, toolName: 'display_diagram' });
-      writer.write({ type: 'tool-input-delta', toolCallId, inputTextDelta: xml });
+      // Stream the XML as JSON input so it matches the tool schema exactly
+      writer.write({ type: 'tool-input-delta', toolCallId, inputTextDelta: JSON.stringify({ xml }) });
+      // Input must match the tool schema (only xml field, no extra fields like fromCache)
       writer.write({ type: 'tool-input-available', toolCallId, toolName: 'display_diagram', input: { xml } });
+      // Include tool output so the message is complete for follow-up conversations
+      writer.write({ type: 'tool-output-available', toolCallId, output: 'Successfully displayed the diagram.' });
       writer.write({ type: 'finish' });
     },
   });
@@ -30,7 +38,12 @@ function createCachedStreamResponse(xml: string): Response {
 
 export async function POST(req: Request) {
   try {
-    const { messages, xml } = await req.json();
+    const { messages, xml, lastGeneratedXml } = await req.json();
+
+    // Basic validation for demo app
+    if (!messages || !Array.isArray(messages) || messages.length === 0) {
+      return Response.json({ error: 'Invalid messages' }, { status: 400 });
+    }
 
     // === CACHE CHECK START ===
     const isFirstMessage = messages.length === 1;
@@ -74,6 +87,7 @@ parameters: {
 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
+- When using edit_diagram: If the current diagram XML is provided in the user message context, use it as the source of truth for constructing search patterns. If no XML is provided, you can use your memory of the diagram structure.
 
 Core capabilities:
 - Generate valid, well-formed XML strings for draw.io diagrams
@@ -120,37 +134,32 @@ When using edit_diagram tool:
     // Extract file parts (images) from the last message
     const fileParts = lastMessage.parts?.filter((part: any) => part.type === 'file') || [];
 
-    const formattedTextContent = `
-Current diagram XML:
-"""xml
-${xml || ''}
-"""
-User input:
+    // Check diagram state - use formatted XML for reliable comparison
+    const hasDiagram = xml && !isMinimalDiagram(xml);
+    const noHistory = !lastGeneratedXml || lastGeneratedXml.trim() === '';
+    const formattedXml = hasDiagram ? formatXML(xml) : '';
+    const formattedLastGenXml = lastGeneratedXml ? formatXML(lastGeneratedXml) : '';
+    const userModified = hasDiagram && formattedLastGenXml && formattedXml !== formattedLastGenXml;
+
+    // Build context based on diagram state
+    let diagramContext = '';
+    if (hasDiagram && noHistory) {
+      // No history (e.g., cached response) - include XML directly
+      diagramContext = `\n\n[Current diagram XML - use this as source of truth for edits:]\n\`\`\`xml\n${xml}\n\`\`\``;
+    } else if (userModified) {
+      // User modified - include XML
+      diagramContext = `\n\n[User modified the diagram. Current XML:]\n\`\`\`xml\n${xml}\n\`\`\``;
+    }
+    // If unchanged and has history, agent can use memory (no XML sent = save tokens)
+
+    const formattedTextContent = `User input:
 """md
 ${lastMessageText}
-"""`;
+"""${diagramContext}`;
 
     // Convert UIMessages to ModelMessages and add system message
     const modelMessages = convertToModelMessages(messages);
 
-    // Log messages with empty content for debugging (helps identify root cause)
-    const emptyMessages = modelMessages.filter((msg: any) =>
-      !msg.content || !Array.isArray(msg.content) || msg.content.length === 0
-    );
-    if (emptyMessages.length > 0) {
-      console.warn('[Chat API] Messages with empty content detected:',
-        JSON.stringify(emptyMessages.map((m: any) => ({ role: m.role, contentLength: m.content?.length })))
-      );
-      console.warn('[Chat API] Original UI messages structure:',
-        JSON.stringify(messages.map((m: any) => ({
-          id: m.id,
-          role: m.role,
-          partsCount: m.parts?.length,
-          partTypes: m.parts?.map((p: any) => p.type)
-        })))
-      );
-    }
-
     // Filter out messages with empty content arrays (Bedrock API rejects these)
     // This is a safety measure - ideally convertToModelMessages should handle all cases
     let enhancedMessages = modelMessages.filter((msg: any) =>
@@ -217,13 +226,8 @@ ${lastMessageText}
       messages: [systemMessageWithCache, ...enhancedMessages],
       ...(providerOptions && { providerOptions }),
       ...(headers && { headers }),
-      onFinish: ({ usage, providerMetadata }) => {
-        console.log('[Cache] Usage:', JSON.stringify({
-          inputTokens: usage?.inputTokens,
-          outputTokens: usage?.outputTokens,
-          cachedInputTokens: usage?.cachedInputTokens,
-        }, null, 2));
-        console.log('[Cache] Provider metadata:', JSON.stringify(providerMetadata, null, 2));
+      onFinish: ({ usage }) => {
+        console.log('[API] Tokens:', usage?.inputTokens, 'in /', usage?.outputTokens, 'out, cached:', usage?.cachedInputTokens);
       },
       tools: {
         // Client-side tool that will be executed on the client
@@ -261,6 +265,7 @@ IMPORTANT: Keep edits concise:
         },
       },
       temperature: 0,
+      maxSteps: 5, // Allow model to continue after server-side tool execution
     });
 
     // Error handler function to provide detailed error messages

components/chat-message-display.tsx

@@ -157,7 +157,9 @@ 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">
-                                {output || (toolName === "display_diagram"
+                                {typeof output === "object" && output !== null
+                                    ? (output as any).message || JSON.stringify(output)
+                                    : output || (toolName === "display_diagram"
                                         ? "Diagram generated"
                                         : toolName === "edit_diagram"
                                             ? "Diagram edited"
@@ -165,7 +167,9 @@ export function ChatMessageDisplay({
                             </div>
                         ) : state === "output-error" ? (
                             <div className="text-red-600">
-                                {output || (toolName === "display_diagram"
+                                {typeof output === "object" && output !== null
+                                    ? (output as any).message || JSON.stringify(output)
+                                    : output || (toolName === "display_diagram"
                                         ? "Error generating diagram"
                                         : toolName === "edit_diagram"
                                             ? "Error editing diagram"

components/chat-panel.tsx

@@ -19,6 +19,7 @@ import { ChatInput } from "@/components/chat-input";
 import { ChatMessageDisplay } from "./chat-message-display";
 import { useDiagram } from "@/contexts/diagram-context";
 import { replaceNodes, formatXML } from "@/lib/utils";
+import { CACHED_TOOL_PREFIX } from "@/app/api/chat/route";
 import { ButtonWithTooltip } from "@/components/button-with-tooltip";
 
 interface ChatPanelProps {
@@ -33,6 +34,8 @@ export default function ChatPanel({ isVisible, onToggleVisibility }: ChatPanelPr
         resolverRef,
         chartXML,
         clearDiagram,
+        getLastAgentGeneratedXml,
+        markAgentDiagramPending,
     } = useDiagram();
 
     const onFetchChart = () => {
@@ -73,7 +76,15 @@ export default function ChatPanel({ isVisible, onToggleVisibility }: ChatPanelPr
             }),
             async onToolCall({ toolCall }) {
                 if (toolCall.toolName === "display_diagram") {
-                    // Diagram is handled streamingly in the ChatMessageDisplay component
+                    // Check if this is a cached response by looking at the toolCallId prefix
+                    const isCached = toolCall.toolCallId.startsWith(CACHED_TOOL_PREFIX);
+
+                    // Only mark as pending if agent actually generated it (not cached)
+                    // This ensures lastAgentGeneratedXml stays empty for cached responses
+                    if (!isCached) {
+                        markAgentDiagramPending();
+                    }
+
                     addToolResult({
                         tool: "display_diagram",
                         toolCallId: toolCall.toolCallId,
@@ -96,6 +107,9 @@ export default function ChatPanel({ isVisible, onToggleVisibility }: ChatPanelPr
                         // Load the edited diagram
                         onDisplayChart(editedXml);
 
+                        // Mark that an agent diagram is pending - the next export will update lastAgentGeneratedXml
+                        markAgentDiagramPending();
+
                         addToolResult({
                             tool: "edit_diagram",
                             toolCallId: toolCall.toolCallId,
@@ -134,10 +148,6 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
         }
     }, [messages]);
 
-    // Debug: Log status changes
-    useEffect(() => {
-        console.log('[ChatPanel] Status changed to:', status);
-    }, [status]);
 
     const onFormSubmit = async (e: React.FormEvent<HTMLFormElement>) => {
         e.preventDefault();
@@ -171,11 +181,14 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
                     }
                 }
 
+                const lastGenXml = getLastAgentGeneratedXml();
+
                 sendMessage(
                     { parts },
                     {
                         body: {
                             xml: chartXml,
+                            lastGeneratedXml: lastGenXml,
                         },
                     }
                 );

contexts/diagram-context.tsx

@@ -2,12 +2,16 @@
 
 import React, { createContext, useContext, useRef, useState } from "react";
 import type { DrawIoEmbedRef } from "react-drawio";
-import { extractDiagramXML } from "../lib/utils";
+import { extractDiagramXML, formatXML } from "../lib/utils";
 
 interface DiagramContextType {
     chartXML: string;
     latestSvg: string;
     diagramHistory: { svg: string; xml: string }[];
+    lastAgentGeneratedXml: string;
+    getLastAgentGeneratedXml: () => string;
+    setLastAgentGeneratedXml: (xml: string) => void;
+    markAgentDiagramPending: () => void;
     loadDiagram: (chart: string) => void;
     handleExport: () => void;
     resolverRef: React.Ref<((value: string) => void) | null>;
@@ -24,9 +28,25 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
     const [diagramHistory, setDiagramHistory] = useState<
         { svg: string; xml: string }[]
     >([]);
+    const [lastAgentGeneratedXml, setLastAgentGeneratedXmlState] = useState<string>("");
+    const lastAgentGeneratedXmlRef = useRef<string>("");
+    const agentDiagramPendingRef = useRef<boolean>(false);
     const drawioRef = useRef<DrawIoEmbedRef | null>(null);
     const resolverRef = useRef<((value: string) => void) | null>(null);
 
+    // Wrapper to keep ref and state in sync
+    const setLastAgentGeneratedXml = (xml: string) => {
+        lastAgentGeneratedXmlRef.current = xml;
+        setLastAgentGeneratedXmlState(xml);
+    };
+
+    // Getter that returns the ref value (always up-to-date, even in async contexts)
+    const getLastAgentGeneratedXml = () => lastAgentGeneratedXmlRef.current;
+
+    const markAgentDiagramPending = () => {
+        agentDiagramPendingRef.current = true;
+    };
+
     const handleExport = () => {
         if (drawioRef.current) {
             drawioRef.current.exportDiagram({
@@ -54,6 +74,15 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
                 xml: extractedXML,
             },
         ]);
+
+        // If agent just generated a diagram, update lastAgentGeneratedXml with the exported XML
+        // This ensures we compare apples-to-apples (both formatted the same way)
+        if (agentDiagramPendingRef.current) {
+            const formatted = formatXML(extractedXML);
+            setLastAgentGeneratedXml(formatted);
+            agentDiagramPendingRef.current = false;
+        }
+
         if (resolverRef.current) {
             resolverRef.current(extractedXML);
             resolverRef.current = null;
@@ -66,6 +95,7 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
         setChartXML(emptyDiagram);
         setLatestSvg("");
         setDiagramHistory([]);
+        setLastAgentGeneratedXml("");
     };
 
     return (
@@ -74,6 +104,10 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
                 chartXML,
                 latestSvg,
                 diagramHistory,
+                lastAgentGeneratedXml,
+                getLastAgentGeneratedXml,
+                setLastAgentGeneratedXml,
+                markAgentDiagramPending,
                 loadDiagram,
                 handleExport,
                 resolverRef,