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,7 +17,7 @@ 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 }) => {
@@ -36,6 +40,11 @@ export async function POST(req: Request) {
   try {
     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;
     const isEmptyDiagram = !xml || xml.trim() === '' || isMinimalDiagram(xml);
@@ -125,10 +134,12 @@ When using edit_diagram tool:
     // Extract file parts (images) from the last message
     const fileParts = lastMessage.parts?.filter((part: any) => part.type === 'file') || [];
 
-    // Check diagram state
+    // Check diagram state - use formatted XML for reliable comparison
     const hasDiagram = xml && !isMinimalDiagram(xml);
     const noHistory = !lastGeneratedXml || lastGeneratedXml.trim() === '';
-    const userModified = hasDiagram && lastGeneratedXml && xml !== lastGeneratedXml;
+    const formattedXml = hasDiagram ? formatXML(xml) : '';
+    const formattedLastGenXml = lastGeneratedXml ? formatXML(lastGeneratedXml) : '';
+    const userModified = hasDiagram && formattedLastGenXml && formattedXml !== formattedLastGenXml;
 
     // Build context based on diagram state
     let diagramContext = '';
@@ -149,28 +160,6 @@ ${lastMessageText}
     // Convert UIMessages to ModelMessages and add system message
     const modelMessages = convertToModelMessages(messages);
 
-    // Debug: Log the full structure of model messages to diagnose Bedrock API errors
-    console.log('[Debug] Model messages structure:');
-    console.log('[Debug] Full messages JSON:', JSON.stringify(modelMessages, null, 2));
-
-    // 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) =>
@@ -237,19 +226,8 @@ ${lastMessageText}
       messages: [systemMessageWithCache, ...enhancedMessages],
       ...(providerOptions && { providerOptions }),
       ...(headers && { headers }),
-      onStepFinish: ({ stepType, toolCalls, toolResults }) => {
-        console.log('[Step] Type:', stepType);
-        console.log('[Step] Tool calls:', toolCalls?.map(t => t.toolName));
-        console.log('[Step] Tool results:', toolResults?.length);
-      },
-      onFinish: ({ usage, providerMetadata, steps }) => {
-        console.log('[Finish] Total steps:', steps?.length);
-        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

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 {
@@ -76,7 +77,7 @@ export default function ChatPanel({ isVisible, onToggleVisibility }: ChatPanelPr
             async onToolCall({ toolCall }) {
                 if (toolCall.toolName === "display_diagram") {
                     // Check if this is a cached response by looking at the toolCallId prefix
-                    const isCached = toolCall.toolCallId.startsWith('cached-');
+                    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
@@ -85,8 +86,9 @@ export default function ChatPanel({ isVisible, onToggleVisibility }: ChatPanelPr
                     }
 
                     addToolResult({
+                        tool: "display_diagram",
                         toolCallId: toolCall.toolCallId,
-                        result: "Successfully displayed the diagram.",
+                        output: "Successfully displayed the diagram.",
                     });
                 } else if (toolCall.toolName === "edit_diagram") {
                     const { edits } = toolCall.input as {
@@ -109,8 +111,9 @@ export default function ChatPanel({ isVisible, onToggleVisibility }: ChatPanelPr
                         markAgentDiagramPending();
 
                         addToolResult({
+                            tool: "edit_diagram",
                             toolCallId: toolCall.toolCallId,
-                            result: `Successfully applied ${edits.length} edit(s) to the diagram.`,
+                            output: `Successfully applied ${edits.length} edit(s) to the diagram.`,
                         });
                     } catch (error) {
                         console.error("Edit diagram failed:", error);
@@ -119,8 +122,9 @@ export default function ChatPanel({ isVisible, onToggleVisibility }: ChatPanelPr
 
                         // Provide detailed error with current diagram XML
                         addToolResult({
+                            tool: "edit_diagram",
                             toolCallId: toolCall.toolCallId,
-                            result: `Edit failed: ${errorMessage}
+                            output: `Edit failed: ${errorMessage}
 
 Current diagram XML:
 \`\`\`xml
@@ -144,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();
@@ -182,9 +182,6 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
                 }
 
                 const lastGenXml = getLastAgentGeneratedXml();
-                console.log('[ChatPanel] Sending message with xml length:', chartXml.length);
-                console.log('[ChatPanel] lastGeneratedXml length:', lastGenXml.length);
-                console.log('[ChatPanel] Are they equal:', chartXml === lastGenXml);
 
                 sendMessage(
                     { parts },

contexts/diagram-context.tsx

@@ -44,7 +44,6 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
     const getLastAgentGeneratedXml = () => lastAgentGeneratedXmlRef.current;
 
     const markAgentDiagramPending = () => {
-        console.log('[DiagramContext] markAgentDiagramPending called');
         agentDiagramPendingRef.current = true;
     };
 
@@ -80,7 +79,6 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
         // This ensures we compare apples-to-apples (both formatted the same way)
         if (agentDiagramPendingRef.current) {
             const formatted = formatXML(extractedXML);
-            console.log('[DiagramContext] Setting lastAgentGeneratedXml from export, length:', formatted.length);
             setLastAgentGeneratedXml(formatted);
             agentDiagramPendingRef.current = false;
         }

lib/utils.ts

@@ -190,15 +190,9 @@ export function replaceXMLParts(
   let result = formatXML(xmlContent);
   let lastProcessedIndex = 0;
 
-  console.log('[replaceXMLParts] Input XML length:', xmlContent.length);
-  console.log('[replaceXMLParts] Formatted XML length:', result.length);
-  console.log('[replaceXMLParts] Number of edits:', searchReplacePairs.length);
-
   for (const { search, replace } of searchReplacePairs) {
     // Also format the search content for consistency
     const formattedSearch = formatXML(search);
-    console.log('[replaceXMLParts] Search pattern (first 200):', search.substring(0, 200));
-    console.log('[replaceXMLParts] Formatted search (first 200):', formattedSearch.substring(0, 200));
     const searchLines = formattedSearch.split('\n');
 
     // Split into lines for exact line matching
@@ -282,9 +276,6 @@ export function replaceXMLParts(
     }
 
     if (!matchFound) {
-      console.log('[replaceXMLParts] SEARCH FAILED!');
-      console.log('[replaceXMLParts] Current XML content:\n', result);
-      console.log('[replaceXMLParts] Search pattern:\n', formattedSearch);
       throw new Error(`Search pattern not found in the diagram. The pattern may not exist in the current structure.`);
     }