From 06621764b2a1041f3614b8cad86d5d1e33028fd0 Mon Sep 17 00:00:00 2001
From: "dayuan.jiang" <jdy.toh@gmail.com>
Date: Mon, 12 Oct 2026 08:17:11 +0900
Subject: [PATCH] refactor(chat): drop redesign leftovers and old-format
 migrations

- stores/versions-store.ts, hooks/use-session-manager.ts,
  lib/session-storage.ts: remove the diagramHistory conversion layer,
  the localStorage chat migration and the empty templates store
- stores/canvas-store.ts, lib/drawio/editor-bridge.ts: remove the
  pages list nothing read
- components/chat/chat-engine.tsx: use the context's chartXMLRef and
  latestSvgRef instead of local copies, subscribe to one settings field,
  stop prefetching quota limits (lib/use-quota-manager.tsx takes the
  values from the 429 body), readDiagramFile delegates to the MCP
  loader (which now strips XML prologs and accepts <svg:svg>)
- contexts/diagram-context.tsx: one requestTaggedExport path for
  thumbnails, validation, fetch and save; fetchDiagramXml replaces
  handleExport plus the exported resolver map
- hooks/use-validate-diagram.ts and app/api/validate-diagram: plain
  fetch returning JSON instead of a streamed object nobody consumed
- hooks/use-diagram-tool-handlers.ts: drop the DEBUG blocks; canvas
  vars inheritance lives in placeOnPage (lib/diagram-pages.ts)
- lobby-hero, icon-button, drawio-frame: remove unreachable branches
  and unused props
---
 app/api/config/route.ts                       |   3 -
 app/api/validate-diagram/route.ts             |  48 ++--
 components/canvas/canvas-stage.tsx            |   3 +-
 components/canvas/drawio-frame.tsx            |  32 +--
 components/chat/chat-engine.tsx               | 124 ++-------
 components/chat/lobby-hero.tsx                |  31 +--
 components/workspace/icon-button.tsx          |  16 +-
 contexts/diagram-context.tsx                  | 263 ++++++++----------
 hooks/use-diagram-tool-handlers.ts            |  91 +-----
 hooks/use-session-manager.ts                  |  22 +-
 hooks/use-validate-diagram.ts                 | 154 +++-------
 lib/diagram-diff.ts                           |  25 +-
 lib/diagram-pages.ts                          |  52 ++--
 lib/drawio/editor-bridge.ts                   |  13 -
 lib/session-storage.ts                        | 116 +-------
 lib/use-quota-manager.tsx                     |  41 ++-
 packages/mcp-server/src/load-diagram.ts       |  16 +-
 .../mcp-server/tests/load-diagram.test.ts     |  34 +++
 stores/canvas-store.ts                        |   7 -
 stores/versions-store.ts                      |  20 +-
 tests/e2e/settings-providers.spec.ts          |   9 +-
 tests/e2e/workspace.spec.ts                   |  23 +-
 tests/unit/diagram-context.test.tsx           |  40 ++-
 tests/unit/diagram-pages.test.ts              |  19 +-
 tests/unit/redesign-helpers.test.ts           |  25 +-
 tests/unit/use-session-manager.test.tsx       |   1 -
 tests/unit/validate-diagram-route.test.ts     |  49 ++--
 27 files changed, 411 insertions(+), 866 deletions(-)

diff --git a/app/api/config/route.ts b/app/api/config/route.ts
index f4b5bc4e..f4dac222 100644
--- a/app/api/config/route.ts
+++ b/app/api/config/route.ts
@@ -3,8 +3,5 @@ import { NextResponse } from "next/server"
 export async function GET() {
     return NextResponse.json({
         accessCodeRequired: !!process.env.ACCESS_CODE_LIST,
-        dailyRequestLimit: Number(process.env.DAILY_REQUEST_LIMIT) || 0,
-        dailyTokenLimit: Number(process.env.DAILY_TOKEN_LIMIT) || 0,
-        tpmLimit: Number(process.env.TPM_LIMIT) || 0,
     })
 }
diff --git a/app/api/validate-diagram/route.ts b/app/api/validate-diagram/route.ts
index 3fba6628..c91e6377 100644
--- a/app/api/validate-diagram/route.ts
+++ b/app/api/validate-diagram/route.ts
@@ -1,9 +1,9 @@
 /**
  * API endpoint for VLM-based diagram validation.
- * Accepts a PNG image and streams validation results using useObject-compatible format.
+ * Accepts a PNG image and answers with the validation result as JSON.
  */
 
-import { Output, streamText } from "ai"
+import { generateText, Output } from "ai"
 import { checkAccessCode, rejectCrossSite } from "@/lib/access-code"
 import { getValidationModel } from "@/lib/ai-providers"
 import {
@@ -35,13 +35,6 @@ const DEFAULT_VALID_RESULT: ValidationResult = {
     suggestions: [],
 }
 
-/** A fixed result in the text format useObject reads */
-function createStreamingResponse(result: ValidationResult): Response {
-    return new Response(JSON.stringify(result), {
-        headers: { "Content-Type": "text/plain; charset=utf-8" },
-    })
-}
-
 export async function POST(req: Request): Promise<Response> {
     const crossSite = rejectCrossSite(req)
     if (crossSite) return crossSite
@@ -53,7 +46,7 @@ export async function POST(req: Request): Promise<Response> {
         // Check if VLM validation is enabled (default: true)
         const enableValidation = process.env.ENABLE_VLM_VALIDATION !== "false"
         if (!enableValidation) {
-            return createStreamingResponse(DEFAULT_VALID_RESULT)
+            return Response.json(DEFAULT_VALID_RESULT)
         }
 
         const body: ValidateDiagramRequest = await req.json()
@@ -123,7 +116,7 @@ export async function POST(req: Request): Promise<Response> {
                 error,
             )
             // Return valid if no vision model is configured
-            return createStreamingResponse(DEFAULT_VALID_RESULT)
+            return Response.json(DEFAULT_VALID_RESULT)
         }
 
         // Parse timeout with validation (minimum 1000ms, default 10000ms)
@@ -133,8 +126,7 @@ export async function POST(req: Request): Promise<Response> {
                 parseInt(process.env.VALIDATION_TIMEOUT || "10000", 10),
             ) || 10000
 
-        // Stream the VLM response for useObject consumption
-        const result = streamText({
+        const { output, totalUsage } = await generateText({
             model,
             output: Output.object({ schema: ValidationResultSchema }),
             system: VALIDATION_SYSTEM_PROMPT,
@@ -155,23 +147,19 @@ export async function POST(req: Request): Promise<Response> {
             ],
             maxOutputTokens: 1024,
             abortSignal: AbortSignal.timeout(timeout),
-            onFinish: ({ output, totalUsage }) => {
-                if (countsQuota && totalUsage) {
-                    recordTokenUsage(
-                        userId,
-                        (totalUsage.inputTokens || 0) +
-                            (totalUsage.outputTokens || 0),
-                    )
-                }
-                if (sessionId && output) {
-                    console.log(
-                        `[validate-diagram] Session ${sessionId}: valid=${output.valid}, issues=${output.issues?.length ?? 0}`,
-                    )
-                }
-            },
         })
-
-        return result.toTextStreamResponse()
+        if (countsQuota) {
+            await recordTokenUsage(
+                userId,
+                (totalUsage.inputTokens || 0) + (totalUsage.outputTokens || 0),
+            )
+        }
+        if (sessionId) {
+            console.log(
+                `[validate-diagram] Session ${sessionId}: valid=${output.valid}, issues=${output.issues?.length ?? 0}`,
+            )
+        }
+        return Response.json(output)
     } catch (error) {
         // Log with session context if available
         const errorMessage =
@@ -179,6 +167,6 @@ export async function POST(req: Request): Promise<Response> {
         console.error("[validate-diagram] Error:", errorMessage)
 
         // On error, return valid to not block the user
-        return createStreamingResponse(DEFAULT_VALID_RESULT)
+        return Response.json(DEFAULT_VALID_RESULT)
     }
 }
diff --git a/components/canvas/canvas-stage.tsx b/components/canvas/canvas-stage.tsx
index 2cb1617e..8b40dcb0 100644
--- a/components/canvas/canvas-stage.tsx
+++ b/components/canvas/canvas-stage.tsx
@@ -136,8 +136,7 @@ export function CanvasStage({
                     key={src}
                     ref={drawioRef}
                     src={src}
-                    configCss={DRAWIO_CSS}
-                    editorCss={DRAWIO_CSS}
+                    css={DRAWIO_CSS}
                     dark={isDark}
                     onLoad={onDrawioLoad}
                     onAutoSave={handleDiagramAutoSave}
diff --git a/components/canvas/drawio-frame.tsx b/components/canvas/drawio-frame.tsx
index df47a3fc..177094d8 100644
--- a/components/canvas/drawio-frame.tsx
+++ b/components/canvas/drawio-frame.tsx
@@ -30,10 +30,11 @@ export interface DrawioFrameHandle {
 
 interface DrawioFrameProps {
     src: string
-    /** CSS sent with the configure message (works cross-origin too) */
-    configCss: string
-    /** CSS injected once the editor is reachable directly (same origin) */
-    editorCss: string
+    /**
+     * Theme CSS: sent with the configure message (works cross-origin too)
+     * and, same origin, injected as a <style> so it updates without a reload
+     */
+    css: string
     /** draw.io's dark mode, switched in place when the editor is reachable */
     dark: boolean
     onLoad?: () => void
@@ -52,25 +53,14 @@ const STYLE_ID = "next-ai-drawio-theme"
  */
 export const DrawioFrame = forwardRef<DrawioFrameHandle, DrawioFrameProps>(
     function DrawioFrame(
-        {
-            src,
-            configCss,
-            editorCss,
-            dark,
-            onLoad,
-            onAutoSave,
-            onExport,
-            onEditorAccess,
-        },
+        { src, css, dark, onLoad, onAutoSave, onExport, onEditorAccess },
         ref,
     ) {
         const iframeRef = useRef<HTMLIFrameElement>(null)
         const editorUiRef = useRef<any>(null)
         const attachedRef = useRef(false)
-        const configCssRef = useRef(configCss)
-        configCssRef.current = configCss
-        const editorCssRef = useRef(editorCss)
-        editorCssRef.current = editorCss
+        const cssRef = useRef(css)
+        cssRef.current = css
         const darkRef = useRef(dark)
         darkRef.current = dark
         // Handlers change on every render; the message listener reads refs
@@ -111,7 +101,7 @@ export const DrawioFrame = forwardRef<DrawioFrameHandle, DrawioFrameProps>(
                     style.id = STYLE_ID
                     doc.head.appendChild(style)
                 }
-                style.textContent = editorCssRef.current
+                style.textContent = cssRef.current
             } catch {
                 // Cross-origin: the configure message already applied it
             }
@@ -119,7 +109,7 @@ export const DrawioFrame = forwardRef<DrawioFrameHandle, DrawioFrameProps>(
 
         useEffect(() => {
             applyThemeCss()
-        }, [editorCss])
+        }, [css])
 
         // Same-origin only: draw.io's own dark mode, without a reload
         const applyDarkMode = () => {
@@ -209,7 +199,7 @@ export const DrawioFrame = forwardRef<DrawioFrameHandle, DrawioFrameProps>(
                         send("configure", {
                             config: {
                                 ...DRAWIO_CONFIG,
-                                css: configCssRef.current,
+                                css: cssRef.current,
                             },
                         })
                         break
diff --git a/components/chat/chat-engine.tsx b/components/chat/chat-engine.tsx
index 2d37756c..4e3b719c 100644
--- a/components/chat/chat-engine.tsx
+++ b/components/chat/chat-engine.tsx
@@ -23,10 +23,7 @@ import { flushSync } from "react-dom"
 import { toast } from "sonner"
 import type { ValidationState } from "@/components/chat/ValidationCard"
 import { useDiagram } from "@/contexts/diagram-context"
-import {
-    keepFileVars,
-    useDiagramToolHandlers,
-} from "@/hooks/use-diagram-tool-handlers"
+import { useDiagramToolHandlers } from "@/hooks/use-diagram-tool-handlers"
 import { useDictionary } from "@/hooks/use-dictionary"
 import {
     getSelectedAIConfig,
@@ -34,7 +31,6 @@ import {
     useModelConfig,
 } from "@/hooks/use-model-config"
 import {
-    getSessionVersions,
     type SessionData,
     useSessionManager,
 } from "@/hooks/use-session-manager"
@@ -121,31 +117,14 @@ function newLangfuseSessionId() {
     return `session-${Date.now()}-${Math.random().toString(36).slice(2, 9)}`
 }
 
-const XML_PROLOG =
-    /^(?:\s|<\?[\s\S]*?\?>|<!--[\s\S]*?-->|<!DOCTYPE[^>[]*(?:\[[\s\S]*?\])?\s*>)*/i
-
 /**
  * Read the diagram document out of a .drawio, .xml or .drawio.svg file, with
  * its pages inflated: previews and the model read it before draw.io saves
  * it again. A broken or cut off file gives null.
  */
 async function readDiagramFile(file: File): Promise<string | null> {
-    const raw = await file.text()
-    // What XML allows before the root element: the declaration, comments,
-    // a DOCTYPE (the parser below expects the root first)
-    let text = raw.replace(XML_PROLOG, "")
-    // Editable SVG: the diagram is stored in the root's `content` attribute.
-    // Parsed whole: its DOCTYPE may define entities the content uses.
-    if (/^<(?:[\w.-]+:)?svg[\s>]/i.test(text)) {
-        const doc = new DOMParser().parseFromString(raw, "image/svg+xml")
-        text = (doc.documentElement?.getAttribute("content") ?? "").replace(
-            XML_PROLOG,
-            "",
-        )
-    }
-    const loaded = parseDrawioFileContent(text)
-    if (!loaded.ok) return null
-    return withUniquePageIds(loaded.xml)
+    const loaded = parseDrawioFileContent(await file.text())
+    return loaded.ok ? withUniquePageIds(loaded.xml) : null
 }
 
 /**
@@ -267,11 +246,10 @@ export function ChatEngineProvider({
 }) {
     const {
         loadDiagram: onDisplayChart,
-        handleExport: onExport,
-        exportResolversRef,
+        fetchDiagramXml,
         chartXML,
-        chartXMLRef: liveChartXMLRef,
-        latestSvg,
+        chartXMLRef,
+        latestSvgRef,
         clearDiagram,
         getThumbnailSvg,
         captureValidationPng,
@@ -284,26 +262,10 @@ export function ChatEngineProvider({
     const searchParams = useSearchParams()
     const urlSessionId = searchParams.get("session")
 
-    const settings = useSettingsStore()
+    const vlmValidationEnabled = useSettingsStore((s) => s.vlmValidationEnabled)
     const openSettings = useUiStore((s) => s.openSettings)
     const focusComposer = useUiStore((s) => s.focusComposer)
 
-    const onFetchChart = () => {
-        // Waits for the reply to its own export, by its tag
-        const tag = onExport()
-        return Promise.race([
-            new Promise<string>((resolve) => {
-                if (tag) exportResolversRef.current[tag] = resolve
-            }),
-            new Promise<string>((_, reject) => {
-                setTimeout(() => {
-                    delete exportResolversRef.current[tag]
-                    reject(new Error("Chart export timed out after 10 seconds"))
-                }, 10000)
-            }),
-        ])
-    }
-
     const { files, pdfData, handleFileChange, setFiles } = useFileProcessor({
         tooLong: dict.errors.fileTooLong,
         readFailed: dict.errors.fileReadFailed,
@@ -318,9 +280,6 @@ export function ChatEngineProvider({
     const sessionManager = useSessionManager({ initialSessionId: urlSessionId })
 
     const [input, setInputState] = useState("")
-    const [dailyRequestLimit, setDailyRequestLimit] = useState(0)
-    const [dailyTokenLimit, setDailyTokenLimit] = useState(0)
-    const [tpmLimit, setTpmLimit] = useState(0)
 
     // Restore the unsent input from sessionStorage
     useEffect(() => {
@@ -340,21 +299,7 @@ export function ChatEngineProvider({
         else sessionStorage.removeItem(SESSION_STORAGE_INPUT_KEY)
     }, [])
 
-    useEffect(() => {
-        fetch(getApiEndpoint("/api/config"))
-            .then((res) => res.json())
-            .then((data) => {
-                setDailyRequestLimit(data.dailyRequestLimit || 0)
-                setDailyTokenLimit(data.dailyTokenLimit || 0)
-                setTpmLimit(data.tpmLimit || 0)
-            })
-            .catch(() => {})
-    }, [])
-
     const quotaManager = useQuotaManager({
-        dailyRequestLimit,
-        dailyTokenLimit,
-        tpmLimit,
         onConfigModel: () => openSettings("models"),
     })
 
@@ -392,20 +337,12 @@ export function ChatEngineProvider({
     const hasRestoredRef = useRef(false)
     const [isRestored, setIsRestored] = useState(false)
 
-    // Latest chartXML for callbacks (avoids stale closures)
-    const chartXMLRef = useRef(chartXML)
     // Session loaded without a diagram (prevents thumbnail contamination)
     const justLoadedSessionIdRef = useRef<string | null>(null)
     useEffect(() => {
-        chartXMLRef.current = chartXML
         if (chartXML) justLoadedSessionIdRef.current = null
     }, [chartXML])
 
-    const latestSvgRef = useRef(latestSvg)
-    useEffect(() => {
-        latestSvgRef.current = latestSvg
-    }, [latestSvg])
-
     const autoRetryCountRef = useRef(0)
     const continuationRetryCountRef = useRef(0)
     // Partial XML when output was truncated; non-empty means continuation mode
@@ -458,13 +395,13 @@ export function ChatEngineProvider({
         processedToolCallsRef,
         validationRetryCountRef,
         // A preview undone just before the tool call is in this one already
-        chartXMLRef: liveChartXMLRef,
+        chartXMLRef,
         turnPageIdRef,
         onDisplayChart,
-        onFetchChart,
+        onFetchChart: fetchDiagramXml,
         captureValidationPng,
         validateDiagram: validateWithFallback,
-        enableVlmValidation: settings.vlmValidationEnabled,
+        enableVlmValidation: vlmValidationEnabled,
         sessionId: langfuseSessionId,
         watchStop: () => {
             const stopsBefore = stopCountRef.current
@@ -832,13 +769,7 @@ export function ChatEngineProvider({
     const messagesBeforeShownRef = useRef<UIMessage[] | null>(null)
 
     const syncUIWithSession = useCallback(
-        (
-            data:
-                | (SessionData & {
-                      diagramHistory?: { svg: string; xml: string }[]
-                  })
-                | null,
-        ) => {
+        (data: SessionData | null) => {
             // An answer still coming in belongs to the chat that was on
             // screen (browser back or forward): it stops, before this chat's
             // messages are put in its place, so none of it lands here
@@ -877,24 +808,15 @@ export function ChatEngineProvider({
                 )
                 setMessages(data.messages as unknown as UIMessage[])
                 xmlSnapshotsRef.current = new Map(data.xmlSnapshots)
-                if (data.diagramXml) {
-                    onDisplayChart(data.diagramXml, true)
-                    chartXMLRef.current = data.diagramXml
-                } else {
-                    clearDiagram()
-                    chartXMLRef.current = ""
-                }
+                if (data.diagramXml) onDisplayChart(data.diagramXml, true)
+                else clearDiagram()
                 if (!isRealDiagram(data.diagramXml)) latestSvgRef.current = ""
-                useVersionsStore
-                    .getState()
-                    .setVersions(getSessionVersions(data))
+                useVersionsStore.getState().setVersions(data.versions ?? [])
             } else {
                 loadedMessageIdsRef.current = new Set()
                 setMessages([])
                 xmlSnapshotsRef.current.clear()
                 clearDiagram()
-                chartXMLRef.current = ""
-                latestSvgRef.current = ""
                 useVersionsStore.getState().clear()
             }
         },
@@ -927,7 +849,6 @@ export function ChatEngineProvider({
             if (isRealDiagram(currentDiagramXml) && options.withThumbnail) {
                 const freshThumb = await getThumbnailSvg()
                 if (freshThumb) {
-                    latestSvgRef.current = freshThumb
                     thumbnailDataUrl = freshThumb
                 } else if (latestSvgRef.current) {
                     thumbnailDataUrl = latestSvgRef.current
@@ -1256,7 +1177,7 @@ export function ChatEngineProvider({
         selectedCells: MessageSelection = [],
         generation = getChatGeneration(),
     ): Promise<boolean> => {
-        const chartXml = formatXML(await onFetchChart())
+        const chartXml = formatXML(await fetchDiagramXml())
         if (getChatGeneration() !== generation || leavingRef.current) {
             return false
         }
@@ -1339,12 +1260,8 @@ export function ChatEngineProvider({
                 })
                 if (prepared.ok) {
                     const pageId = viewedPageId()
-                    const canvasXml = chartXMLRef.current || ""
                     onDisplayChart(
-                        keepFileVars(
-                            placeOnPage(prepared.xml, canvasXml, pageId),
-                            canvasXml,
-                        ),
+                        placeOnPage(prepared.xml, chartXMLRef.current, pageId),
                         true,
                         "commit",
                         { toolCallId, pageId },
@@ -1473,11 +1390,6 @@ export function ChatEngineProvider({
         }
     }
 
-    const restoreDiagramFromSnapshot = (savedXml: string) => {
-        onDisplayChart(savedXml, true)
-        chartXMLRef.current = savedXml
-    }
-
     // Snapshots of the messages after this one go with them
     const dropLaterSnapshots = (messageIndex: number) => {
         for (const key of xmlSnapshotsRef.current.keys()) {
@@ -1495,7 +1407,7 @@ export function ChatEngineProvider({
             return false
         }
         const previousXml = getPreviousXml(index)
-        restoreDiagramFromSnapshot(savedXml)
+        onDisplayChart(savedXml, true)
         dropLaterSnapshots(index)
         useVersionsStore.getState().removeFromTurn(index)
 
@@ -1563,7 +1475,7 @@ export function ChatEngineProvider({
         // Like submit: the diagram export comes first
         preparingSendRef.current = true
         try {
-            const chartXml = formatXML(await onFetchChart())
+            const chartXml = formatXML(await fetchDiagramXml())
             // Another chat came on screen meanwhile (browser back/forward),
             // or this one is being left
             if (getChatGeneration() !== generation || leavingRef.current) {
diff --git a/components/chat/lobby-hero.tsx b/components/chat/lobby-hero.tsx
index 542612fd..2a4e863f 100644
--- a/components/chat/lobby-hero.tsx
+++ b/components/chat/lobby-hero.tsx
@@ -1,4 +1,4 @@
-import { FileText, FileUp, History, Pencil } from "lucide-react"
+import { FileUp, History, Pencil } from "lucide-react"
 import type React from "react"
 import { useDictionary } from "@/hooks/use-dictionary"
 import { getAssetUrl } from "@/lib/base-path"
@@ -12,7 +12,7 @@ interface Example {
     title: string
     description: string
     /** Picture of the result (or of the input, for "replicate" examples) */
-    image?: string
+    image: string
     prompt: string
     /** Bundled file sent with the prompt; matches the cached example */
     file?: { path: string; name: string; type: string }
@@ -34,7 +34,6 @@ function ExampleTile({
     title,
     description,
     picture,
-    plain,
     testId,
 }: {
     onClick: () => void
@@ -42,8 +41,6 @@ function ExampleTile({
     /** Shown as the tooltip: the title alone names the example */
     description: string
     picture: React.ReactNode
-    /** An icon on a tinted tile instead of a diagram on paper */
-    plain?: boolean
     testId?: string
 }) {
     return (
@@ -54,12 +51,7 @@ function ExampleTile({
             className="group/example flex min-w-0 flex-col gap-2 rounded-xl border border-border bg-card p-1.5 pb-2.5 text-left transition-colors hover:border-foreground/25"
             data-testid={testId}
         >
-            <span
-                className={cn(
-                    "flex h-16 w-full items-center justify-center overflow-hidden rounded-lg",
-                    plain ? "bg-muted" : "sheet-light bg-white",
-                )}
-            >
+            <span className="flex h-16 w-full items-center justify-center overflow-hidden rounded-lg sheet-light bg-white">
                 {picture}
             </span>
             <span className="px-1 text-[13px] leading-snug font-medium text-foreground">
@@ -220,18 +212,13 @@ export function LobbyHero({
                             description={example.description}
                             testId={`example-${example.key}`}
                             picture={
-                                example.image ? (
-                                    // biome-ignore lint/performance/noImgElement: static preview
-                                    <img
-                                        src={getAssetUrl(example.image)}
-                                        alt=""
-                                        className="max-h-14 max-w-[88%] object-contain transition-transform duration-300 group-hover/example:scale-[1.06]"
-                                    />
-                                ) : (
-                                    <FileText className="size-5 text-muted-foreground" />
-                                )
+                                // biome-ignore lint/performance/noImgElement: static preview
+                                <img
+                                    src={getAssetUrl(example.image)}
+                                    alt=""
+                                    className="max-h-14 max-w-[88%] object-contain transition-transform duration-300 group-hover/example:scale-[1.06]"
+                                />
                             }
-                            plain={!example.image}
                         />
                     ),
                 )}
diff --git a/components/workspace/icon-button.tsx b/components/workspace/icon-button.tsx
index eb70f076..e9a4a299 100644
--- a/components/workspace/icon-button.tsx
+++ b/components/workspace/icon-button.tsx
@@ -13,9 +13,7 @@ interface IconButtonProps extends React.ComponentProps<"button"> {
     label: string
     /** Shown next to the label in the tooltip, e.g. "⌘B" */
     shortcut?: string
-    active?: boolean
     tooltipSide?: "top" | "bottom" | "left" | "right"
-    size?: "sm" | "md"
 }
 
 /** Square icon button with a tooltip; the tooltip text is also its label */
@@ -24,9 +22,7 @@ export const IconButton = forwardRef<HTMLButtonElement, IconButtonProps>(
         {
             label,
             shortcut,
-            active = false,
             tooltipSide = "bottom",
-            size = "md",
             className,
             children,
             ...props
@@ -40,17 +36,12 @@ export const IconButton = forwardRef<HTMLButtonElement, IconButtonProps>(
                         ref={ref}
                         type="button"
                         aria-label={label}
-                        aria-pressed={active || undefined}
                         className={cn(
-                            "inline-flex shrink-0 items-center justify-center rounded-lg text-muted-foreground transition-colors",
+                            "inline-flex size-8 shrink-0 items-center justify-center rounded-lg text-muted-foreground transition-colors",
                             "hover:bg-accent hover:text-foreground",
                             "focus-visible:outline-2 focus-visible:outline-offset-1 focus-visible:outline-ring",
                             "disabled:pointer-events-none disabled:opacity-35",
                             "[&_svg]:size-4 [&_svg]:stroke-[1.75]",
-                            size === "md" ? "size-8" : "size-7",
-                            // A tinted tile: solid ink is for the main action
-                            active &&
-                                "bg-interactive-active text-foreground hover:bg-interactive-active [&_svg]:stroke-2",
                             className,
                         )}
                         {...props}
@@ -73,8 +64,3 @@ export const IconButton = forwardRef<HTMLButtonElement, IconButtonProps>(
         )
     },
 )
-
-/** Thin vertical divider between toolbar groups */
-export function ToolbarDivider() {
-    return <span className="mx-1 h-5 w-px shrink-0 bg-border" aria-hidden />
-}
diff --git a/contexts/diagram-context.tsx b/contexts/diagram-context.tsx
index b9850708..d7bfab6e 100644
--- a/contexts/diagram-context.tsx
+++ b/contexts/diagram-context.tsx
@@ -64,20 +64,16 @@ interface DiagramContextType {
     chartXML: string
     // chartXML right away, before the re-render (loadDiagram sets both)
     chartXMLRef: React.MutableRefObject<string>
-    latestSvg: string
+    // The diagram's latest SVG export, a fallback thumbnail for saves
+    latestSvgRef: React.MutableRefObject<string>
     loadDiagram: (
         chart: string,
         skipValidation?: boolean,
         mode?: LoadMode,
         meta?: { toolCallId?: string; pageId?: string | null },
     ) => string | null
-    // Returns the export's tag (empty when draw.io is not there yet)
-    handleExport: () => string
-    // Pending exports by tag; a plain export's resolver gets the XML of the
-    // page on screen
-    exportResolversRef: React.MutableRefObject<
-        Record<string, (data: string, xml?: string) => void>
-    >
+    /** The XML of the page on screen, exported from draw.io */
+    fetchDiagramXml: () => Promise<string>
     drawioRef: React.MutableRefObject<DrawioFrameHandle | null>
     handleDiagramExport: (data: DrawioExportEvent) => void
     handleDiagramAutoSave: (data: { xml?: string }) => void
@@ -87,7 +83,7 @@ interface DiagramContextType {
         format: ExportFormat,
         sessionId?: string,
         successMessage?: string,
-    ) => void
+    ) => Promise<void>
     getThumbnailSvg: () => Promise<string | null>
     getVersionSvg: () => Promise<string | null>
     captureValidationPng: () => Promise<string | null>
@@ -108,11 +104,17 @@ const DiagramContext = createContext<DiagramContextType | undefined>(undefined)
 // echoes the request back in the export event, so each result reaches its
 // own caller. Tags end in a request number, so a late result never answers
 // a newer request.
-type ExportTag = "thumbnail" | "validation" | "version"
+type ExportTag = "thumbnail" | "validation" | "version" | "fetch" | "save"
+
+/** What draw.io answers: the data URL asked for, and the whole document */
+interface ExportResult {
+    data: string
+    xml: string
+}
 
 export function DiagramProvider({ children }: { children: React.ReactNode }) {
     const [chartXML, setChartXML] = useState<string>("")
-    const [latestSvg, setLatestSvg] = useState<string>("")
+    const latestSvgRef = useRef("")
     const [isDrawioReady, setIsDrawioReady] = useState(false)
     const hasCalledOnLoadRef = useRef(false)
     const mountedRef = useRef(true)
@@ -125,7 +127,7 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
     const drawioRef = useRef<DrawioFrameHandle | null>(null)
     // Pending exports, keyed by their export tag
     const exportResolversRef = useRef<
-        Record<string, (data: string, xml?: string) => void>
+        Record<string, (result: ExportResult) => void>
     >({})
     const exportSeqRef = useRef(0)
     // Track latest chartXML for restoration after remount
@@ -184,26 +186,15 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
         setChartXML(xml)
     }
 
-    const handleExport = () => {
-        if (!drawioRef.current) return ""
-        // Export of the current state (for the chat and edit_diagram)
-        const tag = `fetch-${++exportSeqRef.current}`
-        drawioRef.current.exportDiagram({
-            format: "xmlsvg",
-            message: tag,
-        })
-        return tag
-    }
-
     // Export with a tag in `message` (draw.io echoes it back in the export
     // event) and wait for that result. Resolves to null on timeout, which is
     // expected occasionally.
     // (Reads refs only, so it keeps one identity)
     const requestTaggedExport = useCallback(
         (tag: ExportTag, format: "xmlsvg" | "svg" | "png", timeoutMs: number) =>
-            new Promise<string | null>((resolve) => {
+            new Promise<ExportResult | null>((resolve) => {
                 const id = `${tag}-${++exportSeqRef.current}`
-                const finish = (value: string | null) => {
+                const finish = (value: ExportResult | null) => {
                     clearTimeout(timer)
                     delete exportResolversRef.current[id]
                     resolve(value)
@@ -216,20 +207,18 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
     )
 
     // Get current diagram as SVG for thumbnail (used by session storage).
-    // One identity: the chat's auto-save depends on it, and each thumbnail
-    // renders this provider again (latestSvg), which would otherwise start
-    // the next save
+    // One identity: the chat's auto-save effect depends on it
     const getThumbnailSvg = useCallback(async (): Promise<string | null> => {
         if (!drawioRef.current) return null
         // Don't export if diagram is empty
         if (!isRealDiagram(chartXMLRef.current)) return null
 
         // xmlsvg exports return an SVG data URL
-        const svgData = await requestTaggedExport("thumbnail", "xmlsvg", 3000)
-        if (svgData?.startsWith("data:image/svg")) {
-            // Update latestSvg so it's available for future saves
-            setLatestSvg(svgData)
-            return svgData
+        const result = await requestTaggedExport("thumbnail", "xmlsvg", 3000)
+        if (result?.data.startsWith("data:image/svg")) {
+            // Kept for later saves that take no fresh thumbnail
+            latestSvgRef.current = result.data
+            return result.data
         }
         return null
     }, [requestTaggedExport])
@@ -237,8 +226,8 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
     // Plain SVG (no embedded diagram data) for version cards
     const getVersionSvg = async (): Promise<string | null> => {
         if (!drawioRef.current) return null
-        const svgData = await requestTaggedExport("version", "svg", 4000)
-        return svgData?.startsWith("data:image/svg") ? svgData : null
+        const result = await requestTaggedExport("version", "svg", 4000)
+        return result?.data.startsWith("data:image/svg") ? result.data : null
     }
 
     // Capture current diagram as PNG for VLM validation
@@ -247,9 +236,25 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
         // Don't export if diagram is empty
         if (!isRealDiagram(chartXMLRef.current)) return null
 
-        const pngData = await requestTaggedExport("validation", "png", 5000)
+        const result = await requestTaggedExport("validation", "png", 5000)
         // PNG data should be a base64 data URL
-        return pngData?.startsWith("data:image/png") ? pngData : null
+        return result?.data.startsWith("data:image/png") ? result.data : null
+    }
+
+    // The page the user is viewing (the first one with an external draw.io,
+    // which cannot tell us), for the chat and edit_diagram.
+    // Not written to chartXML: exports don't change the diagram, and the
+    // export's xml has compressed <diagram> payloads that would break
+    // edit_diagram/display_diagram. Autosave keeps chartXML up to date with
+    // the full uncompressed multi-page document (#879).
+    const fetchDiagramXml = async (): Promise<string> => {
+        const result = await requestTaggedExport("fetch", "xmlsvg", 10000)
+        if (!result) throw new Error("Chart export timed out after 10 seconds")
+        latestSvgRef.current = result.data
+        return extractDiagramXML(
+            result.data,
+            useCanvasStore.getState().currentPageId,
+        )
     }
 
     const loadDiagram = (
@@ -335,34 +340,11 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
         resetPreview()
     }
 
+    // Each export result goes to the caller that asked for it, by its tag
     const handleDiagramExport = (data: DrawioExportEvent) => {
-        // Thumbnail, version, validation PNG and file save exports go only
-        // to their own caller
         const tag = data.message?.message
-        if (/^(thumbnail|validation|version|save)-/.test(tag ?? "")) {
-            exportResolversRef.current[tag as string]?.(data.data, data.xml)
-            return
-        }
-
-        // Don't write chartXML here: exports don't change the diagram, and
-        // data.xml from xmlsvg exports has compressed <diagram> payloads that
-        // would break edit_diagram/display_diagram. Autosave keeps chartXML
-        // up to date with the full uncompressed multi-page document (#879).
-        // The chat gets the page the user is viewing (the first one with
-        // an external draw.io, which cannot tell us)
-        const extractedXML = extractDiagramXML(
-            data.data,
-            useCanvasStore.getState().currentPageId,
-        )
-        setLatestSvg(data.data)
-
-        // The chat's own export (onFetchChart), not another one in flight
-        const resolve =
-            tag !== undefined ? exportResolversRef.current[tag] : undefined
-        if (resolve) {
-            delete exportResolversRef.current[tag as string]
-            resolve(extractedXML)
-        }
+        if (!tag) return
+        exportResolversRef.current[tag]?.({ data: data.data, xml: data.xml })
     }
 
     // The frame registers this callback once per iframe mount, so it must
@@ -383,10 +365,10 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
         const emptyDiagram = BLANK_MXFILE
         // Skip validation for trusted internal template (loadDiagram also sets chartXML)
         loadDiagram(emptyDiagram, true)
-        setLatestSvg("")
+        latestSvgRef.current = ""
     }
 
-    const saveDiagramToFile = (
+    const saveDiagramToFile = async (
         filename: string,
         format: ExportFormat,
         sessionId?: string,
@@ -401,85 +383,79 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
         const drawioFormat =
             format === "drawio" || format === "xmlsvg" ? "xmlsvg" : format
 
-        // Each save has its own tag, so two at once never swap results
-        const tag = `save-${++exportSeqRef.current}`
-        exportResolversRef.current[tag] = (
-            exportData: string,
-            fullDiagramXML?: string,
-        ) => {
-            delete exportResolversRef.current[tag]
-            let fileContent: string | Blob
-            let mimeType: string
-            let extension: string
+        // Each save has its own tag, so two at once never swap results. A
+        // large PNG takes a while to render.
+        const result = await requestTaggedExport("save", drawioFormat, 30000)
+        if (!result) {
+            console.warn("Diagram export for saving timed out")
+            return
+        }
+        const { data: exportData, xml: fullDiagramXML } = result
+        let fileContent: string | Blob
+        let mimeType: string
+        let extension: string
 
-            if (format === "drawio") {
-                // Prefer the complete document from the export event so all pages are saved.
-                const xml = fullDiagramXML?.trim()
-                    ? fullDiagramXML
-                    : extractDiagramXML(exportData)
-                fileContent =
-                    normalizeToMxfile(xml, {
-                        pageId: "page-1",
-                        pageName: "Page-1",
-                    }) ?? xml
-                mimeType = "application/xml"
-                extension = ".drawio"
-            } else if (format === "png") {
-                // PNG data comes as base64 data URL
-                fileContent = exportData
-                mimeType = "image/png"
-                extension = ".png"
-            } else if (format === "xmlsvg") {
-                // Editable SVG: pass data URL directly (like PNG)
-                fileContent = exportData
-                mimeType = "image/svg+xml"
-                extension = ".drawio.svg"
-            } else {
-                // SVG format (view-only)
-                fileContent = exportData
-                mimeType = "image/svg+xml"
-                extension = ".svg"
-            }
-
-            // Log save event to Langfuse (flags the trace)
-            logSaveToLangfuse(filename, format, sessionId)
-
-            // Handle download
-            let url: string
-            if (
-                typeof fileContent === "string" &&
-                fileContent.startsWith("data:")
-            ) {
-                // Already a data URL (PNG)
-                url = fileContent
-            } else {
-                const blob = new Blob([fileContent], { type: mimeType })
-                url = URL.createObjectURL(blob)
-            }
-
-            const a = document.createElement("a")
-            a.href = url
-            a.download = `${filename}${extension}`
-            document.body.appendChild(a)
-            a.click()
-            document.body.removeChild(a)
-
-            // Show success toast after download is initiated
-            if (successMessage) {
-                toast.success(successMessage, { duration: 2500 })
-            }
-
-            // Delay URL revocation to ensure download completes
-            if (!url.startsWith("data:")) {
-                setTimeout(() => URL.revokeObjectURL(url), 100)
-            }
+        if (format === "drawio") {
+            // Prefer the complete document from the export event so all pages are saved.
+            const xml = fullDiagramXML?.trim()
+                ? fullDiagramXML
+                : extractDiagramXML(exportData)
+            fileContent =
+                normalizeToMxfile(xml, {
+                    pageId: "page-1",
+                    pageName: "Page-1",
+                }) ?? xml
+            mimeType = "application/xml"
+            extension = ".drawio"
+        } else if (format === "png") {
+            // PNG data comes as base64 data URL
+            fileContent = exportData
+            mimeType = "image/png"
+            extension = ".png"
+        } else if (format === "xmlsvg") {
+            // Editable SVG: pass data URL directly (like PNG)
+            fileContent = exportData
+            mimeType = "image/svg+xml"
+            extension = ".drawio.svg"
+        } else {
+            // SVG format (view-only)
+            fileContent = exportData
+            mimeType = "image/svg+xml"
+            extension = ".svg"
         }
 
-        // Export diagram - callback will be handled in handleDiagramExport
-        drawioRef.current.exportDiagram({
-            format: drawioFormat,
-            message: tag,
-        })
+        // Log save event to Langfuse (flags the trace)
+        logSaveToLangfuse(filename, format, sessionId)
+
+        // Handle download
+        let url: string
+        if (
+            typeof fileContent === "string" &&
+            fileContent.startsWith("data:")
+        ) {
+            // Already a data URL (PNG)
+            url = fileContent
+        } else {
+            const blob = new Blob([fileContent], { type: mimeType })
+            url = URL.createObjectURL(blob)
+        }
+
+        const a = document.createElement("a")
+        a.href = url
+        a.download = `${filename}${extension}`
+        document.body.appendChild(a)
+        a.click()
+        document.body.removeChild(a)
+
+        // Show success toast after download is initiated
+        if (successMessage) {
+            toast.success(successMessage, { duration: 2500 })
+        }
+
+        // Delay URL revocation to ensure download completes
+        if (!url.startsWith("data:")) {
+            setTimeout(() => URL.revokeObjectURL(url), 100)
+        }
     }
 
     // Log save event to Langfuse (just flags the trace, doesn't send content)
@@ -510,10 +486,9 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
             value={{
                 chartXML,
                 chartXMLRef,
-                latestSvg,
+                latestSvgRef,
                 loadDiagram,
-                handleExport,
-                exportResolversRef,
+                fetchDiagramXml,
                 drawioRef,
                 handleDiagramExport,
                 handleDiagramAutoSave,
diff --git a/hooks/use-diagram-tool-handlers.ts b/hooks/use-diagram-tool-handlers.ts
index e44ec0da..b6290fc9 100644
--- a/hooks/use-diagram-tool-handlers.ts
+++ b/hooks/use-diagram-tool-handlers.ts
@@ -14,32 +14,10 @@ import { editDiagram } from "@/packages/mcp-server/src/edit-diagram.ts"
 import { prepareNewDiagram } from "@/packages/mcp-server/src/new-diagram.ts"
 import { hasCells } from "@/packages/mcp-server/src/pages.ts"
 
-const DEBUG = process.env.NODE_ENV === "development"
-
 // The page a drawing is wrapped in; on the canvas it replaces the page the
 // model works on, so the canvas page keeps its own id and name
 const NEW_PAGE = { pageId: "page-1", pageName: "Page-1" }
 
-/**
- * A new diagram written without file variables (%name% placeholders) keeps
- * the canvas file's, as replacing the page in draw.io does
- */
-export function keepFileVars(xml: string, canvasXml: string): string {
-    const varsOf = (doc: Document) =>
-        doc.documentElement?.nodeName === "mxfile"
-            ? doc.documentElement.getAttribute("vars")
-            : null
-    const parser = new DOMParser()
-    const canvasVars = varsOf(parser.parseFromString(canvasXml, "text/xml"))
-    if (!canvasVars) return xml
-    const doc = parser.parseFromString(xml, "text/xml")
-    if (doc.documentElement?.nodeName !== "mxfile" || varsOf(doc) !== null) {
-        return xml
-    }
-    doc.documentElement.setAttribute("vars", canvasVars)
-    return new XMLSerializer().serializeToString(doc)
-}
-
 interface ToolCall {
     toolCallId: string
     toolName: string
@@ -157,12 +135,6 @@ export function useDiagramToolHandlers({
         { toolCall }: { toolCall: ToolCall },
         addToolOutput: AddToolOutputFn,
     ) => {
-        if (DEBUG) {
-            console.log(
-                `[onToolCall] Tool: ${toolCall.toolName}, CallId: ${toolCall.toolCallId}`,
-            )
-        }
-
         processedToolCallsRef.current.add(toolCall.toolCallId)
         // Only display_diagram, edit_diagram and a completing append_diagram
         // put their result on the canvas. Other tools (get_shape_library,
@@ -204,12 +176,9 @@ export function useDiagramToolHandlers({
 
     // A drawn page goes onto the model's page of the canvas file, which
     // keeps its other pages; a drawing that brings several pages replaces
-    // the file, and gets its variables
+    // the file. Either way it keeps the file's variables unless it sets its own.
     const onAiPage = (drawnXml: string) =>
-        keepFileVars(
-            placeOnPage(drawnXml, chartXMLRef.current, aiPageId()),
-            chartXMLRef.current,
-        )
+        placeOnPage(drawnXml, chartXMLRef.current, aiPageId())
 
     const takeOriginals = (): string | undefined => {
         const [originalXml] = editDiagramOriginalXmlRef.current.values()
@@ -228,20 +197,8 @@ export function useDiagramToolHandlers({
     ) => {
         const { xml } = toolCall.input as { xml: string }
 
-        // DEBUG: Log raw input to diagnose false truncation detection
-        if (DEBUG) {
-            console.log(
-                "[display_diagram] XML ending (last 100 chars):",
-                xml.slice(-100),
-            )
-            console.log("[display_diagram] XML length:", xml.length)
-        }
-
         // Check if XML is truncated (incomplete mxCell indicates truncated output)
         const isTruncated = !isMxCellXmlComplete(xml)
-        if (DEBUG) {
-            console.log("[display_diagram] isTruncated:", isTruncated)
-        }
 
         if (isTruncated) {
             // Store the partial XML for continuation via append_diagram
@@ -298,11 +255,6 @@ NEXT STEP: Call append_diagram with the continuation XML.
             // keeps the diagram from before this failed call
             if (originalXml) onDisplayChart(originalXml, true, "revert")
             // Return error to model - sendAutomaticallyWhen will trigger retry
-            if (DEBUG) {
-                console.log(
-                    "[display_diagram] Adding tool output with state: output-error",
-                )
-            }
             addToolOutput({
                 tool: "display_diagram",
                 toolCallId: toolCall.toolCallId,
@@ -318,12 +270,6 @@ ${finalXml}
             })
         } else {
             // Success - diagram will be rendered by chat-message-display
-            if (DEBUG) {
-                console.log(
-                    "[display_diagram] Success! Checking if VLM validation is enabled...",
-                )
-            }
-
             // VLM validation after successful display
             if (
                 enableVlmValidation &&
@@ -348,12 +294,6 @@ ${finalXml}
                     if (stopped?.()) {
                         updateValidationState(toolCall.toolCallId, "skipped")
                     } else if (capturedPngData) {
-                        if (DEBUG) {
-                            console.log(
-                                "[display_diagram] Captured PNG for validation",
-                            )
-                        }
-
                         // Each retry is a new tool call, so count attempts
                         // per user turn (the chat resets it when the user sends)
                         const attempt = validationRetryCountRef.current + 1
@@ -379,12 +319,6 @@ ${finalXml}
                             if (attempt < MAX_VALIDATION_RETRIES) {
                                 const feedback =
                                     formatValidationFeedback(result)
-                                if (DEBUG) {
-                                    console.log(
-                                        `[display_diagram] Validation failed (attempt ${attempt}/${MAX_VALIDATION_RETRIES}):`,
-                                        result.issues,
-                                    )
-                                }
 
                                 // Notify UI of validation failure (include the image)
                                 updateValidationState(
@@ -407,11 +341,6 @@ ${finalXml}
                                 return
                             } else {
                                 // Last attempt - accept the diagram with warning
-                                if (DEBUG) {
-                                    console.log(
-                                        "[display_diagram] Max validation retries reached, accepting diagram",
-                                    )
-                                }
                                 // Notify UI that we're accepting with issues (include the image)
                                 updateValidationState(
                                     toolCall.toolCallId,
@@ -427,12 +356,6 @@ ${finalXml}
                                 return
                             }
                         } else {
-                            if (DEBUG) {
-                                console.log(
-                                    "[display_diagram] Validation passed!",
-                                )
-                            }
-
                             // Notify UI of success (include the image)
                             // Use "success_with_warnings" if valid but has issues
                             const hasWarnings = result.issues.length > 0
@@ -474,21 +397,11 @@ ${finalXml}
                 }
             }
 
-            if (DEBUG) {
-                console.log(
-                    "[display_diagram] Adding tool output with state: output-available",
-                )
-            }
             addToolOutput({
                 tool: "display_diagram",
                 toolCallId: toolCall.toolCallId,
                 output: "Successfully displayed the diagram.",
             })
-            if (DEBUG) {
-                console.log(
-                    "[display_diagram] Tool output added. Diagram should be visible now.",
-                )
-            }
         }
     }
 
diff --git a/hooks/use-session-manager.ts b/hooks/use-session-manager.ts
index c5f79484..74eafce8 100644
--- a/hooks/use-session-manager.ts
+++ b/hooks/use-session-manager.ts
@@ -12,7 +12,6 @@ import {
     getAllSessionMetadata,
     getSession,
     isIndexedDBAvailable,
-    migrateFromLocalStorage,
     readSessionCount,
     renameSession as renameSessionInDB,
     type SessionMetadata,
@@ -20,10 +19,7 @@ import {
     saveSession,
 } from "@/lib/session-storage"
 import { STORAGE_KEYS } from "@/lib/storage"
-import {
-    type DiagramVersion,
-    versionsFromLegacyHistory,
-} from "@/stores/versions-store"
+import type { DiagramVersion } from "@/stores/versions-store"
 
 export interface SessionData {
     messages: StoredMessage[]
@@ -35,13 +31,6 @@ export interface SessionData {
     title?: string
 }
 
-/** Versions of a stored session, converting the old history format */
-export function getSessionVersions(
-    session: Pick<ChatSession, "versions" | "diagramHistory">,
-): DiagramVersion[] {
-    return session.versions ?? versionsFromLegacyHistory(session.diagramHistory)
-}
-
 // Taken right before a save's data is read: the chat on screen then, and
 // the order of the reads
 export interface SaveTicket {
@@ -192,9 +181,6 @@ export function useSessionManager(
             setIsAvailable(true)
 
             try {
-                // Run migration first (one-time conversion from localStorage)
-                await migrateFromLocalStorage()
-
                 // Load sessions list
                 const metadata = await getAllSessionMetadata()
                 setSessions(metadata)
@@ -334,7 +320,7 @@ export function useSessionManager(
                 xmlSnapshots: session.xmlSnapshots,
                 diagramXml: session.diagramXml,
                 thumbnailDataUrl: session.thumbnailDataUrl,
-                versions: getSessionVersions(session),
+                versions: session.versions ?? [],
             }
         },
         [currentSessionId, changeChat],
@@ -417,9 +403,7 @@ export function useSessionManager(
                     thumbnailDataUrl:
                         data.thumbnailDataUrl ??
                         currentSession.thumbnailDataUrl,
-                    versions:
-                        data.versions ?? getSessionVersions(currentSession),
-                    diagramHistory: undefined,
+                    versions: data.versions ?? currentSession.versions ?? [],
                     updatedAt: Date.now(),
                     // Update title if it's still default and we have messages
                     title:
diff --git a/hooks/use-validate-diagram.ts b/hooks/use-validate-diagram.ts
index d3402400..faec6189 100644
--- a/hooks/use-validate-diagram.ts
+++ b/hooks/use-validate-diagram.ts
@@ -1,10 +1,9 @@
 "use client"
 
 /**
- * Hook for VLM-based diagram validation using AI SDK's useObject.
+ * Hook for VLM-based diagram validation: one request, one JSON result.
  */
 
-import { experimental_useObject as useObject } from "@ai-sdk/react"
 import { useCallback, useRef } from "react"
 import { getSelectedAIConfig } from "@/hooks/use-model-config"
 import { getApiEndpoint } from "@/lib/base-path"
@@ -13,8 +12,6 @@ import {
     ValidationResultSchema,
 } from "@/lib/validation-schema"
 
-export type { ValidationResult }
-
 // Default valid result for fallback cases
 const DEFAULT_VALID_RESULT: ValidationResult = {
     valid: true,
@@ -22,111 +19,42 @@ const DEFAULT_VALID_RESULT: ValidationResult = {
     suggestions: [],
 }
 
-interface UseValidateDiagramOptions {
-    onSuccess?: (result: ValidationResult) => void
-    onError?: (error: Error) => void
-}
-
-// Track pending validation promises for imperative API
-type PendingValidation = {
-    resolve: (result: ValidationResult) => void
-    reject: (error: Error) => void
-}
-
-export function useValidateDiagram(options: UseValidateDiagramOptions = {}) {
-    const { onSuccess, onError } = options
-    const pendingValidationRef = useRef<PendingValidation | null>(null)
-
-    const { object, submit, isLoading, error, stop } = useObject({
-        api: getApiEndpoint("/api/validate-diagram"),
-        schema: ValidationResultSchema,
-        // Resolved per request so a changed access code is picked up
-        headers: () => ({ "x-access-code": getSelectedAIConfig().accessCode }),
-        onFinish: ({
-            object,
-            error: finishError,
-        }: {
-            object: ValidationResult | undefined
-            error: Error | undefined
-        }) => {
-            if (finishError) {
-                console.error(
-                    "[useValidateDiagram] Validation error:",
-                    finishError,
-                )
-                onError?.(finishError)
-                pendingValidationRef.current?.reject(finishError)
-                pendingValidationRef.current = null
-                return
-            }
-
-            if (object) {
-                const result = object as ValidationResult
-                onSuccess?.(result)
-                pendingValidationRef.current?.resolve(result)
-                pendingValidationRef.current = null
-            }
-        },
-        onError: (err: Error) => {
-            console.error("[useValidateDiagram] Stream error:", err)
-            onError?.(err)
-            pendingValidationRef.current?.reject(err)
-            pendingValidationRef.current = null
-        },
-    })
+export function useValidateDiagram() {
+    // The running check's controller; a new check or cancel() aborts it
+    const controllerRef = useRef<AbortController | null>(null)
 
     /**
-     * Validate a diagram image.
-     * Returns a promise that resolves with the validation result.
-     */
-    const validate = useCallback(
-        async (
-            imageData: string,
-            sessionId?: string,
-        ): Promise<ValidationResult> => {
-            // Reject any pending validation to prevent promise leaks
-            if (pendingValidationRef.current) {
-                pendingValidationRef.current.reject(
-                    new Error("Validation superseded by new request"),
-                )
-                pendingValidationRef.current = null
-            }
-
-            return new Promise((resolve, reject) => {
-                // Store the promise handlers
-                pendingValidationRef.current = { resolve, reject }
-
-                // Submit the validation request
-                submit({ imageData, sessionId })
-            })
-        },
-        [submit],
-    )
-
-    /**
-     * End a running check (the user pressed Stop): its promise rejects with
-     * an AbortError, so the tool handler can finish at once.
-     */
-    const cancel = useCallback(() => {
-        const pending = pendingValidationRef.current
-        if (!pending) return
-        pendingValidationRef.current = null
-        stop()
-        pending.reject(new DOMException("Validation cancelled", "AbortError"))
-    }, [stop])
-
-    /**
-     * Validate with fallback - returns default valid result on error.
-     * Use this to avoid blocking the user on validation failures.
-     * A cancelled check is passed on as its AbortError.
+     * Validate a diagram image, with fallback: a failed check gives the
+     * default valid result so it never blocks the user. A cancelled check
+     * is passed on as its AbortError.
      */
     const validateWithFallback = useCallback(
         async (
             imageData: string,
             sessionId?: string,
         ): Promise<ValidationResult> => {
+            controllerRef.current?.abort()
+            const controller = new AbortController()
+            controllerRef.current = controller
             try {
-                return await validate(imageData, sessionId)
+                const res = await fetch(
+                    getApiEndpoint("/api/validate-diagram"),
+                    {
+                        method: "POST",
+                        headers: {
+                            "Content-Type": "application/json",
+                            // Read per request so a changed access code is
+                            // picked up
+                            "x-access-code": getSelectedAIConfig().accessCode,
+                        },
+                        body: JSON.stringify({ imageData, sessionId }),
+                        signal: controller.signal,
+                    },
+                )
+                if (!res.ok) {
+                    throw new Error(`Validation request failed: ${res.status}`)
+                }
+                return ValidationResultSchema.parse(await res.json())
             } catch (error) {
                 if ((error as Error)?.name === "AbortError") throw error
                 console.warn(
@@ -134,21 +62,23 @@ export function useValidateDiagram(options: UseValidateDiagramOptions = {}) {
                     error,
                 )
                 return DEFAULT_VALID_RESULT
+            } finally {
+                if (controllerRef.current === controller) {
+                    controllerRef.current = null
+                }
             }
         },
-        [validate],
+        [],
     )
 
-    return {
-        // Validation functions
-        validate,
-        validateWithFallback,
-        stop,
-        cancel,
+    /**
+     * End a running check (the user pressed Stop): its promise rejects with
+     * an AbortError, so the tool handler can finish at once.
+     */
+    const cancel = useCallback(() => {
+        controllerRef.current?.abort()
+        controllerRef.current = null
+    }, [])
 
-        // State
-        isValidating: isLoading,
-        partialResult: object as ValidationResult | undefined,
-        error,
-    }
+    return { validateWithFallback, cancel }
 }
diff --git a/lib/diagram-diff.ts b/lib/diagram-diff.ts
index cabd5cf3..01e8fdb8 100644
--- a/lib/diagram-diff.ts
+++ b/lib/diagram-diff.ts
@@ -1,4 +1,9 @@
-import { modelOfPage, pageElement } from "@/lib/diagram-pages"
+import {
+    fileVars,
+    modelOfPage,
+    pageElement,
+    parseXml,
+} from "@/lib/diagram-pages"
 
 /** What changed on one page between two versions of a diagram */
 export interface ChangeSummary {
@@ -32,18 +37,12 @@ interface CellInfo {
     signature: string
 }
 
-function parse(xml: string): Document | null {
-    if (!xml?.trim()) return null
-    const doc = new DOMParser().parseFromString(xml, "text/xml")
-    return doc.querySelector("parsererror") ? null : doc
-}
-
 /**
  * A page's mxGraphModel element, whatever wrapper the XML has: the page
  * with this id, else the first
  */
 function pageModel(xml: string, pageId?: string | null): Element | null {
-    const doc = parse(xml)
+    const doc = parseXml(xml)
     if (!doc) return null
     const diagram = pageElement(doc, pageId)
     if (diagram) return modelOfPage(diagram)
@@ -298,14 +297,10 @@ export function isSameDocument(a: string, b: string): boolean {
             model: modelOfPage(d),
         }))
     }
-    const docA = parse(a)
-    const docB = parse(b)
+    const docA = parseXml(a)
+    const docB = parseXml(b)
     if (!docA || !docB) return a === b
-    const varsOf = (doc: Document) =>
-        doc.documentElement.nodeName === "mxfile"
-            ? doc.documentElement.getAttribute("vars")
-            : null
-    if (!sameFileVars(varsOf(docA), varsOf(docB))) return false
+    if (!sameFileVars(fileVars(docA), fileVars(docB))) return false
     const pagesA = pagesOf(docA)
     const pagesB = pagesOf(docB)
     if (pagesA.length !== pagesB.length) return false
diff --git a/lib/diagram-pages.ts b/lib/diagram-pages.ts
index 3a885568..1f47e664 100644
--- a/lib/diagram-pages.ts
+++ b/lib/diagram-pages.ts
@@ -12,12 +12,20 @@ import {
     type PageSelector,
 } from "@/packages/mcp-server/src/pages.ts"
 
-function parse(xml: string): Document | null {
+/** The parsed XML, or null when it is empty or does not parse */
+export function parseXml(xml: string): Document | null {
     if (!xml?.trim()) return null
     const doc = new DOMParser().parseFromString(xml, "text/xml")
     return doc.querySelector("parsererror") ? null : doc
 }
 
+/** The file variables (%name% placeholders) an <mxfile> sets, else null */
+export function fileVars(doc: Document): string | null {
+    return doc.documentElement.nodeName === "mxfile"
+        ? doc.documentElement.getAttribute("vars")
+        : null
+}
+
 /** The <diagram> the AI works on: the one with this id, else the first */
 export function pageElement(
     doc: Document,
@@ -40,7 +48,7 @@ export function modelOfPage(diagram: Element): Element | null {
     if (model) return model
     const inflated = decompressPageContent(diagram.textContent || "")
     if (!inflated) return null
-    return parse(inflated)?.querySelector("mxGraphModel") ?? null
+    return parseXml(inflated)?.querySelector("mxGraphModel") ?? null
 }
 
 /**
@@ -52,7 +60,7 @@ export function pageSelectorFor(
     pageId: string | null | undefined,
 ): PageSelector {
     if (!pageId) return {}
-    const doc = parse(xml)
+    const doc = parseXml(xml)
     const found = doc
         ? Array.from(doc.getElementsByTagName("diagram")).some(
               (diagram) => diagram.getAttribute("id") === pageId,
@@ -69,7 +77,7 @@ export function pageModelXml(
     xml: string,
     pageId: string | null | undefined,
 ): string | null {
-    const doc = parse(xml)
+    const doc = parseXml(xml)
     if (!doc) return null
     const root = doc.documentElement
     if (root.nodeName === "mxGraphModel") return xml
@@ -90,15 +98,15 @@ export function replacePageModel(
     pageId: string | null | undefined,
     modelXml: string,
 ): string {
-    const model = parse(modelXml)
+    const model = parseXml(modelXml)
     if (model?.documentElement.nodeName !== "mxGraphModel") {
         throw new Error("replacePageModel needs an <mxGraphModel>")
     }
-    let doc = parse(fileXml)
-    if (doc?.documentElement.nodeName !== "mxfile") doc = parse(BLANK_MXFILE)
+    let doc = parseXml(fileXml)
+    if (doc?.documentElement.nodeName !== "mxfile") doc = parseXml(BLANK_MXFILE)
     let diagram = pageElement(doc as Document, pageId)
     if (!diagram) {
-        doc = parse(BLANK_MXFILE)
+        doc = parseXml(BLANK_MXFILE)
         diagram = pageElement(doc as Document, null)
     }
     const target = diagram as Element
@@ -112,27 +120,37 @@ export function replacePageModel(
 /**
  * Where a drawn diagram goes. One page's worth (bare cells wrapped by
  * prepareNewDiagram, or a one-page file) replaces the AI's page of the
- * canvas file; file variables the drawn file sets replace the canvas
- * file's. A file with several pages replaces the whole document: the model
- * wrote the pages on purpose.
+ * canvas file. A file with several pages replaces the whole document: the
+ * model wrote the pages on purpose. Either way, file variables the drawn
+ * file sets replace the canvas file's, and a drawn file without any keeps
+ * the canvas file's (as replacing the page in draw.io does).
  */
 export function placeOnPage(
     drawnXml: string,
     canvasXml: string,
     pageId: string | null | undefined,
 ): string {
-    const doc = parse(drawnXml)
+    const doc = parseXml(drawnXml)
     const drawnFile =
         doc?.documentElement.nodeName === "mxfile" ? doc.documentElement : null
-    if (drawnFile && drawnFile.getElementsByTagName("diagram").length > 1) {
-        return drawnXml
+    const severalPages =
+        (drawnFile?.getElementsByTagName("diagram").length ?? 0) > 1
+    const model = severalPages ? null : pageModelXml(drawnXml, null)
+    if (!model) {
+        // The drawn document as it is, with the canvas file's variables
+        // when it sets none itself
+        const canvas = parseXml(canvasXml)
+        const canvasVars = canvas ? fileVars(canvas) : null
+        if (!doc || !drawnFile || fileVars(doc) !== null || !canvasVars) {
+            return drawnXml
+        }
+        drawnFile.setAttribute("vars", canvasVars)
+        return new XMLSerializer().serializeToString(doc)
     }
-    const model = pageModelXml(drawnXml, null)
-    if (!model) return drawnXml
     const placed = replacePageModel(canvasXml, pageId, model)
     const vars = drawnFile?.getAttribute("vars")
     if (vars === null || vars === undefined) return placed
-    const result = parse(placed) as Document
+    const result = parseXml(placed) as Document
     result.documentElement.setAttribute("vars", vars)
     return new XMLSerializer().serializeToString(result)
 }
diff --git a/lib/drawio/editor-bridge.ts b/lib/drawio/editor-bridge.ts
index 89f67506..6af32de3 100644
--- a/lib/drawio/editor-bridge.ts
+++ b/lib/drawio/editor-bridge.ts
@@ -600,18 +600,6 @@ function readSelectionRect() {
     }
 }
 
-// ---------------------------------------------------------------------------
-// Pages
-// ---------------------------------------------------------------------------
-
-function readPages() {
-    const pages = (ui?.pages ?? []) as any[]
-    return pages.map((page, index) => ({
-        id: String(page.getId?.() ?? index),
-        name: String(page.getName?.() ?? `Page-${index + 1}`),
-    }))
-}
-
 // ---------------------------------------------------------------------------
 // Keeping the store in sync
 // ---------------------------------------------------------------------------
@@ -624,7 +612,6 @@ function syncAll() {
     if (!ui) return
     const g = graph()
     useCanvasStore.getState().set({
-        pages: readPages(),
         currentPageId: ui.currentPage?.getId
             ? String(ui.currentPage.getId())
             : null,
diff --git a/lib/session-storage.ts b/lib/session-storage.ts
index abead739..f149dcce 100644
--- a/lib/session-storage.ts
+++ b/lib/session-storage.ts
@@ -2,13 +2,11 @@ import { type DBSchema, type IDBPDatabase, openDB } from "idb"
 import { nanoid } from "nanoid"
 import { toast } from "sonner"
 import type { DiagramVersion } from "@/stores/versions-store"
-import type { Template } from "./template-storage"
 
 // Constants
 const DB_NAME = "next-ai-drawio"
 const DB_VERSION = 2
 const STORE_NAME = "sessions"
-const MIGRATION_FLAG = "next-ai-drawio-migrated-to-idb"
 const MAX_SESSIONS = 50
 
 // Types
@@ -21,8 +19,6 @@ export interface ChatSession {
     xmlSnapshots: [number, string][]
     diagramXml: string
     thumbnailDataUrl?: string // Small PNG preview of the diagram
-    // Old version list (before AI versions existed); read once, then replaced
-    diagramHistory?: { svg: string; xml: string }[]
     // Diagram states produced by AI changes, shown as cards in the chat
     versions?: DiagramVersion[]
 }
@@ -51,16 +47,6 @@ interface ChatSessionDB extends DBSchema {
         value: ChatSession
         indexes: { "by-updated": number }
     }
-    templates: {
-        key: string
-        value: Template
-        indexes: {
-            "by-updated": number
-            "by-pinned": number
-            "by-run-count": number
-            "by-last-used": number
-        }
-    }
 }
 
 // Database singleton
@@ -77,24 +63,6 @@ async function getDB(): Promise<IDBPDatabase<ChatSessionDB>> {
                     })
                     store.createIndex("by-updated", "updatedAt")
                 }
-                // Version 2: templates store (added by template-storage.ts)
-                // Note: We also need to include this here to ensure the upgrade
-                // callback properly handles all migrations when opening from this file
-                if (oldVersion < 2) {
-                    // Check if templates store already exists (created by template-storage.ts)
-                    if (!db.objectStoreNames.contains("templates")) {
-                        const templateStore = db.createObjectStore(
-                            "templates",
-                            {
-                                keyPath: "id",
-                            },
-                        )
-                        templateStore.createIndex("by-updated", "updatedAt")
-                        templateStore.createIndex("by-pinned", "pinned")
-                        templateStore.createIndex("by-run-count", "runCount")
-                        templateStore.createIndex("by-last-used", "lastUsedAt")
-                    }
-                }
             },
             blocked() {
                 // An older tab keeps the DB open, so the upgrade has to wait
@@ -307,16 +275,13 @@ export function sanitizeMessage(message: unknown): StoredMessage | null {
     const role = msg.role as string
     if (!["user", "assistant", "system"].includes(role)) return null
 
-    // Extract parts, removing streaming state artifacts
     let parts: Array<{ type: string; [key: string]: unknown }> = []
     if (Array.isArray(msg.parts)) {
-        parts = msg.parts.map((part: unknown) => {
-            if (!part || typeof part !== "object") return { type: "unknown" }
-            const p = part as Record<string, unknown>
-            // Remove streaming-related fields
-            const { isStreaming, streamingState, ...cleanPart } = p
-            return cleanPart as { type: string; [key: string]: unknown }
-        })
+        parts = msg.parts.map((part: unknown) =>
+            part && typeof part === "object"
+                ? (part as { type: string; [key: string]: unknown })
+                : { type: "unknown" },
+        )
     }
 
     return {
@@ -332,74 +297,3 @@ export function sanitizeMessages(messages: unknown[]): StoredMessage[] {
         .map(sanitizeMessage)
         .filter((m): m is StoredMessage => m !== null)
 }
-
-// Migration from localStorage
-export async function migrateFromLocalStorage(): Promise<string | null> {
-    if (typeof window === "undefined") return null
-    if (!isIndexedDBAvailable()) return null
-
-    // Check if already migrated
-    if (localStorage.getItem(MIGRATION_FLAG)) return null
-
-    try {
-        const savedMessages = localStorage.getItem("next-ai-draw-io-messages")
-        const savedSnapshots = localStorage.getItem(
-            "next-ai-draw-io-xml-snapshots",
-        )
-        const savedXml = localStorage.getItem("next-ai-draw-io-diagram-xml")
-
-        let newSessionId: string | null = null
-        let migrationSucceeded = false
-
-        if (savedMessages) {
-            const messages = JSON.parse(savedMessages)
-            if (Array.isArray(messages) && messages.length > 0) {
-                const sanitized = sanitizeMessages(messages)
-                const session: ChatSession = {
-                    id: nanoid(),
-                    title: extractTitle(sanitized),
-                    createdAt: Date.now(),
-                    updatedAt: Date.now(),
-                    messages: sanitized,
-                    xmlSnapshots: savedSnapshots
-                        ? JSON.parse(savedSnapshots)
-                        : [],
-                    diagramXml: savedXml || "",
-                }
-                const saved = await saveSession(session)
-                if (saved) {
-                    // Verify the session was actually written
-                    const verified = await getSession(session.id)
-                    if (verified) {
-                        newSessionId = session.id
-                        migrationSucceeded = true
-                    }
-                }
-            } else {
-                // Empty array or invalid data - nothing to migrate, mark as success
-                migrationSucceeded = true
-            }
-        } else {
-            // No data to migrate - mark as success
-            migrationSucceeded = true
-        }
-
-        // Only clean up old data if migration succeeded
-        if (migrationSucceeded) {
-            localStorage.setItem(MIGRATION_FLAG, "true")
-            localStorage.removeItem("next-ai-draw-io-messages")
-            localStorage.removeItem("next-ai-draw-io-xml-snapshots")
-            localStorage.removeItem("next-ai-draw-io-diagram-xml")
-        } else {
-            console.warn(
-                "Migration to IndexedDB failed - keeping localStorage data for retry",
-            )
-        }
-
-        return newSessionId
-    } catch (error) {
-        console.error("Migration failed:", error)
-        // Don't mark as migrated - allow retry on next load
-        return null
-    }
-}
diff --git a/lib/use-quota-manager.tsx b/lib/use-quota-manager.tsx
index fd705c67..9cdefca5 100644
--- a/lib/use-quota-manager.tsx
+++ b/lib/use-quota-manager.tsx
@@ -7,34 +7,30 @@ import { useDictionary } from "@/hooks/use-dictionary"
 import { formatMessage } from "@/lib/i18n/utils"
 
 export interface QuotaConfig {
-    dailyRequestLimit: number
-    dailyTokenLimit: number
-    tpmLimit: number
     onConfigModel?: () => void
 }
 
 /**
  * Hook for displaying quota limit toasts.
- * Server-side handles actual quota enforcement via DynamoDB.
- * This hook only provides UI feedback when limits are exceeded.
+ * Server-side handles actual quota enforcement via DynamoDB, and its 429
+ * responses carry the used and limit values shown here.
  */
 export function useQuotaManager(config: QuotaConfig): {
-    showQuotaLimitToast: (used?: number, limit?: number) => void
-    showTokenLimitToast: (used?: number, limit?: number) => void
-    showTPMLimitToast: (limit?: number) => void
+    showQuotaLimitToast: (used: number, limit: number) => void
+    showTokenLimitToast: (used: number, limit: number) => void
+    showTPMLimitToast: (limit: number) => void
 } {
-    const { dailyRequestLimit, dailyTokenLimit, tpmLimit, onConfigModel } =
-        config
+    const { onConfigModel } = config
     const dict = useDictionary()
 
     // Show quota limit toast (request-based)
     const showQuotaLimitToast = useCallback(
-        (used?: number, limit?: number) => {
+        (used: number, limit: number) => {
             toast.custom(
                 (t) => (
                     <QuotaLimitToast
-                        used={used ?? dailyRequestLimit}
-                        limit={limit ?? dailyRequestLimit}
+                        used={used}
+                        limit={limit}
                         onDismiss={() => toast.dismiss(t)}
                         onConfigModel={onConfigModel}
                     />
@@ -42,18 +38,18 @@ export function useQuotaManager(config: QuotaConfig): {
                 { duration: 15000 },
             )
         },
-        [dailyRequestLimit, onConfigModel],
+        [onConfigModel],
     )
 
     // Show token limit toast
     const showTokenLimitToast = useCallback(
-        (used?: number, limit?: number) => {
+        (used: number, limit: number) => {
             toast.custom(
                 (t) => (
                     <QuotaLimitToast
                         type="token"
-                        used={used ?? dailyTokenLimit}
-                        limit={limit ?? dailyTokenLimit}
+                        used={used}
+                        limit={limit}
                         onDismiss={() => toast.dismiss(t)}
                         onConfigModel={onConfigModel}
                     />
@@ -61,24 +57,21 @@ export function useQuotaManager(config: QuotaConfig): {
                 { duration: 15000 },
             )
         },
-        [dailyTokenLimit, onConfigModel],
+        [onConfigModel],
     )
 
     // Show TPM limit toast
     const showTPMLimitToast = useCallback(
-        (limit?: number) => {
-            const effectiveLimit = limit ?? tpmLimit
+        (limit: number) => {
             const limitDisplay =
-                effectiveLimit >= 1000
-                    ? `${effectiveLimit / 1000}k`
-                    : String(effectiveLimit)
+                limit >= 1000 ? `${limit / 1000}k` : String(limit)
             const message = formatMessage(dict.quota.tpmMessageDetailed, {
                 limit: limitDisplay,
                 seconds: 60,
             })
             toast.error(message, { duration: 8000 })
         },
-        [tpmLimit, dict],
+        [dict],
     )
 
     return {
diff --git a/packages/mcp-server/src/load-diagram.ts b/packages/mcp-server/src/load-diagram.ts
index ab395762..e4ce179f 100644
--- a/packages/mcp-server/src/load-diagram.ts
+++ b/packages/mcp-server/src/load-diagram.ts
@@ -46,7 +46,8 @@ export function decompressPageContent(compressed: string): string | null {
 /**
  * Read the diagram XML a .drawio.svg (Editable SVG) carries in its root
  * element's content attribute. Returns undefined when the text is not an
- * SVG, and "" when the SVG has no embedded diagram.
+ * SVG, and "" when the SVG has no embedded diagram. The root may carry a
+ * namespace prefix (<svg:svg>).
  *
  * Uses getAttributeNode().value rather than getAttribute(): linkedom's
  * getAttribute re-escapes <, > and & in XML documents, which would turn
@@ -55,10 +56,16 @@ export function decompressPageContent(compressed: string): string | null {
 export function extractSvgEmbeddedXml(text: string): string | undefined {
     const doc = new DOMParser().parseFromString(text, "text/xml")
     const root = doc.documentElement
-    if (root?.tagName !== "svg") return undefined
+    // (linkedom keeps the prefix in localName too)
+    if (root?.tagName.split(":").pop() !== "svg") return undefined
     return (root.getAttributeNode("content")?.value ?? "").trim()
 }
 
+// What XML allows before the root element: the declaration, comments, a
+// DOCTYPE (the root checks below expect the root first)
+const XML_PROLOG =
+    /^(?:\s|<\?[\s\S]*?\?>|<!--[\s\S]*?-->|<!DOCTYPE[^>[]*(?:\[[\s\S]*?\])?\s*>)*/i
+
 /**
  * Parse the content of a .drawio file into the canonical session shape:
  * an <mxfile> whose every page holds plain <mxGraphModel> XML. Accepts a
@@ -66,7 +73,7 @@ export function extractSvgEmbeddedXml(text: string): string | undefined {
  * the diagram embedded, and decompresses any compressed pages.
  */
 export function parseDrawioFileContent(content: string): LoadResult {
-    let trimmed = content.trim()
+    let trimmed = content.replace(XML_PROLOG, "").trim()
     if (!trimmed) return { ok: false, error: "File is empty." }
 
     if (isMxGraphModel(trimmed)) {
@@ -78,7 +85,8 @@ export function parseDrawioFileContent(content: string): LoadResult {
         trimmed = normalized
     }
     if (!isMxFile(trimmed)) {
-        const embedded = extractSvgEmbeddedXml(trimmed)
+        // Parsed whole: the SVG's DOCTYPE may define entities its content uses
+        const embedded = extractSvgEmbeddedXml(content)
         if (embedded === undefined) {
             return {
                 ok: false,
diff --git a/packages/mcp-server/tests/load-diagram.test.ts b/packages/mcp-server/tests/load-diagram.test.ts
index 1dc12a2e..065bc555 100644
--- a/packages/mcp-server/tests/load-diagram.test.ts
+++ b/packages/mcp-server/tests/load-diagram.test.ts
@@ -190,12 +190,46 @@ const SVG_PROLOG = `<?xml version="1.0" encoding="UTF-8"?>
 <!DOCTYPE svg PUBLIC "-//W3C//DTD SVG 1.1//EN" "http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd">
 `
 
+describe("parseDrawioFileContent with a prolog before the root", () => {
+    it("accepts a declaration, a comment and a DOCTYPE before <mxfile>", () => {
+        const prolog = `<?xml version="1.0" encoding="UTF-8"?>\n<!-- Architecture, v2 -->\n<!DOCTYPE mxfile>\n`
+        expect(parseDrawioFileContent(prolog + PLAIN_MXFILE)).toEqual({
+            ok: true,
+            xml: PLAIN_MXFILE,
+            hadCompressedPages: false,
+        })
+    })
+
+    it("accepts a comment before a bare <mxGraphModel>", () => {
+        const r = parseDrawioFileContent(`<!-- exported -->\n${MODEL_XML}`)
+        expect(r.ok).toBe(true)
+        if (r.ok) expect(r.xml).toContain('value="Hello"')
+    })
+
+    it("treats a file with nothing after its prolog as empty", () => {
+        expect(parseDrawioFileContent(`<?xml version="1.0"?>\n`)).toEqual({
+            ok: false,
+            error: "File is empty.",
+        })
+    })
+})
+
 describe("extractSvgEmbeddedXml", () => {
     it("returns undefined for non-SVG text", () => {
         expect(extractSvgEmbeddedXml(PLAIN_MXFILE)).toBeUndefined()
         expect(extractSvgEmbeddedXml("hello")).toBeUndefined()
     })
 
+    it("accepts a root element with a namespace prefix", () => {
+        const svg = `<svg:svg xmlns:svg="http://www.w3.org/2000/svg" content="${escapeAttr(PLAIN_MXFILE)}"></svg:svg>`
+        expect(extractSvgEmbeddedXml(svg)).toBe(PLAIN_MXFILE)
+        expect(parseDrawioFileContent(svg)).toEqual({
+            ok: true,
+            xml: PLAIN_MXFILE,
+            hadCompressedPages: false,
+        })
+    })
+
     it("returns an empty string for an SVG without content", () => {
         expect(extractSvgEmbeddedXml("<svg><rect/></svg>")).toBe("")
     })
diff --git a/stores/canvas-store.ts b/stores/canvas-store.ts
index 7f6ed600..49763ac9 100644
--- a/stores/canvas-store.ts
+++ b/stores/canvas-store.ts
@@ -14,15 +14,9 @@ interface CanvasRect {
     height: number
 }
 
-interface CanvasPage {
-    id: string
-    name: string
-}
-
 interface CanvasState {
     /** The page can call the draw.io editor directly (same-origin iframe) */
     hasEditor: boolean
-    pages: CanvasPage[]
     currentPageId: string | null
     selection: SelectedCell[]
     selectionRect: CanvasRect | null
@@ -35,7 +29,6 @@ interface CanvasState {
 
 export const useCanvasStore = create<CanvasState>((set) => ({
     hasEditor: false,
-    pages: [],
     currentPageId: null,
     selection: [],
     selectionRect: null,
diff --git a/stores/versions-store.ts b/stores/versions-store.ts
index 86623f05..dca6ed92 100644
--- a/stores/versions-store.ts
+++ b/stores/versions-store.ts
@@ -1,7 +1,6 @@
 import { nanoid } from "nanoid"
 import { create } from "zustand"
 import type { ChangeSummary } from "@/lib/diagram-diff"
-import { EMPTY_SUMMARY } from "@/lib/diagram-diff"
 
 /** One diagram state produced by an AI change, shown as a card in the chat */
 export interface DiagramVersion {
@@ -71,10 +70,8 @@ export const useVersionsStore = create<VersionsState>((set) => ({
     setVersions: (versions) =>
         set({
             // A missing thumbnail could not be made: no "rendering" forever.
-            // Versions saved before they had numbers count from 1.
-            versions: versions.map((v, i) => ({
+            versions: versions.map((v) => ({
                 ...v,
-                number: v.number ?? i + 1,
                 svg: v.svg ?? "",
             })),
             undoneVersionId: versions.find((v) => v.undone)?.id ?? null,
@@ -141,18 +138,3 @@ export const useVersionsStore = create<VersionsState>((set) => ({
     clear: () =>
         set({ versions: [], undoneVersionId: null, onCanvasVersionId: null }),
 }))
-
-/** Convert the old `diagramHistory` list ({svg, xml}) saved by earlier versions */
-export function versionsFromLegacyHistory(
-    history: { svg: string; xml: string }[] | undefined,
-): DiagramVersion[] {
-    if (!history?.length) return []
-    return history.slice(-MAX_VERSIONS).map((entry, index) => ({
-        id: `legacy-${index}`,
-        number: index + 1,
-        xml: entry.xml,
-        svg: entry.svg,
-        turnIndex: -1,
-        summary: { ...EMPTY_SUMMARY },
-    }))
-}
diff --git a/tests/e2e/settings-providers.spec.ts b/tests/e2e/settings-providers.spec.ts
index 69087532..d23ed594 100644
--- a/tests/e2e/settings-providers.spec.ts
+++ b/tests/e2e/settings-providers.spec.ts
@@ -219,14 +219,7 @@ test("an access code error opens the general tab with the code field", async ({
     page,
 }) => {
     await page.route("**/api/config", (route) =>
-        route.fulfill({
-            json: {
-                accessCodeRequired: true,
-                dailyRequestLimit: 0,
-                dailyTokenLimit: 0,
-                tpmLimit: 0,
-            },
-        }),
+        route.fulfill({ json: { accessCodeRequired: true } }),
     )
     await page.route("**/api/chat", (route) =>
         route.fulfill({
diff --git a/tests/e2e/workspace.spec.ts b/tests/e2e/workspace.spec.ts
index 73a88ed5..bd1f8c6d 100644
--- a/tests/e2e/workspace.spec.ts
+++ b/tests/e2e/workspace.spec.ts
@@ -2987,9 +2987,25 @@ test.describe("Edge cases", () => {
         ).toBeVisible()
     })
 
-    test("compare closes when another chat with the same old version ids comes on screen", async ({
+    test("compare closes when another chat with the same version ids comes on screen", async ({
         page,
     }) => {
+        // Both chats' versions share their ids
+        const version = (id: string, number: number) => ({
+            id,
+            number,
+            xml: `<mxGraphModel><root/></mxGraphModel>`,
+            svg: "",
+            turnIndex: 0,
+            summary: {
+                shapesAdded: 0,
+                shapesRemoved: 0,
+                shapesChanged: 0,
+                edgesAdded: 0,
+                edgesRemoved: 0,
+                edgesChanged: 0,
+            },
+        })
         const legacy = (id: string, label: string) => ({
             id,
             title: `Chat ${id}`,
@@ -3004,10 +3020,7 @@ test.describe("Edge cases", () => {
             ],
             xmlSnapshots: [],
             diagramXml: `<mxfile><diagram name="P" id="p"><mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/>${boxCell("a", label, 40)}</root></mxGraphModel></diagram></mxfile>`,
-            diagramHistory: [
-                { svg: "", xml: `<mxGraphModel><root/></mxGraphModel>` },
-                { svg: "", xml: `<mxGraphModel><root/></mxGraphModel>` },
-            ],
+            versions: [version("v-1", 1), version("v-2", 2)],
         })
         await openApp(page)
         await page.evaluate(
diff --git a/tests/unit/diagram-context.test.tsx b/tests/unit/diagram-context.test.tsx
index a251b507..380a99c0 100644
--- a/tests/unit/diagram-context.test.tsx
+++ b/tests/unit/diagram-context.test.tsx
@@ -18,17 +18,18 @@ function setup() {
         exportDiagram: (r: any) => requests.push(r),
         load: vi.fn(),
     } as any
-    // draw.io's reply to a request: it echoes the request in `message`
+    // draw.io's reply to a request: it echoes the request in `message`.
+    // Async, so what the caller does with the result runs before it returns
     const reply = (request: { message: string }, data: string, xml = "") =>
-        act(() =>
+        act(async () => {
             result.current.handleDiagramExport({
                 event: "export",
                 data,
                 xml,
                 format: "xmlsvg",
                 message: request,
-            } as any),
-        )
+            } as any)
+        })
     return { result, requests, reply }
 }
 
@@ -53,25 +54,20 @@ afterEach(() => {
 })
 
 describe("exports in flight at the same time", () => {
-    it("give the chat's export only its own reply", () => {
+    it("give the chat's export only its own reply", async () => {
         const { result, requests, reply } = setup()
         // An edit's history export is still on its way when the chat exports
-        act(() => {
-            result.current.handleExport()
+        const older = result.current.fetchDiagramXml()
+        const current = result.current.fetchDiagramXml()
+        let currentDone = false
+        current.then(() => {
+            currentDone = true
         })
-        let tag = ""
-        const got: string[] = []
-        act(() => {
-            tag = result.current.handleExport()
-            result.current.exportResolversRef.current[tag] = (xml) =>
-                got.push(xml)
-        })
-        reply(requests[0], svgOf("older"))
-        expect(got).toEqual([])
-        reply(requests[1], svgOf("current"))
-        expect(got).toHaveLength(1)
-        expect(got[0]).toContain('value="current"')
-        expect(result.current.exportResolversRef.current[tag]).toBeUndefined()
+        await reply(requests[0], svgOf("older"))
+        expect(await older).toContain('value="older"')
+        expect(currentDone).toBe(false)
+        await reply(requests[1], svgOf("current"))
+        expect(await current).toContain('value="current"')
     })
 
     it("save each file with its own result", async () => {
@@ -101,8 +97,8 @@ describe("exports in flight at the same time", () => {
             result.current.saveDiagramToFile("pic", "png")
         })
         // The PNG answers first
-        reply(requests[1], "data:image/png;base64,iVBORw0KGgo=")
-        reply(requests[0], svgOf("doc"), twoPages)
+        await reply(requests[1], "data:image/png;base64,iVBORw0KGgo=")
+        await reply(requests[0], svgOf("doc"), twoPages)
 
         expect(saved.map((s) => s.name)).toEqual(["pic.png", "doc.drawio"])
         expect(saved[0].href).toMatch(/^data:image\/png/)
diff --git a/tests/unit/diagram-pages.test.ts b/tests/unit/diagram-pages.test.ts
index c5f26d44..ae6e8da2 100644
--- a/tests/unit/diagram-pages.test.ts
+++ b/tests/unit/diagram-pages.test.ts
@@ -113,7 +113,24 @@ describe("placeOnPage", () => {
 
     it("lets a drawing with several pages replace the document", () => {
         const drawn = `<mxfile>${page("x", box("X"))}${page("y", box("Y"))}</mxfile>`
-        expect(placeOnPage(drawn, TWO_PAGES, "b")).toBe(drawn)
+        expect(pagesOf(placeOnPage(drawn, TWO_PAGES, "b"))).toEqual([
+            { id: "x", name: "Page x", cells: ["X"] },
+            { id: "y", name: "Page y", cells: ["Y"] },
+        ])
+        // A canvas without variables leaves it as it is
+        const plainCanvas = `<mxfile>${page("a", box("A"))}</mxfile>`
+        expect(placeOnPage(drawn, plainCanvas, "a")).toBe(drawn)
+    })
+
+    it("gives a drawing with several pages the canvas's variables unless it has its own", () => {
+        const drawn = `<mxfile>${page("x", box("X"))}${page("y", box("Y"))}</mxfile>`
+        expect(placeOnPage(drawn, TWO_PAGES, "b")).toContain(
+            'vars="{&quot;x&quot;:&quot;1&quot;}"',
+        )
+        const withVars = `<mxfile vars="{&quot;team&quot;:&quot;New&quot;}">${page("x", box("X"))}${page("y", box("Y"))}</mxfile>`
+        const result = placeOnPage(withVars, TWO_PAGES, "b")
+        expect(result).toContain('vars="{&quot;team&quot;:&quot;New&quot;}"')
+        expect(result).not.toContain("&quot;x&quot;")
     })
 
     it("draws on the blank file when the canvas is empty", () => {
diff --git a/tests/unit/redesign-helpers.test.ts b/tests/unit/redesign-helpers.test.ts
index 58ce8e60..3ce3ee88 100644
--- a/tests/unit/redesign-helpers.test.ts
+++ b/tests/unit/redesign-helpers.test.ts
@@ -7,11 +7,7 @@ import jaDict from "@/lib/i18n/dictionaries/ja.json"
 import zhDict from "@/lib/i18n/dictionaries/zh.json"
 import zhHantDict from "@/lib/i18n/dictionaries/zh-Hant.json"
 import { describeChanges, describeTotals } from "@/lib/version-text"
-import {
-    MAX_VERSIONS,
-    useVersionsStore,
-    versionsFromLegacyHistory,
-} from "@/stores/versions-store"
+import { MAX_VERSIONS, useVersionsStore } from "@/stores/versions-store"
 
 describe("withPageDefaults", () => {
     it("adds page and grid when the model has none", () => {
@@ -157,7 +153,7 @@ describe("versions store", () => {
         expect(versions.at(-1)?.number).toBe(MAX_VERSIONS + 5)
     })
 
-    it("numbers on from the last version kept, and numbers saved ones", () => {
+    it("numbers on from the last version kept", () => {
         add(0)
         add(2)
         add(4)
@@ -166,14 +162,6 @@ describe("versions store", () => {
         expect(
             useVersionsStore.getState().versions.map((v) => v.number),
         ).toEqual([1, 2, 3])
-        const { versions } = useVersionsStore.getState()
-        // Saved before versions had numbers
-        useVersionsStore
-            .getState()
-            .setVersions(versions.map(({ number, ...v }) => v as any))
-        expect(
-            useVersionsStore.getState().versions.map((v) => v.number),
-        ).toEqual([1, 2, 3])
     })
 
     it("drops versions from a turn that is re-run, none left on the canvas", () => {
@@ -199,13 +187,4 @@ describe("versions store", () => {
         expect(state.undoneVersionId).toBeNull()
         expect(state.versions[0].undone).toBeUndefined()
     })
-
-    it("converts the old history list", () => {
-        const versions = versionsFromLegacyHistory([
-            { svg: "data:image/svg+xml;base64,AA", xml: "<a/>" },
-        ])
-        expect(versions).toHaveLength(1)
-        expect(versions[0]).toMatchObject({ xml: "<a/>", turnIndex: -1 })
-        expect(versionsFromLegacyHistory(undefined)).toEqual([])
-    })
 })
diff --git a/tests/unit/use-session-manager.test.tsx b/tests/unit/use-session-manager.test.tsx
index c697db6f..c3c034af 100644
--- a/tests/unit/use-session-manager.test.tsx
+++ b/tests/unit/use-session-manager.test.tsx
@@ -17,7 +17,6 @@ vi.mock("@/lib/session-storage", async (importActual) => {
         createEmptySession: actual.createEmptySession,
         extractTitle: actual.extractTitle,
         isIndexedDBAvailable: () => true,
-        migrateFromLocalStorage: async () => null,
         readSessionCount: async () => stored.size,
         enforceSessionLimit: async () => {},
         getSession: async (id: string) => {
diff --git a/tests/unit/validate-diagram-route.test.ts b/tests/unit/validate-diagram-route.test.ts
index 4f3df7d1..9775ad66 100644
--- a/tests/unit/validate-diagram-route.test.ts
+++ b/tests/unit/validate-diagram-route.test.ts
@@ -1,5 +1,4 @@
 // @vitest-environment node
-import { simulateReadableStream } from "ai"
 import { MockLanguageModelV3 } from "ai/test"
 import { afterEach, describe, expect, it, vi } from "vitest"
 import { POST as validateDiagram } from "@/app/api/validate-diagram/route"
@@ -16,36 +15,19 @@ const RESULT = {
     suggestions: ["Move box B to the right"],
 }
 
-// A vision model that answers with the JSON in a few text chunks
+// A vision model that answers with the result as JSON text
 vi.mock("@/lib/ai-providers", () => ({
     getValidationModel: () =>
         new MockLanguageModelV3({
-            doStream: (async () => {
-                const json = JSON.stringify(RESULT)
-                return {
-                    stream: simulateReadableStream({
-                        chunks: [
-                            { type: "text-start", id: "t" },
-                            ...[json.slice(0, 20), json.slice(20)].map(
-                                (delta) => ({
-                                    type: "text-delta",
-                                    id: "t",
-                                    delta,
-                                }),
-                            ),
-                            { type: "text-end", id: "t" },
-                            {
-                                type: "finish",
-                                finishReason: { unified: "stop", raw: "stop" },
-                                usage: {
-                                    inputTokens: { total: 1 },
-                                    outputTokens: { total: 1 },
-                                },
-                            },
-                        ],
-                    }),
-                }
-            }) as any,
+            doGenerate: (async () => ({
+                content: [{ type: "text", text: JSON.stringify(RESULT) }],
+                finishReason: { unified: "stop", raw: "stop" },
+                usage: {
+                    inputTokens: { total: 1 },
+                    outputTokens: { total: 1 },
+                },
+                warnings: [],
+            })) as any,
         }),
 }))
 
@@ -110,24 +92,25 @@ describe("the quota", () => {
         // get its check, and the check does not count as a chat
         quota.enabled = true
         const res = await validateDiagram(post())
-        expect(JSON.parse(await res.text())).toEqual(RESULT)
+        expect(await res.json()).toEqual(RESULT)
         expect(quota.checks).toHaveLength(1)
         expect(quota.checks[0].increment).toBe(0)
         expect(quota.checks[0].limits.requests).toBe(0)
-        await vi.waitFor(() => expect(quota.recorded).toEqual([2]))
+        expect(quota.recorded).toEqual([2])
     })
 })
 
 describe("POST /api/validate-diagram", () => {
-    it("streams the model's result as JSON text for useObject", async () => {
+    it("answers the model's result as JSON", async () => {
         const res = await validateDiagram(post())
-        expect(JSON.parse(await res.text())).toEqual(RESULT)
+        expect(res.headers.get("content-type")).toContain("application/json")
+        expect(await res.json()).toEqual(RESULT)
     })
 
     it("answers valid when the check is turned off", async () => {
         process.env.ENABLE_VLM_VALIDATION = "false"
         const res = await validateDiagram(post())
-        expect(JSON.parse(await res.text())).toEqual({
+        expect(await res.json()).toEqual({
             valid: true,
             issues: [],
             suggestions: [],