i/next-ai-draw-io
1
0
Fork 0
mirror of https://github.com/DayuanJiang/next-ai-draw-io.git synced 2026-10-12 12:39:52 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
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.
This commit is contained in:
dayuan.jiang
2026-10-12 12:36:45 +09:00
parent 60105ddb51
commit 92e785bcb5
16 changed files with 578 additions and 439 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/api/chat/route.ts
+51 -187
View File
@@ -2,12 +2,14 @@ import {
APICallError,
convertToModelMessages,
InvalidToolInputError,
isFileUIPart,
isTextUIPart,
pruneMessages,
safeValidateUIMessages,
stepCountIs,
streamText,
} from "ai"
import { jsonrepair } from "jsonrepair"
import path from "path"
import { z } from "zod"
import { checkAccessCode, rejectCrossSite } from "@/lib/access-code"
import {
CACHE_POINT,
@@ -25,6 +27,8 @@ import {
formatSelectionContext,
validateFileParts,
} from "@/lib/chat-helpers"
import type { ChatUIMessage } from "@/lib/chat-message"
import { chatTools } from "@/lib/chat-tools"
import { withDeprecatedParamsFallback } from "@/lib/deprecated-params"
import {
checkAndIncrementRequest,
@@ -52,14 +56,6 @@ import { getSystemPrompt } from "@/lib/system-prompts"
import { normalizeBaseUrl } from "@/lib/types/model-config"
import { getUserIdFromRequest } from "@/lib/user-id"
import { foldCells } from "@/packages/mcp-server/src/compact-cells.ts"
import {
getShapeLibrary,
SHAPE_LIBRARY_LIST,
} from "@/packages/mcp-server/src/shape-library.ts"
import {
STYLE_CLASS_EXAMPLE,
SWIMLANE_EXAMPLE,
} from "@/packages/mcp-server/src/xml-examples.ts"
// No explicit cap: a reasoning model can spend minutes planning before it emits
// the tool call, so take whatever the host allows. Vercel's own default is 300s,
@@ -99,7 +95,19 @@ async function handleChatRequest(req: Request): Promise<Response> {
if (accessDenied) return accessDenied
const body = await req.json()
const { messages, xml, previousXml, sessionId } = body
// The messages in the SDK's UI shape. The route writes its own system
// messages, so none may come from the client.
const validated = await safeValidateUIMessages<ChatUIMessage>({
messages: body.messages,
})
if (!validated.success) {
return Response.json(
{ error: validated.error.message },
{ status: 400 },
)
}
const messages = validated.data.filter((m) => m.role !== "system")
const { xml, previousXml, sessionId } = body
const customSystemMessage =
typeof body.customSystemMessage === "string"
? body.customSystemMessage.slice(0, 5000)
@@ -118,9 +126,8 @@ async function handleChatRequest(req: Request): Promise<Response> {
// Find the last USER message, not just the last message (which could be assistant in multi-step tool flows)
const lastUserMessage = [...messages]
.reverse()
.find((m: any) => m.role === "user")
const userInputText =
lastUserMessage?.parts?.find((p: any) => p.type === "text")?.text || ""
.find((m) => m.role === "user")
const userInputText = lastUserMessage?.parts.find(isTextUIPart)?.text || ""
// Update Langfuse trace with input, session, and user
setTraceInput({
@@ -318,9 +325,7 @@ async function handleChatRequest(req: Request): Promise<Response> {
: systemMessage
// Extract file parts (images) from the last user message
const fileParts =
lastUserMessage?.parts?.filter((part: any) => part.type === "file") ||
[]
const fileParts = lastUserMessage?.parts.filter(isFileUIPart) ?? []
// Note: we used to pre-emptively reject images for models we guessed were
// text-only (by name matching). That heuristic misfired on newer models
@@ -333,12 +338,16 @@ async function handleChatRequest(req: Request): Promise<Response> {
${userInputText}
"""`
// Convert UIMessages to ModelMessages and add system message. A tool
// call that never got its result (the user stopped while it ran) is
// left out: the SDK would refuse this and every later request of the
// chat (MissingToolResultsError)
const modelMessages = await convertToModelMessages(messages, {
ignoreIncompleteToolCalls: true,
// Convert UIMessages to ModelMessages. A tool call that never got its
// result (the user stopped while it ran) is left out: the SDK would
// refuse this and every later request of the chat
// (MissingToolResultsError). Messages left without content (a notice
// the chat shows is a data part, which the model never sees) are
// dropped: Bedrock rejects empty content.
const modelMessages = pruneMessages({
messages: await convertToModelMessages(messages, {
ignoreIncompleteToolCalls: true,
}),
})
// DEBUG_LLM_PAYLOAD=true logs the incoming message structure
@@ -349,17 +358,10 @@ ${userInputText}
)
}
// Filter out messages with empty content arrays (Bedrock API rejects these)
// This is a safety measure - ideally convertToModelMessages should handle all cases
let enhancedMessages = modelMessages.filter(
(msg: any) =>
msg.content && Array.isArray(msg.content) && msg.content.length > 0,
)
// Filter out tool-calls with invalid inputs (from failed repair or interrupted streaming)
// and their results. Bedrock API rejects messages where toolUse.input is not a valid
// JSON object, and every provider rejects a tool result whose call is gone.
enhancedMessages = dropInvalidToolCalls(enhancedMessages)
let enhancedMessages = dropInvalidToolCalls(modelMessages)
// DEBUG_LLM_PAYLOAD=true logs what is sent to the model
if (DEBUG_LLM_PAYLOAD) {
@@ -471,13 +473,10 @@ IMPORTANT: The "Current diagram XML" is the SINGLE SOURCE OF TRUTH for what's on
const allMessages = [...systemMessages, ...enhancedMessages]
// Set by onAbort, which records the finished steps' tokens itself
let stopped = false
const result = streamText({
model,
// The system messages carry cache points, so they go in messages.
// A client's own system messages have string content and were
// dropped by the empty-content filter above.
// The system messages carry cache points, so they go in messages
// (the client's own system messages were dropped above)
allowSystemInMessages: true,
abortSignal: req.signal,
// Must be sent: unset means the provider's own default, and Bedrock's is
@@ -524,166 +523,31 @@ IMPORTANT: The "Current diagram XML" is the SINGLE SOURCE OF TRUTH for what's on
},
messages: allMessages,
...(providerOptions && { providerOptions }), // This now includes all reasoning configs
// Langfuse telemetry config (returns undefined if not configured)
...(getTelemetryConfig({ sessionId: validSessionId, userId }) && {
experimental_telemetry: getTelemetryConfig({
sessionId: validSessionId,
userId,
}),
// Langfuse telemetry (undefined when not configured)
experimental_telemetry: getTelemetryConfig({
sessionId: validSessionId,
userId,
}),
onFinish: ({ text, totalUsage }) => {
// The quota counts each step as it ends: a later step that fails or
// is stopped has still used these tokens (onFinish runs only for a
// completed answer). Awaited, so the write is done before the host
// may freeze the function after the answer ends.
onStepFinish: async ({ usage }) => {
if (countsQuota) {
await recordTokenUsage(userId, usage.totalTokens ?? 0)
}
},
onFinish: ({ text }) => {
// AI SDK 6 telemetry auto-reports token usage on its spans
setTraceOutput(text)
// Record token usage for server-side quota tracking (if enabled)
// Use totalUsage (cumulative across all steps) instead of usage (final step only)
// inputTokens already includes cache reads and writes in AI SDK 6
if (countsQuota && totalUsage && !stopped) {
const totalTokens =
(totalUsage.inputTokens || 0) +
(totalUsage.outputTokens || 0)
recordTokenUsage(userId, totalTokens)
}
},
// onFinish is skipped when the stream fails or is aborted, so end the trace here
onError: ({ error }) => {
console.error(error) // what AI SDK does without an onError
endTrace()
},
onAbort: ({ steps }) => {
stopped = true
endTrace()
// Stopped (or disconnected) after some steps finished: their
// tokens were used, or stopping every request after a costly
// first step would get around the token limits
if (countsQuota) {
const tokens = steps.reduce(
(sum, step) =>
sum +
(step.usage.inputTokens || 0) +
(step.usage.outputTokens || 0),
0,
)
if (tokens > 0) recordTokenUsage(userId, tokens)
}
},
tools: {
// Client-side tool that will be executed on the client
display_diagram: {
description: `Display a diagram on draw.io. Pass ONLY the mxStyle definitions and the mxCell elements - wrapper tags and root cells are added automatically.
VALIDATION RULES (XML will be rejected if violated):
1. Generate ONLY mxStyle definitions and mxCell elements - NO wrapper tags (<mxfile>, <mxGraphModel>, <root>)
2. Do NOT include root cells (id="0" or id="1") - they are added automatically
3. All mxCell elements must be siblings - never nested
4. Every mxCell needs a unique id (start from "2")
5. parent defaults to "1"; write parent="<container-id>" only for shapes inside a container
6. Escape special chars in values: &lt; &gt; &amp; &quot;
A shape is one self-closing mxCell with x, y, w and h; an edge is one with source and target (a cell with source or target is always an edge). vertex="1", edge="1", parent="1" and the mxGeometry element are added automatically. Write parent only for a shape inside a container, and an mxGeometry element only for edge waypoints or for a separate label cell placed on an edge: <mxCell id="9" value="yes" style="edgeLabel;" parent="<edge id>" connectable="0"><mxGeometry x="-0.5" relative="1" as="geometry"/></mxCell>. An edge's own text simply goes in its value.
Example (generate ONLY this - no wrapper tags):
${SWIMLANE_EXAMPLE}
Styles: define a style used by several cells ONCE with <mxStyle name="..." value="..."/> before the cells and use the name in style like a CSS class; overrides after the name win. Name only styles that two or more cells share; names must not be draw.io's own style names (text, ellipse, rhombus, swimlane, label, image, blue, green, red, gray, yellow, orange, purple, pink). A definition applies to the call it is in. html=1 and whiteSpace=wrap are added automatically, never write them. Labels are HTML: use &lt;br&gt; for a line break, never \\n; a literal < or > is written &amp;lt; or &amp;gt;.
${STYLE_CLASS_EXAMPLE}
Notes:
- For AWS diagrams, use **AWS 2025 icons**.
- For animated connectors, add "flowAnimation=1" to edge style.
`,
inputSchema: z.object({
xml: z
.string()
.describe("XML string to be displayed on draw.io"),
}),
},
edit_diagram: {
description: `Edit the current diagram by ID-based operations (update/add/delete cells).
Operations:
- update: Replace an existing cell by its id. Provide cell_id and complete new_xml.
- add: Add a new cell. Provide cell_id (new unique id) and new_xml.
- delete: Remove a cell. Cascade is automatic: children AND edges (source/target) are auto-deleted. Only specify ONE cell_id.
For update/add, new_xml is the complete mxCell in the compact form (a shape with x, y, w, h; an edge with source and target), with its complete style (named styles are not available in edit_diagram).
⚠️ JSON ESCAPING: Every " inside new_xml MUST be escaped as \\". Example: id=\\"5\\" value=\\"Label\\"
Example - Add a rectangle:
{"operations": [{"operation": "add", "cell_id": "rect-1", "new_xml": "<mxCell id=\\"rect-1\\" value=\\"Hello\\" style=\\"rounded=1;\\" x=\\"100\\" y=\\"100\\" w=\\"120\\" h=\\"60\\"/>"}]}
Example - Delete container (children & edges auto-deleted):
{"operations": [{"operation": "delete", "cell_id": "2"}]}`,
inputSchema: z.object({
operations: z
.array(
z.object({
operation: z
.enum(["update", "add", "delete"])
.describe(
"Operation to perform: add, update, or delete",
),
cell_id: z
.string()
.describe(
"The id of the mxCell. Must match the id attribute in new_xml.",
),
new_xml: z
.string()
.optional()
.describe(
"Complete mxCell XML element (required for update/add)",
),
}),
)
.describe("Array of operations to apply"),
}),
},
append_diagram: {
description: `Continue generating diagram XML when previous display_diagram output was truncated due to length limits.
WHEN TO USE: Only call this tool after display_diagram was truncated (you'll see an error message about truncation).
CRITICAL INSTRUCTIONS:
1. Do NOT include any wrapper tags - just continue the mxCell elements
2. Continue from EXACTLY where your previous output stopped
3. Complete the remaining mxCell elements
4. If still truncated, call append_diagram again with the next fragment
Example: If previous output ended with '<mxCell id="x" style="rounded=1', continue with ';" x="40" y="40" w="120" h="60"/>' and complete the remaining elements.`,
inputSchema: z.object({
xml: z
.string()
.describe(
"Continuation XML fragment to append (NO wrapper tags)",
),
}),
},
get_shape_library: {
description: `Get draw.io shape/icon library documentation with style syntax and shape names.
Available libraries:
${SHAPE_LIBRARY_LIST}
Call this tool to get shape names and usage syntax for a specific library.`,
inputSchema: z.object({
library: z
.string()
.describe(
"Library name (e.g., 'aws4', 'kubernetes', 'flowchart')",
),
}),
execute: async ({ library }) => {
// Only known library names reach the file system
const result = await getShapeLibrary(
library,
path.join(process.cwd(), "docs/shape-libraries"),
)
return result.ok ? result.text : result.error
},
},
},
onAbort: () => endTrace(),
tools: chatTools,
...(process.env.TEMPERATURE !== undefined && {
temperature: parseFloat(process.env.TEMPERATURE),
}),
app/api/validate-diagram/route.ts
+11 -5
View File
@@ -11,6 +11,7 @@ import {
isQuotaEnabled,
recordTokenUsage,
} from "@/lib/dynamo-quota-manager"
import { getTelemetryConfig } from "@/lib/langfuse"
import { getUserIdFromRequest } from "@/lib/user-id"
import { VALIDATION_SYSTEM_PROMPT } from "@/lib/validation-prompts"
import {
@@ -146,13 +147,18 @@ export async function POST(req: Request): Promise<Response> {
},
],
maxOutputTokens: 1024,
abortSignal: AbortSignal.timeout(timeout),
// Ends with the request: the user pressed Stop, or left
abortSignal: req.signal,
timeout,
// Its usage and cost show in Langfuse next to the chat's
experimental_telemetry: getTelemetryConfig({
functionId: "validate-diagram",
sessionId,
userId,
}),
})
if (countsQuota) {
await recordTokenUsage(
userId,
(totalUsage.inputTokens || 0) + (totalUsage.outputTokens || 0),
)
await recordTokenUsage(userId, totalUsage.totalTokens ?? 0)
}
if (sessionId) {
console.log(
components/chat/chat-engine.tsx
+102 -92
View File
@@ -4,8 +4,8 @@ import { useChat } from "@ai-sdk/react"
import {
type ChatStatus,
DefaultChatTransport,
getToolName,
isToolUIPart,
type UIMessage,
} from "ai"
import { usePathname, useRouter, useSearchParams } from "next/navigation"
import type React from "react"
@@ -19,7 +19,6 @@ import {
useRef,
useState,
} from "react"
import { flushSync } from "react-dom"
import { toast } from "sonner"
import type { ValidationState } from "@/components/chat/ValidationCard"
import { useDiagram } from "@/contexts/diagram-context"
@@ -37,6 +36,12 @@ import {
import { useValidateDiagram } from "@/hooks/use-validate-diagram"
import { getApiEndpoint } from "@/lib/base-path"
import { findCachedResponse } from "@/lib/cached-responses"
import {
type ChatToolName,
type ChatUIMessage,
type MessageSelection,
noticeMessage,
} from "@/lib/chat-message"
import { buildChatHeaders } from "@/lib/chat-request"
import { pageModelXml, placeOnPage } from "@/lib/diagram-pages"
import { EMPTY_DRAWIO_DOCUMENT } from "@/lib/drawio/drawio-config"
@@ -76,7 +81,7 @@ const SAVE_DEBOUNCE_MS = 1000
* Only checks the LAST tool part (most recent tool call), not all tool parts.
* A call the user stopped is not one to retry.
*/
function hasToolErrors(messages: UIMessage[]): boolean {
function hasToolErrors(messages: ChatUIMessage[]): boolean {
const lastMessage = messages[messages.length - 1]
if (lastMessage?.role !== "assistant") return false
const lastToolPart = lastMessage.parts.filter(isToolUIPart).at(-1)
@@ -103,14 +108,9 @@ function viewedPageId(): string | null {
return useCanvasStore.getState().currentPageId
}
// Shapes sent with a user message (also kept in its metadata)
type MessageSelection = Pick<SelectedCell, "id" | "label">[]
function selectionOf(message: UIMessage): MessageSelection {
const meta = message.metadata as
| { selectedCells?: MessageSelection }
| undefined
return Array.isArray(meta?.selectedCells) ? meta.selectedCells : []
function selectionOf(message: ChatUIMessage): MessageSelection {
const selected = message.metadata?.selectedCells
return Array.isArray(selected) ? selected : []
}
function newLangfuseSessionId() {
@@ -169,15 +169,13 @@ function withUniquePageIds(xml: string): string {
class FileReadError extends Error {}
export interface ChatEngine {
messages: UIMessage[]
messages: ChatUIMessage[]
status: ChatStatus
error: Error | undefined
isBusy: boolean
/** The chat on screen is being left (saved, then another one shown) */
isLeaving: boolean
/** A message is being sent again; the chat is not empty meanwhile */
isResending: boolean
setMessages: (messages: UIMessage[]) => void
setMessages: (messages: ChatUIMessage[]) => void
stop: () => void
input: string
@@ -361,9 +359,6 @@ export function ChatEngineProvider({
const [isLeaving, setIsLeaving] = useState(false)
// The user's provider the current turn was sent to, for its errors
const requestProviderIdRef = useRef<string | undefined>(undefined)
// A message is being sent again (retry, regenerate, edit), until the
// SDK has put it back
const [isResending, setIsResending] = useState(false)
// Presses of Stop: a check that began before one still knows of it after
// the next message clears stoppedRef
const stopCountRef = useRef(0)
@@ -429,9 +424,40 @@ export function ChatEngineProvider({
onDisplayChart(original, true, "revert")
}
// Requests made in the turn so far. The first carries the diagram the
// message was sent with (exported then, or the snapshot it is sent
// again from). Each later one, an automatic retry after a tool error,
// carries the canvas as it is then: a drawing is committed before its
// screenshot check can fail, and the model must go on from that
// drawing, not from the one before the turn.
const turnRequestsRef = useRef(0)
const transport = useMemo(
() => new DefaultChatTransport({ api: getApiEndpoint("/api/chat") }),
[],
() =>
new DefaultChatTransport<ChatUIMessage>({
api: getApiEndpoint("/api/chat"),
prepareSendMessagesRequest: ({
id,
messages,
body,
trigger,
messageId,
}) => ({
body: {
...body,
...(turnRequestsRef.current++ > 0 && {
xml: getPageXml(
chartXMLRef.current,
turnPageIdRef.current,
),
}),
id,
messages,
trigger,
messageId,
},
}),
}),
[chartXMLRef],
)
const {
@@ -442,19 +468,19 @@ export function ChatEngineProvider({
error,
setMessages,
stop,
} = useChat({
} = useChat<ChatUIMessage>({
transport,
onToolCall: async ({ toolCall }) => {
await handleToolCall({ toolCall }, addToolOutputSafely)
},
onFinish: ({ message, isAbort, isError }) => {
onFinish: ({ message, isAbort, isError, finishReason }) => {
// Stopped, failed, or the reply just ended (a closed connection):
// tool calls still streaming never reach the tool handler. Mark
// them handled so a later render of the stream does not draw
// their preview again.
const cut = (message.parts as any[]).filter(
(part) => part.state === "input-streaming" && part.toolCallId,
)
const cut = message.parts
.filter(isToolUIPart)
.filter((part) => part.state === "input-streaming")
for (const part of cut) {
processedToolCallsRef.current.add(part.toolCallId)
}
@@ -463,18 +489,16 @@ export function ChatEngineProvider({
// or the page going, aborts a reply). Not once another chat is
// on screen: the call is not among its messages.
if (isAbort && getChatGeneration() === turnChatRef.current) {
for (const part of cut) {
addToolOutputSafely({
tool: part.type.replace("tool-", ""),
toolCallId: part.toolCallId,
state: "output-error",
errorText: "Stopped by user",
})
}
for (const part of cut) stopToolCall(part)
}
if (isAbort || isError || cut.length === 0) return
// Ended with no error: undo the preview, and say so, as an
// error does (with Retry)
if (isAbort || isError) return
// The output budget ran out before the model called a tool (its
// thinking took it all): nothing was drawn
const outOfTokens =
finishReason === "length" && !message.parts.some(isToolUIPart)
if (cut.length === 0 && !outOfTokens) return
// Ended with no error and no diagram: undo the preview, and say
// so, as an error does (with Retry)
const [originalXml] = editDiagramOriginalXmlRef.current.values()
if (originalXml) onDisplayChart(originalXml, true, "revert")
editDiagramOriginalXmlRef.current.clear()
@@ -483,16 +507,11 @@ export function ChatEngineProvider({
if (getChatGeneration() !== turnChatRef.current) return
setMessages((currentMessages) => [
...currentMessages,
{
id: `error-${Date.now()}`,
role: "system" as const,
parts: [
{
type: "text" as const,
text: dict.errors.answerEndedEarly,
},
],
},
noticeMessage({
text: outOfTokens
? dict.errors.llm.output_truncated
: dict.errors.answerEndedEarly,
}),
])
},
onError: (error) => {
@@ -561,19 +580,15 @@ export function ChatEngineProvider({
// the model selected may have changed since, e.g. in another tab
const providerId = requestProviderIdRef.current
// A system message, so it can be cleared with the conversation
// In the chat, so it is cleared with the conversation; with a
// button that opens model settings, on the page of the provider
// that refused
setMessages((currentMessages) => [
...currentMessages,
{
id: `error-${Date.now()}`,
role: "system" as const,
parts: [{ type: "text" as const, text }],
// The message shows a button that opens model settings,
// on the page of the provider that refused
...(openModelConfig && {
metadata: { openModelConfig: true, providerId },
}),
},
noticeMessage({
text,
...(openModelConfig && { openModelConfig, providerId }),
}),
])
if (isAccessCodeError) openSettings("general")
@@ -603,9 +618,19 @@ export function ChatEngineProvider({
)
}
// A call the user stopped before it reached the tool handler. The SDK
// finds the call by its id; the tool name is for the types only.
const stopToolCall = (part: Parameters<typeof getToolName>[0]) =>
addToolOutputSafely({
tool: getToolName(part) as ChatToolName,
toolCallId: part.toolCallId,
state: "output-error",
errorText: "Stopped by user",
})
// After each request: send the next one (retry, continue a cut off
// drawing) or end the turn
const shouldSendAgain = (messages: UIMessage[]) => {
const shouldSendAgain = (messages: ChatUIMessage[]) => {
// The user stopped: a tool result that arrives later (a VLM
// check still running) must not start a new request
if (stoppedRef.current) return false
@@ -657,9 +682,6 @@ export function ChatEngineProvider({
}
const isBusy = status === "streaming" || status === "submitted"
useEffect(() => {
if (messages.length > 0) setIsResending(false)
}, [messages.length])
// isBusy of the latest render, for code that runs after an await
const busyRef = useRef(isBusy)
busyRef.current = isBusy
@@ -771,7 +793,7 @@ export function ChatEngineProvider({
const [shownChats, setShownChats] = useState(0)
// The messages shown when another chat was put on screen: useChat can
// show that chat's own up to 150 ms later than the rest of its state
const messagesBeforeShownRef = useRef<UIMessage[] | null>(null)
const messagesBeforeShownRef = useRef<ChatUIMessage[] | null>(null)
const syncUIWithSession = useCallback(
(data: SessionData | null) => {
@@ -811,7 +833,7 @@ export function ChatEngineProvider({
loadedMessageIdsRef.current = new Set(
(data.messages as { id: string }[]).map((m) => m.id),
)
setMessages(data.messages as unknown as UIMessage[])
setMessages(data.messages as unknown as ChatUIMessage[])
xmlSnapshotsRef.current = new Map(data.xmlSnapshots)
if (data.diagramXml) onDisplayChart(data.diagramXml, true)
else clearDiagram()
@@ -1118,12 +1140,15 @@ export function ChatEngineProvider({
: ""
}
// messageId: the user message sent again (regenerate, edit, retry); the
// SDK puts the new one in its place and drops the answers after it
const sendChatMessage = (
parts: any,
xml: string,
previousXml: string,
turnIndex: number,
selectedCells: MessageSelection,
messageId?: string,
) => {
autoRetryCountRef.current = 0
continuationRetryCountRef.current = 0
@@ -1161,9 +1186,11 @@ export function ChatEngineProvider({
maxOutputTokens,
}),
}
turnRequestsRef.current = 0
// The selection stays with the message, for regenerate and retry
sendMessage(
{
...(messageId !== undefined && { messageId }),
parts,
...(selected.length > 0 && {
metadata: { selectedCells: selected },
@@ -1416,18 +1443,13 @@ export function ChatEngineProvider({
dropLaterSnapshots(index)
useVersionsStore.getState().removeFromTurn(index)
// Remove the user message onwards (sendMessage re-adds it). For a
// moment there may be no message: that is no empty chat.
flushSync(() => {
setIsResending(true)
setMessages(messages.slice(0, index))
})
sendChatMessage(
parts,
getPageXml(savedXml, viewedPageId()),
previousXml,
index,
selectionOf(messages[index]),
messages[index].id,
)
return true
@@ -1493,16 +1515,13 @@ export function ChatEngineProvider({
chartXMLRef.current || chartXml,
)
}
flushSync(() => {
setIsResending(true)
setMessages(messagesRef.current.slice(0, index))
})
sendChatMessage(
message.parts,
chartXml,
getPreviousXml(index),
index,
selectionOf(message),
message.id,
)
} catch (error) {
console.error("Error fetching chart data:", error)
@@ -1521,25 +1540,17 @@ export function ChatEngineProvider({
const lastMessage = messages[messages.length - 1]
// Calls the tool handler already took can still show as streaming:
// the messages update at most every 150 ms (useChat throttle)
const toolParts = lastMessage?.parts?.filter(
(part: any) =>
part.type?.startsWith("tool-") &&
part.state === "input-streaming" &&
!processedToolCallsRef.current.has(part.toolCallId),
)
for (const part of (toolParts ?? []) as any[]) {
if (part.toolCallId) {
addToolOutputSafely({
tool: part.type.replace("tool-", ""),
toolCallId: part.toolCallId,
state: "output-error",
errorText: "Stopped by user",
})
}
}
const streaming = (lastMessage?.parts ?? [])
.filter(isToolUIPart)
.filter(
(part) =>
part.state === "input-streaming" &&
!processedToolCallsRef.current.has(part.toolCallId),
)
for (const part of streaming) stopToolCall(part)
stop()
abandonContinuation()
}, [messages, addToolOutputSafely, stop, cancelValidation])
}, [messages, stopToolCall, stop, cancelValidation])
// ---------------------------------------------------------------------
// Session actions
@@ -1849,8 +1860,7 @@ export function ChatEngineProvider({
error,
isBusy,
isLeaving,
isResending,
setMessages: setMessages as (messages: UIMessage[]) => void,
setMessages: setMessages as (messages: ChatUIMessage[]) => void,
stop: handleStop,
input,
setInput,
components/chat/message-list.tsx
+23 -28
View File
@@ -1,4 +1,4 @@
import type { UIMessage } from "ai"
import { isToolUIPart, type UIMessage } from "ai"
import {
AlertCircle,
BookmarkPlus,
@@ -24,6 +24,7 @@ import { ValidationCard } from "@/components/chat/ValidationCard"
import Image from "@/components/image-with-basepath"
import { useDictionary } from "@/hooks/use-dictionary"
import { getApiEndpoint } from "@/lib/base-path"
import { noticeOf } from "@/lib/chat-message"
import { useCopy } from "@/lib/clipboard"
import { cn } from "@/lib/utils"
import { useUiStore } from "@/stores/ui-store"
@@ -369,7 +370,7 @@ function AssistantMessage({
parts.forEach((part, partIndex) => {
if (part.type === "reasoning")
blocks.push({ kind: "reasoning", index: partIndex, texts: [] })
else if (part.type?.startsWith("tool-"))
else if (isToolUIPart(part))
blocks.push({ kind: "tool", index: partIndex, texts: [] })
else if (part.type === "text") {
const text = (part as { text: string }).text
@@ -568,9 +569,10 @@ export function MessageList() {
}
}, [messages])
const lastAssistantIndex = messages
.map((m) => m.role)
.lastIndexOf("assistant")
// A notice (a failed request) is an assistant message without an answer
const lastAssistantIndex = messages.findLastIndex(
(m) => m.role === "assistant" && !noticeOf(m),
)
const lastUserIndex = messages.map((m) => m.role).lastIndexOf("user")
return (
@@ -579,6 +581,7 @@ export function MessageList() {
const isRestored = engine.loadedMessageIdsRef.current.has(
message.id,
)
const notice = noticeOf(message)
return (
<div
key={message.id}
@@ -590,33 +593,18 @@ export function MessageList() {
index={index}
isLast={index === lastUserIndex}
/>
) : message.role === "assistant" ? (
<AssistantMessage
message={message}
index={index}
isLastMessage={index === messages.length - 1}
isLastAssistant={index === lastAssistantIndex}
/>
) : (
) : notice ? (
<ErrorNotice
text={getMessageTextContent(message)}
openModelConfig={
!!(
message.metadata as
| { openModelConfig?: boolean }
| undefined
)?.openModelConfig
}
onOpenModelConfig={() => {
const { providerId } =
(message.metadata as
| { providerId?: string }
| undefined) ?? {}
text={notice.text}
openModelConfig={!!notice.openModelConfig}
onOpenModelConfig={() =>
openSettings(
"models",
providerId ? { providerId } : "list",
notice.providerId
? { providerId: notice.providerId }
: "list",
)
}}
}
onRetry={
index === messages.length - 1 &&
!engine.isBusy &&
@@ -625,6 +613,13 @@ export function MessageList() {
: undefined
}
/>
) : (
<AssistantMessage
message={message}
index={index}
isLastMessage={index === messages.length - 1}
isLastAssistant={index === lastAssistantIndex}
/>
)}
</div>
)
components/chat/tool-activity.tsx
+3 -2
View File
@@ -1,3 +1,4 @@
import { getToolName } from "ai"
import {
AlertCircle,
Check,
@@ -164,7 +165,7 @@ function ToolRow({
tone === "error" && "border-destructive/25",
)}
data-testid="tool-row"
data-tool-name={part.type?.replace("tool-", "")}
data-tool-name={getToolName(part)}
data-tool-state={part.state}
>
<button
@@ -246,7 +247,7 @@ export function ToolActivity({
const t = dict.versions
const engine = useChatEngine()
const versions = useVersionsStore((s) => s.versions)
const toolName = part.type?.replace("tool-", "")
const toolName = getToolName(part)
const { state, errorText, output } = part
const versionIndex = versions.findIndex(
components/chat/types.ts
+12 -8
View File
@@ -1,15 +1,19 @@
import type { DynamicToolUIPart, ToolUIPart, UITools } from "ai"
import type { DiagramOperation } from "@/packages/mcp-server/src/diagram-operations.ts"
export type { DiagramOperation }
export interface ToolPartLike {
type: string
toolCallId: string
state?: string
input?: {
xml?: string
operations?: DiagramOperation[]
} & Record<string, unknown>
/**
* A tool part as the UI reads it, in any state: input, output and errorText
* are each there or not, instead of one shape per state.
*/
export type ToolPartLike = (ToolUIPart<UITools> | DynamicToolUIPart) & {
input?:
| ({
xml?: string
operations?: DiagramOperation[]
} & Record<string, unknown>)
| undefined
output?: string
errorText?: string
}
components/workspace/workspace.tsx
-1
View File
@@ -150,7 +150,6 @@ function WorkspaceLayout() {
const showHero =
engine.isRestored &&
engine.messages.length === 0 &&
!engine.isResending &&
!isRealDiagram(chartXML) &&
!heroDismissed
hooks/use-diagram-tool-handlers.ts
+4 -19
View File
@@ -1,3 +1,4 @@
import type { ChatAddToolOutputFunction } from "ai"
import type { RefObject } from "react"
import type { DiagramOperation } from "@/components/chat/types"
import type {
@@ -5,6 +6,7 @@ import type {
ValidationStatus,
} from "@/components/chat/ValidationCard"
import type { LoadMode } from "@/contexts/diagram-context"
import type { ChatUIMessage } from "@/lib/chat-message"
import { pageModelXml, pageSelectorFor, placeOnPage } from "@/lib/diagram-pages"
import type { ValidationResult } from "@/lib/diagram-validator"
import { formatValidationFeedback } from "@/lib/diagram-validator"
@@ -24,25 +26,8 @@ interface ToolCall {
input: unknown
}
type AddToolOutputSuccess = {
tool: string
toolCallId: string
state?: "output-available"
output: string
errorText?: undefined
}
type AddToolOutputError = {
tool: string
toolCallId: string
state: "output-error"
output?: undefined
errorText: string
}
type AddToolOutputParams = AddToolOutputSuccess | AddToolOutputError
type AddToolOutputFn = (params: AddToolOutputParams) => void
// useChat's addToolOutput, typed with the chat's tools
type AddToolOutputFn = ChatAddToolOutputFunction<ChatUIMessage>
const MAX_VALIDATION_RETRIES = 3
lib/chat-message.ts
+44
View File
@@ -0,0 +1,44 @@
import type { InferUITools, UIMessage } from "ai"
import type { chatTools } from "@/lib/chat-tools"
import type { SelectedCell } from "@/stores/canvas-store"
/** Shapes sent with a user message (also kept in its metadata) */
export type MessageSelection = Pick<SelectedCell, "id" | "label">[]
export interface ChatMessageMetadata {
selectedCells?: MessageSelection
}
/**
* A failed request or an answer that drew nothing, shown in the chat with
* what to do: an assistant message with this as its only part. The server
* ignores data parts, so it never reaches the model.
*/
export interface ChatNotice {
text: string
/** Show a button that opens model settings, on the provider that refused */
openModelConfig?: boolean
providerId?: string
}
export type ChatToolName = keyof InferUITools<typeof chatTools>
export type ChatUIMessage = UIMessage<
ChatMessageMetadata,
{ notice: ChatNotice },
InferUITools<typeof chatTools>
>
export function noticeMessage(notice: ChatNotice): ChatUIMessage {
return {
id: `notice-${Date.now()}`,
role: "assistant",
parts: [{ type: "data-notice", data: notice }],
}
}
/** The notice a message carries, if it is one */
export function noticeOf(message: ChatUIMessage): ChatNotice | undefined {
const part = message.parts.find((p) => p.type === "data-notice")
return part?.data
}
lib/chat-tools.ts
+138
View File
@@ -0,0 +1,138 @@
import { type InferUITools, tool } from "ai"
import path from "path"
import { z } from "zod"
import {
getShapeLibrary,
SHAPE_LIBRARY_LIST,
} from "@/packages/mcp-server/src/shape-library.ts"
import {
STYLE_CLASS_EXAMPLE,
SWIMLANE_EXAMPLE,
} from "@/packages/mcp-server/src/xml-examples.ts"
/**
* The chat's tools. display_diagram, edit_diagram and append_diagram run in
* the browser: the chat engine's tool handler puts their result on the
* canvas and reports it back as text, so they have an output schema and no
* execute. get_shape_library runs on the server.
*/
export const chatTools = {
display_diagram: tool({
description: `Display a diagram on draw.io. Pass ONLY the mxStyle definitions and the mxCell elements - wrapper tags and root cells are added automatically.
VALIDATION RULES (XML will be rejected if violated):
1. Generate ONLY mxStyle definitions and mxCell elements - NO wrapper tags (<mxfile>, <mxGraphModel>, <root>)
2. Do NOT include root cells (id="0" or id="1") - they are added automatically
3. All mxCell elements must be siblings - never nested
4. Every mxCell needs a unique id (start from "2")
5. parent defaults to "1"; write parent="<container-id>" only for shapes inside a container
6. Escape special chars in values: &lt; &gt; &amp; &quot;
A shape is one self-closing mxCell with x, y, w and h; an edge is one with source and target (a cell with source or target is always an edge). vertex="1", edge="1", parent="1" and the mxGeometry element are added automatically. Write parent only for a shape inside a container, and an mxGeometry element only for edge waypoints or for a separate label cell placed on an edge: <mxCell id="9" value="yes" style="edgeLabel;" parent="<edge id>" connectable="0"><mxGeometry x="-0.5" relative="1" as="geometry"/></mxCell>. An edge's own text simply goes in its value.
Example (generate ONLY this - no wrapper tags):
${SWIMLANE_EXAMPLE}
Styles: define a style used by several cells ONCE with <mxStyle name="..." value="..."/> before the cells and use the name in style like a CSS class; overrides after the name win. Name only styles that two or more cells share; names must not be draw.io's own style names (text, ellipse, rhombus, swimlane, label, image, blue, green, red, gray, yellow, orange, purple, pink). A definition applies to the call it is in. html=1 and whiteSpace=wrap are added automatically, never write them. Labels are HTML: use &lt;br&gt; for a line break, never \\n; a literal < or > is written &amp;lt; or &amp;gt;.
${STYLE_CLASS_EXAMPLE}
Notes:
- For AWS diagrams, use **AWS 2025 icons**.
- For animated connectors, add "flowAnimation=1" to edge style.
`,
inputSchema: z.object({
xml: z.string().describe("XML string to be displayed on draw.io"),
}),
outputSchema: z.string(),
}),
edit_diagram: tool({
description: `Edit the current diagram by ID-based operations (update/add/delete cells).
Operations:
- update: Replace an existing cell by its id. Provide cell_id and complete new_xml.
- add: Add a new cell. Provide cell_id (new unique id) and new_xml.
- delete: Remove a cell. Cascade is automatic: children AND edges (source/target) are auto-deleted. Only specify ONE cell_id.
For update/add, new_xml is the complete mxCell in the compact form (a shape with x, y, w, h; an edge with source and target), with its complete style (named styles are not available in edit_diagram).
⚠️ JSON ESCAPING: Every " inside new_xml MUST be escaped as \\". Example: id=\\"5\\" value=\\"Label\\"
Example - Add a rectangle:
{"operations": [{"operation": "add", "cell_id": "rect-1", "new_xml": "<mxCell id=\\"rect-1\\" value=\\"Hello\\" style=\\"rounded=1;\\" x=\\"100\\" y=\\"100\\" w=\\"120\\" h=\\"60\\"/>"}]}
Example - Delete container (children & edges auto-deleted):
{"operations": [{"operation": "delete", "cell_id": "2"}]}`,
inputSchema: z.object({
operations: z
.array(
z.object({
operation: z
.enum(["update", "add", "delete"])
.describe(
"Operation to perform: add, update, or delete",
),
cell_id: z
.string()
.describe(
"The id of the mxCell. Must match the id attribute in new_xml.",
),
new_xml: z
.string()
.optional()
.describe(
"Complete mxCell XML element (required for update/add)",
),
}),
)
.describe("Array of operations to apply"),
}),
outputSchema: z.string(),
}),
append_diagram: tool({
description: `Continue generating diagram XML when previous display_diagram output was truncated due to length limits.
WHEN TO USE: Only call this tool after display_diagram was truncated (you'll see an error message about truncation).
CRITICAL INSTRUCTIONS:
1. Do NOT include any wrapper tags - just continue the mxCell elements
2. Continue from EXACTLY where your previous output stopped
3. Complete the remaining mxCell elements
4. If still truncated, call append_diagram again with the next fragment
Example: If previous output ended with '<mxCell id="x" style="rounded=1', continue with ';" x="40" y="40" w="120" h="60"/>' and complete the remaining elements.`,
inputSchema: z.object({
xml: z
.string()
.describe(
"Continuation XML fragment to append (NO wrapper tags)",
),
}),
outputSchema: z.string(),
}),
get_shape_library: tool({
description: `Get draw.io shape/icon library documentation with style syntax and shape names.
Available libraries:
${SHAPE_LIBRARY_LIST}
Call this tool to get shape names and usage syntax for a specific library.`,
inputSchema: z.object({
library: z
.string()
.describe(
"Library name (e.g., 'aws4', 'kubernetes', 'flowchart')",
),
}),
execute: async ({ library }) => {
// Only known library names reach the file system
const result = await getShapeLibrary(
library,
path.join(process.cwd(), "docs/shape-libraries"),
)
return result.ok ? result.text : result.error
},
}),
}
/** The tool names and their input and output types, for the UI messages */
export type ChatTools = InferUITools<typeof chatTools>
lib/langfuse.ts
+8 -4
View File
@@ -1,6 +1,7 @@
import { LangfuseClient } from "@langfuse/client"
import { observe, updateActiveTrace } from "@langfuse/tracing"
import * as api from "@opentelemetry/api"
import type { TelemetrySettings } from "ai"
import { randomUUID } from "crypto"
// Singleton LangfuseClient instance for direct API calls
@@ -115,20 +116,23 @@ export function endTrace() {
}
}
// Get telemetry config for streamText
// Telemetry settings for streamText and generateText
export function getTelemetryConfig(params: {
sessionId?: string
userId?: string
}) {
/** Names the trace of a model call outside the chat's observe() span */
functionId?: string
}): TelemetrySettings | undefined {
if (!isLangfuseEnabled()) return undefined
return {
isEnabled: true,
functionId: params.functionId,
recordInputs: true,
recordOutputs: true,
metadata: {
sessionId: params.sessionId,
userId: params.userId,
...(params.sessionId && { sessionId: params.sessionId }),
...(params.userId && { userId: params.userId }),
},
}
}
lib/test-model-connection.ts
+8 -3
View File
@@ -199,13 +199,16 @@ export async function testModelConnection(
tools: { ping: PING_TOOL },
maxOutputTokens: 1024,
maxRetries: 0,
abortSignal: AbortSignal.timeout(TEST_TIMEOUT_MS),
// Ends with the request: the user closed the dialog or the page
abortSignal: req.signal,
timeout: TEST_TIMEOUT_MS,
})
let calledTool = false
let finishReason: string | undefined
for await (const part of result.fullStream) {
if (part.type === "error") throw part.error
// The timeout ends the stream with an abort part, not an error
// The timeout (or the request's end) ends the stream with an
// abort part, not an error
if (part.type === "abort") {
const timeout = new Error(
`The model did not answer within ${TEST_TIMEOUT_MS / 1000} s.`,
@@ -213,7 +216,9 @@ export async function testModelConnection(
timeout.name = "TimeoutError"
throw timeout
}
if (part.type === "tool-call") {
// A call of a tool that does not exist, or with unreadable
// input, arrives marked invalid: no proof of tool support
if (part.type === "tool-call" && !part.invalid) {
calledTool = true
break
}
tests/unit/chat-route-abort.test.ts
+101 -79
View File
@@ -28,60 +28,89 @@ const sse = (chunks: object[], end = true) =>
chunks.map((c) => `data: ${JSON.stringify(c)}\n\n`).join("") +
(end ? "data: [DONE]\n\n" : "")
// Step 1: the model asks for a shape library (a tool run on the server) and
// reports its usage, so the chat makes a second request
const firstStep = () =>
new Response(
sse([
{
id: "c1",
choices: [
{
index: 0,
delta: {
role: "assistant",
tool_calls: [
{
index: 0,
id: "call_1",
type: "function",
function: {
name: "get_shape_library",
arguments: '{"library":"aws4"}',
},
},
],
},
finish_reason: null,
},
],
},
{
id: "c1",
choices: [{ index: 0, delta: {}, finish_reason: "tool_calls" }],
usage: { prompt_tokens: 1200, completion_tokens: 30 },
},
]),
{ headers: { "content-type": "text/event-stream" } },
)
const send = (signal?: AbortSignal) =>
chat(
new Request("http://localhost/api/chat", {
method: "POST",
signal,
headers: {
"Content-Type": "application/json",
"x-forwarded-for": "203.0.113.7",
// The server's own network: counted
"x-ai-provider": "glm",
"x-ai-base-url": "http://127.0.0.1:9000/v1",
"x-ai-api-key": "dummy",
"x-ai-model": "glm-5",
},
body: JSON.stringify({
messages: [
{
id: "u1",
role: "user",
parts: [{ type: "text", text: "Draw AWS" }],
},
],
xml: "",
}),
}),
)
const drain = async (res: Response) => {
const reader = res.body?.getReader()
while (
reader &&
!(await reader.read().catch(() => ({ done: true }))).done
) {
// drain
}
}
describe("a request stopped after a finished step", () => {
it("counts that step's tokens", async () => {
// Step 1 asks for a shape library (run on the server) and reports
// its usage; step 2 never ends, and the user stops
// Step 2 never ends, and the user stops
let call = 0
vi.stubGlobal(
"fetch",
vi.fn(async (_url: string, init?: RequestInit) => {
call++
if (call === 1) {
return new Response(
sse([
{
id: "c1",
choices: [
{
index: 0,
delta: {
role: "assistant",
tool_calls: [
{
index: 0,
id: "call_1",
type: "function",
function: {
name: "get_shape_library",
arguments:
'{"library":"aws4"}',
},
},
],
},
finish_reason: null,
},
],
},
{
id: "c1",
choices: [
{
index: 0,
delta: {},
finish_reason: "tool_calls",
},
],
usage: {
prompt_tokens: 1200,
completion_tokens: 30,
},
},
]),
{ headers: { "content-type": "text/event-stream" } },
)
}
if (call === 1) return firstStep()
// Never ends, until the request is aborted (as fetch does)
const body = new ReadableStream({
start(controller) {
@@ -98,42 +127,35 @@ describe("a request stopped after a finished step", () => {
}),
)
const stop = new AbortController()
const res = await chat(
new Request("http://localhost/api/chat", {
method: "POST",
signal: stop.signal,
headers: {
"Content-Type": "application/json",
"x-forwarded-for": "203.0.113.7",
// The server's own network: counted
"x-ai-provider": "glm",
"x-ai-base-url": "http://127.0.0.1:9000/v1",
"x-ai-api-key": "dummy",
"x-ai-model": "glm-5",
},
body: JSON.stringify({
messages: [
{
id: "u1",
role: "user",
parts: [{ type: "text", text: "Draw AWS" }],
},
],
xml: "",
}),
}),
)
const reader = res.body?.getReader()
const res = await send(stop.signal)
const draining = drain(res)
// Read until the second step has started
await vi.waitFor(() => expect(call).toBe(2), { timeout: 3000 })
stop.abort()
// The answer stream ends; the SDK handles the stop as it is read
while (
reader &&
!(await reader.read().catch(() => ({ done: true }))).done
) {
// drain
}
await draining
await vi.waitFor(() => expect(quota.recorded).toEqual([1230]))
})
})
describe("a request that fails after a finished step", () => {
it("counts that step's tokens", async () => {
// Step 2 is refused by the provider (a status the SDK does not
// retry): the stream ends with an error
let call = 0
vi.stubGlobal(
"fetch",
vi.fn(async () => {
call++
if (call === 1) return firstStep()
return new Response('{"error":{"message":"bad request"}}', {
status: 400,
headers: { "content-type": "application/json" },
})
}),
)
await drain(await send())
expect(call).toBe(2)
expect(quota.recorded).toEqual([1230])
})
})
tests/unit/chat-route-quota.test.ts
+15
View File
@@ -205,6 +205,21 @@ describe("request checks", () => {
expect(await res.text()).toMatch(/Invalid AWS region/)
expect(fetch).not.toHaveBeenCalled()
})
it("refuses messages that are not UI messages", async () => {
const res = await chat(
new Request("http://localhost/api/chat", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({
messages: [{ id: "u1", role: "user", text: "no parts" }],
xml: "",
}),
}),
)
expect(res.status).toBe(400)
expect(fetch).not.toHaveBeenCalled()
})
})
describe("server model allowlist", () => {
tests/unit/validate-diagram-route.test.ts
+40 -11
View File
@@ -15,19 +15,32 @@ const RESULT = {
suggestions: ["Move box B to the right"],
}
// A vision model that answers with the result as JSON text
// A vision model that answers with the result as JSON text; told to wait,
// it answers only once its call is aborted
const model = vi.hoisted(() => ({ waitForAbort: false }))
vi.mock("@/lib/ai-providers", () => ({
getValidationModel: () =>
new MockLanguageModelV3({
doGenerate: (async () => ({
content: [{ type: "text", text: JSON.stringify(RESULT) }],
finishReason: { unified: "stop", raw: "stop" },
usage: {
inputTokens: { total: 1 },
outputTokens: { total: 1 },
},
warnings: [],
})) as any,
doGenerate: (async (options: { abortSignal?: AbortSignal }) => {
if (model.waitForAbort) {
const signal = options.abortSignal
await new Promise((_, reject) => {
if (signal?.aborted) reject(signal.reason)
signal?.addEventListener("abort", () =>
reject(signal.reason),
)
})
}
return {
content: [{ type: "text", text: JSON.stringify(RESULT) }],
finishReason: { unified: "stop", raw: "stop" },
usage: {
inputTokens: { total: 1 },
outputTokens: { total: 1 },
},
warnings: [],
}
}) as any,
}),
}))
@@ -61,9 +74,10 @@ vi.mock("@/lib/dynamo-quota-manager", () => ({
},
}))
const post = () =>
const post = (signal?: AbortSignal) =>
new Request("http://localhost/api/validate-diagram", {
method: "POST",
signal,
headers: {
"Content-Type": "application/json",
"x-forwarded-for": "203.0.113.7",
@@ -72,6 +86,7 @@ const post = () =>
})
afterEach(() => {
model.waitForAbort = false
quota.enabled = false
quota.allowed = true
quota.checks = []
@@ -106,6 +121,20 @@ describe("POST /api/validate-diagram", () => {
expect(await res.json()).toEqual(RESULT)
})
it("ends its model call when the request is aborted", async () => {
// The user pressed Stop: the browser aborts the request, and the
// model call must not run on (and be paid for)
model.waitForAbort = true
const stop = new AbortController()
const pending = validateDiagram(post(stop.signal))
stop.abort()
expect(await (await pending).json()).toEqual({
valid: true,
issues: [],
suggestions: [],
})
})
it("answers valid when the check is turned off", async () => {
vi.stubEnv("ENABLE_VLM_VALIDATION", "false")
const res = await validateDiagram(post())
tests/unit/validate-model-route.test.ts
+18
View File
@@ -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", () => {