feat: improve diagram edit tracking and cache handling

- Track lastGeneratedXml to detect user modifications
- Only send XML context when needed (saves tokens)
- Fix cached response streaming to include tool output
- Add debug logging for model messages and steps
- Enable multi-step tool execution with maxSteps: 5

app/api/chat/route.ts

@@ -19,8 +19,12 @@ function createCachedStreamResponse(xml: string): Response {
     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 +34,7 @@ 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();
 
     // === CACHE CHECK START ===
     const isFirstMessage = messages.length === 1;
@@ -74,6 +78,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,19 +125,34 @@ 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 = `
+    // Check diagram state
-Current diagram XML:
+    const hasDiagram = xml && !isMinimalDiagram(xml);
-"""xml
+    const noHistory = !lastGeneratedXml || lastGeneratedXml.trim() === '';
-${xml || ''}
+    const userModified = hasDiagram && lastGeneratedXml && xml !== lastGeneratedXml;
-"""
+
-User input:
+    // 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);
 
+    // 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
@@ -217,7 +237,13 @@ ${lastMessageText}
       messages: [systemMessageWithCache, ...enhancedMessages],
       ...(providerOptions && { providerOptions }),
       ...(headers && { headers }),
-      onFinish: ({ usage, providerMetadata }) => {
+      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,
@@ -261,6 +287,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

@@ -33,6 +33,8 @@ export default function ChatPanel({ isVisible, onToggleVisibility }: ChatPanelPr
         resolverRef,
         chartXML,
         clearDiagram,
+        getLastAgentGeneratedXml,
+        markAgentDiagramPending,
     } = useDiagram();
 
     const onFetchChart = () => {
@@ -73,11 +75,18 @@ 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-');
+
+                    // 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,
-                        output: "Successfully displayed the diagram.",
+                        result: "Successfully displayed the diagram.",
                     });
                 } else if (toolCall.toolName === "edit_diagram") {
                     const { edits } = toolCall.input as {
@@ -96,10 +105,12 @@ 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,
-                            output: `Successfully applied ${edits.length} edit(s) to the diagram.`,
+                            result: `Successfully applied ${edits.length} edit(s) to the diagram.`,
                         });
                     } catch (error) {
                         console.error("Edit diagram failed:", error);
@@ -108,9 +119,8 @@ export default function ChatPanel({ isVisible, onToggleVisibility }: ChatPanelPr
 
                         // Provide detailed error with current diagram XML
                         addToolResult({
-                            tool: "edit_diagram",
                             toolCallId: toolCall.toolCallId,
-                            output: `Edit failed: ${errorMessage}
+                            result: `Edit failed: ${errorMessage}
 
 Current diagram XML:
 \`\`\`xml
@@ -171,11 +181,17 @@ 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 },
                     {
                         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,26 @@ 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 = () => {
+        console.log('[DiagramContext] markAgentDiagramPending called');
+        agentDiagramPendingRef.current = true;
+    };
+
     const handleExport = () => {
         if (drawioRef.current) {
             drawioRef.current.exportDiagram({
@@ -54,6 +75,16 @@ 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);
+            console.log('[DiagramContext] Setting lastAgentGeneratedXml from export, length:', formatted.length);
+            setLastAgentGeneratedXml(formatted);
+            agentDiagramPendingRef.current = false;
+        }
+
         if (resolverRef.current) {
             resolverRef.current(extractedXML);
             resolverRef.current = null;
@@ -66,6 +97,7 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
         setChartXML(emptyDiagram);
         setLatestSvg("");
         setDiagramHistory([]);
+        setLastAgentGeneratedXml("");
     };
 
     return (
@@ -74,6 +106,10 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
                 chartXML,
                 latestSvg,
                 diagramHistory,
+                lastAgentGeneratedXml,
+                getLastAgentGeneratedXml,
+                setLastAgentGeneratedXml,
+                markAgentDiagramPending,
                 loadDiagram,
                 handleExport,
                 resolverRef,

lib/utils.ts

@@ -190,9 +190,15 @@ 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
@@ -276,6 +282,9 @@ 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.`);
     }