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.
2026-10-12 12:39:52 +08:00 · 2026-10-12 12:36:45 +09:00
26 changed files with 2285 additions and 1908 deletions
@@ -12,13 +12,13 @@ jobs:
    runs-on: ubuntu-latest
    steps:
      - name: Checkout
-        uses: actions/checkout@v7
+        uses: actions/checkout@v6
        with:
          ref: ${{ github.event.pull_request.head.sha }}
          token: ${{ secrets.GITHUB_TOKEN }}

      - name: Setup Node.js
-        uses: actions/setup-node@v7
+        uses: actions/setup-node@v6
        with:
          node-version: '24'

@@ -23,14 +23,14 @@ jobs:
      
    steps:
      - name: Checkout repository
-        uses: actions/checkout@v7
+        uses: actions/checkout@v6

      - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v4
+        uses: docker/setup-buildx-action@v3

      - name: Log in to GitHub Container Registry
        if: github.event_name != 'pull_request'
-        uses: docker/login-action@v4
+        uses: docker/login-action@v3
        with:
          registry: ${{ env.REGISTRY }}
          username: ${{ github.actor }}
@@ -38,7 +38,7 @@ jobs:

      - name: Extract metadata (tags, labels)
        id: meta
-        uses: docker/metadata-action@v6
+        uses: docker/metadata-action@v5
        with:
          images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
          tags: |
@@ -51,7 +51,7 @@ jobs:
            type=raw,value=latest,enable={{is_default_branch}}

      - name: Build and push Docker image
-        uses: docker/build-push-action@v7
+        uses: docker/build-push-action@v6
        with:
          context: .
          push: ${{ github.event_name != 'pull_request' }}
@@ -68,7 +68,7 @@ jobs:
      # Push to AWS ECR for App Runner auto-deploy
      - name: Configure AWS credentials
        if: github.event_name != 'pull_request' && github.ref == 'refs/heads/main'
-        uses: aws-actions/configure-aws-credentials@v6
+        uses: aws-actions/configure-aws-credentials@v5
        with:
          aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }}
          aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
@@ -26,10 +26,10 @@ jobs:
    runs-on: ${{ matrix.os }}
    steps:
      - name: Checkout code
-        uses: actions/checkout@v7
+        uses: actions/checkout@v6

      - name: Setup Node.js
-        uses: actions/setup-node@v7
+        uses: actions/setup-node@v6
        with:
          node-version: 24
          cache: "npm"
@@ -50,10 +50,10 @@ jobs:
    runs-on: windows-latest
    steps:
      - name: Checkout code
-        uses: actions/checkout@v7
+        uses: actions/checkout@v6

      - name: Setup Node.js
-        uses: actions/setup-node@v7
+        uses: actions/setup-node@v6
        with:
          node-version: 24
          cache: "npm"
@@ -69,7 +69,7 @@ jobs:
          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}

      - name: Upload unsigned artifacts for signing
-        uses: actions/upload-artifact@v7
+        uses: actions/upload-artifact@v6
        id: upload-unsigned
        with:
          name: windows-unsigned
@@ -77,7 +77,7 @@ jobs:
          retention-days: 1

      - name: Sign with SignPath
-        uses: signpath/github-action-submit-signing-request@v3
+        uses: signpath/github-action-submit-signing-request@v2
        with:
          api-token: ${{ secrets.SIGNPATH_API_TOKEN }}
          organization-id: '880a211d-2cd3-4e7b-8d04-3d1f8eb39df5'
@@ -92,7 +92,7 @@ jobs:
          output-artifact-directory: release-signed

      - name: Upload signed artifacts to release
-        uses: softprops/action-gh-release@v3
+        uses: softprops/action-gh-release@v2
        with:
          files: release-signed/*.exe
        env:
@@ -39,10 +39,10 @@ jobs:
        working-directory: packages/mcp-server
    steps:
      - name: Checkout
-        uses: actions/checkout@v7
+        uses: actions/checkout@v6

      - name: Setup Node.js
-        uses: actions/setup-node@v7
+        uses: actions/setup-node@v6
        with:
          node-version: 24
          cache: "npm"
@@ -18,10 +18,10 @@ jobs:
    name: Lint & Unit Tests
    runs-on: ubuntu-latest
    steps:
-      - uses: actions/checkout@v7
+      - uses: actions/checkout@v6

      - name: Setup Node.js
-        uses: actions/setup-node@v7
+        uses: actions/setup-node@v6
        with:
          node-version: "24"
          cache: "npm"
@@ -46,15 +46,15 @@ jobs:
    name: MCP Server & Shell
    runs-on: ubuntu-latest
    steps:
-      - uses: actions/checkout@v7
+      - uses: actions/checkout@v6
        with:
          # HEAD^1 is the base branch for the version bump reminder below
          fetch-depth: 2

      - name: Setup Node.js
-        uses: actions/setup-node@v7
+        uses: actions/setup-node@v6
        with:
-          node-version: "24"
+          node-version: "20"
          cache: "npm"
          cache-dependency-path: |
            package-lock.json
@@ -69,7 +69,7 @@ jobs:
        run: npm --prefix packages/mcp-server ci

      - name: Cache Playwright browsers
-        uses: actions/cache@v6
+        uses: actions/cache@v5
        id: playwright-cache
        with:
          path: ~/.cache/ms-playwright
@@ -129,7 +129,7 @@ jobs:
          fi

      - name: Upload test results
-        uses: actions/upload-artifact@v7
+        uses: actions/upload-artifact@v6
        if: always()
        with:
          name: playwright-report-mcp-shell
@@ -146,10 +146,10 @@ jobs:
      matrix:
        shard: [1, 2, 3, 4, 5, 6]
    steps:
-      - uses: actions/checkout@v7
+      - uses: actions/checkout@v6

      - name: Setup Node.js
-        uses: actions/setup-node@v7
+        uses: actions/setup-node@v6
        with:
          node-version: "24"
          cache: "npm"
@@ -158,7 +158,7 @@ jobs:
        run: npm ci

      - name: Cache Playwright browsers
-        uses: actions/cache@v6
+        uses: actions/cache@v5
        id: playwright-cache
        with:
          path: ~/.cache/ms-playwright
@@ -181,7 +181,7 @@ jobs:
          CI: true

      - name: Upload test results
-        uses: actions/upload-artifact@v7
+        uses: actions/upload-artifact@v6
        if: always()
        with:
          name: playwright-report-${{ matrix.shard }}
@@ -19,7 +19,7 @@ jobs:
    if: github.event_name != 'issue_comment'
    runs-on: ubuntu-latest
    steps:
-      - uses: actions/stale@v11
+      - uses: actions/stale@v9
        with:
          only-issue-labels: "wait for confrim"
          days-before-issue-stale: 15
@@ -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),
        }),
@@ -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(
@@ -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,
@@ -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>
                )
@@ -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(
@@ -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
 }
@@ -150,7 +150,6 @@ function WorkspaceLayout() {
    const showHero =
        engine.isRestored &&
        engine.messages.length === 0 &&
-        !engine.isResending &&
        !isRealDiagram(chartXML) &&
        !heroDismissed

@@ -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

@@ -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
+}
@@ -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>
@@ -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 }),
        },
    }
 }
@@ -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
            }
@@ -44,11 +44,11 @@
        "@aihubmix/ai-sdk-provider": "^2.2.1",
        "@aws-sdk/client-dynamodb": "^3.957.0",
        "@aws-sdk/credential-providers": "^3.943.0",
-        "@extractus/article-extractor": "^9.0.0",
+        "@extractus/article-extractor": "^8.0.18",
        "@formatjs/intl-localematcher": "^0.9.0",
-        "@langfuse/client": "^5.0.0",
-        "@langfuse/otel": "^5.0.0",
-        "@langfuse/tracing": "^5.0.0",
+        "@langfuse/client": "^4.4.9",
+        "@langfuse/otel": "^4.4.4",
+        "@langfuse/tracing": "^4.4.9",
        "@next/third-parties": "^16.0.6",
        "@openrouter/ai-sdk-provider": "^2.10.0",
        "@opentelemetry/api": "^1.9.0",
@@ -70,12 +70,12 @@
        "cmdk": "^1.1.1",
        "idb": "^8.0.3",
        "jsonrepair": "^3.13.1",
-        "lucide-react": "^1.0.0",
-        "nanoid": "^6.0.0",
+        "lucide-react": "^0.577.0",
+        "nanoid": "^5.0.0",
        "negotiator": "^1.0.0",
        "next": "^16.0.7",
        "ollama-ai-provider-v2": "^3.6.0",
-        "pako": "^3.0.0",
+        "pako": "^2.1.0",
        "prism-react-renderer": "^2.4.1",
        "react": "^19.1.2",
        "react-dom": "^19.1.2",
@@ -111,24 +111,24 @@
        "@testing-library/react": "^16.3.1",
        "@types/negotiator": "^0.6.4",
        "@types/node": "^24.0.0",
-        "@types/pako": "^3.0.0",
+        "@types/pako": "^2.0.3",
        "@types/react": "^19",
        "@types/react-dom": "^19",
        "@types/turndown": "^5.0.6",
-        "@vitejs/plugin-react": "^6.0.0",
+        "@vitejs/plugin-react": "^5.1.2",
        "cross-env": "^10.1.0",
-        "dotenv": "^18.0.0",
+        "dotenv": "^16.6.1",
        "electron": "^41.10.6",
        "electron-builder": "^26.0.12",
        "esbuild": "^0.28.0",
        "husky": "^9.1.7",
-        "jsdom": "^30.0.0",
-        "lint-staged": "^17.0.0",
+        "jsdom": "^27.4.0",
+        "lint-staged": "^16.2.7",
        "postcss": "^8.5.6",
        "tailwindcss": "^4",
-        "typescript": "^7.0.0",
+        "typescript": "^5",
        "vite-tsconfig-paths": "^6.0.3",
-        "vitest": "^5.0.0",
+        "vitest": "^4.0.16",
        "wrangler": "^4.60.0"
    },
    "overrides": {
@@ -44,16 +44,16 @@
        "@modelcontextprotocol/sdk": "^1.31.0",
        "linkedom": "^0.18.0",
        "open": "^11.0.0",
-        "pako": "^3.0.0",
+        "pako": "^2.2.0",
        "saxes": "^6.0.0",
        "zod": "^4.0.0"
    },
    "devDependencies": {
        "@types/node": "^24.0.0",
-        "@types/pako": "^3.0.0",
+        "@types/pako": "^2.0.4",
        "tsx": "^4.19.0",
-        "typescript": "^7.0.0",
-        "vitest": "^5.0.0"
+        "typescript": "^5",
+        "vitest": "^4.1.8"
    },
    "engines": {
        "node": ">=18"
@@ -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])
+    })
+})
@@ -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", () => {
@@ -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())
@@ -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", () => {