From 92e785bcb56e378ab700ab59589aa3a9c1c2fca3 Mon Sep 17 00:00:00 2001 From: "dayuan.jiang" Date: Mon, 12 Oct 2026 11:40:33 +0900 Subject: [PATCH] fix(chat): use the AI SDK's own mechanisms for quota, retries, stops and notices Server - Record quota tokens in onStepFinish and await the write, so a step that fails or is stopped later still counts; drop the onFinish/onAbort copies - validate-diagram and the connection test end their model call with the request (req.signal) and use the SDK's timeout option; validate-diagram reports to Langfuse (functionId "validate-diagram") - The connection test no longer counts an invalid tool call as support - Validate incoming messages with safeValidateUIMessages (400 otherwise), drop client system messages explicitly, prune empty model messages with pruneMessages - Chat tools move to lib/chat-tools.ts with an output schema for the client-side tools Client - Every automatic retry sends the canvas as it is then (prepareSendMessagesRequest): a drawing is committed before its screenshot check can fail, so the model must go on from it - A notice when the output budget ran out before any tool call (finishReason "length") - Error notices are assistant messages with a data-notice part instead of role "system" messages (lib/chat-message.ts: ChatUIMessage with metadata, data part and InferUITools types) - regenerate, edit and retry use sendMessage({ messageId }); isResending and flushSync go - isToolUIPart/getToolName instead of slicing "tool-"; ToolPartLike derives from the SDK's part types; addToolOutput typed with ChatAddToolOutputFunction Tests: step failure after a finished step, abort of the screenshot check, invalid tool call in the connection test, malformed messages. --- app/api/chat/route.ts | 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} >