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

## Problem

A review of how the app uses the Vercel AI SDK (two Claude reviewers and Codex, each finding checked against the SDK's type definitions and source) found five behaviours that are wrong and a set of places where the app re-implements what the SDK already provides:

1. Quota tokens were recorded in `onFinish`, which the SDK does not call when a later step fails; only a user stop was handled (`onAbort`), and the write was not awaited.
2. An automatic retry after a tool error re-sent the request body as it was when the message was sent, so the model was told the "authoritative" canvas XML was the diagram from before the turn, although `display_diagram` commits a drawing before its screenshot check can fail.
3. The screenshot check and the connection test ran their model call on `AbortSignal.timeout` only, so pressing Stop (or closing the dialog) did not stop a paid model call.
4. The client ignored `finishReason`; a reasoning model that spent the whole output budget before calling a tool left the user with no diagram and no message.
5. The connection test counted an `invalid` tool call (a tool that does not exist, unreadable input) as proof of tool support.

## Change

- `onStepFinish` records each step's `usage.totalTokens` and awaits the write; the copies in `onFinish` and `onAbort` are gone.
- `DefaultChatTransport.prepareSendMessagesRequest` sends the current canvas page XML on every request after the first of a turn; the first keeps the exported (or snapshot) XML.
- `validate-diagram` and the connection test pass `req.signal` and use the SDK's `timeout` option; `validate-diagram` reports to Langfuse under `functionId: "validate-diagram"`.
- The connection test only counts a valid tool call.
- The client's `onFinish` shows a notice (with Retry) when `finishReason === "length"` and no tool was called.
- The server validates incoming messages with `safeValidateUIMessages` (400 otherwise), drops client system messages explicitly and prunes empty model messages with `pruneMessages`.
- The chat's tools live in `lib/chat-tools.ts` with an `outputSchema` for the client-side ones; `lib/chat-message.ts` defines `ChatUIMessage` (metadata, a `data-notice` part, `InferUITools`) so `useChat`, `addToolOutput` and the tool handlers are typed.
- Error notices are assistant messages with a `data-notice` part instead of role `system` UI messages.
- Regenerate, edit and retry use `sendMessage({ messageId })`; `isResending` and `flushSync` are gone.
- `isToolUIPart` / `getToolName` replace slicing `"tool-"` off part types; `ToolPartLike` derives from the SDK's part types.

Net: 578 insertions, 439 deletions over 16 files.

## Testing

- New unit tests: a step that fails after a finished step still counts that step's tokens; the screenshot check's model call ends when the request is aborted; an invalid tool call in the connection test gives the warning; malformed messages get a 400. The first and third confirm SDK behaviour the review relied on.
- `tsc`, `biome ci` (0 errors), root `vitest run` 723 passed on the branch rebased onto #986, `next build`: clean
- Playwright, 10 chat-related specs on the dev server: 199 passed, 0 failed, both before and after the rebase onto #986

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")
-    const userInputText =
-        lastUserMessage?.parts?.find((p: any) => p.type === "text")?.text || ""
+        .find((m) => m.role === "user")
+    const userInputText = lastUserMessage?.parts.find(isTextUIPart)?.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 =
-        lastUserMessage?.parts?.filter((part: any) => part.type === "file") ||
-        []
+    const fileParts = lastUserMessage?.parts.filter(isFileUIPart) ?? []
 
     // 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
-    // call that never got its result (the user stopped while it ran) is
-    // left out: the SDK would refuse this and every later request of the
-    // chat (MissingToolResultsError)
-    const modelMessages = await convertToModelMessages(messages, {
-        ignoreIncompleteToolCalls: true,
+    // Convert UIMessages to ModelMessages. A tool call that never got its
+    // result (the user stopped while it ran) is left out: the SDK would
+    // refuse this and every later request of the chat
+    // (MissingToolResultsError). Messages left without content (a notice
+    // the chat shows is a data part, which the model never sees) are
+    // 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.
-        // A client's own system messages have string content and were
-        // dropped by the empty-content filter above.
+        // The system messages carry cache points, so they go in messages
+        // (the client's own system messages were dropped 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)
-        ...(getTelemetryConfig({ sessionId: validSessionId, userId }) && {
-            experimental_telemetry: getTelemetryConfig({
-                sessionId: validSessionId,
-                userId,
-            }),
+        // Langfuse telemetry (undefined when not configured)
+        experimental_telemetry: getTelemetryConfig({
+            sessionId: validSessionId,
+            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 }) => {
-            stopped = true
-            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
-                },
-            },
-        },
+        onAbort: () => endTrace(),
+        tools: chatTools,
         ...(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(
-                userId,
-                (totalUsage.inputTokens || 0) + (totalUsage.outputTokens || 0),
-            )
+            await recordTokenUsage(userId, totalUsage.totalTokens ?? 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)
-type MessageSelection = Pick<SelectedCell, "id" | "label">[]
-
-function selectionOf(message: UIMessage): MessageSelection {
-    const meta = message.metadata as
-        | { selectedCells?: MessageSelection }
-        | undefined
-    return Array.isArray(meta?.selectedCells) ? meta.selectedCells : []
+function selectionOf(message: ChatUIMessage): MessageSelection {
+    const selected = message.metadata?.selectedCells
+    return Array.isArray(selected) ? selected : []
 }
 
 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 */
-    isResending: boolean
-    setMessages: (messages: UIMessage[]) => void
+    setMessages: (messages: ChatUIMessage[]) => 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(
-                (part) => part.state === "input-streaming" && part.toolCallId,
-            )
+            const cut = message.parts
+                .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) {
-                    addToolOutputSafely({
-                        tool: part.type.replace("tool-", ""),
-                        toolCallId: part.toolCallId,
-                        state: "output-error",
-                        errorText: "Stopped by user",
-                    })
-                }
+                for (const part of cut) stopToolCall(part)
             }
-            if (isAbort || isError || cut.length === 0) return
-            // Ended with no error: undo the preview, and say so, as an
-            // error does (with Retry)
+            if (isAbort || isError) return
+            // The output budget ran out before the model called a tool (its
+            // 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,
-                {
-                    id: `error-${Date.now()}`,
-                    role: "system" as const,
-                    parts: [
-                        {
-                            type: "text" as const,
-                            text: dict.errors.answerEndedEarly,
-                        },
-                    ],
-                },
+                noticeMessage({
+                    text: outOfTokens
+                        ? dict.errors.llm.output_truncated
+                        : 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,
-                {
-                    id: `error-${Date.now()}`,
-                    role: "system" as const,
-                    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 },
-                    }),
-                },
+                noticeMessage({
+                    text,
+                    ...(openModelConfig && { openModelConfig, 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(
-            (part: any) =>
-                part.type?.startsWith("tool-") &&
-                part.state === "input-streaming" &&
-                !processedToolCallsRef.current.has(part.toolCallId),
-        )
-        for (const part of (toolParts ?? []) as any[]) {
-            if (part.toolCallId) {
-                addToolOutputSafely({
-                    tool: part.type.replace("tool-", ""),
-                    toolCallId: part.toolCallId,
-                    state: "output-error",
-                    errorText: "Stopped by user",
-                })
-            }
-        }
+        const streaming = (lastMessage?.parts ?? [])
+            .filter(isToolUIPart)
+            .filter(
+                (part) =>
+                    part.state === "input-streaming" &&
+                    !processedToolCallsRef.current.has(part.toolCallId),
+            )
+        for (const part of streaming) stopToolCall(part)
         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: UIMessage[]) => void,
+        setMessages: setMessages as (messages: ChatUIMessage[]) => 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
-        .map((m) => m.role)
-        .lastIndexOf("assistant")
+    // A notice (a failed request) is an assistant message without an answer
+    const lastAssistantIndex = messages.findLastIndex(
+        (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" ? (
-                            <AssistantMessage
-                                message={message}
-                                index={index}
-                                isLastMessage={index === messages.length - 1}
-                                isLastAssistant={index === lastAssistantIndex}
-                            />
-                        ) : (
+                        ) : notice ? (
                             <ErrorNotice
-                                text={getMessageTextContent(message)}
-                                openModelConfig={
-                                    !!(
-                                        message.metadata as
-                                            | { openModelConfig?: boolean }
-                                            | undefined
-                                    )?.openModelConfig
-                                }
-                                onOpenModelConfig={() => {
-                                    const { providerId } =
-                                        (message.metadata as
-                                            | { providerId?: string }
-                                            | undefined) ?? {}
+                                text={notice.text}
+                                openModelConfig={!!notice.openModelConfig}
+                                onOpenModelConfig={() =>
                                     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
-    toolCallId: string
-    state?: string
-    input?: {
-        xml?: string
-        operations?: DiagramOperation[]
-    } & Record<string, unknown>
+/**
+ * A tool part as the UI reads it, in any state: input, output and errorText
+ * are each there or not, instead of one shape per state.
+ */
+export type ToolPartLike = (ToolUIPart<UITools> | DynamicToolUIPart) & {
+    input?:
+        | ({
+              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 = {
-    tool: string
-    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
+// useChat's addToolOutput, typed with the chat's tools
+type AddToolOutputFn = ChatAddToolOutputFunction<ChatUIMessage>
 
 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,
-            userId: params.userId,
+            ...(params.sessionId && { sessionId: params.sessionId }),
+            ...(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
-        // its usage; step 2 never ends, and the user stops
+        // 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" } },
-                    )
-                }
+                if (call === 1) return firstStep()
                 // 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(
-            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()
+        const res = await send(stop.signal)
+        const draining = drain(res)
         // 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 draining
         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 () => ({
-                content: [{ type: "text", text: JSON.stringify(RESULT) }],
-                finishReason: { unified: "stop", raw: "stop" },
-                usage: {
-                    inputTokens: { total: 1 },
-                    outputTokens: { total: 1 },
-                },
-                warnings: [],
-            })) as any,
+            doGenerate: (async (options: { abortSignal?: AbortSignal }) => {
+                if (model.waitForAbort) {
+                    const signal = options.abortSignal
+                    await new Promise((_, reject) => {
+                        if (signal?.aborted) reject(signal.reason)
+                        signal?.addEventListener("abort", () =>
+                            reject(signal.reason),
+                        )
+                    })
+                }
+                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", () => {