Files
next-ai-draw-io/hooks/use-diagram-tool-handlers.ts
dayuan.jiang a3814f702d feat(diagram-engine): flowcharts, swimlanes, sequence diagrams and mind maps
Extends the declarative engine past cloud architecture. The tool routing was
divided by icon library — AWS through the engine, everything else hand-written
XML — which is the wrong axis. What matters is the LAYOUT SHAPE.

Measured first: a six-step approval flow declared in its natural order comes out
as one column, because the layout only arranges what nesting tells it to and
never looked at the arrows. That forces the arrow from the decision to its second
branch to jump over the first branch.

graph.ts computes what the layout should have looked at: layer assignment by
longest path, cycle breaking so a loop is drawn without setting the order, and
barycentre sweeping to cut edge crossings. It emits ordinary container
operations, so layout, routing and round-tripping are unchanged — reaching zero
arrows-through-boxes on a 14-node pipeline and zero crossings on a bipartite
graph whose declared order forces three.

Three new container kinds, each because one layout rule cannot serve them all:

  pool     — swimlanes. Lanes are real cells and each step is parented to its
             band, so dragging a step to another role records the change.
  sequence — participants across the top, one lifeline cell per participant so
             head and line stay together on a drag. Messages bypass the router:
             a message's height IS its order.
  radial   — mind maps and org charts. Children are a flat list and the
             hierarchy comes from the links, because a branch is a box and a box
             cannot hold children.

Flowchart box shapes (diamond, stadium, parallelogram, document) so a reader can
tell a branch from a step.

Two bugs the new tests caught: the duplicate-link guard blocked a sequence
diagram from having two messages between the same pair, and the fallback message
numbering was shared across containers, pushing a second diagram's messages off
its own lifelines.

523 unit tests and 17 diagram e2e tests pass. Every kind verified round-trip
stable to a fixed point, and rendered in a real browser — draw.io keeps the
lifeline shape and the lane markers.
2026-08-09 13:49:11 +09:00

713 lines
26 KiB
TypeScript

import type { MutableRefObject } from "react"
import { useRef } from "react"
import type { DiagramOperation } from "@/components/chat/types"
import type {
ValidationState,
ValidationStatus,
} from "@/components/chat/ValidationCard"
import type { Operation } from "@/lib/diagram-engine"
import {
drawGraph,
type GraphEdge,
type GraphNode,
restructureDiagram,
} from "@/lib/diagram-engine"
import type { ValidationResult } from "@/lib/diagram-validator"
import { formatValidationFeedback } from "@/lib/diagram-validator"
import { isMxCellXmlComplete, isRealDiagram, wrapWithMxFile } from "@/lib/utils"
const DEBUG = process.env.NODE_ENV === "development"
interface ToolCall {
toolCallId: string
toolName: string
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
const MAX_VALIDATION_RETRIES = 3
// Type for the validation function passed from useValidateDiagram hook
type ValidateDiagramFn = (
imageData: string,
sessionId?: string,
) => Promise<ValidationResult>
interface UseDiagramToolHandlersParams {
partialXmlRef: MutableRefObject<string>
editDiagramOriginalXmlRef: MutableRefObject<Map<string, string>>
chartXMLRef: MutableRefObject<string>
onDisplayChart: (xml: string, skipValidation?: boolean) => string | null
onFetchChart: (saveToHistory?: boolean) => Promise<string>
onExport: () => void
captureValidationPng?: () => Promise<string | null>
validateDiagram?: ValidateDiagramFn
enableVlmValidation?: boolean
sessionId?: string
onValidationStateChange?: (
toolCallId: string,
state: ValidationState,
) => void
}
/**
* Hook that creates the onToolCall handler for diagram-related tools.
* Handles display_diagram, edit_diagram, and append_diagram tools.
*
* Note: addToolOutput is passed at call time (not hook init) because
* it comes from useChat which creates a circular dependency.
*/
export function useDiagramToolHandlers({
partialXmlRef,
editDiagramOriginalXmlRef,
chartXMLRef,
onDisplayChart,
onFetchChart,
onExport,
captureValidationPng,
validateDiagram,
enableVlmValidation = true,
sessionId,
onValidationStateChange,
}: UseDiagramToolHandlersParams) {
// Track validation retry count per tool call
const validationRetryCountRef = useRef<Map<string, number>>(new Map())
// Helper to update validation state
const updateValidationState = (
toolCallId: string,
status: ValidationStatus,
options?: {
attempt?: number
maxAttempts?: number
result?: ValidationResult
error?: string
imageData?: string
},
) => {
if (onValidationStateChange) {
onValidationStateChange(toolCallId, {
status,
...options,
})
}
}
const handleToolCall = async (
{ toolCall }: { toolCall: ToolCall },
addToolOutput: AddToolOutputFn,
) => {
if (DEBUG) {
console.log(
`[onToolCall] Tool: ${toolCall.toolName}, CallId: ${toolCall.toolCallId}`,
)
}
if (toolCall.toolName === "display_diagram") {
await handleDisplayDiagram(toolCall, addToolOutput)
} else if (toolCall.toolName === "edit_diagram") {
await handleEditDiagram(toolCall, addToolOutput)
} else if (toolCall.toolName === "append_diagram") {
handleAppendDiagram(toolCall, addToolOutput)
} else if (toolCall.toolName === "restructure_diagram") {
await handleRestructureDiagram(toolCall, addToolOutput)
} else if (toolCall.toolName === "draw_graph") {
await handleDrawGraph(toolCall, addToolOutput)
}
}
const handleDisplayDiagram = async (
toolCall: ToolCall,
addToolOutput: AddToolOutputFn,
) => {
const { xml } = toolCall.input as { xml: string }
// DEBUG: Log raw input to diagnose false truncation detection
if (DEBUG) {
console.log(
"[display_diagram] XML ending (last 100 chars):",
xml.slice(-100),
)
console.log("[display_diagram] XML length:", xml.length)
}
// Check if XML is truncated (incomplete mxCell indicates truncated output)
const isTruncated = !isMxCellXmlComplete(xml)
if (DEBUG) {
console.log("[display_diagram] isTruncated:", isTruncated)
}
if (isTruncated) {
// Store the partial XML for continuation via append_diagram
partialXmlRef.current = xml
// Tell LLM to use append_diagram to continue
const partialEnding = partialXmlRef.current.slice(-500)
addToolOutput({
tool: "display_diagram",
toolCallId: toolCall.toolCallId,
state: "output-error",
errorText: `Output was truncated due to length limits. Use the append_diagram tool to continue.
Your output ended with:
\`\`\`
${partialEnding}
\`\`\`
NEXT STEP: Call append_diagram with the continuation XML.
- Do NOT include wrapper tags or root cells (id="0", id="1")
- Start from EXACTLY where you stopped
- Complete all remaining mxCell elements`,
})
return
}
// Complete XML received - use it directly
// (continuation is now handled via append_diagram tool)
const finalXml = xml
partialXmlRef.current = "" // Reset any partial from previous truncation
// Wrap raw XML with full mxfile structure for draw.io
const fullXml = wrapWithMxFile(finalXml)
// loadDiagram validates and returns error if invalid
const validationError = onDisplayChart(fullXml)
if (validationError) {
console.warn("[display_diagram] Validation error:", validationError)
// Return error to model - sendAutomaticallyWhen will trigger retry
if (DEBUG) {
console.log(
"[display_diagram] Adding tool output with state: output-error",
)
}
addToolOutput({
tool: "display_diagram",
toolCallId: toolCall.toolCallId,
state: "output-error",
errorText: `${validationError}
Please fix the XML issues and call display_diagram again with corrected XML.
Your failed XML:
\`\`\`xml
${finalXml}
\`\`\``,
})
} else {
// Success - diagram will be rendered by chat-message-display
if (DEBUG) {
console.log(
"[display_diagram] Success! Checking if VLM validation is enabled...",
)
}
// VLM validation after successful display
if (
enableVlmValidation &&
captureValidationPng &&
validateDiagram
) {
let capturedPngData: string | null = null
try {
// Notify UI that we're starting capture
updateValidationState(toolCall.toolCallId, "capturing")
// Small delay (100ms) to allow diagram rendering to complete before capture.
// This is a best-effort heuristic and may need adjustment for complex diagrams or slower devices.
await new Promise((resolve) => setTimeout(resolve, 100))
capturedPngData = await captureValidationPng()
if (capturedPngData) {
if (DEBUG) {
console.log(
"[display_diagram] Captured PNG for validation",
)
}
const retryCount =
validationRetryCountRef.current.get(
toolCall.toolCallId,
) || 0
// Notify UI that we're validating (include the image)
updateValidationState(
toolCall.toolCallId,
"validating",
{
attempt: retryCount + 1,
maxAttempts: MAX_VALIDATION_RETRIES,
imageData: capturedPngData,
},
)
const result = await validateDiagram(
capturedPngData,
sessionId,
)
if (!result.valid) {
if (retryCount < MAX_VALIDATION_RETRIES) {
validationRetryCountRef.current.set(
toolCall.toolCallId,
retryCount + 1,
)
const feedback =
formatValidationFeedback(result)
if (DEBUG) {
console.log(
`[display_diagram] Validation failed (attempt ${retryCount + 1}/${MAX_VALIDATION_RETRIES}):`,
result.issues,
)
}
// Notify UI of validation failure (include the image)
updateValidationState(
toolCall.toolCallId,
"failed",
{
attempt: retryCount + 1,
maxAttempts: MAX_VALIDATION_RETRIES,
result,
imageData: capturedPngData,
},
)
addToolOutput({
tool: "display_diagram",
toolCallId: toolCall.toolCallId,
state: "output-error",
errorText: `[Validation attempt ${retryCount + 1}/${MAX_VALIDATION_RETRIES}]\n${feedback}`,
})
return
} else {
// Max retries reached - accept the diagram with warning
if (DEBUG) {
console.log(
"[display_diagram] Max validation retries reached, accepting diagram",
)
}
validationRetryCountRef.current.delete(
toolCall.toolCallId,
)
// Notify UI that we're accepting with issues (include the image)
updateValidationState(
toolCall.toolCallId,
"skipped",
{ result, imageData: capturedPngData },
)
addToolOutput({
tool: "display_diagram",
toolCallId: toolCall.toolCallId,
output: "Diagram displayed (validation issues noted but max retries reached).",
})
return
}
} else {
// Validation passed - clean up retry count
validationRetryCountRef.current.delete(
toolCall.toolCallId,
)
if (DEBUG) {
console.log(
"[display_diagram] Validation passed!",
)
}
// Notify UI of success (include the image)
// Use "success_with_warnings" if valid but has issues
const hasWarnings = result.issues.length > 0
updateValidationState(
toolCall.toolCallId,
hasWarnings
? "success_with_warnings"
: "success",
{ result, imageData: capturedPngData },
)
}
} else {
// PNG capture failed - skip validation
updateValidationState(toolCall.toolCallId, "skipped")
}
} catch (error) {
// VLM validation error - log but don't block the user
console.warn(
"[display_diagram] VLM validation error:",
error,
)
updateValidationState(toolCall.toolCallId, "error", {
error:
error instanceof Error
? error.message
: "Validation failed",
imageData: capturedPngData || undefined,
})
}
}
if (DEBUG) {
console.log(
"[display_diagram] Adding tool output with state: output-available",
)
}
addToolOutput({
tool: "display_diagram",
toolCallId: toolCall.toolCallId,
output: "Successfully displayed the diagram.",
})
if (DEBUG) {
console.log(
"[display_diagram] Tool output added. Diagram should be visible now.",
)
}
}
}
const handleEditDiagram = async (
toolCall: ToolCall,
addToolOutput: AddToolOutputFn,
) => {
const { operations } = toolCall.input as {
operations: DiagramOperation[]
}
let currentXml = ""
try {
// Use the original XML captured during streaming (shared with chat-message-display)
// This ensures we apply operations to the same base XML that streaming used
const originalXml = editDiagramOriginalXmlRef.current.get(
toolCall.toolCallId,
)
if (originalXml) {
currentXml = originalXml
} else {
// Fallback: use chartXML from ref if streaming didn't capture original
const cachedXML = chartXMLRef.current
if (cachedXML) {
currentXml = cachedXML
} else {
// Last resort: export from iframe
currentXml = await onFetchChart(false)
}
}
const { applyDiagramOperations } = await import("@/lib/utils")
const { result: editedXml, errors } = applyDiagramOperations(
currentXml,
operations,
)
// Check for operation errors
if (errors.length > 0) {
const errorMessages = errors
.map(
(e) =>
`- ${e.type} on cell_id="${e.cellId}": ${e.message}`,
)
.join("\n")
addToolOutput({
tool: "edit_diagram",
toolCallId: toolCall.toolCallId,
state: "output-error",
errorText: `Some operations failed:\n${errorMessages}
Current diagram XML:
\`\`\`xml
${currentXml}
\`\`\`
Please check the cell IDs and retry.`,
})
// Clean up the shared original XML ref
editDiagramOriginalXmlRef.current.delete(toolCall.toolCallId)
return
}
// loadDiagram validates and returns error if invalid
const validationError = onDisplayChart(editedXml)
if (validationError) {
console.warn(
"[edit_diagram] Validation error:",
validationError,
)
addToolOutput({
tool: "edit_diagram",
toolCallId: toolCall.toolCallId,
state: "output-error",
errorText: `Edit produced invalid XML: ${validationError}
Current diagram XML:
\`\`\`xml
${currentXml}
\`\`\`
Please fix the operations to avoid structural issues.`,
})
// Clean up the shared original XML ref
editDiagramOriginalXmlRef.current.delete(toolCall.toolCallId)
return
}
onExport()
addToolOutput({
tool: "edit_diagram",
toolCallId: toolCall.toolCallId,
output: `Successfully applied ${operations.length} operation(s) to the diagram.`,
})
// Clean up the shared original XML ref
editDiagramOriginalXmlRef.current.delete(toolCall.toolCallId)
} catch (error) {
console.error("[edit_diagram] Failed:", error)
const errorMessage =
error instanceof Error ? error.message : String(error)
addToolOutput({
tool: "edit_diagram",
toolCallId: toolCall.toolCallId,
state: "output-error",
errorText: `Edit failed: ${errorMessage}
Current diagram XML:
\`\`\`xml
${currentXml || "No XML available"}
\`\`\`
Please check cell IDs and retry, or use display_diagram to regenerate.`,
})
// Clean up the shared original XML ref even on error
editDiagramOriginalXmlRef.current.delete(toolCall.toolCallId)
}
}
const handleAppendDiagram = (
toolCall: ToolCall,
addToolOutput: AddToolOutputFn,
) => {
const { xml } = toolCall.input as { xml: string }
// Detect if LLM incorrectly started fresh instead of continuing
// LLM should only output bare mxCells now, so wrapper tags indicate error
const trimmed = xml.trim()
const isFreshStart =
trimmed.startsWith("<mxGraphModel") ||
trimmed.startsWith("<root") ||
trimmed.startsWith("<mxfile") ||
trimmed.startsWith('<mxCell id="0"') ||
trimmed.startsWith('<mxCell id="1"')
if (isFreshStart) {
addToolOutput({
tool: "append_diagram",
toolCallId: toolCall.toolCallId,
state: "output-error",
errorText: `ERROR: You started fresh with wrapper tags. Do NOT include wrapper tags or root cells (id="0", id="1").
Continue from EXACTLY where the partial ended:
\`\`\`
${partialXmlRef.current.slice(-500)}
\`\`\`
Start your continuation with the NEXT character after where it stopped.`,
})
return
}
// Append to accumulated XML
partialXmlRef.current += xml
// Check if XML is now complete (last mxCell is complete)
const isComplete = isMxCellXmlComplete(partialXmlRef.current)
if (isComplete) {
// Wrap and display the complete diagram
const finalXml = partialXmlRef.current
partialXmlRef.current = "" // Reset
const fullXml = wrapWithMxFile(finalXml)
const validationError = onDisplayChart(fullXml)
if (validationError) {
addToolOutput({
tool: "append_diagram",
toolCallId: toolCall.toolCallId,
state: "output-error",
errorText: `Validation error after assembly: ${validationError}
Assembled XML:
\`\`\`xml
${finalXml.substring(0, 2000)}...
\`\`\`
Please use display_diagram with corrected XML.`,
})
} else {
addToolOutput({
tool: "append_diagram",
toolCallId: toolCall.toolCallId,
output: "Diagram assembly complete and displayed successfully.",
})
}
} else {
// Still incomplete - signal to continue
addToolOutput({
tool: "append_diagram",
toolCallId: toolCall.toolCallId,
state: "output-error",
errorText: `XML still incomplete (mxCell not closed). Call append_diagram again to continue.
Current ending:
\`\`\`
${partialXmlRef.current.slice(-500)}
\`\`\`
Continue from EXACTLY where you stopped.`,
})
}
}
/**
* Structural editing. The model sends operations against the tree; the engine
* re-derives that tree from whatever is on the canvas right now — including anything
* the user moved or recoloured by hand — applies the operations, recomputes every
* coordinate, and returns new XML.
*
* Nothing about the tree is stored between calls, so there is no second copy of the
* state to drift out of sync with the canvas.
*/
const handleRestructureDiagram = async (
toolCall: ToolCall,
addToolOutput: AddToolOutputFn,
) => {
const { operations } = toolCall.input as { operations: Operation[] }
// Read the live canvas, not the last thing we generated: the user may have
// edited it since.
let currentXml = ""
try {
currentXml = await onFetchChart(false)
} catch {
currentXml = chartXMLRef.current ?? ""
}
if (!isRealDiagram(currentXml)) currentXml = ""
const result = restructureDiagram(currentXml, operations)
if (result.errors.length > 0 || !result.xml) {
addToolOutput({
tool: "restructure_diagram",
toolCallId: toolCall.toolCallId,
state: "output-error",
errorText: `Could not apply the operations:
${result.errors.map((e) => `- ${e}`).join("\n")}
Structure as it stands:
${result.outline}
Fix the operations and call restructure_diagram again.`,
})
return
}
const loadError = onDisplayChart(result.xml)
if (loadError) {
addToolOutput({
tool: "restructure_diagram",
toolCallId: toolCall.toolCallId,
state: "output-error",
errorText: `The diagram was built but draw.io rejected it: ${loadError}`,
})
return
}
// Report the outline rather than the XML: it is what the model needs to name ids
// in the next call, at a fraction of the tokens.
const notes = result.warnings.length
? `\n\nNotes:\n${result.warnings.map((w) => `- ${w}`).join("\n")}`
: ""
addToolOutput({
tool: "restructure_diagram",
toolCallId: toolCall.toolCallId,
output: `Diagram updated.\n\n${result.outline}${notes}`,
})
}
/**
* draw_graph: build a flowchart from nodes and arrows, with no positions given.
*
* The current canvas is deliberately NOT read. Which row a node belongs in depends on
* every arrow in the graph, so one new edge can move half the diagram — there is no
* meaningful way to merge a new graph into an existing layout. The model edits afterwards
* through restructure_diagram, which does read the canvas.
*/
const handleDrawGraph = async (
toolCall: ToolCall,
addToolOutput: AddToolOutputFn,
) => {
const { nodes, edges, title, flow } = toolCall.input as {
nodes: GraphNode[]
edges: GraphEdge[]
title?: string
flow?: "col" | "row"
}
const result = drawGraph(nodes ?? [], edges ?? [], { title, flow })
if (result.errors.length > 0 || !result.xml) {
addToolOutput({
tool: "draw_graph",
toolCallId: toolCall.toolCallId,
state: "output-error",
errorText: `Could not draw the graph:
${result.errors.map((e) => `- ${e}`).join("\n")}
Fix the nodes or edges and call draw_graph again.`,
})
return
}
const loadError = onDisplayChart(result.xml)
if (loadError) {
addToolOutput({
tool: "draw_graph",
toolCallId: toolCall.toolCallId,
state: "output-error",
errorText: `The diagram was built but draw.io rejected it: ${loadError}`,
})
return
}
const notes = result.warnings.length
? `\n\nNotes:\n${result.warnings.map((w) => `- ${w}`).join("\n")}`
: ""
addToolOutput({
tool: "draw_graph",
toolCallId: toolCall.toolCallId,
output: `Diagram updated.\n\n${result.outline}${notes}`,
})
}
return { handleToolCall }
}