fix: improve LLM diagram context awareness and image preview (#202)

- Add replaceHistoricalToolInputs to replace XML in tool calls with placeholders
- Send both previousXml and current xml so LLM can understand user's manual edits
- Update system message to mark current XML as authoritative source of truth
- Fix React StrictMode issue with blob URL cleanup in FilePreviewList
- Add unoptimized prop to Image components for blob URLs

app/api/chat/route.ts

@@ -66,6 +66,35 @@ function isMinimalDiagram(xml: string): boolean {
     return !stripped.includes('id="2"')
 }
 
+// Helper function to replace historical tool call XML with placeholders
+// This reduces token usage and forces LLM to rely on the current diagram XML (source of truth)
+function replaceHistoricalToolInputs(messages: any[]): any[] {
+    return messages.map((msg) => {
+        if (msg.role !== "assistant" || !Array.isArray(msg.content)) {
+            return msg
+        }
+        const replacedContent = msg.content.map((part: any) => {
+            if (part.type === "tool-call") {
+                const toolName = part.toolName
+                if (
+                    toolName === "display_diagram" ||
+                    toolName === "edit_diagram"
+                ) {
+                    return {
+                        ...part,
+                        input: {
+                            placeholder:
+                                "[XML content replaced - see current diagram XML in system context]",
+                        },
+                    }
+                }
+            }
+            return part
+        })
+        return { ...msg, content: replacedContent }
+    })
+}
+
 // Helper function to fix tool call inputs for Bedrock API
 // Bedrock requires toolUse.input to be a JSON object, not a string
 function fixToolCallInputs(messages: any[]): any[] {
@@ -144,7 +173,7 @@ async function handleChatRequest(req: Request): Promise<Response> {
         }
     }
 
-    const { messages, xml, sessionId } = await req.json()
+    const { messages, xml, previousXml, sessionId } = await req.json()
 
     // Get user IP for Langfuse tracking
     const forwardedFor = req.headers.get("x-forwarded-for")
@@ -242,9 +271,12 @@ ${lastMessageText}
     // Fix tool call inputs for Bedrock API (requires JSON objects, not strings)
     const fixedMessages = fixToolCallInputs(modelMessages)
 
+    // Replace historical tool call XML with placeholders to reduce tokens and avoid confusion
+    const placeholderMessages = replaceHistoricalToolInputs(fixedMessages)
+
     // Filter out messages with empty content arrays (Bedrock API rejects these)
     // This is a safety measure - ideally convertToModelMessages should handle all cases
-    let enhancedMessages = fixedMessages.filter(
+    let enhancedMessages = placeholderMessages.filter(
         (msg: any) =>
             msg.content && Array.isArray(msg.content) && msg.content.length > 0,
     )
@@ -308,10 +340,10 @@ ${lastMessageText}
                 },
             }),
         },
-        // Cache breakpoint 2: Current diagram XML context
+        // Cache breakpoint 2: Previous and Current diagram XML context
         {
             role: "system" as const,
-            content: `Current diagram XML:\n"""xml\n${xml || ""}\n"""\nWhen using edit_diagram, COPY search patterns exactly from this XML - attribute order matters!`,
+            content: `${previousXml ? `Previous diagram XML (before user's last message):\n"""xml\n${previousXml}\n"""\n\n` : ""}Current diagram XML (AUTHORITATIVE - the source of truth):\n"""xml\n${xml || ""}\n"""\n\nIMPORTANT: The "Current diagram XML" is the SINGLE SOURCE OF TRUTH for what's on the canvas right now. The user can manually add, delete, or modify shapes directly in draw.io. Always count and describe elements based on the CURRENT XML, not on what you previously generated. If both previous and current XML are shown, compare them to understand what the user changed. When using edit_diagram, COPY search patterns exactly from the CURRENT XML - attribute order matters!`,
             ...(shouldCache && {
                 providerOptions: {
                     bedrock: { cachePoint: { type: "default" } },

components/chat-panel.tsx

@@ -764,6 +764,15 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
                     }
                 }
 
+                // Get previous XML from the last snapshot (before this message)
+                const snapshotKeys = Array.from(
+                    xmlSnapshotsRef.current.keys(),
+                ).sort((a, b) => b - a)
+                const previousXml =
+                    snapshotKeys.length > 0
+                        ? xmlSnapshotsRef.current.get(snapshotKeys[0]) || ""
+                        : ""
+
                 // Save XML snapshot for this message (will be at index = current messages.length)
                 const messageIndex = messages.length
                 xmlSnapshotsRef.current.set(messageIndex, chartXml)
@@ -805,6 +814,7 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
                     {
                         body: {
                             xml: chartXml,
+                            previousXml,
                             sessionId,
                         },
                         headers: {
@@ -869,6 +879,15 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
             return
         }
 
+        // Get previous XML (snapshot before the one being regenerated)
+        const snapshotKeys = Array.from(xmlSnapshotsRef.current.keys())
+            .filter((k) => k < userMessageIndex)
+            .sort((a, b) => b - a)
+        const previousXml =
+            snapshotKeys.length > 0
+                ? xmlSnapshotsRef.current.get(snapshotKeys[0]) || ""
+                : ""
+
         // Restore the diagram to the saved state (skip validation for trusted snapshots)
         onDisplayChart(savedXml, true)
 
@@ -923,6 +942,7 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
             {
                 body: {
                     xml: savedXml,
+                    previousXml,
                     sessionId,
                 },
                 headers: {
@@ -956,6 +976,15 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
             return
         }
 
+        // Get previous XML (snapshot before the one being edited)
+        const snapshotKeys = Array.from(xmlSnapshotsRef.current.keys())
+            .filter((k) => k < messageIndex)
+            .sort((a, b) => b - a)
+        const previousXml =
+            snapshotKeys.length > 0
+                ? xmlSnapshotsRef.current.get(snapshotKeys[0]) || ""
+                : ""
+
         // Restore the diagram to the saved state (skip validation for trusted snapshots)
         onDisplayChart(savedXml, true)
 
@@ -1018,6 +1047,7 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
             {
                 body: {
                     xml: savedXml,
+                    previousXml,
                     sessionId,
                 },
                 headers: {

components/file-preview-list.tsx

@@ -48,6 +48,8 @@ export function FilePreviewList({ files, onRemoveFile }: FilePreviewListProps) {
             imageUrlsRef.current.forEach((url) => {
                 URL.revokeObjectURL(url)
             })
+            // Clear the ref so StrictMode remount creates fresh URLs
+            imageUrlsRef.current = new Map()
         }
     }, [])
 
@@ -83,6 +85,7 @@ export function FilePreviewList({ files, onRemoveFile }: FilePreviewListProps) {
                                         width={80}
                                         height={80}
                                         className="object-cover w-full h-full"
+                                        unoptimized
                                     />
                                 ) : (
                                     <div className="flex items-center justify-center h-full text-xs text-center p-1">
@@ -124,6 +127,7 @@ export function FilePreviewList({ files, onRemoveFile }: FilePreviewListProps) {
                             height={900}
                             className="object-contain max-w-full max-h-[90vh] w-auto h-auto"
                             onClick={(e) => e.stopPropagation()}
+                            unoptimized
                         />
                     </div>
                 </div>