i/next-ai-draw-io
1
0
Fork 0
mirror of https://github.com/DayuanJiang/next-ai-draw-io.git synced 2026-10-12 20:49:53 +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
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
}