fix(chat): use the AI SDK's own mechanisms for quota, retries, stops and notices

Server
- Record quota tokens in onStepFinish and await the write, so a step that
  fails or is stopped later still counts; drop the onFinish/onAbort copies
- validate-diagram and the connection test end their model call with the
  request (req.signal) and use the SDK's timeout option; validate-diagram
  reports to Langfuse (functionId "validate-diagram")
- The connection test no longer counts an invalid tool call as support
- Validate incoming messages with safeValidateUIMessages (400 otherwise),
  drop client system messages explicitly, prune empty model messages with
  pruneMessages
- Chat tools move to lib/chat-tools.ts with an output schema for the
  client-side tools

Client
- Every automatic retry sends the canvas as it is then
  (prepareSendMessagesRequest): a drawing is committed before its
  screenshot check can fail, so the model must go on from it
- A notice when the output budget ran out before any tool call
  (finishReason "length")
- Error notices are assistant messages with a data-notice part instead of
  role "system" messages (lib/chat-message.ts: ChatUIMessage with metadata,
  data part and InferUITools types)
- regenerate, edit and retry use sendMessage({ messageId }); isResending
  and flushSync go
- isToolUIPart/getToolName instead of slicing "tool-"; ToolPartLike derives
  from the SDK's part types; addToolOutput typed with ChatAddToolOutputFunction

Tests: step failure after a finished step, abort of the screenshot check,
invalid tool call in the connection test, malformed messages.

app/api/chat/route.ts

@@ -2,12 +2,14 @@ import {
     APICallError,
     convertToModelMessages,
     InvalidToolInputError,
+    isFileUIPart,
+    isTextUIPart,
+    pruneMessages,
+    safeValidateUIMessages,
     stepCountIs,
     streamText,
 } from "ai"
 import { jsonrepair } from "jsonrepair"
-import path from "path"
-import { z } from "zod"
 import { checkAccessCode, rejectCrossSite } from "@/lib/access-code"
 import {
     CACHE_POINT,
@@ -25,6 +27,8 @@ import {
     formatSelectionContext,
     validateFileParts,
 } from "@/lib/chat-helpers"
+import type { ChatUIMessage } from "@/lib/chat-message"
+import { chatTools } from "@/lib/chat-tools"
 import { withDeprecatedParamsFallback } from "@/lib/deprecated-params"
 import {
     checkAndIncrementRequest,
@@ -52,14 +56,6 @@ import { getSystemPrompt } from "@/lib/system-prompts"
 import { normalizeBaseUrl } from "@/lib/types/model-config"
 import { getUserIdFromRequest } from "@/lib/user-id"
 import { foldCells } from "@/packages/mcp-server/src/compact-cells.ts"
-import {
-    getShapeLibrary,
-    SHAPE_LIBRARY_LIST,
-} from "@/packages/mcp-server/src/shape-library.ts"
-import {
-    STYLE_CLASS_EXAMPLE,
-    SWIMLANE_EXAMPLE,
-} from "@/packages/mcp-server/src/xml-examples.ts"
 
 // No explicit cap: a reasoning model can spend minutes planning before it emits
 // the tool call, so take whatever the host allows. Vercel's own default is 300s,
@@ -99,7 +95,19 @@ async function handleChatRequest(req: Request): Promise<Response> {
     if (accessDenied) return accessDenied
 
     const body = await req.json()
-    const { messages, xml, previousXml, sessionId } = body
+    // The messages in the SDK's UI shape. The route writes its own system
+    // messages, so none may come from the client.
+    const validated = await safeValidateUIMessages<ChatUIMessage>({
+        messages: body.messages,
+    })
+    if (!validated.success) {
+        return Response.json(
+            { error: validated.error.message },
+            { status: 400 },
+        )
+    }
+    const messages = validated.data.filter((m) => m.role !== "system")
+    const { xml, previousXml, sessionId } = body
     const customSystemMessage =
         typeof body.customSystemMessage === "string"
             ? body.customSystemMessage.slice(0, 5000)
@@ -118,9 +126,8 @@ async function handleChatRequest(req: Request): Promise<Response> {
     // Find the last USER message, not just the last message (which could be assistant in multi-step tool flows)
     const lastUserMessage = [...messages]
         .reverse()
-        .find((m: any) => m.role === "user")
+        .find((m) => m.role === "user")
-    const userInputText =
+    const userInputText = lastUserMessage?.parts.find(isTextUIPart)?.text || ""
-        lastUserMessage?.parts?.find((p: any) => p.type === "text")?.text || ""
 
     // Update Langfuse trace with input, session, and user
     setTraceInput({
@@ -318,9 +325,7 @@ async function handleChatRequest(req: Request): Promise<Response> {
         : systemMessage
 
     // Extract file parts (images) from the last user message
-    const fileParts =
+    const fileParts = lastUserMessage?.parts.filter(isFileUIPart) ?? []
-        lastUserMessage?.parts?.filter((part: any) => part.type === "file") ||
-        []
 
     // Note: we used to pre-emptively reject images for models we guessed were
     // text-only (by name matching). That heuristic misfired on newer models
@@ -333,12 +338,16 @@ async function handleChatRequest(req: Request): Promise<Response> {
 ${userInputText}
 """`
 
-    // Convert UIMessages to ModelMessages and add system message. A tool
+    // Convert UIMessages to ModelMessages. A tool call that never got its
-    // call that never got its result (the user stopped while it ran) is
+    // result (the user stopped while it ran) is left out: the SDK would
-    // left out: the SDK would refuse this and every later request of the
+    // refuse this and every later request of the chat
-    // chat (MissingToolResultsError)
+    // (MissingToolResultsError). Messages left without content (a notice
-    const modelMessages = await convertToModelMessages(messages, {
+    // the chat shows is a data part, which the model never sees) are
-        ignoreIncompleteToolCalls: true,
+    // dropped: Bedrock rejects empty content.
+    const modelMessages = pruneMessages({
+        messages: await convertToModelMessages(messages, {
+            ignoreIncompleteToolCalls: true,
+        }),
     })
 
     // DEBUG_LLM_PAYLOAD=true logs the incoming message structure
@@ -349,17 +358,10 @@ ${userInputText}
         )
     }
 
-    // Filter out messages with empty content arrays (Bedrock API rejects these)
-    // This is a safety measure - ideally convertToModelMessages should handle all cases
-    let enhancedMessages = modelMessages.filter(
-        (msg: any) =>
-            msg.content && Array.isArray(msg.content) && msg.content.length > 0,
-    )
-
     // Filter out tool-calls with invalid inputs (from failed repair or interrupted streaming)
     // and their results. Bedrock API rejects messages where toolUse.input is not a valid
     // JSON object, and every provider rejects a tool result whose call is gone.
-    enhancedMessages = dropInvalidToolCalls(enhancedMessages)
+    let enhancedMessages = dropInvalidToolCalls(modelMessages)
 
     // DEBUG_LLM_PAYLOAD=true logs what is sent to the model
     if (DEBUG_LLM_PAYLOAD) {
@@ -471,13 +473,10 @@ IMPORTANT: The "Current diagram XML" is the SINGLE SOURCE OF TRUTH for what's on
 
     const allMessages = [...systemMessages, ...enhancedMessages]
 
-    // Set by onAbort, which records the finished steps' tokens itself
-    let stopped = false
     const result = streamText({
         model,
-        // The system messages carry cache points, so they go in messages.
+        // The system messages carry cache points, so they go in messages
-        // A client's own system messages have string content and were
+        // (the client's own system messages were dropped above)
-        // dropped by the empty-content filter above.
         allowSystemInMessages: true,
         abortSignal: req.signal,
         // Must be sent: unset means the provider's own default, and Bedrock's is
@@ -524,166 +523,31 @@ IMPORTANT: The "Current diagram XML" is the SINGLE SOURCE OF TRUTH for what's on
         },
         messages: allMessages,
         ...(providerOptions && { providerOptions }), // This now includes all reasoning configs
-        // Langfuse telemetry config (returns undefined if not configured)
+        // Langfuse telemetry (undefined when not configured)
-        ...(getTelemetryConfig({ sessionId: validSessionId, userId }) && {
+        experimental_telemetry: getTelemetryConfig({
-            experimental_telemetry: getTelemetryConfig({
+            sessionId: validSessionId,
-                sessionId: validSessionId,
+            userId,
-                userId,
-            }),
         }),
-        onFinish: ({ text, totalUsage }) => {
+        // The quota counts each step as it ends: a later step that fails or
+        // is stopped has still used these tokens (onFinish runs only for a
+        // completed answer). Awaited, so the write is done before the host
+        // may freeze the function after the answer ends.
+        onStepFinish: async ({ usage }) => {
+            if (countsQuota) {
+                await recordTokenUsage(userId, usage.totalTokens ?? 0)
+            }
+        },
+        onFinish: ({ text }) => {
             // AI SDK 6 telemetry auto-reports token usage on its spans
             setTraceOutput(text)
-
-            // Record token usage for server-side quota tracking (if enabled)
-            // Use totalUsage (cumulative across all steps) instead of usage (final step only)
-            // inputTokens already includes cache reads and writes in AI SDK 6
-            if (countsQuota && totalUsage && !stopped) {
-                const totalTokens =
-                    (totalUsage.inputTokens || 0) +
-                    (totalUsage.outputTokens || 0)
-                recordTokenUsage(userId, totalTokens)
-            }
         },
         // onFinish is skipped when the stream fails or is aborted, so end the trace here
         onError: ({ error }) => {
             console.error(error) // what AI SDK does without an onError
             endTrace()
         },
-        onAbort: ({ steps }) => {
+        onAbort: () => endTrace(),
-            stopped = true
+        tools: chatTools,
-            endTrace()
-            // Stopped (or disconnected) after some steps finished: their
-            // tokens were used, or stopping every request after a costly
-            // first step would get around the token limits
-            if (countsQuota) {
-                const tokens = steps.reduce(
-                    (sum, step) =>
-                        sum +
-                        (step.usage.inputTokens || 0) +
-                        (step.usage.outputTokens || 0),
-                    0,
-                )
-                if (tokens > 0) recordTokenUsage(userId, tokens)
-            }
-        },
-        tools: {
-            // Client-side tool that will be executed on the client
-            display_diagram: {
-                description: `Display a diagram on draw.io. Pass ONLY the mxStyle definitions and the mxCell elements - wrapper tags and root cells are added automatically.
-
-VALIDATION RULES (XML will be rejected if violated):
-1. Generate ONLY mxStyle definitions and mxCell elements - NO wrapper tags (<mxfile>, <mxGraphModel>, <root>)
-2. Do NOT include root cells (id="0" or id="1") - they are added automatically
-3. All mxCell elements must be siblings - never nested
-4. Every mxCell needs a unique id (start from "2")
-5. parent defaults to "1"; write parent="<container-id>" only for shapes inside a container
-6. Escape special chars in values: &lt; &gt; &amp; &quot;
-
-A shape is one self-closing mxCell with x, y, w and h; an edge is one with source and target (a cell with source or target is always an edge). vertex="1", edge="1", parent="1" and the mxGeometry element are added automatically. Write parent only for a shape inside a container, and an mxGeometry element only for edge waypoints or for a separate label cell placed on an edge: <mxCell id="9" value="yes" style="edgeLabel;" parent="<edge id>" connectable="0"><mxGeometry x="-0.5" relative="1" as="geometry"/></mxCell>. An edge's own text simply goes in its value.
-
-Example (generate ONLY this - no wrapper tags):
-${SWIMLANE_EXAMPLE}
-
-Styles: define a style used by several cells ONCE with <mxStyle name="..." value="..."/> before the cells and use the name in style like a CSS class; overrides after the name win. Name only styles that two or more cells share; names must not be draw.io's own style names (text, ellipse, rhombus, swimlane, label, image, blue, green, red, gray, yellow, orange, purple, pink). A definition applies to the call it is in. html=1 and whiteSpace=wrap are added automatically, never write them. Labels are HTML: use &lt;br&gt; for a line break, never \\n; a literal < or > is written &amp;lt; or &amp;gt;.
-${STYLE_CLASS_EXAMPLE}
-
-Notes:
-- For AWS diagrams, use **AWS 2025 icons**.
-- For animated connectors, add "flowAnimation=1" to edge style.
-`,
-                inputSchema: z.object({
-                    xml: z
-                        .string()
-                        .describe("XML string to be displayed on draw.io"),
-                }),
-            },
-            edit_diagram: {
-                description: `Edit the current diagram by ID-based operations (update/add/delete cells).
-
-Operations:
-- update: Replace an existing cell by its id. Provide cell_id and complete new_xml.
-- add: Add a new cell. Provide cell_id (new unique id) and new_xml.
-- delete: Remove a cell. Cascade is automatic: children AND edges (source/target) are auto-deleted. Only specify ONE cell_id.
-
-For update/add, new_xml is the complete mxCell in the compact form (a shape with x, y, w, h; an edge with source and target), with its complete style (named styles are not available in edit_diagram).
-
-⚠️ JSON ESCAPING: Every " inside new_xml MUST be escaped as \\". Example: id=\\"5\\" value=\\"Label\\"
-
-Example - Add a rectangle:
-{"operations": [{"operation": "add", "cell_id": "rect-1", "new_xml": "<mxCell id=\\"rect-1\\" value=\\"Hello\\" style=\\"rounded=1;\\" x=\\"100\\" y=\\"100\\" w=\\"120\\" h=\\"60\\"/>"}]}
-
-Example - Delete container (children & edges auto-deleted):
-{"operations": [{"operation": "delete", "cell_id": "2"}]}`,
-                inputSchema: z.object({
-                    operations: z
-                        .array(
-                            z.object({
-                                operation: z
-                                    .enum(["update", "add", "delete"])
-                                    .describe(
-                                        "Operation to perform: add, update, or delete",
-                                    ),
-                                cell_id: z
-                                    .string()
-                                    .describe(
-                                        "The id of the mxCell. Must match the id attribute in new_xml.",
-                                    ),
-                                new_xml: z
-                                    .string()
-                                    .optional()
-                                    .describe(
-                                        "Complete mxCell XML element (required for update/add)",
-                                    ),
-                            }),
-                        )
-                        .describe("Array of operations to apply"),
-                }),
-            },
-            append_diagram: {
-                description: `Continue generating diagram XML when previous display_diagram output was truncated due to length limits.
-
-WHEN TO USE: Only call this tool after display_diagram was truncated (you'll see an error message about truncation).
-
-CRITICAL INSTRUCTIONS:
-1. Do NOT include any wrapper tags - just continue the mxCell elements
-2. Continue from EXACTLY where your previous output stopped
-3. Complete the remaining mxCell elements
-4. If still truncated, call append_diagram again with the next fragment
-
-Example: If previous output ended with '<mxCell id="x" style="rounded=1', continue with ';" x="40" y="40" w="120" h="60"/>' and complete the remaining elements.`,
-                inputSchema: z.object({
-                    xml: z
-                        .string()
-                        .describe(
-                            "Continuation XML fragment to append (NO wrapper tags)",
-                        ),
-                }),
-            },
-            get_shape_library: {
-                description: `Get draw.io shape/icon library documentation with style syntax and shape names.
-
-Available libraries:
-${SHAPE_LIBRARY_LIST}
-
-Call this tool to get shape names and usage syntax for a specific library.`,
-                inputSchema: z.object({
-                    library: z
-                        .string()
-                        .describe(
-                            "Library name (e.g., 'aws4', 'kubernetes', 'flowchart')",
-                        ),
-                }),
-                execute: async ({ library }) => {
-                    // Only known library names reach the file system
-                    const result = await getShapeLibrary(
-                        library,
-                        path.join(process.cwd(), "docs/shape-libraries"),
-                    )
-                    return result.ok ? result.text : result.error
-                },
-            },
-        },
         ...(process.env.TEMPERATURE !== undefined && {
             temperature: parseFloat(process.env.TEMPERATURE),
         }),

app/api/validate-diagram/route.ts

@@ -11,6 +11,7 @@ import {
     isQuotaEnabled,
     recordTokenUsage,
 } from "@/lib/dynamo-quota-manager"
+import { getTelemetryConfig } from "@/lib/langfuse"
 import { getUserIdFromRequest } from "@/lib/user-id"
 import { VALIDATION_SYSTEM_PROMPT } from "@/lib/validation-prompts"
 import {
@@ -146,13 +147,18 @@ export async function POST(req: Request): Promise<Response> {
                 },
             ],
             maxOutputTokens: 1024,
-            abortSignal: AbortSignal.timeout(timeout),
+            // Ends with the request: the user pressed Stop, or left
+            abortSignal: req.signal,
+            timeout,
+            // Its usage and cost show in Langfuse next to the chat's
+            experimental_telemetry: getTelemetryConfig({
+                functionId: "validate-diagram",
+                sessionId,
+                userId,
+            }),
         })
         if (countsQuota) {
-            await recordTokenUsage(
+            await recordTokenUsage(userId, totalUsage.totalTokens ?? 0)
-                userId,
-                (totalUsage.inputTokens || 0) + (totalUsage.outputTokens || 0),
-            )
         }
         if (sessionId) {
             console.log(

components/chat/chat-engine.tsx

@@ -4,8 +4,8 @@ import { useChat } from "@ai-sdk/react"
 import {
     type ChatStatus,
     DefaultChatTransport,
+    getToolName,
     isToolUIPart,
-    type UIMessage,
 } from "ai"
 import { usePathname, useRouter, useSearchParams } from "next/navigation"
 import type React from "react"
@@ -19,7 +19,6 @@ import {
     useRef,
     useState,
 } from "react"
-import { flushSync } from "react-dom"
 import { toast } from "sonner"
 import type { ValidationState } from "@/components/chat/ValidationCard"
 import { useDiagram } from "@/contexts/diagram-context"
@@ -37,6 +36,12 @@ import {
 import { useValidateDiagram } from "@/hooks/use-validate-diagram"
 import { getApiEndpoint } from "@/lib/base-path"
 import { findCachedResponse } from "@/lib/cached-responses"
+import {
+    type ChatToolName,
+    type ChatUIMessage,
+    type MessageSelection,
+    noticeMessage,
+} from "@/lib/chat-message"
 import { buildChatHeaders } from "@/lib/chat-request"
 import { pageModelXml, placeOnPage } from "@/lib/diagram-pages"
 import { EMPTY_DRAWIO_DOCUMENT } from "@/lib/drawio/drawio-config"
@@ -76,7 +81,7 @@ const SAVE_DEBOUNCE_MS = 1000
  * Only checks the LAST tool part (most recent tool call), not all tool parts.
  * A call the user stopped is not one to retry.
  */
-function hasToolErrors(messages: UIMessage[]): boolean {
+function hasToolErrors(messages: ChatUIMessage[]): boolean {
     const lastMessage = messages[messages.length - 1]
     if (lastMessage?.role !== "assistant") return false
     const lastToolPart = lastMessage.parts.filter(isToolUIPart).at(-1)
@@ -103,14 +108,9 @@ function viewedPageId(): string | null {
     return useCanvasStore.getState().currentPageId
 }
 
-// Shapes sent with a user message (also kept in its metadata)
+function selectionOf(message: ChatUIMessage): MessageSelection {
-type MessageSelection = Pick<SelectedCell, "id" | "label">[]
+    const selected = message.metadata?.selectedCells
-
+    return Array.isArray(selected) ? selected : []
-function selectionOf(message: UIMessage): MessageSelection {
-    const meta = message.metadata as
-        | { selectedCells?: MessageSelection }
-        | undefined
-    return Array.isArray(meta?.selectedCells) ? meta.selectedCells : []
 }
 
 function newLangfuseSessionId() {
@@ -169,15 +169,13 @@ function withUniquePageIds(xml: string): string {
 class FileReadError extends Error {}
 
 export interface ChatEngine {
-    messages: UIMessage[]
+    messages: ChatUIMessage[]
     status: ChatStatus
     error: Error | undefined
     isBusy: boolean
     /** The chat on screen is being left (saved, then another one shown) */
     isLeaving: boolean
-    /** A message is being sent again; the chat is not empty meanwhile */
+    setMessages: (messages: ChatUIMessage[]) => void
-    isResending: boolean
-    setMessages: (messages: UIMessage[]) => void
     stop: () => void
 
     input: string
@@ -361,9 +359,6 @@ export function ChatEngineProvider({
     const [isLeaving, setIsLeaving] = useState(false)
     // The user's provider the current turn was sent to, for its errors
     const requestProviderIdRef = useRef<string | undefined>(undefined)
-    // A message is being sent again (retry, regenerate, edit), until the
-    // SDK has put it back
-    const [isResending, setIsResending] = useState(false)
     // Presses of Stop: a check that began before one still knows of it after
     // the next message clears stoppedRef
     const stopCountRef = useRef(0)
@@ -429,9 +424,40 @@ export function ChatEngineProvider({
         onDisplayChart(original, true, "revert")
     }
 
+    // Requests made in the turn so far. The first carries the diagram the
+    // message was sent with (exported then, or the snapshot it is sent
+    // again from). Each later one, an automatic retry after a tool error,
+    // carries the canvas as it is then: a drawing is committed before its
+    // screenshot check can fail, and the model must go on from that
+    // drawing, not from the one before the turn.
+    const turnRequestsRef = useRef(0)
     const transport = useMemo(
-        () => new DefaultChatTransport({ api: getApiEndpoint("/api/chat") }),
+        () =>
-        [],
+            new DefaultChatTransport<ChatUIMessage>({
+                api: getApiEndpoint("/api/chat"),
+                prepareSendMessagesRequest: ({
+                    id,
+                    messages,
+                    body,
+                    trigger,
+                    messageId,
+                }) => ({
+                    body: {
+                        ...body,
+                        ...(turnRequestsRef.current++ > 0 && {
+                            xml: getPageXml(
+                                chartXMLRef.current,
+                                turnPageIdRef.current,
+                            ),
+                        }),
+                        id,
+                        messages,
+                        trigger,
+                        messageId,
+                    },
+                }),
+            }),
+        [chartXMLRef],
     )
 
     const {
@@ -442,19 +468,19 @@ export function ChatEngineProvider({
         error,
         setMessages,
         stop,
-    } = useChat({
+    } = useChat<ChatUIMessage>({
         transport,
         onToolCall: async ({ toolCall }) => {
             await handleToolCall({ toolCall }, addToolOutputSafely)
         },
-        onFinish: ({ message, isAbort, isError }) => {
+        onFinish: ({ message, isAbort, isError, finishReason }) => {
             // Stopped, failed, or the reply just ended (a closed connection):
             // tool calls still streaming never reach the tool handler. Mark
             // them handled so a later render of the stream does not draw
             // their preview again.
-            const cut = (message.parts as any[]).filter(
+            const cut = message.parts
-                (part) => part.state === "input-streaming" && part.toolCallId,
+                .filter(isToolUIPart)
-            )
+                .filter((part) => part.state === "input-streaming")
             for (const part of cut) {
                 processedToolCallsRef.current.add(part.toolCallId)
             }
@@ -463,18 +489,16 @@ export function ChatEngineProvider({
             // or the page going, aborts a reply). Not once another chat is
             // on screen: the call is not among its messages.
             if (isAbort && getChatGeneration() === turnChatRef.current) {
-                for (const part of cut) {
+                for (const part of cut) stopToolCall(part)
-                    addToolOutputSafely({
-                        tool: part.type.replace("tool-", ""),
-                        toolCallId: part.toolCallId,
-                        state: "output-error",
-                        errorText: "Stopped by user",
-                    })
-                }
             }
-            if (isAbort || isError || cut.length === 0) return
+            if (isAbort || isError) return
-            // Ended with no error: undo the preview, and say so, as an
+            // The output budget ran out before the model called a tool (its
-            // error does (with Retry)
+            // thinking took it all): nothing was drawn
+            const outOfTokens =
+                finishReason === "length" && !message.parts.some(isToolUIPart)
+            if (cut.length === 0 && !outOfTokens) return
+            // Ended with no error and no diagram: undo the preview, and say
+            // so, as an error does (with Retry)
             const [originalXml] = editDiagramOriginalXmlRef.current.values()
             if (originalXml) onDisplayChart(originalXml, true, "revert")
             editDiagramOriginalXmlRef.current.clear()
@@ -483,16 +507,11 @@ export function ChatEngineProvider({
             if (getChatGeneration() !== turnChatRef.current) return
             setMessages((currentMessages) => [
                 ...currentMessages,
-                {
+                noticeMessage({
-                    id: `error-${Date.now()}`,
+                    text: outOfTokens
-                    role: "system" as const,
+                        ? dict.errors.llm.output_truncated
-                    parts: [
+                        : dict.errors.answerEndedEarly,
-                        {
+                }),
-                            type: "text" as const,
-                            text: dict.errors.answerEndedEarly,
-                        },
-                    ],
-                },
             ])
         },
         onError: (error) => {
@@ -561,19 +580,15 @@ export function ChatEngineProvider({
             // the model selected may have changed since, e.g. in another tab
             const providerId = requestProviderIdRef.current
 
-            // A system message, so it can be cleared with the conversation
+            // In the chat, so it is cleared with the conversation; with a
+            // button that opens model settings, on the page of the provider
+            // that refused
             setMessages((currentMessages) => [
                 ...currentMessages,
-                {
+                noticeMessage({
-                    id: `error-${Date.now()}`,
+                    text,
-                    role: "system" as const,
+                    ...(openModelConfig && { openModelConfig, providerId }),
-                    parts: [{ type: "text" as const, text }],
+                }),
-                    // The message shows a button that opens model settings,
-                    // on the page of the provider that refused
-                    ...(openModelConfig && {
-                        metadata: { openModelConfig: true, providerId },
-                    }),
-                },
             ])
 
             if (isAccessCodeError) openSettings("general")
@@ -603,9 +618,19 @@ export function ChatEngineProvider({
         )
     }
 
+    // A call the user stopped before it reached the tool handler. The SDK
+    // finds the call by its id; the tool name is for the types only.
+    const stopToolCall = (part: Parameters<typeof getToolName>[0]) =>
+        addToolOutputSafely({
+            tool: getToolName(part) as ChatToolName,
+            toolCallId: part.toolCallId,
+            state: "output-error",
+            errorText: "Stopped by user",
+        })
+
     // After each request: send the next one (retry, continue a cut off
     // drawing) or end the turn
-    const shouldSendAgain = (messages: UIMessage[]) => {
+    const shouldSendAgain = (messages: ChatUIMessage[]) => {
         // The user stopped: a tool result that arrives later (a VLM
         // check still running) must not start a new request
         if (stoppedRef.current) return false
@@ -657,9 +682,6 @@ export function ChatEngineProvider({
     }
 
     const isBusy = status === "streaming" || status === "submitted"
-    useEffect(() => {
-        if (messages.length > 0) setIsResending(false)
-    }, [messages.length])
     // isBusy of the latest render, for code that runs after an await
     const busyRef = useRef(isBusy)
     busyRef.current = isBusy
@@ -771,7 +793,7 @@ export function ChatEngineProvider({
     const [shownChats, setShownChats] = useState(0)
     // The messages shown when another chat was put on screen: useChat can
     // show that chat's own up to 150 ms later than the rest of its state
-    const messagesBeforeShownRef = useRef<UIMessage[] | null>(null)
+    const messagesBeforeShownRef = useRef<ChatUIMessage[] | null>(null)
 
     const syncUIWithSession = useCallback(
         (data: SessionData | null) => {
@@ -811,7 +833,7 @@ export function ChatEngineProvider({
                 loadedMessageIdsRef.current = new Set(
                     (data.messages as { id: string }[]).map((m) => m.id),
                 )
-                setMessages(data.messages as unknown as UIMessage[])
+                setMessages(data.messages as unknown as ChatUIMessage[])
                 xmlSnapshotsRef.current = new Map(data.xmlSnapshots)
                 if (data.diagramXml) onDisplayChart(data.diagramXml, true)
                 else clearDiagram()
@@ -1118,12 +1140,15 @@ export function ChatEngineProvider({
             : ""
     }
 
+    // messageId: the user message sent again (regenerate, edit, retry); the
+    // SDK puts the new one in its place and drops the answers after it
     const sendChatMessage = (
         parts: any,
         xml: string,
         previousXml: string,
         turnIndex: number,
         selectedCells: MessageSelection,
+        messageId?: string,
     ) => {
         autoRetryCountRef.current = 0
         continuationRetryCountRef.current = 0
@@ -1161,9 +1186,11 @@ export function ChatEngineProvider({
                 maxOutputTokens,
             }),
         }
+        turnRequestsRef.current = 0
         // The selection stays with the message, for regenerate and retry
         sendMessage(
             {
+                ...(messageId !== undefined && { messageId }),
                 parts,
                 ...(selected.length > 0 && {
                     metadata: { selectedCells: selected },
@@ -1416,18 +1443,13 @@ export function ChatEngineProvider({
         dropLaterSnapshots(index)
         useVersionsStore.getState().removeFromTurn(index)
 
-        // Remove the user message onwards (sendMessage re-adds it). For a
-        // moment there may be no message: that is no empty chat.
-        flushSync(() => {
-            setIsResending(true)
-            setMessages(messages.slice(0, index))
-        })
         sendChatMessage(
             parts,
             getPageXml(savedXml, viewedPageId()),
             previousXml,
             index,
             selectionOf(messages[index]),
+            messages[index].id,
         )
 
         return true
@@ -1493,16 +1515,13 @@ export function ChatEngineProvider({
                     chartXMLRef.current || chartXml,
                 )
             }
-            flushSync(() => {
-                setIsResending(true)
-                setMessages(messagesRef.current.slice(0, index))
-            })
             sendChatMessage(
                 message.parts,
                 chartXml,
                 getPreviousXml(index),
                 index,
                 selectionOf(message),
+                message.id,
             )
         } catch (error) {
             console.error("Error fetching chart data:", error)
@@ -1521,25 +1540,17 @@ export function ChatEngineProvider({
         const lastMessage = messages[messages.length - 1]
         // Calls the tool handler already took can still show as streaming:
         // the messages update at most every 150 ms (useChat throttle)
-        const toolParts = lastMessage?.parts?.filter(
+        const streaming = (lastMessage?.parts ?? [])
-            (part: any) =>
+            .filter(isToolUIPart)
-                part.type?.startsWith("tool-") &&
+            .filter(
-                part.state === "input-streaming" &&
+                (part) =>
-                !processedToolCallsRef.current.has(part.toolCallId),
+                    part.state === "input-streaming" &&
-        )
+                    !processedToolCallsRef.current.has(part.toolCallId),
-        for (const part of (toolParts ?? []) as any[]) {
+            )
-            if (part.toolCallId) {
+        for (const part of streaming) stopToolCall(part)
-                addToolOutputSafely({
-                    tool: part.type.replace("tool-", ""),
-                    toolCallId: part.toolCallId,
-                    state: "output-error",
-                    errorText: "Stopped by user",
-                })
-            }
-        }
         stop()
         abandonContinuation()
-    }, [messages, addToolOutputSafely, stop, cancelValidation])
+    }, [messages, stopToolCall, stop, cancelValidation])
 
     // ---------------------------------------------------------------------
     // Session actions
@@ -1849,8 +1860,7 @@ export function ChatEngineProvider({
         error,
         isBusy,
         isLeaving,
-        isResending,
+        setMessages: setMessages as (messages: ChatUIMessage[]) => void,
-        setMessages: setMessages as (messages: UIMessage[]) => void,
         stop: handleStop,
         input,
         setInput,

components/chat/message-list.tsx

@@ -1,4 +1,4 @@
-import type { UIMessage } from "ai"
+import { isToolUIPart, type UIMessage } from "ai"
 import {
     AlertCircle,
     BookmarkPlus,
@@ -24,6 +24,7 @@ import { ValidationCard } from "@/components/chat/ValidationCard"
 import Image from "@/components/image-with-basepath"
 import { useDictionary } from "@/hooks/use-dictionary"
 import { getApiEndpoint } from "@/lib/base-path"
+import { noticeOf } from "@/lib/chat-message"
 import { useCopy } from "@/lib/clipboard"
 import { cn } from "@/lib/utils"
 import { useUiStore } from "@/stores/ui-store"
@@ -369,7 +370,7 @@ function AssistantMessage({
     parts.forEach((part, partIndex) => {
         if (part.type === "reasoning")
             blocks.push({ kind: "reasoning", index: partIndex, texts: [] })
-        else if (part.type?.startsWith("tool-"))
+        else if (isToolUIPart(part))
             blocks.push({ kind: "tool", index: partIndex, texts: [] })
         else if (part.type === "text") {
             const text = (part as { text: string }).text
@@ -568,9 +569,10 @@ export function MessageList() {
         }
     }, [messages])
 
-    const lastAssistantIndex = messages
+    // A notice (a failed request) is an assistant message without an answer
-        .map((m) => m.role)
+    const lastAssistantIndex = messages.findLastIndex(
-        .lastIndexOf("assistant")
+        (m) => m.role === "assistant" && !noticeOf(m),
+    )
     const lastUserIndex = messages.map((m) => m.role).lastIndexOf("user")
 
     return (
@@ -579,6 +581,7 @@ export function MessageList() {
                 const isRestored = engine.loadedMessageIdsRef.current.has(
                     message.id,
                 )
+                const notice = noticeOf(message)
                 return (
                     <div
                         key={message.id}
@@ -590,33 +593,18 @@ export function MessageList() {
                                 index={index}
                                 isLast={index === lastUserIndex}
                             />
-                        ) : message.role === "assistant" ? (
+                        ) : notice ? (
-                            <AssistantMessage
-                                message={message}
-                                index={index}
-                                isLastMessage={index === messages.length - 1}
-                                isLastAssistant={index === lastAssistantIndex}
-                            />
-                        ) : (
                             <ErrorNotice
-                                text={getMessageTextContent(message)}
+                                text={notice.text}
-                                openModelConfig={
+                                openModelConfig={!!notice.openModelConfig}
-                                    !!(
+                                onOpenModelConfig={() =>
-                                        message.metadata as
-                                            | { openModelConfig?: boolean }
-                                            | undefined
-                                    )?.openModelConfig
-                                }
-                                onOpenModelConfig={() => {
-                                    const { providerId } =
-                                        (message.metadata as
-                                            | { providerId?: string }
-                                            | undefined) ?? {}
                                     openSettings(
                                         "models",
-                                        providerId ? { providerId } : "list",
+                                        notice.providerId
+                                            ? { providerId: notice.providerId }
+                                            : "list",
                                     )
-                                }}
+                                }
                                 onRetry={
                                     index === messages.length - 1 &&
                                     !engine.isBusy &&
@@ -625,6 +613,13 @@ export function MessageList() {
                                         : undefined
                                 }
                             />
+                        ) : (
+                            <AssistantMessage
+                                message={message}
+                                index={index}
+                                isLastMessage={index === messages.length - 1}
+                                isLastAssistant={index === lastAssistantIndex}
+                            />
                         )}
                     </div>
                 )

components/chat/tool-activity.tsx

@@ -1,3 +1,4 @@
+import { getToolName } from "ai"
 import {
     AlertCircle,
     Check,
@@ -164,7 +165,7 @@ function ToolRow({
                 tone === "error" && "border-destructive/25",
             )}
             data-testid="tool-row"
-            data-tool-name={part.type?.replace("tool-", "")}
+            data-tool-name={getToolName(part)}
             data-tool-state={part.state}
         >
             <button
@@ -246,7 +247,7 @@ export function ToolActivity({
     const t = dict.versions
     const engine = useChatEngine()
     const versions = useVersionsStore((s) => s.versions)
-    const toolName = part.type?.replace("tool-", "")
+    const toolName = getToolName(part)
     const { state, errorText, output } = part
 
     const versionIndex = versions.findIndex(

components/chat/types.ts

@@ -1,15 +1,19 @@
+import type { DynamicToolUIPart, ToolUIPart, UITools } from "ai"
 import type { DiagramOperation } from "@/packages/mcp-server/src/diagram-operations.ts"
 
 export type { DiagramOperation }
 
-export interface ToolPartLike {
+/**
-    type: string
+ * A tool part as the UI reads it, in any state: input, output and errorText
-    toolCallId: string
+ * are each there or not, instead of one shape per state.
-    state?: string
+ */
-    input?: {
+export type ToolPartLike = (ToolUIPart<UITools> | DynamicToolUIPart) & {
-        xml?: string
+    input?:
-        operations?: DiagramOperation[]
+        | ({
-    } & Record<string, unknown>
+              xml?: string
+              operations?: DiagramOperation[]
+          } & Record<string, unknown>)
+        | undefined
     output?: string
     errorText?: string
 }

components/workspace/workspace.tsx

@@ -150,7 +150,6 @@ function WorkspaceLayout() {
     const showHero =
         engine.isRestored &&
         engine.messages.length === 0 &&
-        !engine.isResending &&
         !isRealDiagram(chartXML) &&
         !heroDismissed
 

hooks/use-diagram-tool-handlers.ts

@@ -1,3 +1,4 @@
+import type { ChatAddToolOutputFunction } from "ai"
 import type { RefObject } from "react"
 import type { DiagramOperation } from "@/components/chat/types"
 import type {
@@ -5,6 +6,7 @@ import type {
     ValidationStatus,
 } from "@/components/chat/ValidationCard"
 import type { LoadMode } from "@/contexts/diagram-context"
+import type { ChatUIMessage } from "@/lib/chat-message"
 import { pageModelXml, pageSelectorFor, placeOnPage } from "@/lib/diagram-pages"
 import type { ValidationResult } from "@/lib/diagram-validator"
 import { formatValidationFeedback } from "@/lib/diagram-validator"
@@ -24,25 +26,8 @@ interface ToolCall {
     input: unknown
 }
 
-type AddToolOutputSuccess = {
+// useChat's addToolOutput, typed with the chat's tools
-    tool: string
+type AddToolOutputFn = ChatAddToolOutputFunction<ChatUIMessage>
-    toolCallId: string
-    state?: "output-available"
-    output: string
-    errorText?: undefined
-}
-
-type AddToolOutputError = {
-    tool: string
-    toolCallId: string
-    state: "output-error"
-    output?: undefined
-    errorText: string
-}
-
-type AddToolOutputParams = AddToolOutputSuccess | AddToolOutputError
-
-type AddToolOutputFn = (params: AddToolOutputParams) => void
 
 const MAX_VALIDATION_RETRIES = 3
 

lib/chat-message.ts

@@ -0,0 +1,44 @@
+import type { InferUITools, UIMessage } from "ai"
+import type { chatTools } from "@/lib/chat-tools"
+import type { SelectedCell } from "@/stores/canvas-store"
+
+/** Shapes sent with a user message (also kept in its metadata) */
+export type MessageSelection = Pick<SelectedCell, "id" | "label">[]
+
+export interface ChatMessageMetadata {
+    selectedCells?: MessageSelection
+}
+
+/**
+ * A failed request or an answer that drew nothing, shown in the chat with
+ * what to do: an assistant message with this as its only part. The server
+ * ignores data parts, so it never reaches the model.
+ */
+export interface ChatNotice {
+    text: string
+    /** Show a button that opens model settings, on the provider that refused */
+    openModelConfig?: boolean
+    providerId?: string
+}
+
+export type ChatToolName = keyof InferUITools<typeof chatTools>
+
+export type ChatUIMessage = UIMessage<
+    ChatMessageMetadata,
+    { notice: ChatNotice },
+    InferUITools<typeof chatTools>
+>
+
+export function noticeMessage(notice: ChatNotice): ChatUIMessage {
+    return {
+        id: `notice-${Date.now()}`,
+        role: "assistant",
+        parts: [{ type: "data-notice", data: notice }],
+    }
+}
+
+/** The notice a message carries, if it is one */
+export function noticeOf(message: ChatUIMessage): ChatNotice | undefined {
+    const part = message.parts.find((p) => p.type === "data-notice")
+    return part?.data
+}

lib/chat-tools.ts

@@ -0,0 +1,138 @@
+import { type InferUITools, tool } from "ai"
+import path from "path"
+import { z } from "zod"
+import {
+    getShapeLibrary,
+    SHAPE_LIBRARY_LIST,
+} from "@/packages/mcp-server/src/shape-library.ts"
+import {
+    STYLE_CLASS_EXAMPLE,
+    SWIMLANE_EXAMPLE,
+} from "@/packages/mcp-server/src/xml-examples.ts"
+
+/**
+ * The chat's tools. display_diagram, edit_diagram and append_diagram run in
+ * the browser: the chat engine's tool handler puts their result on the
+ * canvas and reports it back as text, so they have an output schema and no
+ * execute. get_shape_library runs on the server.
+ */
+export const chatTools = {
+    display_diagram: tool({
+        description: `Display a diagram on draw.io. Pass ONLY the mxStyle definitions and the mxCell elements - wrapper tags and root cells are added automatically.
+
+VALIDATION RULES (XML will be rejected if violated):
+1. Generate ONLY mxStyle definitions and mxCell elements - NO wrapper tags (<mxfile>, <mxGraphModel>, <root>)
+2. Do NOT include root cells (id="0" or id="1") - they are added automatically
+3. All mxCell elements must be siblings - never nested
+4. Every mxCell needs a unique id (start from "2")
+5. parent defaults to "1"; write parent="<container-id>" only for shapes inside a container
+6. Escape special chars in values: &lt; &gt; &amp; &quot;
+
+A shape is one self-closing mxCell with x, y, w and h; an edge is one with source and target (a cell with source or target is always an edge). vertex="1", edge="1", parent="1" and the mxGeometry element are added automatically. Write parent only for a shape inside a container, and an mxGeometry element only for edge waypoints or for a separate label cell placed on an edge: <mxCell id="9" value="yes" style="edgeLabel;" parent="<edge id>" connectable="0"><mxGeometry x="-0.5" relative="1" as="geometry"/></mxCell>. An edge's own text simply goes in its value.
+
+Example (generate ONLY this - no wrapper tags):
+${SWIMLANE_EXAMPLE}
+
+Styles: define a style used by several cells ONCE with <mxStyle name="..." value="..."/> before the cells and use the name in style like a CSS class; overrides after the name win. Name only styles that two or more cells share; names must not be draw.io's own style names (text, ellipse, rhombus, swimlane, label, image, blue, green, red, gray, yellow, orange, purple, pink). A definition applies to the call it is in. html=1 and whiteSpace=wrap are added automatically, never write them. Labels are HTML: use &lt;br&gt; for a line break, never \\n; a literal < or > is written &amp;lt; or &amp;gt;.
+${STYLE_CLASS_EXAMPLE}
+
+Notes:
+- For AWS diagrams, use **AWS 2025 icons**.
+- For animated connectors, add "flowAnimation=1" to edge style.
+`,
+        inputSchema: z.object({
+            xml: z.string().describe("XML string to be displayed on draw.io"),
+        }),
+        outputSchema: z.string(),
+    }),
+    edit_diagram: tool({
+        description: `Edit the current diagram by ID-based operations (update/add/delete cells).
+
+Operations:
+- update: Replace an existing cell by its id. Provide cell_id and complete new_xml.
+- add: Add a new cell. Provide cell_id (new unique id) and new_xml.
+- delete: Remove a cell. Cascade is automatic: children AND edges (source/target) are auto-deleted. Only specify ONE cell_id.
+
+For update/add, new_xml is the complete mxCell in the compact form (a shape with x, y, w, h; an edge with source and target), with its complete style (named styles are not available in edit_diagram).
+
+⚠️ JSON ESCAPING: Every " inside new_xml MUST be escaped as \\". Example: id=\\"5\\" value=\\"Label\\"
+
+Example - Add a rectangle:
+{"operations": [{"operation": "add", "cell_id": "rect-1", "new_xml": "<mxCell id=\\"rect-1\\" value=\\"Hello\\" style=\\"rounded=1;\\" x=\\"100\\" y=\\"100\\" w=\\"120\\" h=\\"60\\"/>"}]}
+
+Example - Delete container (children & edges auto-deleted):
+{"operations": [{"operation": "delete", "cell_id": "2"}]}`,
+        inputSchema: z.object({
+            operations: z
+                .array(
+                    z.object({
+                        operation: z
+                            .enum(["update", "add", "delete"])
+                            .describe(
+                                "Operation to perform: add, update, or delete",
+                            ),
+                        cell_id: z
+                            .string()
+                            .describe(
+                                "The id of the mxCell. Must match the id attribute in new_xml.",
+                            ),
+                        new_xml: z
+                            .string()
+                            .optional()
+                            .describe(
+                                "Complete mxCell XML element (required for update/add)",
+                            ),
+                    }),
+                )
+                .describe("Array of operations to apply"),
+        }),
+        outputSchema: z.string(),
+    }),
+    append_diagram: tool({
+        description: `Continue generating diagram XML when previous display_diagram output was truncated due to length limits.
+
+WHEN TO USE: Only call this tool after display_diagram was truncated (you'll see an error message about truncation).
+
+CRITICAL INSTRUCTIONS:
+1. Do NOT include any wrapper tags - just continue the mxCell elements
+2. Continue from EXACTLY where your previous output stopped
+3. Complete the remaining mxCell elements
+4. If still truncated, call append_diagram again with the next fragment
+
+Example: If previous output ended with '<mxCell id="x" style="rounded=1', continue with ';" x="40" y="40" w="120" h="60"/>' and complete the remaining elements.`,
+        inputSchema: z.object({
+            xml: z
+                .string()
+                .describe(
+                    "Continuation XML fragment to append (NO wrapper tags)",
+                ),
+        }),
+        outputSchema: z.string(),
+    }),
+    get_shape_library: tool({
+        description: `Get draw.io shape/icon library documentation with style syntax and shape names.
+
+Available libraries:
+${SHAPE_LIBRARY_LIST}
+
+Call this tool to get shape names and usage syntax for a specific library.`,
+        inputSchema: z.object({
+            library: z
+                .string()
+                .describe(
+                    "Library name (e.g., 'aws4', 'kubernetes', 'flowchart')",
+                ),
+        }),
+        execute: async ({ library }) => {
+            // Only known library names reach the file system
+            const result = await getShapeLibrary(
+                library,
+                path.join(process.cwd(), "docs/shape-libraries"),
+            )
+            return result.ok ? result.text : result.error
+        },
+    }),
+}
+
+/** The tool names and their input and output types, for the UI messages */
+export type ChatTools = InferUITools<typeof chatTools>

lib/langfuse.ts

@@ -1,6 +1,7 @@
 import { LangfuseClient } from "@langfuse/client"
 import { observe, updateActiveTrace } from "@langfuse/tracing"
 import * as api from "@opentelemetry/api"
+import type { TelemetrySettings } from "ai"
 import { randomUUID } from "crypto"
 
 // Singleton LangfuseClient instance for direct API calls
@@ -115,20 +116,23 @@ export function endTrace() {
     }
 }
 
-// Get telemetry config for streamText
+// Telemetry settings for streamText and generateText
 export function getTelemetryConfig(params: {
     sessionId?: string
     userId?: string
-}) {
+    /** Names the trace of a model call outside the chat's observe() span */
+    functionId?: string
+}): TelemetrySettings | undefined {
     if (!isLangfuseEnabled()) return undefined
 
     return {
         isEnabled: true,
+        functionId: params.functionId,
         recordInputs: true,
         recordOutputs: true,
         metadata: {
-            sessionId: params.sessionId,
+            ...(params.sessionId && { sessionId: params.sessionId }),
-            userId: params.userId,
+            ...(params.userId && { userId: params.userId }),
         },
     }
 }

lib/test-model-connection.ts

@@ -199,13 +199,16 @@ export async function testModelConnection(
             tools: { ping: PING_TOOL },
             maxOutputTokens: 1024,
             maxRetries: 0,
-            abortSignal: AbortSignal.timeout(TEST_TIMEOUT_MS),
+            // Ends with the request: the user closed the dialog or the page
+            abortSignal: req.signal,
+            timeout: TEST_TIMEOUT_MS,
         })
         let calledTool = false
         let finishReason: string | undefined
         for await (const part of result.fullStream) {
             if (part.type === "error") throw part.error
-            // The timeout ends the stream with an abort part, not an error
+            // The timeout (or the request's end) ends the stream with an
+            // abort part, not an error
             if (part.type === "abort") {
                 const timeout = new Error(
                     `The model did not answer within ${TEST_TIMEOUT_MS / 1000} s.`,
@@ -213,7 +216,9 @@ export async function testModelConnection(
                 timeout.name = "TimeoutError"
                 throw timeout
             }
-            if (part.type === "tool-call") {
+            // A call of a tool that does not exist, or with unreadable
+            // input, arrives marked invalid: no proof of tool support
+            if (part.type === "tool-call" && !part.invalid) {
                 calledTool = true
                 break
             }

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

@@ -28,60 +28,89 @@ const sse = (chunks: object[], end = true) =>
     chunks.map((c) => `data: ${JSON.stringify(c)}\n\n`).join("") +
     (end ? "data: [DONE]\n\n" : "")
 
+// Step 1: the model asks for a shape library (a tool run on the server) and
+// reports its usage, so the chat makes a second request
+const firstStep = () =>
+    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" } },
+    )
+
+const send = (signal?: AbortSignal) =>
+    chat(
+        new Request("http://localhost/api/chat", {
+            method: "POST",
+            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 drain = async (res: Response) => {
+    const reader = res.body?.getReader()
+    while (
+        reader &&
+        !(await reader.read().catch(() => ({ done: true }))).done
+    ) {
+        // drain
+    }
+}
+
 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
+        // Step 2 never ends, and the user stops
-        // 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) {
+                if (call === 1) return firstStep()
-                    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) {
@@ -98,42 +127,35 @@ describe("a request stopped after a finished step", () => {
             }),
         )
         const stop = new AbortController()
-        const res = await chat(
+        const res = await send(stop.signal)
-            new Request("http://localhost/api/chat", {
+        const draining = drain(res)
-                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 (
+        await draining
-            reader &&
-            !(await reader.read().catch(() => ({ done: true }))).done
-        ) {
-            // drain
-        }
         await vi.waitFor(() => expect(quota.recorded).toEqual([1230]))
     })
 })
+
+describe("a request that fails after a finished step", () => {
+    it("counts that step's tokens", async () => {
+        // Step 2 is refused by the provider (a status the SDK does not
+        // retry): the stream ends with an error
+        let call = 0
+        vi.stubGlobal(
+            "fetch",
+            vi.fn(async () => {
+                call++
+                if (call === 1) return firstStep()
+                return new Response('{"error":{"message":"bad request"}}', {
+                    status: 400,
+                    headers: { "content-type": "application/json" },
+                })
+            }),
+        )
+        await drain(await send())
+        expect(call).toBe(2)
+        expect(quota.recorded).toEqual([1230])
+    })
+})

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

@@ -205,6 +205,21 @@ describe("request checks", () => {
         expect(await res.text()).toMatch(/Invalid AWS region/)
         expect(fetch).not.toHaveBeenCalled()
     })
+
+    it("refuses messages that are not UI messages", async () => {
+        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", text: "no parts" }],
+                    xml: "",
+                }),
+            }),
+        )
+        expect(res.status).toBe(400)
+        expect(fetch).not.toHaveBeenCalled()
+    })
 })
 
 describe("server model allowlist", () => {

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

@@ -15,19 +15,32 @@ const RESULT = {
     suggestions: ["Move box B to the right"],
 }
 
-// A vision model that answers with the result as JSON text
+// A vision model that answers with the result as JSON text; told to wait,
+// it answers only once its call is aborted
+const model = vi.hoisted(() => ({ waitForAbort: false }))
 vi.mock("@/lib/ai-providers", () => ({
     getValidationModel: () =>
         new MockLanguageModelV3({
-            doGenerate: (async () => ({
+            doGenerate: (async (options: { abortSignal?: AbortSignal }) => {
-                content: [{ type: "text", text: JSON.stringify(RESULT) }],
+                if (model.waitForAbort) {
-                finishReason: { unified: "stop", raw: "stop" },
+                    const signal = options.abortSignal
-                usage: {
+                    await new Promise((_, reject) => {
-                    inputTokens: { total: 1 },
+                        if (signal?.aborted) reject(signal.reason)
-                    outputTokens: { total: 1 },
+                        signal?.addEventListener("abort", () =>
-                },
+                            reject(signal.reason),
-                warnings: [],
+                        )
-            })) as any,
+                    })
+                }
+                return {
+                    content: [{ type: "text", text: JSON.stringify(RESULT) }],
+                    finishReason: { unified: "stop", raw: "stop" },
+                    usage: {
+                        inputTokens: { total: 1 },
+                        outputTokens: { total: 1 },
+                    },
+                    warnings: [],
+                }
+            }) as any,
         }),
 }))
 
@@ -61,9 +74,10 @@ vi.mock("@/lib/dynamo-quota-manager", () => ({
     },
 }))
 
-const post = () =>
+const post = (signal?: AbortSignal) =>
     new Request("http://localhost/api/validate-diagram", {
         method: "POST",
+        signal,
         headers: {
             "Content-Type": "application/json",
             "x-forwarded-for": "203.0.113.7",
@@ -72,6 +86,7 @@ const post = () =>
     })
 
 afterEach(() => {
+    model.waitForAbort = false
     quota.enabled = false
     quota.allowed = true
     quota.checks = []
@@ -106,6 +121,20 @@ describe("POST /api/validate-diagram", () => {
         expect(await res.json()).toEqual(RESULT)
     })
 
+    it("ends its model call when the request is aborted", async () => {
+        // The user pressed Stop: the browser aborts the request, and the
+        // model call must not run on (and be paid for)
+        model.waitForAbort = true
+        const stop = new AbortController()
+        const pending = validateDiagram(post(stop.signal))
+        stop.abort()
+        expect(await (await pending).json()).toEqual({
+            valid: true,
+            issues: [],
+            suggestions: [],
+        })
+    })
+
     it("answers valid when the check is turned off", async () => {
         vi.stubEnv("ENABLE_VLM_VALIDATION", "false")
         const res = await validateDiagram(post())

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

@@ -141,6 +141,24 @@ describe("POST /api/validate-model", () => {
         expect(data.valid).toBe(true)
         expect(data.warning).toMatch(/without calling a tool/)
     })
+
+    it("warns when the model calls a tool that does not exist", async () => {
+        // The SDK passes the call on, marked invalid
+        streamReply({
+            role: "assistant",
+            tool_calls: [
+                {
+                    index: 0,
+                    id: "call_1",
+                    type: "function",
+                    function: { name: "pong", arguments: "{}" },
+                },
+            ],
+        })
+        const data = await testGlm()
+        expect(data.valid).toBe(true)
+        expect(data.warning).toMatch(/without calling a tool/)
+    })
 })
 
 describe("chat requests to a client base URL", () => {