mirror of
https://github.com/DayuanJiang/next-ai-draw-io.git
synced 2026-10-12 12:39:52 +08:00
Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 dayuan.jiang
|
92e785bcb5 |
+51
-187
@@ -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: < > & "
|
||||
|
||||
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 <br> for a line break, never \\n; a literal < or > is written &lt; or &gt;.
|
||||
${STYLE_CLASS_EXAMPLE}
|
||||
|
||||
Notes:
|
||||
- For AWS diagrams, use **AWS 2025 icons**.
|
||||
- For animated connectors, add "flowAnimation=1" to edge style.
|
||||
`,
|
||||
inputSchema: z.object({
|
||||
xml: z
|
||||
.string()
|
||||
.describe("XML string to be displayed on draw.io"),
|
||||
}),
|
||||
},
|
||||
edit_diagram: {
|
||||
description: `Edit the current diagram by ID-based operations (update/add/delete cells).
|
||||
|
||||
Operations:
|
||||
- update: Replace an existing cell by its id. Provide cell_id and complete new_xml.
|
||||
- add: Add a new cell. Provide cell_id (new unique id) and new_xml.
|
||||
- delete: Remove a cell. Cascade is automatic: children AND edges (source/target) are auto-deleted. Only specify ONE cell_id.
|
||||
|
||||
For update/add, new_xml is the complete mxCell in the compact form (a shape with x, y, w, h; an edge with source and target), with its complete style (named styles are not available in edit_diagram).
|
||||
|
||||
⚠️ JSON ESCAPING: Every " inside new_xml MUST be escaped as \\". Example: id=\\"5\\" value=\\"Label\\"
|
||||
|
||||
Example - Add a rectangle:
|
||||
{"operations": [{"operation": "add", "cell_id": "rect-1", "new_xml": "<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),
|
||||
}),
|
||||
|
||||
@@ -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(
|
||||
|
||||
+102
-92
@@ -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,
|
||||
|
||||
@@ -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>
|
||||
)
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -150,7 +150,6 @@ function WorkspaceLayout() {
|
||||
const showHero =
|
||||
engine.isRestored &&
|
||||
engine.messages.length === 0 &&
|
||||
!engine.isResending &&
|
||||
!isRealDiagram(chartXML) &&
|
||||
!heroDismissed
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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: < > & "
|
||||
|
||||
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 <br> for a line break, never \\n; a literal < or > is written &lt; or &gt;.
|
||||
${STYLE_CLASS_EXAMPLE}
|
||||
|
||||
Notes:
|
||||
- For AWS diagrams, use **AWS 2025 icons**.
|
||||
- For animated connectors, add "flowAnimation=1" to edge style.
|
||||
`,
|
||||
inputSchema: z.object({
|
||||
xml: z.string().describe("XML string to be displayed on draw.io"),
|
||||
}),
|
||||
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>
|
||||
+8
-4
@@ -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 }),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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])
|
||||
})
|
||||
})
|
||||
|
||||
@@ -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", () => {
|
||||
|
||||
@@ -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())
|
||||
|
||||
@@ -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", () => {
|
||||
|
||||
Reference in New Issue
Block a user