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 (#987)

Browse Source 
## 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
This commit is contained in:
Dayuan Jiang
2026-10-12 12:58:52 +09:00
committed by GitHub
parent 60105ddb51
commit 99ced88e5b
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", () => {