From 99ced88e5b45282937d5e871603591936b22d629 Mon Sep 17 00:00:00 2001 From: Dayuan Jiang <34411969+DayuanJiang@users.noreply.github.com> Date: Mon, 12 Oct 2026 12:58:52 +0900 Subject: [PATCH] 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 | 238 +++++----------------- app/api/validate-diagram/route.ts | 16 +- components/chat/chat-engine.tsx | 194 +++++++++--------- components/chat/message-list.tsx | 51 +++-- components/chat/tool-activity.tsx | 5 +- components/chat/types.ts | 20 +- components/workspace/workspace.tsx | 1 - hooks/use-diagram-tool-handlers.ts | 23 +-- lib/chat-message.ts | 44 ++++ lib/chat-tools.ts | 138 +++++++++++++ lib/langfuse.ts | 12 +- lib/test-model-connection.ts | 11 +- tests/unit/chat-route-abort.test.ts | 180 +++++++++------- tests/unit/chat-route-quota.test.ts | 15 ++ tests/unit/validate-diagram-route.test.ts | 51 ++++- tests/unit/validate-model-route.test.ts | 18 ++ 16 files changed, 578 insertions(+), 439 deletions(-) create mode 100644 lib/chat-message.ts create mode 100644 lib/chat-tools.ts diff --git a/app/api/chat/route.ts b/app/api/chat/route.ts index 4ea92508..e63e6f63 100644 --- a/app/api/chat/route.ts +++ b/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 { 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({ + 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 { // 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 { : 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 { ${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 (, , ) -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="" only for shapes inside a container -6. Escape special chars in values: < > & " - -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: . 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 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 <br> for a line break, never \\n; a literal < or > is written &lt; or &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": ""}]} - -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 '' 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), }), diff --git a/app/api/validate-diagram/route.ts b/app/api/validate-diagram/route.ts index c91e6377..c985a021 100644 --- a/app/api/validate-diagram/route.ts +++ b/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 { }, ], 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( diff --git a/components/chat/chat-engine.tsx b/components/chat/chat-engine.tsx index b6e770d3..a83a2eec 100644 --- a/components/chat/chat-engine.tsx +++ b/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[] - -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(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({ + 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({ 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[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(null) + const messagesBeforeShownRef = useRef(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, diff --git a/components/chat/message-list.tsx b/components/chat/message-list.tsx index f62e5e53..b8e85647 100644 --- a/components/chat/message-list.tsx +++ b/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 (
- ) : message.role === "assistant" ? ( - - ) : ( + ) : notice ? ( { - 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 } /> + ) : ( + )}
) diff --git a/components/chat/tool-activity.tsx b/components/chat/tool-activity.tsx index ff6e5f97..68940c2e 100644 --- a/components/chat/tool-activity.tsx +++ b/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} >