fix: older defects (batch C) and the second batch's review

Chats:
- New Chat right after an answer saves that chat once. Saves run one at
  a time and read the chat on screen when their turn comes; a save
  scheduled for a chat that is no longer on screen is dropped. A chat
  whose id was still on its way to the URL no longer comes back after
  New Chat (the next answer went into it).
- Crossing the 768 px breakpoint keeps the chat panel: a streaming answer,
  unsaved messages and attachments stay. The panel gets the sizes of each
  side, and a panel collapsed on desktop opens on mobile.
- The chat's export waits for its own reply: an edit's history export
  still on its way no longer answers it with the older diagram, and two
  file saves at once no longer swap results.
- A second edit in one answer is previewed on the first edit's result.
- Stop also ends a running screenshot check; a chat that cannot be saved
  (storage full) can be left with "Continue without saving".
- Small diagrams with shapes count as diagrams; the tool card no longer
  crashes on malformed operations.

Quota and providers:
- Requests that reach the server's own endpoints count toward the quota:
  EdgeOne (always its own endpoint now), a private base URL whatever key
  header is sent, keyless Ollama without a URL. With the quota on, a
  redirect is followed only to a public address. The output cap applies
  to these requests too.
- Stop records the tokens of the steps that finished; the screenshot check
  counts its tokens without counting a request.
- EdgeOne configured only by AI_PROVIDER works, also in the admin Test,
  which forwards the access code. Azure set up only in the admin panel
  works in chat. The Test sends a Bedrock session token.
- The admin panel's Test of an entry without a URL uses the server's URL
  as the server does (no private address check for it); the admin panel
  no longer writes an Ollama URL.

MCP server:
- Write tools and start_session run one at a time, so two at once never
  drop each other's change; a cancelled call waiting its turn is skipped.
  get_diagram and export_diagram keep the session they started with.
- Export to .drawio first gets the user's latest edits from the browser.
- History thumbnails: one that arrives after the next AI write is
  dropped; a sync reply keeps the image; a version that changed only page
  settings is its own entry.
- A diagram over the 10 MB limit is saved without its image, or the user
  is told to download it (the server now answers 413 instead of cutting
  the connection).
- Labels holding text like id='1' or parent='1' are no longer read as
  attributes (a layer or a parent was deleted). A broken bare
  <mxGraphModel> file is refused.
- After a sync reply the tab no longer sends its autosave copy again.

Desktop and files:
- A newer switch of the same preset is not rolled back by an older one
  that failed. .env values with escaped quotes are read whole.
- MCP saved files: a file that could not be read stays protected while a
  folder without permission hides it, and is saved again once deleted.
- The desktop app reports "no chats" only when the count was read and no
  model settings are stored.

tests/e2e/chat.spec.ts

@@ -1,4 +1,10 @@
-import { expect, getIframe, test } from "./lib/fixtures"
+import {
+    expect,
+    getChatInput,
+    getIframe,
+    sendMessage,
+    test,
+} from "./lib/fixtures"
 
 test.describe("Chat Panel", () => {
     test.beforeEach(async ({ page }) => {
@@ -20,3 +26,82 @@ test.describe("Chat Panel", () => {
         expect(src).toBeTruthy()
     })
 })
+
+test.describe("Crossing the mobile breakpoint", () => {
+    test.beforeEach(async ({ page }) => {
+        // A text answer that arrives in parts over a few seconds
+        await page.addInitScript(() => {
+            const realFetch = window.fetch
+            window.fetch = async (input, init) => {
+                const url = input instanceof Request ? input.url : String(input)
+                if (!url.endsWith("/api/chat")) return realFetch(input, init)
+                const events = [
+                    { type: "start" },
+                    { type: "text-start", id: "t" },
+                    { type: "text-delta", id: "t", delta: "Once upon" },
+                    { type: "text-delta", id: "t", delta: " a time." },
+                    { type: "text-end", id: "t" },
+                    { type: "finish" },
+                ]
+                const body = new ReadableStream({
+                    async start(controller) {
+                        for (const event of events) {
+                            controller.enqueue(
+                                new TextEncoder().encode(
+                                    `data: ${JSON.stringify(event)}\n\n`,
+                                ),
+                            )
+                            await new Promise((r) => setTimeout(r, 1500))
+                        }
+                        controller.enqueue(
+                            new TextEncoder().encode("data: [DONE]\n\n"),
+                        )
+                        controller.close()
+                    },
+                })
+                return new Response(body, {
+                    headers: { "content-type": "text/event-stream" },
+                })
+            }
+        })
+        await page.setViewportSize({ width: 1280, height: 800 })
+        await page.goto("/", { waitUntil: "networkidle" })
+        await getIframe(page).waitFor({ state: "visible", timeout: 30000 })
+    })
+
+    test("keeps the chat and its streaming answer", async ({ page }) => {
+        const chat = page.locator('[data-panel-id="chat-panel"]')
+        await sendMessage(page, "Tell me a story")
+        await expect(page.getByText("Once upon")).toBeVisible({
+            timeout: 10000,
+        })
+
+        await page.setViewportSize({ width: 600, height: 900 })
+        await expect(page.getByText("Tell me a story")).toBeVisible()
+        // Half the height on mobile
+        await expect
+            .poll(async () => (await chat.boundingBox())?.height ?? 0)
+            .toBeCloseTo(450, -1)
+
+        await page.setViewportSize({ width: 1280, height: 800 })
+        // A third of the width on desktop
+        await expect
+            .poll(async () => (await chat.boundingBox())?.width ?? 0)
+            .toBeCloseTo(1280 / 3, -1)
+        await expect(page.getByText("Once upon a time.")).toBeVisible({
+            timeout: 10000,
+        })
+        await expect(page.getByText("Tell me a story")).toBeVisible()
+    })
+
+    test("opens a chat collapsed on desktop", async ({ page }) => {
+        await page.locator("button:has(svg.lucide-panel-right-close)").click()
+        await expect(getChatInput(page)).toBeHidden()
+
+        await page.setViewportSize({ width: 600, height: 900 })
+        await expect(getChatInput(page)).toBeVisible()
+
+        await page.setViewportSize({ width: 1280, height: 800 })
+        await expect(getChatInput(page)).toBeVisible()
+    })
+})

tests/e2e/diagram-content.spec.ts

@@ -409,6 +409,14 @@ const drawReply = (id: string, xml: string) => {
     const call = toolCallEvents(id, "display_diagram", { xml })
     return `${sse([{ type: "start" }, call.start, ...call.deltas, call.done, { type: "finish" }])}data: [DONE]\n\n`
 }
+const textReply = (text: string) =>
+    `${sse([
+        { type: "start" },
+        { type: "text-start", id: "t" },
+        { type: "text-delta", id: "t", delta: text },
+        { type: "text-end", id: "t" },
+        { type: "finish" },
+    ])}data: [DONE]\n\n`
 // SSE comments keep a stream open without sending anything
 const KEEP_OPEN = Array(20).fill(":\n\n")
 
@@ -685,3 +693,45 @@ test("stopping during the screenshot check starts no new request", async ({
     await p.waitForTimeout(5000)
     expect(chatRequests).toBe(1)
 })
+
+test("stopping during the screenshot check lets the next message go at once", async ({
+    page: p,
+}) => {
+    // The check was still running when the user stopped; it held up the
+    // chat until it ended, and its call never got a result
+    await p.addInitScript(() => {
+        localStorage.setItem("next-ai-draw-io-vlm-validation-enabled", "true")
+    })
+    const bodies: Array<{ messages: any[] }> = []
+    await p.route("**/api/chat", async (route) => {
+        bodies.push(route.request().postDataJSON())
+        const n = bodies.length
+        await route.fulfill({
+            status: 200,
+            contentType: "text/event-stream",
+            body:
+                n === 1
+                    ? drawReply("d1", cell("a", "Alpha", 40))
+                    : textReply("Second answer"),
+        })
+    })
+    let checking = false
+    await p.route("**/api/validate-diagram", async (route) => {
+        checking = true
+        // Much longer than this test waits for the second answer
+        await new Promise((r) => setTimeout(r, 30000))
+        await route.fulfill({ status: 200, body: "{}" }).catch(() => {})
+    })
+    await p.goto("/", { waitUntil: "networkidle" })
+    await getIframe(p).waitFor({ state: "visible", timeout: 30000 })
+    await sendMessage(p, "Draw a box")
+    await expect.poll(() => checking, { timeout: 15000 }).toBe(true)
+    await p.getByRole("button", { name: "Stop generation" }).click()
+    await sendMessage(p, "Thanks")
+    await expect(p.getByText("Second answer")).toBeVisible({ timeout: 8000 })
+    // The drawing call had its result when the next message was sent
+    const draw = bodies[1].messages
+        .flatMap((m: any) => m.parts ?? [])
+        .find((part: any) => part.type === "tool-display_diagram")
+    expect(draw?.state).toBe("output-available")
+})

tests/e2e/history-restore.spec.ts

@@ -86,6 +86,40 @@ test.describe("History and Session Restore", () => {
         ).toBeVisible()
     })
 
+    test("new chat can go on without saving when storage is full", async ({
+        page,
+    }) => {
+        // Old chats can only be deleted from the empty chat's list, so the
+        // user must be able to get there
+        await page.route("**/api/chat", async (route) => {
+            await route.fulfill({
+                status: 200,
+                contentType: "text/event-stream",
+                body: createMockSSEResponse(
+                    SINGLE_BOX_XML,
+                    "Created your test diagram.",
+                ),
+            })
+        })
+        await page.goto("/", { waitUntil: "networkidle" })
+        await getIframe(page).waitFor({ state: "visible", timeout: 30000 })
+        await sendMessage(page, "Create a test diagram")
+        await waitForText(page, "Created your test diagram.")
+        await page.evaluate(() => {
+            IDBObjectStore.prototype.put = () => {
+                throw new DOMException("Storage is full", "QuotaExceededError")
+            }
+        })
+        await page.locator('[data-testid="new-chat-button"]').click()
+        await page
+            .getByRole("button", { name: "Continue without saving" })
+            .click({ timeout: 5000 })
+        await expect(
+            page.locator('text="Created your test diagram."'),
+        ).toHaveCount(0, { timeout: 5000 })
+        await expect(page.getByText("Paper to Diagram")).toBeVisible()
+    })
+
     // A diagram drawn by hand, without chat messages: loaded into draw.io
     // directly, then moved with an arrow key, which draw.io reports as an
     // edit like any manual change
@@ -338,3 +372,55 @@ test.describe("History and Session Restore", () => {
         })
     })
 })
+
+/** Number of chats stored in this origin's IndexedDB */
+const countSessions = (page: Page) =>
+    page.evaluate(
+        () =>
+            new Promise<number>((resolve, reject) => {
+                const open = indexedDB.open("next-ai-drawio")
+                open.onerror = () => reject(open.error)
+                open.onsuccess = () => {
+                    const db = open.result
+                    if (!db.objectStoreNames.contains("sessions")) {
+                        db.close()
+                        return resolve(0)
+                    }
+                    const count = db
+                        .transaction("sessions", "readonly")
+                        .objectStore("sessions")
+                        .count()
+                    count.onsuccess = () => {
+                        db.close()
+                        resolve(count.result)
+                    }
+                }
+            }),
+    )
+
+test("new chat right after an answer saves that chat once", async ({
+    page,
+}) => {
+    test.setTimeout(180_000)
+    await page.route("**/api/chat", async (route) => {
+        await route.fulfill({
+            status: 200,
+            contentType: "text/event-stream",
+            body: createMockSSEResponse(SINGLE_BOX_XML, "Drew the box."),
+        })
+    })
+    await page.goto("/", { waitUntil: "networkidle" })
+    await getIframe(page).waitFor({ state: "visible", timeout: 30000 })
+    const newChat = page.locator('[data-testid="new-chat-button"]')
+    // The auto-save runs a second after the answer; New Chat around then
+    // waits for its thumbnail while the auto-save starts
+    for (let run = 1; run <= 10; run++) {
+        await sendMessage(page, `Draw box ${run}`)
+        await waitForText(page, "Drew the box.")
+        await page.waitForTimeout(500 + run * 100)
+        await newChat.click()
+        await expect(page.getByText("Drew the box.")).toHaveCount(0)
+        await page.waitForTimeout(2500)
+        expect(await countSessions(page), `run ${run}`).toBe(run)
+    }
+})

tests/e2e/provider-models.spec.ts

@@ -302,6 +302,35 @@ test("an older test does not end a newer one's spinners", async ({ page }) => {
     await expect(dialog.locator('[title="1.0 s"]')).toHaveCount(1)
 })
 
+test("an older test touches nothing, also when the key came back", async ({
+    page,
+}) => {
+    const releases: Array<() => void> = []
+    await page.route("**/api/validate-model", async (route) => {
+        const n = releases.length
+        await new Promise<void>((r) => releases.push(r))
+        await route.fulfill({
+            status: 200,
+            // The older test's result would say 9.0 s
+            json: { valid: true, responseTime: n === 0 ? 9000 : 1000 },
+        })
+    })
+    const dialog = await openQwenSettings(page, TWO_PROVIDERS)
+    await dialog.getByRole("button", { name: "Test", exact: true }).click()
+    await expect.poll(() => releases.length).toBe(1)
+    // The key changes and comes back, and the user tests again
+    await dialog.locator("#api-key").fill("other-key")
+    await dialog.locator("#api-key").fill("test-key")
+    await dialog.getByRole("button", { name: "Test", exact: true }).click()
+    await expect.poll(() => releases.length).toBe(2)
+    releases[0]()
+    await page.waitForTimeout(500)
+    await expect(dialog.locator(".animate-spin").first()).toBeVisible()
+    await expect(dialog.locator('[title="9.0 s"]')).toHaveCount(0)
+    releases[1]()
+    await expect(dialog.locator('[title="1.0 s"]')).toHaveCount(1)
+})
+
 test("no spinner stays after another tab's change while elsewhere", async ({
     page,
 }) => {

tests/unit/admin-providers.test.ts

@@ -69,15 +69,15 @@ describe("deriveEnvUpdates", () => {
         expect(updates.ADMIN_OPENAI_API_KEY_2).toBe("sk-second")
     })
 
-    it("sends an Ollama key without a URL to Ollama Cloud, like its Test", () => {
-        // Chat sends a server Ollama key to OLLAMA_BASE_URL, or to local
-        // Ollama without one; the Test sends it to Ollama Cloud
-        const cloud = deriveEnvUpdates(
+    it("writes an Ollama URL only when the entry has one", () => {
+        // Without one, the operator's own OLLAMA_BASE_URL (or local Ollama)
+        // stays, also for the AI_PROVIDER=ollama default model
+        const keyOnly = deriveEnvUpdates(
             [provider({ provider: "ollama", apiKey: "ollama-key" })],
             [],
         )
-        expect(cloud.OLLAMA_API_KEY).toBe("ollama-key")
-        expect(cloud.OLLAMA_BASE_URL).toBe("https://ollama.com/api")
+        expect(keyOnly.OLLAMA_API_KEY).toBe("ollama-key")
+        expect(keyOnly.OLLAMA_BASE_URL ?? null).toBeNull()
         const own = deriveEnvUpdates(
             [
                 provider({

tests/unit/admin-test-model.test.ts

@@ -48,6 +48,8 @@ describe("admin Test of an entry without a URL", () => {
         process.env.OPENAI_BASE_URL = "https://operator-proxy.example.com/v1"
         await test({ provider: "openai", apiKey: "panel-key" })
         expect(sent.body.baseUrl).toBe("https://operator-proxy.example.com/v1")
+        // The server's own URL, tested without the rules for typed URLs
+        expect(sent.body.serverBaseUrl).toBe(true)
 
         process.env.AI_GATEWAY_BASE_URL = "https://gateway.example.com/v3/ai"
         await test({ provider: "gateway", apiKey: "k" })
@@ -65,4 +67,15 @@ describe("admin Test of an entry without a URL", () => {
         await test({ provider: "deepseek", apiKey: "k" })
         expect(sent.body.baseUrl).toBeUndefined()
     })
+
+    it("does not use Vertex's variable, which the panel writes itself", async () => {
+        // Before a save it still holds the entry's previous URL
+        process.env.GOOGLE_VERTEX_BASE_URL = "https://old-proxy.example.com"
+        try {
+            await test({ provider: "vertexai", vertexApiKey: "new-key" })
+            expect(sent.body.baseUrl).toBeUndefined()
+        } finally {
+            delete process.env.GOOGLE_VERTEX_BASE_URL
+        }
+    })
 })

tests/unit/ai-providers-credentials.test.ts

@@ -28,6 +28,14 @@ vi.mock("@ai-sdk/openai", () => {
     }
 })
 
+vi.mock("@ai-sdk/azure", () => {
+    const mockModel = { modelId: "test-model" }
+    const mockProviderFn = vi.fn(() => mockModel) as any
+    mockProviderFn.chat = vi.fn(() => mockModel)
+    mockProviderFn.responses = vi.fn(() => mockModel)
+    return { createAzure: vi.fn(() => mockProviderFn) }
+})
+
 vi.mock("@ai-sdk/amazon-bedrock", () => {
     const mockProviderFn = vi.fn(() => ({ modelId: "test-model" }))
     return { createAmazonBedrock: vi.fn(() => mockProviderFn) }
@@ -446,6 +454,43 @@ describe("whose keys a request uses", () => {
         }
     })
 
+    it("runs an Azure entry set up only in the admin panel", async () => {
+        // No AZURE_BASE_URL or AZURE_RESOURCE_NAME: the entry's own
+        // variables hold the key and the resource URL
+        process.env.ADMIN_AZURE_API_KEY = "panel-key"
+        process.env.ADMIN_AZURE_BASE_URL = "https://res.openai.azure.com/openai"
+        try {
+            const { createAzure } = await import("@ai-sdk/azure")
+            expect(() =>
+                getAIModel({
+                    provider: "azure",
+                    modelId: "gpt-4o",
+                    apiKeyEnv: "ADMIN_AZURE_API_KEY",
+                    baseUrlEnv: "ADMIN_AZURE_BASE_URL",
+                }),
+            ).not.toThrow()
+            expect(createAzure).toHaveBeenLastCalledWith(
+                expect.objectContaining({
+                    apiKey: "panel-key",
+                    baseURL: "https://res.openai.azure.com/openai",
+                }),
+            )
+            // Without any URL it still says what is missing
+            delete process.env.ADMIN_AZURE_BASE_URL
+            expect(() =>
+                getAIModel({
+                    provider: "azure",
+                    modelId: "gpt-4o",
+                    apiKeyEnv: "ADMIN_AZURE_API_KEY",
+                    baseUrlEnv: "ADMIN_AZURE_BASE_URL",
+                }),
+            ).toThrow(/AZURE_BASE_URL/)
+        } finally {
+            delete process.env.ADMIN_AZURE_API_KEY
+            delete process.env.ADMIN_AZURE_BASE_URL
+        }
+    })
+
     it("needs a base URL with a user's Azure key", () => {
         // The SDK would otherwise read the server's AZURE_RESOURCE_NAME
         process.env.AZURE_RESOURCE_NAME = "operator-resource"

tests/unit/app-menu.test.ts

@@ -64,4 +64,20 @@ describe("switchPreset", () => {
         await toC
         expect(state.current).toBe("C")
     })
+
+    it("keeps a newer choice of the same preset", async () => {
+        // A, then B, C, and B again while the first restart is pending
+        const first = switchPreset("B").catch(() => {})
+        const second = switchPreset("C").catch(() => {})
+        const third = switchPreset("B")
+        // The first restart fails: the current preset is B again, but it is
+        // the third switch's, which must not be undone
+        state.restarts[0].reject(new Error("timed out"))
+        await new Promise((r) => setTimeout(r, 0))
+        for (const r of state.restarts.slice(1)) r.resolve()
+        await first
+        await second
+        await third
+        expect(state.current).toBe("B")
+    })
 })

tests/unit/chat-message-display-preview.test.tsx

@@ -0,0 +1,68 @@
+import { render } from "@testing-library/react"
+import { describe, expect, it, vi } from "vitest"
+import en from "@/lib/i18n/dictionaries/en.json"
+
+const page = (cells: string) =>
+    `<mxfile><diagram id="p" name="Page-1"><mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/>${cells}</root></mxGraphModel></diagram></mxfile>`
+const box = (id: string) =>
+    `<mxCell id="${id}" value="${id}" vertex="1" parent="1"><mxGeometry x="0" y="0" width="80" height="40" as="geometry"/></mxCell>`
+
+// The first edit's result is loaded (the ref has it); the chartXML state
+// has not caught up yet
+const BEFORE_FIRST_EDIT = page(box("a"))
+const AFTER_FIRST_EDIT = page(box("a") + box("b"))
+
+vi.mock("@/contexts/diagram-context", () => ({
+    useDiagram: () => ({
+        chartXML: BEFORE_FIRST_EDIT,
+        chartXMLRef: { current: AFTER_FIRST_EDIT },
+        loadDiagram: vi.fn(() => null),
+    }),
+}))
+vi.mock("@/hooks/use-dictionary", () => ({ useDictionary: () => en }))
+
+import { ChatMessageDisplay } from "@/components/chat-message-display"
+
+// jsdom has no layout
+Element.prototype.scrollIntoView = () => {}
+
+describe("the streaming preview of a second edit", () => {
+    it("starts from the first edit's result", () => {
+        const editDiagramOriginalXmlRef = { current: new Map<string, string>() }
+        const messages = [
+            {
+                id: "m1",
+                role: "assistant",
+                parts: [
+                    {
+                        type: "tool-edit_diagram",
+                        toolCallId: "edit-2",
+                        state: "input-streaming",
+                        input: {
+                            operations: [
+                                {
+                                    operation: "add",
+                                    cell_id: "c",
+                                    new_xml: box("c"),
+                                },
+                            ],
+                        },
+                    },
+                ],
+            },
+        ] as any
+        render(
+            <ChatMessageDisplay
+                messages={messages}
+                setInput={() => {}}
+                setFiles={() => {}}
+                processedToolCallsRef={{ current: new Set() }}
+                editDiagramOriginalXmlRef={editDiagramOriginalXmlRef}
+                status="streaming"
+            />,
+        )
+        expect(editDiagramOriginalXmlRef.current.get("edit-2")).toBe(
+            AFTER_FIRST_EDIT,
+        )
+    })
+})

tests/unit/chat-route-abort.test.ts

@@ -0,0 +1,139 @@
+// @vitest-environment node
+import { afterEach, describe, expect, it, vi } from "vitest"
+
+const quota = vi.hoisted(() => ({ recorded: [] as number[] }))
+vi.mock("@/lib/dynamo-quota-manager", () => ({
+    isQuotaEnabled: () => true,
+    checkAndIncrementRequest: async () => ({ allowed: true }),
+    recordTokenUsage: async (_ip: string, tokens: number) => {
+        quota.recorded.push(tokens)
+    },
+}))
+
+// No DNS in tests: only loopback addresses are private
+vi.mock("@/lib/ssrf-protection", async (importOriginal) => ({
+    ...(await importOriginal<typeof import("@/lib/ssrf-protection")>()),
+    isPrivateUrl: async (url: string) =>
+        /^https?:\/\/(127\.0\.0\.1|localhost)\b/.test(url),
+}))
+
+import { POST as chat } from "@/app/api/chat/route"
+
+afterEach(() => {
+    quota.recorded = []
+    vi.unstubAllGlobals()
+})
+
+const sse = (chunks: object[], end = true) =>
+    chunks.map((c) => `data: ${JSON.stringify(c)}\n\n`).join("") +
+    (end ? "data: [DONE]\n\n" : "")
+
+describe("a request stopped after a finished step", () => {
+    it("counts that step's tokens", async () => {
+        // Step 1 asks for a shape library (run on the server) and reports
+        // its usage; step 2 never ends, and the user stops
+        let call = 0
+        vi.stubGlobal(
+            "fetch",
+            vi.fn(async (_url: string, init?: RequestInit) => {
+                call++
+                if (call === 1) {
+                    return new Response(
+                        sse([
+                            {
+                                id: "c1",
+                                choices: [
+                                    {
+                                        index: 0,
+                                        delta: {
+                                            role: "assistant",
+                                            tool_calls: [
+                                                {
+                                                    index: 0,
+                                                    id: "call_1",
+                                                    type: "function",
+                                                    function: {
+                                                        name: "get_shape_library",
+                                                        arguments:
+                                                            '{"library":"aws4"}',
+                                                    },
+                                                },
+                                            ],
+                                        },
+                                        finish_reason: null,
+                                    },
+                                ],
+                            },
+                            {
+                                id: "c1",
+                                choices: [
+                                    {
+                                        index: 0,
+                                        delta: {},
+                                        finish_reason: "tool_calls",
+                                    },
+                                ],
+                                usage: {
+                                    prompt_tokens: 1200,
+                                    completion_tokens: 30,
+                                },
+                            },
+                        ]),
+                        { headers: { "content-type": "text/event-stream" } },
+                    )
+                }
+                // Never ends, until the request is aborted (as fetch does)
+                const body = new ReadableStream({
+                    start(controller) {
+                        init?.signal?.addEventListener("abort", () =>
+                            controller.error(
+                                new DOMException("aborted", "AbortError"),
+                            ),
+                        )
+                    },
+                })
+                return new Response(body, {
+                    headers: { "content-type": "text/event-stream" },
+                })
+            }),
+        )
+        const stop = new AbortController()
+        const res = await chat(
+            new Request("http://localhost/api/chat", {
+                method: "POST",
+                signal: stop.signal,
+                headers: {
+                    "Content-Type": "application/json",
+                    "x-forwarded-for": "203.0.113.7",
+                    // The server's own network: counted
+                    "x-ai-provider": "glm",
+                    "x-ai-base-url": "http://127.0.0.1:9000/v1",
+                    "x-ai-api-key": "dummy",
+                    "x-ai-model": "glm-5",
+                },
+                body: JSON.stringify({
+                    messages: [
+                        {
+                            id: "u1",
+                            role: "user",
+                            parts: [{ type: "text", text: "Draw AWS" }],
+                        },
+                    ],
+                    xml: "",
+                }),
+            }),
+        )
+        const reader = res.body?.getReader()
+        // Read until the second step has started
+        await vi.waitFor(() => expect(call).toBe(2), { timeout: 3000 })
+        stop.abort()
+        // The answer stream ends; the SDK handles the stop as it is read
+        while (
+            reader &&
+            !(await reader.read().catch(() => ({ done: true }))).done
+        ) {
+            // drain
+        }
+        await vi.waitFor(() => expect(quota.recorded).toEqual([1230]))
+    })
+})

tests/unit/chat-route-edgeone.test.ts

@@ -77,3 +77,63 @@ describe("EdgeOne as a server model", () => {
         expect(calls[0]?.headers.get("cookie")).toBe("eo_token=t; eo_time=1")
     })
 })
+
+const send = (headers: Record<string, string>) =>
+    chat(
+        new Request("http://localhost/api/chat", {
+            method: "POST",
+            headers: { "Content-Type": "application/json", ...headers },
+            body: JSON.stringify({
+                messages: [
+                    {
+                        id: "u1",
+                        role: "user",
+                        parts: [{ type: "text", text: "Draw two boxes" }],
+                    },
+                ],
+                xml: "",
+            }),
+        }),
+    ).then((r) => r.text())
+
+describe("EdgeOne endpoints", () => {
+    it("works when the deployment names EdgeOne only in AI_PROVIDER", async () => {
+        process.env.AI_PROVIDER = "edgeone"
+        process.env.AI_MODEL = "@tx/deepseek-ai/deepseek-v3-0324"
+        await send({})
+        expect(calls[0]?.url).toBe(
+            "http://localhost/api/edgeai/chat/completions",
+        )
+    })
+
+    it("always calls the site's own function, whatever URL the request names", async () => {
+        // Another host would get the user's EdgeOne cookies
+        await send({
+            "x-ai-provider": "edgeone",
+            "x-ai-model": "@tx/deepseek-ai/deepseek-v3-0324",
+            "x-ai-base-url": "https://elsewhere.example/api/edgeai",
+            cookie: "eo_token=t",
+        })
+        expect(calls[0]?.url).toBe(
+            "http://localhost/api/edgeai/chat/completions",
+        )
+    })
+
+    it("keeps the deployment's base path", async () => {
+        const savedPath = process.env.NEXT_PUBLIC_BASE_PATH
+        process.env.NEXT_PUBLIC_BASE_PATH = "/draw"
+        try {
+            await send({
+                "x-ai-provider": "edgeone",
+                "x-ai-model": "@tx/deepseek-ai/deepseek-v3-0324",
+            })
+            expect(calls[0]?.url).toBe(
+                "http://localhost/draw/api/edgeai/chat/completions",
+            )
+        } finally {
+            if (savedPath === undefined)
+                delete process.env.NEXT_PUBLIC_BASE_PATH
+            else process.env.NEXT_PUBLIC_BASE_PATH = savedPath
+        }
+    })
+})

tests/unit/chat-route-errors.test.ts

@@ -89,6 +89,24 @@ describe("provider error texts in the stream", () => {
         // The SDK retries a refused connection twice, waiting between
     }, 20_000)
 
+    it("shows the server's keyless Ollama error on the web too", async () => {
+        // No key, no money involved; round three hid this text
+        process.env.AI_PROVIDER = "ollama"
+        process.env.AI_MODEL = "llama3"
+        process.env.OLLAMA_BASE_URL = "http://ollama.internal:11434/api"
+        vi.stubGlobal(
+            "fetch",
+            vi.fn(async () => {
+                throw Object.assign(new TypeError("fetch failed"), {
+                    cause: new Error("connect ECONNREFUSED 10.0.0.9:11434"),
+                })
+            }),
+        )
+        expect(await streamedError({})).not.toBe(
+            "The provider returned an error.",
+        )
+    }, 20_000)
+
     it("shows EdgeOne's own daily quota explanation", async () => {
         // The function answers 429, which the SDK retries with a wait;
         // the status does not decide whether the text is shown
@@ -124,3 +142,107 @@ describe("provider error texts in the stream", () => {
         expect(message).toBe("The provider returned an error.")
     })
 })
+
+describe("the output cap", () => {
+    it("holds for the server's own keyless endpoints too", async () => {
+        process.env.MAX_OUTPUT_TOKENS = "8000"
+        const sent: string[] = []
+        vi.stubGlobal(
+            "fetch",
+            vi.fn(async (_url: string, init?: RequestInit) => {
+                sent.push(String(init?.body ?? ""))
+                return new Response("{}", { status: 400 })
+            }),
+        )
+        try {
+            const res = await chat(
+                new Request("http://localhost/api/chat", {
+                    method: "POST",
+                    headers: {
+                        "Content-Type": "application/json",
+                        "x-ai-provider": "ollama",
+                        "x-ai-base-url": "http://127.0.0.1:11434/api",
+                        "x-ai-model": "llama3",
+                        "x-max-output-tokens": "200000",
+                    },
+                    body: JSON.stringify({
+                        messages: [
+                            {
+                                id: "u1",
+                                role: "user",
+                                parts: [{ type: "text", text: "Draw" }],
+                            },
+                        ],
+                        xml: "",
+                    }),
+                }),
+            )
+            await res.text()
+            expect(sent[0]).toContain('"max_output_tokens":8000')
+        } finally {
+            delete process.env.MAX_OUTPUT_TOKENS
+        }
+    })
+})
+
+describe("a tool call that never got its result", () => {
+    it("is left out of the prompt instead of failing every later message", async () => {
+        // Stop while the screenshot check ran left display_diagram without
+        // a result, and the chat was saved like that
+        process.env.AI_PROVIDER = "openai"
+        process.env.AI_MODEL = "gpt-5.5"
+        process.env.OPENAI_API_KEY = "server-key"
+        const sent: string[] = []
+        vi.stubGlobal(
+            "fetch",
+            vi.fn(async (_url: string, init?: RequestInit) => {
+                sent.push(String(init?.body ?? ""))
+                return new Response(
+                    JSON.stringify({ error: { message: "x" } }),
+                    {
+                        status: 400,
+                        headers: { "Content-Type": "application/json" },
+                    },
+                )
+            }),
+        )
+        const res = await chat(
+            new Request("http://localhost/api/chat", {
+                method: "POST",
+                headers: { "Content-Type": "application/json" },
+                body: JSON.stringify({
+                    messages: [
+                        {
+                            id: "u1",
+                            role: "user",
+                            parts: [{ type: "text", text: "Draw a box" }],
+                        },
+                        {
+                            id: "a1",
+                            role: "assistant",
+                            parts: [
+                                {
+                                    type: "tool-display_diagram",
+                                    toolCallId: "call-without-result",
+                                    state: "input-available",
+                                    input: { xml: "<mxCell/>" },
+                                },
+                            ],
+                        },
+                        {
+                            id: "u2",
+                            role: "user",
+                            parts: [{ type: "text", text: "Make it red" }],
+                        },
+                    ],
+                    xml: "",
+                }),
+            }),
+        )
+        await res.text()
+        // The request reached the model, without the unanswered call
+        expect(sent).toHaveLength(1)
+        expect(sent[0]).toContain("Make it red")
+        expect(sent[0]).not.toContain("call-without-result")
+    })
+})

tests/unit/chat-route-quota.test.ts

@@ -140,6 +140,67 @@ describe("chat quota", () => {
         expect(quota.checks).toBe(1)
     })
 
+    it("counts the server's network whatever key header comes along", async () => {
+        // A keyless Ollama or a local SGLang ignores a dummy key
+        for (const headers of [
+            {
+                "x-ai-provider": "ollama",
+                "x-ai-base-url": "http://127.0.0.1:11434/api",
+                "x-ai-api-key": "dummy",
+                "x-ai-model": "llama3.2",
+            },
+            {
+                "x-ai-provider": "openai",
+                "x-ai-base-url": "http://127.0.0.1:30000/v1",
+                "x-ai-api-key": "dummy",
+                "x-ai-model": "m",
+            },
+        ]) {
+            expect((await send(headers)).status).toBe(429)
+        }
+        expect(quota.checks).toBe(2)
+    })
+
+    it("does not count a provider that never uses the base URL header", async () => {
+        // Bedrock on the user's own AWS keys goes to AWS, whatever the
+        // leftover base URL says
+        const res = await send({
+            "x-ai-provider": "bedrock",
+            "x-ai-model": "amazon.nova-lite-v1:0",
+            "x-ai-base-url": "http://127.0.0.1:8080",
+            "x-aws-access-key-id": "id",
+            "x-aws-secret-access-key": "secret",
+            "x-aws-region": "us-east-1",
+        })
+        expect(res.status).not.toBe(429)
+        expect(quota.checks).toBe(0)
+    })
+
+    it("counts EdgeOne even with a base URL header", async () => {
+        const res = await send({
+            "x-ai-provider": "edgeone",
+            "x-ai-model": "@tx/deepseek-ai/deepseek-v3-0324",
+            "x-ai-base-url": "https://this-site.example/api/edgeai",
+        })
+        expect(res.status).toBe(429)
+        expect(quota.checks).toBe(1)
+    })
+
+    it("never counts in the desktop app, where every endpoint is the user's", async () => {
+        process.env.NEXT_AI_DRAWIO_DESKTOP = "1"
+        try {
+            const res = await send({
+                "x-ai-provider": "ollama",
+                "x-ai-base-url": "http://127.0.0.1:11434/api",
+                "x-ai-model": "llama3.2",
+            })
+            expect(res.status).not.toBe(429)
+            expect(quota.checks).toBe(0)
+        } finally {
+            delete process.env.NEXT_AI_DRAWIO_DESKTOP
+        }
+    })
+
     it("does not count Ollama on the user's own server", async () => {
         const res = await send({
             "x-ai-provider": "ollama",

tests/unit/diagram-context.test.tsx

@@ -0,0 +1,119 @@
+import { deflateRawSync } from "node:zlib"
+import { act, renderHook } from "@testing-library/react"
+import type React from "react"
+import { afterEach, describe, expect, it, vi } from "vitest"
+import { DiagramProvider, useDiagram } from "@/contexts/diagram-context"
+
+vi.mock("sonner", () => ({ toast: { success: vi.fn() } }))
+
+// The provider with a stand-in draw.io that records each export request
+function setup() {
+    const { result } = renderHook(() => useDiagram(), {
+        wrapper: ({ children }: { children: React.ReactNode }) => (
+            <DiagramProvider>{children}</DiagramProvider>
+        ),
+    })
+    const requests: { format: string; message: string }[] = []
+    result.current.drawioRef.current = {
+        exportDiagram: (r: any) => requests.push(r),
+        load: vi.fn(),
+    } as any
+    // draw.io's reply to a request: it echoes the request in `message`
+    const reply = (request: { message: string }, data: string, xml = "") =>
+        act(() =>
+            result.current.handleDiagramExport({
+                event: "export",
+                data,
+                xml,
+                format: "xmlsvg",
+                message: request,
+            } as any),
+        )
+    return { result, requests, reply }
+}
+
+// An editable SVG as draw.io exports it: the diagram, compressed, in its
+// content attribute
+const svgOf = (label: string) => {
+    const model = `<mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/><mxCell id="2" value="${label}" vertex="1" parent="1"/></root></mxGraphModel>`
+    const packed = deflateRawSync(
+        Buffer.from(encodeURIComponent(model)),
+    ).toString("base64")
+    const content = `<mxfile><diagram id="p">${packed}</diagram></mxfile>`
+        .replaceAll("&", "&amp;")
+        .replaceAll("<", "&lt;")
+        .replaceAll(">", "&gt;")
+        .replaceAll('"', "&quot;")
+    const svg = `<svg xmlns="http://www.w3.org/2000/svg" content="${content}"></svg>`
+    return `data:image/svg+xml;base64,${btoa(svg)}`
+}
+
+afterEach(() => {
+    vi.restoreAllMocks()
+})
+
+describe("exports in flight at the same time", () => {
+    it("give the chat's export only its own reply", () => {
+        const { result, requests, reply } = setup()
+        // An edit's history export is still on its way when the chat exports
+        act(() => {
+            result.current.handleExport()
+        })
+        let tag = ""
+        const got: string[] = []
+        act(() => {
+            tag = result.current.handleExportWithoutHistory()
+            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()
+    })
+
+    it("save each file with its own result", async () => {
+        const { result, requests, reply } = setup()
+        const saved: { name: string; href: string }[] = []
+        vi.spyOn(HTMLAnchorElement.prototype, "click").mockImplementation(
+            function (this: HTMLAnchorElement) {
+                saved.push({ name: this.download, href: this.href })
+            },
+        )
+        const blobs = new Map<string, Blob>()
+        URL.createObjectURL = vi.fn((blob: Blob) => {
+            const url = `blob:test-${blobs.size}`
+            blobs.set(url, blob)
+            return url
+        })
+        URL.revokeObjectURL = vi.fn()
+        vi.stubGlobal(
+            "fetch",
+            vi.fn(async () => new Response("{}")),
+        )
+
+        const twoPages =
+            '<mxfile><diagram id="a" name="A"><mxGraphModel><root/></mxGraphModel></diagram><diagram id="b" name="B"><mxGraphModel><root/></mxGraphModel></diagram></mxfile>'
+        act(() => {
+            result.current.saveDiagramToFile("doc", "drawio")
+            result.current.saveDiagramToFile("pic", "png")
+        })
+        // The PNG answers first
+        reply(requests[1], "data:image/png;base64,iVBORw0KGgo=")
+        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/)
+        const file = blobs.get(saved[1].href)
+        const text = await new Promise<string>((resolve) => {
+            const reader = new FileReader()
+            reader.onload = () => resolve(String(reader.result))
+            reader.readAsText(file as Blob)
+        })
+        expect(text).toContain('name="A"')
+        expect(text).toContain('name="B"')
+        vi.unstubAllGlobals()
+    })
+})

tests/unit/env-loader.test.ts

@@ -25,6 +25,9 @@ const KEYS = [
     "T_HASH",
     "T_AFTER",
     "T_JOINED",
+    "T_ESC_HASH",
+    "T_ESC_INNER",
+    "T_ESC_COMMENT",
 ]
 afterEach(() => {
     for (const k of KEYS) delete process.env[k]
@@ -72,4 +75,21 @@ describe("loadEnvFile", () => {
         expect(process.env.T_JOINED).toBe(`"a"b`)
         expect(process.env.T_HASH).toBe("http://host/#/x")
     })
+
+    it("does not end a quoted value at an escaped quote, like dotenv", () => {
+        dir.path = mkdtempSync(join(tmpdir(), "env-loader-"))
+        writeFileSync(
+            join(dir.path, ".env"),
+            [
+                'T_ESC_HASH="abc\\" #def"',
+                'T_ESC_INNER="a # \\"b\\""',
+                'T_ESC_COMMENT="x\\"y" # c',
+            ].join("\n"),
+        )
+        loadEnvFile()
+        // Expected values from dotenv 16.6.1, which keeps the backslashes
+        expect(process.env.T_ESC_HASH).toBe('abc\\" #def')
+        expect(process.env.T_ESC_INNER).toBe('a # \\"b\\"')
+        expect(process.env.T_ESC_COMMENT).toBe('x\\"y')
+    })
 })

tests/unit/mcp-preview-recovery.test.ts

@@ -0,0 +1,372 @@
+import { readFileSync } from "node:fs"
+import { join } from "node:path"
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"
+
+// The MCP preview page's script, run in this document with a stubbed
+// server (fetch) and draw.io iframe (its postMessage), so the tab's side of
+// a recreated session can be driven step by step
+const dir = join(process.cwd(), "packages/mcp-server/src/preview")
+const DRAWIO = "https://embed.diagrams.net"
+const html = readFileSync(join(dir, "index.html"), "utf8")
+    .replace("{{CSS}}", "")
+    .replace("{{SESSION_BADGE}}", "")
+    .replaceAll("{{DISABLED}}", "")
+    .replace("{{DRAWIO_URL}}", "about:blank")
+    .replace("{{SESSION_JSON}}", '"mcp-test"')
+    .replace("{{ORIGIN_JSON}}", JSON.stringify(DRAWIO))
+const scripts = [...html.matchAll(/<script>([\s\S]*?)<\/script>/g)].map((m) =>
+    m[1].replace("{{SCRIPT}}", ""),
+)
+const preview = readFileSync(join(dir, "preview.js"), "utf8")
+
+const pageListeners: Array<[string, EventListener]> = []
+
+type Answer = { status: number; body: unknown }
+interface Call {
+    url: string
+    method: string
+    body: any
+    answer: (a: Answer) => void
+    fail: () => void
+}
+
+function openPage() {
+    document.body.innerHTML = html.replace(/<script>[\s\S]*?<\/script>/g, "")
+    const toDrawio: any[] = []
+    const iframe = document.getElementById("drawio") as HTMLIFrameElement
+    Object.defineProperty(iframe, "contentWindow", {
+        value: { postMessage: (m: string) => toDrawio.push(JSON.parse(m)) },
+    })
+    // Every request waits until the test answers it
+    const calls: Call[] = []
+    vi.stubGlobal(
+        "fetch",
+        vi.fn(
+            (url: string, init?: RequestInit) =>
+                new Promise((resolve, reject) => {
+                    calls.push({
+                        url,
+                        method: init?.method ?? "GET",
+                        body: init?.body ? JSON.parse(String(init.body)) : null,
+                        answer: ({ status, body }) =>
+                            resolve(
+                                new Response(JSON.stringify(body), { status }),
+                            ),
+                        fail: () => reject(new TypeError("Failed to fetch")),
+                    })
+                }),
+        ),
+    )
+    // The page's window listeners, removed after the test
+    const addListener = window.addEventListener
+    window.addEventListener = ((type: string, listener: any, options?: any) => {
+        pageListeners.push([type, listener])
+        addListener.call(window, type, listener, options)
+    }) as typeof window.addEventListener
+    const run = new Function(
+        `${scripts.join("\n")}\n${preview}\nreturn { poll, read: () => ({ stateId, currentVersion, lastXml, latestXml }) }`,
+    )
+    let created: unknown
+    try {
+        created = run()
+    } finally {
+        window.addEventListener = addListener
+    }
+    const page = created as {
+        poll: () => Promise<void>
+        read: () => {
+            stateId: string | null
+            currentVersion: number
+            lastXml: string | null
+            latestXml: string | null
+        }
+    }
+    const fromDrawio = (msg: object) =>
+        window.dispatchEvent(
+            new MessageEvent("message", {
+                data: JSON.stringify(msg),
+                origin: DRAWIO,
+            }),
+        )
+    const settle = () => new Promise((r) => setTimeout(r, 0))
+    const next = (method: string) => {
+        const call = calls.find((c) => c.method === method)
+        if (!call) throw new Error(`no pending ${method}`)
+        calls.splice(calls.indexOf(call), 1)
+        return call
+    }
+    return { page, toDrawio, calls, fromDrawio, settle, next }
+}
+
+const state = (
+    stateId: string,
+    version: number,
+    xml: string,
+    blank = false,
+) => ({
+    status: 200,
+    body: { stateId, version, xml, blank, syncRequested: false },
+})
+
+/** A tab in step with state S1 at version 2, showing diagram A */
+async function inStep() {
+    const t = openPage()
+    t.next("GET").answer(state("S1", 2, "<mxfile>A</mxfile>"))
+    await t.settle()
+    t.fromDrawio({ event: "init" })
+    await t.settle()
+    expect(t.page.read().lastXml).toBe("<mxfile>A</mxfile>")
+    return t
+}
+
+beforeEach(() => {
+    vi.useFakeTimers({ toFake: ["setInterval"] })
+})
+afterEach(() => {
+    for (const [type, listener] of pageListeners.splice(0)) {
+        window.removeEventListener(type, listener)
+    }
+    vi.useRealTimers()
+    vi.unstubAllGlobals()
+})
+
+describe("MCP preview after the server recreated its session", () => {
+    it("keeps edits made while the server was down", async () => {
+        const t = await inStep()
+        // The user edits; the server is gone
+        t.fromDrawio({ event: "autosave", xml: "<mxfile>B</mxfile>" })
+        // draw.io's SVG export answers, then the push fails
+        t.fromDrawio({ event: "export", data: "<svg/>" })
+        await t.settle()
+        t.next("POST").fail()
+        await t.settle()
+        // A new process recovered the file, which holds A
+        const poll = t.page.poll()
+        t.next("GET").answer(state("S2", 1, "<mxfile>A</mxfile>"))
+        await poll
+        const push = t.next("POST")
+        expect(push.body).toMatchObject({
+            xml: "<mxfile>B</mxfile>",
+            stateId: "S2",
+            baseVersion: 1,
+        })
+    })
+
+    it("shows the server's diagram and keeps the tab's in History", async () => {
+        const t = await inStep()
+        const poll = t.page.poll()
+        t.next("GET").answer(state("S2", 1, "<mxfile>C</mxfile>"))
+        await poll
+        expect(t.toDrawio.at(-1)).toMatchObject({
+            action: "load",
+            xml: "<mxfile>C</mxfile>",
+        })
+        expect(t.next("POST").body).toMatchObject({
+            xml: "<mxfile>A</mxfile>",
+            source: "recover",
+            stateId: "S2",
+        })
+    })
+
+    it("sends an edit of the replaced canvas to History, until draw.io loaded", async () => {
+        const t = await inStep()
+        const poll = t.page.poll()
+        t.next("GET").answer(state("S2", 1, "<mxfile>C</mxfile>"))
+        await poll
+        t.next("POST") // the tab's copy, to History
+        // An autosave the old canvas sent before the load
+        t.fromDrawio({ event: "autosave", xml: "<mxfile>A edited</mxfile>" })
+        await t.settle()
+        expect(t.next("POST").body).toMatchObject({
+            xml: "<mxfile>A edited</mxfile>",
+            source: "recover",
+        })
+        // After the load, edits are edits again
+        t.fromDrawio({ event: "load" })
+        t.fromDrawio({ event: "autosave", xml: "<mxfile>C edited</mxfile>" })
+        t.fromDrawio({ event: "export", data: "<svg/>" })
+        await t.settle()
+        expect(t.next("POST").body).toMatchObject({
+            xml: "<mxfile>C edited</mxfile>",
+            source: "edit",
+        })
+    })
+
+    it("ignores the late answer to an old state's push", async () => {
+        const t = await inStep()
+        t.fromDrawio({ event: "autosave", xml: "<mxfile>B</mxfile>" })
+        t.fromDrawio({ event: "export", data: "<svg/>" })
+        await t.settle()
+        const oldPush = t.next("POST")
+        const poll = t.page.poll()
+        t.next("GET").answer(state("S2", 1, "<mxfile>B</mxfile>"))
+        await poll
+        oldPush.answer({ status: 200, body: { success: true, version: 21 } })
+        await t.settle()
+        expect(t.page.read()).toMatchObject({
+            stateId: "S2",
+            currentVersion: 1,
+        })
+    })
+
+    it("drops a poll answer older than one already handled", async () => {
+        const t = await inStep()
+        const first = t.page.poll()
+        const firstGet = t.next("GET")
+        const second = t.page.poll()
+        t.next("GET").answer(state("S2", 1, "<mxfile>A</mxfile>"))
+        await second
+        // The answer from before the restart comes last
+        firstGet.answer(state("S1", 3, "<mxfile>old</mxfile>"))
+        await first
+        expect(t.page.read().stateId).toBe("S2")
+    })
+
+    it("saves an undo made while a push was on its way", async () => {
+        const t = await inStep()
+        t.fromDrawio({ event: "autosave", xml: "<mxfile>B</mxfile>" })
+        t.fromDrawio({ event: "export", data: "<svg/>" })
+        await t.settle()
+        const pushB = t.next("POST")
+        // Undo back to A: equal to the last saved diagram, so no push
+        t.fromDrawio({ event: "autosave", xml: "<mxfile>A</mxfile>" })
+        await t.settle()
+        expect(t.calls.filter((c) => c.method === "POST")).toHaveLength(0)
+        pushB.answer({ status: 200, body: { success: true, version: 3 } })
+        await t.settle()
+        await t.settle()
+        expect(t.next("POST").body.xml).toBe("<mxfile>A</mxfile>")
+    })
+
+    it("sends nothing more after a sync reply", async () => {
+        const t = await inStep()
+        const poll = t.page.poll()
+        t.next("GET").answer({
+            status: 200,
+            body: {
+                ...state("S1", 2, "<mxfile>A</mxfile>").body,
+                syncRequested: true,
+            },
+        })
+        await poll
+        const request = t.toDrawio.at(-1)
+        expect(request).toMatchObject({ action: "export", format: "xml" })
+        // draw.io's export of the canvas, formatted unlike its autosave
+        t.fromDrawio({
+            event: "export",
+            format: "xml",
+            xml: '<mxfile host="drawio">A</mxfile>',
+            message: request,
+        })
+        await t.settle()
+        const sync = t.next("POST")
+        expect(sync.body.source).toBe("sync")
+        sync.answer({ status: 200, body: { success: true, version: 3 } })
+        await t.settle()
+        await t.settle()
+        expect(t.calls.filter((c) => c.method === "POST")).toHaveLength(0)
+    })
+})
+
+describe("MCP preview thumbnails and downloads", () => {
+    /** The tab loads the server write B at version 3 and asks for its image */
+    async function loadedB() {
+        const t = await inStep()
+        const poll = t.page.poll()
+        t.next("GET").answer(state("S1", 3, "<mxfile>B</mxfile>"))
+        await poll
+        await new Promise((r) => setTimeout(r, 600))
+        const request = t.toDrawio.at(-1)
+        expect(request).toMatchObject({ action: "export", format: "svg" })
+        return { t, n: request.thumbExport as number }
+    }
+    const thumbnailPosts = (t: ReturnType<typeof openPage>) =>
+        t.calls.filter((c) => c.url === "/api/history-svg")
+
+    it("sends the image with the state and version it shows", async () => {
+        const { t, n } = await loadedB()
+        t.fromDrawio({
+            event: "export",
+            data: "<svg/>",
+            message: { thumbExport: n },
+        })
+        await t.settle()
+        expect(thumbnailPosts(t).map((c) => c.body)).toEqual([
+            expect.objectContaining({ stateId: "S1", version: 3 }),
+        ])
+    })
+
+    it("drops the reply to an older thumbnail export", async () => {
+        const { t, n } = await loadedB()
+        t.fromDrawio({
+            event: "export",
+            data: "<svg/>",
+            message: { thumbExport: n - 1 },
+        })
+        await t.settle()
+        expect(thumbnailPosts(t)).toHaveLength(0)
+    })
+
+    it("drops the image when the user changed the canvas since the load", async () => {
+        const { t, n } = await loadedB()
+        t.fromDrawio({ event: "autosave", xml: "<mxfile>B edited</mxfile>" })
+        t.fromDrawio({
+            event: "export",
+            data: "<svg/>",
+            message: { thumbExport: n },
+        })
+        await t.settle()
+        expect(thumbnailPosts(t)).toHaveLength(0)
+    })
+
+    it("downloads the canvas with an edit the server did not get", async () => {
+        const t = await inStep()
+        t.fromDrawio({ event: "autosave", xml: "<mxfile>B</mxfile>" })
+        t.fromDrawio({ event: "export", data: "<svg/>" })
+        await t.settle()
+        t.next("POST").fail()
+        await t.settle()
+        let saved: Blob | undefined
+        URL.createObjectURL = vi.fn((blob: Blob) => {
+            saved = blob
+            return "blob:test"
+        })
+        URL.revokeObjectURL = vi.fn()
+        ;(document.getElementById("save-format") as HTMLSelectElement).value =
+            "drawio"
+        document.getElementById("save-confirm-btn")?.click()
+        const text = await new Promise<string>((resolve) => {
+            const reader = new FileReader()
+            reader.onload = () => resolve(String(reader.result))
+            reader.readAsText(saved as Blob)
+        })
+        expect(text).toBe("<mxfile>B</mxfile>")
+    })
+})
+
+describe("MCP preview with a diagram over the size limit", () => {
+    it("retries without the image, then tells the user", async () => {
+        const t = await inStep()
+        t.fromDrawio({ event: "autosave", xml: "<mxfile>huge</mxfile>" })
+        t.fromDrawio({ event: "export", data: "<svg/>" })
+        await t.settle()
+        const first = t.next("POST")
+        expect(first.body.svg).not.toBe("")
+        first.answer({ status: 413, body: { error: "Payload too large" } })
+        await t.settle()
+        const retry = t.next("POST")
+        expect(retry.body).toMatchObject({
+            xml: "<mxfile>huge</mxfile>",
+            svg: "",
+        })
+        expect(
+            document.getElementById("notice")?.classList.contains("open"),
+        ).toBe(false)
+        retry.answer({ status: 413, body: { error: "Payload too large" } })
+        await t.settle()
+        expect(t.calls.filter((c) => c.method === "POST")).toHaveLength(0)
+        expect(document.getElementById("notice")?.textContent).toContain(
+            "too large",
+        )
+    })
+})

tests/unit/ssrf-protection.test.ts

@@ -1,5 +1,5 @@
-import { beforeEach, describe, expect, it, vi } from "vitest"
-import { isPrivateUrl } from "@/lib/ssrf-protection"
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"
+import { isPrivateUrl, redirectGuardedFetch } from "@/lib/ssrf-protection"
 
 // Mock DNS so tests are deterministic and never hit the network.
 const lookupMock = vi.hoisted(() => vi.fn())
@@ -77,3 +77,74 @@ describe("isPrivateUrl", () => {
         )
     })
 })
+
+describe("redirectGuardedFetch with the quota on", () => {
+    const answers = (map: Record<string, Response>) =>
+        vi.fn(
+            async (url: string) =>
+                map[String(url)] ?? new Response("?", { status: 404 }),
+        )
+
+    beforeEach(() => {
+        lookupMock.mockReset()
+        // Hosts ending in .example are public
+        lookupMock.mockImplementation(async (host: string) =>
+            host.endsWith(".example")
+                ? [{ address: "93.184.216.34", family: 4 }]
+                : [],
+        )
+        process.env.DYNAMODB_QUOTA_TABLE = "quota"
+        delete process.env.ALLOW_PRIVATE_URLS
+    })
+
+    afterEach(() => {
+        delete process.env.DYNAMODB_QUOTA_TABLE
+        vi.unstubAllGlobals()
+    })
+
+    it("follows a redirect to a public address", async () => {
+        // A user's own proxy that moves http to https
+        vi.stubGlobal(
+            "fetch",
+            answers({
+                "http://proxy.example/v1/chat": new Response(null, {
+                    status: 308,
+                    headers: { location: "https://proxy.example/v1/chat" },
+                }),
+                "https://proxy.example/v1/chat": new Response("ok"),
+            }),
+        )
+        const guarded = redirectGuardedFetch()
+        expect(guarded).toBeDefined()
+        const res = await guarded?.("http://proxy.example/v1/chat", {
+            method: "POST",
+            body: "{}",
+        })
+        expect(await res?.text()).toBe("ok")
+    })
+
+    it("refuses a redirect to the server's own network", async () => {
+        // It would be counted as a public endpoint while using the server's
+        vi.stubGlobal(
+            "fetch",
+            answers({
+                "https://public.example/api/chat": new Response(null, {
+                    status: 307,
+                    headers: { location: "http://127.0.0.1:11434/api/chat" },
+                }),
+            }),
+        )
+        await expect(
+            redirectGuardedFetch()?.("https://public.example/api/chat", {
+                method: "POST",
+                body: "{}",
+            }),
+        ).rejects.toThrow(/private addresses/)
+        expect(fetch).toHaveBeenCalledTimes(1)
+    })
+
+    it("is not used without the quota", () => {
+        delete process.env.DYNAMODB_QUOTA_TABLE
+        expect(redirectGuardedFetch()).toBeUndefined()
+    })
+})

tests/unit/tool-call-card.test.tsx

@@ -0,0 +1,43 @@
+import { cleanup, render } from "@testing-library/react"
+import { afterEach, describe, expect, it } from "vitest"
+import { ToolCallCard } from "@/components/chat/ToolCallCard"
+
+afterEach(cleanup)
+
+const dict = {
+    tools: { complete: "Complete" },
+    chat: { copied: "Copied", failedToCopy: "Failed", copyResponse: "Copy" },
+}
+
+describe("ToolCallCard", () => {
+    it("shows streamed operations without crashing on broken entries", () => {
+        // A partly streamed or invalid edit_diagram input
+        const operations = [
+            null,
+            { operation: {} },
+            { operation: "add", cell_id: {} },
+            { operation: "add", cell_id: "2", new_xml: {} },
+            { operation: "update", cell_id: "3", new_xml: '<mxCell id="3"/>' },
+        ]
+        const { container } = render(
+            <ToolCallCard
+                part={
+                    {
+                        type: "tool-edit_diagram",
+                        toolCallId: "t1",
+                        state: "input-streaming",
+                        input: { operations },
+                    } as any
+                }
+                expandedTools={{ t1: true }}
+                setExpandedTools={() => {}}
+                onCopy={() => {}}
+                copiedToolCallId={null}
+                copyFailedToolCallId={null}
+                dict={dict}
+            />,
+        )
+        expect(container.textContent).toContain("cell_id: 3")
+        expect(container.textContent).toContain('<mxCell id="3"/>')
+    })
+})

tests/unit/use-diagram-tool-handlers.test.tsx

@@ -45,6 +45,77 @@ function setup(partialXml: string) {
     return { refs, onDisplayChart, addToolOutput, append }
 }
 
+describe("the screenshot check and Stop", () => {
+    const draw = async (opts: {
+        isStopped: () => boolean
+        validateDiagram: () => Promise<any>
+    }) => {
+        const onValidationStateChange = vi.fn()
+        const { result } = renderHook(() =>
+            useDiagramToolHandlers({
+                partialXmlRef: { current: "" },
+                editDiagramOriginalXmlRef: { current: new Map() },
+                processedToolCallsRef: { current: new Set() },
+                validationRetryCountRef: { current: 0 },
+                chartXMLRef: { current: "" },
+                onDisplayChart: () => null,
+                onFetchChart: async () => "",
+                onExport: () => {},
+                enableVlmValidation: true,
+                captureValidationPng: async () => "data:image/png;base64,AA",
+                validateDiagram: opts.validateDiagram,
+                isStopped: opts.isStopped,
+                onValidationStateChange,
+            }),
+        )
+        const addToolOutput = vi.fn()
+        await result.current.handleToolCall(
+            {
+                toolCall: {
+                    toolCallId: "d1",
+                    toolName: "display_diagram",
+                    input: { xml: box("2") },
+                },
+            },
+            addToolOutput,
+        )
+        return { addToolOutput, onValidationStateChange }
+    }
+
+    it("skips a check that had not started when the user stopped", async () => {
+        const validateDiagram = vi.fn(async () => ({
+            valid: true,
+            issues: [],
+            suggestions: [],
+        }))
+        const { addToolOutput, onValidationStateChange } = await draw({
+            isStopped: () => true,
+            validateDiagram,
+        })
+        expect(validateDiagram).not.toHaveBeenCalled()
+        expect(onValidationStateChange.mock.lastCall?.[1].status).toBe(
+            "skipped",
+        )
+        expect(addToolOutput.mock.lastCall?.[0].output).toMatch(
+            /Successfully displayed/,
+        )
+    })
+
+    it("ends with the diagram's result when Stop cancels a running check", async () => {
+        const { addToolOutput, onValidationStateChange } = await draw({
+            isStopped: () => false,
+            validateDiagram: async () => {
+                throw new DOMException("Validation cancelled", "AbortError")
+            },
+        })
+        expect(onValidationStateChange.mock.lastCall?.[1].status).toBe(
+            "skipped",
+        )
+        expect(addToolOutput).toHaveBeenCalledTimes(1)
+        expect(addToolOutput.mock.lastCall?.[0].state).toBeUndefined()
+    })
+})
+
 describe("append_diagram and the stored previews", () => {
     it("takes the stored originals when it draws the completed diagram", async () => {
         // Otherwise the preview code later loads the failed edit's original

tests/unit/use-session-manager.test.tsx

@@ -0,0 +1,170 @@
+import { act, renderHook, waitFor } from "@testing-library/react"
+import { beforeEach, describe, expect, it, vi } from "vitest"
+
+const stored = new Map<string, any>()
+// Each write waits for the test to let it finish
+let pendingWrites: (() => void)[] = []
+// Reads wait for this
+let readGate: Promise<void> = Promise.resolve()
+
+vi.mock("@/hooks/use-dictionary", () => ({
+    useDictionary: () => ({ errors: { sessionSaveFailed: "Not saved" } }),
+}))
+
+vi.mock("@/lib/session-storage", async (importActual) => {
+    const actual = await importActual<typeof import("@/lib/session-storage")>()
+    return {
+        createEmptySession: actual.createEmptySession,
+        extractTitle: actual.extractTitle,
+        isIndexedDBAvailable: () => true,
+        migrateFromLocalStorage: async () => null,
+        readSessionCount: async () => stored.size,
+        enforceSessionLimit: async () => {},
+        getSession: async (id: string) => {
+            await readGate
+            return stored.get(id) ?? null
+        },
+        deleteSession: async (id: string) => {
+            stored.delete(id)
+        },
+        getAllSessionMetadata: async () =>
+            [...stored.values()].map((s) => ({ id: s.id, title: s.title })),
+        saveSession: vi.fn(
+            (session: any) =>
+                new Promise<boolean>((resolve) => {
+                    pendingWrites.push(() => {
+                        stored.set(session.id, session)
+                        resolve(true)
+                    })
+                }),
+        ),
+    }
+})
+
+import { useSessionManager } from "@/hooks/use-session-manager"
+
+const data = {
+    messages: [
+        {
+            id: "m1",
+            role: "user" as const,
+            parts: [{ type: "text", text: "Draw a cat" }],
+        },
+    ],
+    xmlSnapshots: [] as [number, string][],
+    diagramXml: "",
+}
+
+// Let every write waiting now (and those it leads to) finish
+async function finishWrites() {
+    for (let i = 0; i < 10; i++) {
+        await act(async () => {
+            const writes = pendingWrites
+            pendingWrites = []
+            for (const finish of writes) finish()
+            await new Promise((r) => setTimeout(r, 0))
+        })
+    }
+}
+
+async function setup() {
+    const hook = renderHook(() => useSessionManager())
+    await waitFor(() => expect(hook.result.current.isLoading).toBe(false))
+    return hook
+}
+
+describe("saving the chat on screen", () => {
+    beforeEach(() => {
+        stored.clear()
+        pendingWrites = []
+    })
+
+    it("creates one session when two saves of a new chat overlap", async () => {
+        const { result } = await setup()
+        let saves!: Promise<boolean[]>
+        act(() => {
+            saves = Promise.all([
+                result.current.saveCurrentSession(data),
+                result.current.saveCurrentSession(data),
+            ])
+        })
+        await finishWrites()
+        expect(await saves).toEqual([true, true])
+        expect(stored.size).toBe(1)
+        expect(result.current.currentSessionId).toBe([...stored.keys()][0])
+    })
+
+    it("drops a save scheduled before New Chat", async () => {
+        const { result } = await setup()
+        const scheduled = result.current.getChatGeneration()
+        act(() => result.current.clearCurrentSession())
+        let save!: Promise<boolean>
+        act(() => {
+            save = result.current.saveCurrentSession(data, scheduled)
+        })
+        await finishWrites()
+        expect(await save).toBe(true)
+        expect(stored.size).toBe(0)
+    })
+
+    it("drops a save of the old chat waiting behind New Chat's save", async () => {
+        const { result } = await setup()
+        // The auto-save is scheduled, then New Chat saves and clears
+        const scheduled = result.current.getChatGeneration()
+        let newChatSave!: Promise<boolean>
+        let autoSave!: Promise<boolean>
+        act(() => {
+            newChatSave = result.current.saveCurrentSession(data)
+            autoSave = result.current.saveCurrentSession(data, scheduled)
+        })
+        await act(async () => {
+            await waitFor(() => expect(pendingWrites).toHaveLength(1))
+            pendingWrites.shift()?.()
+            await newChatSave
+            result.current.clearCurrentSession()
+        })
+        await finishWrites()
+        await autoSave
+        expect(stored.size).toBe(1)
+        expect(result.current.currentSessionId).toBeNull()
+    })
+
+    it("keeps the new blank chat when a save of the old one ends later", async () => {
+        const { result } = await setup()
+        let save!: Promise<boolean>
+        act(() => {
+            save = result.current.saveCurrentSession(data)
+        })
+        await waitFor(() => expect(pendingWrites).toHaveLength(1))
+        // New Chat while the write runs
+        act(() => result.current.clearCurrentSession())
+        await finishWrites()
+        await save
+        expect(stored.size).toBe(1)
+        expect(result.current.currentSessionId).toBeNull()
+        expect(result.current.currentSession).toBeNull()
+    })
+
+    it("keeps New Chat when the URL's chat finishes loading after it", async () => {
+        stored.set("s1", { id: "s1", title: "Old", messages: [] })
+        const hook = renderHook(
+            ({ id }: { id: string | null }) =>
+                useSessionManager({ initialSessionId: id }),
+            { initialProps: { id: null as string | null } },
+        )
+        await waitFor(() => expect(hook.result.current.isLoading).toBe(false))
+        // The new chat's id reaches the URL; reading it takes a moment
+        let release!: () => void
+        readGate = new Promise((r) => {
+            release = r
+        })
+        hook.rerender({ id: "s1" })
+        act(() => hook.result.current.clearCurrentSession())
+        await act(async () => {
+            release()
+            await new Promise((r) => setTimeout(r, 0))
+        })
+        readGate = Promise.resolve()
+        expect(hook.result.current.currentSessionId).toBeNull()
+    })
+})

tests/unit/utils.test.ts

@@ -1,5 +1,37 @@
 import { describe, expect, it } from "vitest"
-import { cn, extractCompleteMxCells, isMxCellXmlComplete } from "@/lib/utils"
+import {
+    cn,
+    extractCompleteMxCells,
+    isMxCellXmlComplete,
+    isRealDiagram,
+} from "@/lib/utils"
+import { BLANK_MXFILE } from "@/packages/mcp-server/src/pages.ts"
+
+describe("isRealDiagram", () => {
+    it("counts a small diagram with one shape", () => {
+        // 234 characters: valid, shown, and saved with its chat
+        const xml =
+            '<mxfile><diagram id="p"><mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/><mxCell id="2" value="A" vertex="1" parent="1"><mxGeometry width="80" height="30" as="geometry"/></mxCell></root></mxGraphModel></diagram></mxfile>'
+        expect(xml.length).toBeLessThan(300)
+        expect(isRealDiagram(xml)).toBe(true)
+    })
+
+    it("does not count a blank page or nothing", () => {
+        expect(isRealDiagram(BLANK_MXFILE)).toBe(false)
+        expect(isRealDiagram("")).toBe(false)
+        expect(isRealDiagram(null)).toBe(false)
+    })
+
+    it("still counts a longer document of empty named pages", () => {
+        // Pages and page settings are worth keeping, as before
+        const pages = Array.from(
+            { length: 3 },
+            (_, i) =>
+                `<diagram id="p${i}" name="Page ${i}"><mxGraphModel pageWidth="1600" pageHeight="900"><root><mxCell id="0"/><mxCell id="1" parent="0"/></root></mxGraphModel></diagram>`,
+        ).join("")
+        expect(isRealDiagram(`<mxfile>${pages}</mxfile>`)).toBe(true)
+    })
+})
 
 describe("isMxCellXmlComplete", () => {
     it("returns false for empty/null input", () => {

tests/unit/validate-diagram-route.test.ts

@@ -49,15 +49,73 @@ vi.mock("@/lib/ai-providers", () => ({
         }),
 }))
 
+// The quota, off unless a test turns it on
+const quota = vi.hoisted(() => ({
+    enabled: false,
+    allowed: true,
+    checks: [] as Array<{ limits: any; increment?: number }>,
+    recorded: [] as number[],
+}))
+vi.mock("@/lib/dynamo-quota-manager", () => ({
+    isQuotaEnabled: () => quota.enabled,
+    checkAndIncrementRequest: async (
+        _ip: string,
+        limits: unknown,
+        increment?: number,
+    ) => {
+        quota.checks.push({ limits, increment })
+        return quota.allowed
+            ? { allowed: true }
+            : {
+                  allowed: false,
+                  type: "token",
+                  error: "Daily token limit exceeded",
+                  used: 10,
+                  limit: 10,
+              }
+    },
+    recordTokenUsage: async (_ip: string, tokens: number) => {
+        quota.recorded.push(tokens)
+    },
+}))
+
 const post = () =>
     new Request("http://localhost/api/validate-diagram", {
         method: "POST",
-        headers: { "Content-Type": "application/json" },
+        headers: {
+            "Content-Type": "application/json",
+            "x-forwarded-for": "203.0.113.7",
+        },
         body: JSON.stringify({ imageData: "data:image/png;base64,AAAA" }),
     })
 
 afterEach(() => {
     delete process.env.ENABLE_VLM_VALIDATION
+    quota.enabled = false
+    quota.allowed = true
+    quota.checks = []
+    quota.recorded = []
+})
+
+describe("the quota", () => {
+    it("refuses a check once the daily tokens are used up", async () => {
+        quota.enabled = true
+        quota.allowed = false
+        const res = await validateDiagram(post())
+        expect(res.status).toBe(429)
+    })
+
+    it("applies the token limits only, and records the tokens", async () => {
+        // The request limit is for chats: the day's last chat must still
+        // 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(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]))
+    })
 })
 
 describe("POST /api/validate-diagram", () => {

tests/unit/validate-model-bedrock.test.ts

@@ -0,0 +1,34 @@
+// @vitest-environment node
+import { describe, expect, it, vi } from "vitest"
+
+// What the Test hands to the provider factory
+const factory = vi.hoisted(() => ({ overrides: null as any }))
+vi.mock("@/lib/ai-providers", async (importOriginal) => ({
+    ...(await importOriginal<typeof import("@/lib/ai-providers")>()),
+    getAIModel: (overrides: unknown) => {
+        factory.overrides = overrides
+        throw new Error("stop here")
+    },
+}))
+
+import { POST as validateModel } from "@/app/api/validate-model/route"
+
+describe("testing Bedrock", () => {
+    it("passes temporary credentials' session token on, as the chat does", async () => {
+        await validateModel(
+            new Request("http://localhost/api/validate-model", {
+                method: "POST",
+                headers: { "Content-Type": "application/json" },
+                body: JSON.stringify({
+                    provider: "bedrock",
+                    modelId: "amazon.nova-lite-v1:0",
+                    awsAccessKeyId: "ASIA-temporary",
+                    awsSecretAccessKey: "secret",
+                    awsRegion: "us-east-1",
+                    awsSessionToken: "session-token",
+                }),
+            }),
+        )
+        expect(factory.overrides.awsSessionToken).toBe("session-token")
+    })
+})

tests/unit/validate-model-route.test.ts

@@ -8,10 +8,11 @@ import { getAIModel } from "@/lib/ai-providers"
 // No saved admin providers
 vi.mock("@/lib/admin/settings", () => ({ loadSettings: () => ({}) }))
 
-// Treat every URL as public so no test hits DNS
+// Every URL is public (no DNS in tests), unless a test says otherwise
+const privateUrls = vi.hoisted(() => ({ all: false }))
 vi.mock("@/lib/ssrf-protection", async (importOriginal) => ({
     ...(await importOriginal<typeof import("@/lib/ssrf-protection")>()),
-    isPrivateUrl: async () => false,
+    isPrivateUrl: async () => privateUrls.all,
 }))
 
 afterEach(() => {
@@ -163,6 +164,120 @@ describe("chat requests to a client base URL", () => {
     })
 })
 
+describe("testing EdgeOne", () => {
+    // The request validate-model sends to the EdgeOne function
+    const capture = () => {
+        const calls: Array<{ url: string; headers: Headers }> = []
+        vi.stubGlobal(
+            "fetch",
+            vi.fn(async (url: string, init?: RequestInit) => {
+                calls.push({
+                    url: String(url),
+                    headers: new Headers(init?.headers),
+                })
+                throw new Error("no network in tests")
+            }),
+        )
+        return calls
+    }
+
+    it("calls the site's own function, also without a base URL", async () => {
+        // The admin panel's Test sends none; a relative one cannot be fetched
+        const calls = capture()
+        await validateModel(
+            new Request("http://localhost/api/validate-model", {
+                method: "POST",
+                headers: {
+                    "Content-Type": "application/json",
+                    origin: "https://draw.example",
+                },
+                body: JSON.stringify({
+                    provider: "edgeone",
+                    modelId: "@tx/deepseek-ai/deepseek-v3-0324",
+                    baseUrl: "https://elsewhere.example/api/edgeai",
+                }),
+            }),
+        )
+        expect(calls[0]?.url).toBe(
+            "https://draw.example/api/edgeai/chat/completions",
+        )
+    })
+
+    it("passes the admin's access code and cookies on", async () => {
+        // The EdgeOne function checks the access code too
+        process.env.ADMIN_PASSWORD = "admin-pw"
+        try {
+            const calls = capture()
+            await testModel(
+                new Request("http://localhost/api/admin/test-model", {
+                    method: "POST",
+                    headers: {
+                        "Content-Type": "application/json",
+                        "x-admin-password": "admin-pw",
+                        "x-access-code": "visitor-code",
+                        cookie: "eo_token=t",
+                        origin: "https://draw.example",
+                    },
+                    body: JSON.stringify({
+                        provider: {
+                            id: "p1",
+                            provider: "edgeone",
+                            models: ["@tx/deepseek-ai/deepseek-v3-0324"],
+                        },
+                        modelId: "@tx/deepseek-ai/deepseek-v3-0324",
+                    }),
+                }),
+            )
+            expect(calls[0]?.url).toBe(
+                "https://draw.example/api/edgeai/chat/completions",
+            )
+            expect(calls[0]?.headers.get("x-access-code")).toBe("visitor-code")
+            expect(calls[0]?.headers.get("cookie")).toBe("eo_token=t")
+        } finally {
+            delete process.env.ADMIN_PASSWORD
+        }
+    })
+})
+
+describe("the admin Test of the server's own base URL", () => {
+    const test = (headers: Record<string, string>) =>
+        validateModel(
+            new Request("http://localhost/api/validate-model", {
+                method: "POST",
+                headers: { "Content-Type": "application/json", ...headers },
+                body: JSON.stringify({
+                    provider: "openai",
+                    apiKey: "panel-key",
+                    modelId: "gpt-5.5",
+                    baseUrl: "http://10.0.0.5:8000/v1",
+                    serverBaseUrl: true,
+                }),
+            }),
+        )
+
+    it("tests it as chat uses it: an internal address is allowed", async () => {
+        // ALLOW_PRIVATE_URLS=false guards URLs users type, not the server's
+        process.env.ALLOW_PRIVATE_URLS = "false"
+        process.env.OPENAI_BASE_URL = "http://10.0.0.5:8000/v1"
+        process.env.ADMIN_PASSWORD = "admin-pw"
+        privateUrls.all = true
+        try {
+            streamReply({ role: "assistant", content: "OK" })
+            const admin = await (
+                await test({ "x-admin-password": "admin-pw" })
+            ).json()
+            expect(admin.valid).toBe(true)
+            // Anyone else claiming it is still refused
+            const other = await test({})
+            expect(other.status).toBe(400)
+        } finally {
+            privateUrls.all = false
+            delete process.env.OPENAI_BASE_URL
+            delete process.env.ADMIN_PASSWORD
+        }
+    })
+})
+
 describe("the admin panel's Test button", () => {
     it("works when access codes are set", async () => {
         // The admin password stands in for the visitor access code