i/next-ai-draw-io
1
0
Fork 0
mirror of https://github.com/DayuanJiang/next-ai-draw-io.git synced 2026-01-02 22:32:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor: extract diagram tool handlers to dedicated hook

Browse Source 
- Create useDiagramToolHandlers hook for display_diagram, edit_diagram, append_diagram
- Remove ~300 lines from chat-panel.tsx
- Remove unused stopRef
- Gate debug console.log statements with DEBUG constant
This commit is contained in:
Dayuan Jiang
2025-12-24 10:15:41 +09:00
parent 1fab261cd0
commit 4dbda5ba8e
2 changed files with 397 additions and 312 deletions
Download Patch File Download Diff File Expand all files Collapse all files

326
components/chat-panel.tsx
View File

@@ -27,6 +27,7 @@ import {
TooltipTrigger,
} from "@/components/ui/tooltip"
import { useDiagram } from "@/contexts/diagram-context"
import { useDiagramToolHandlers } from "@/hooks/use-diagram-tool-handlers"
import { useDictionary } from "@/hooks/use-dictionary"
import { getSelectedAIConfig, useModelConfig } from "@/hooks/use-model-config"
import { getApiEndpoint } from "@/lib/base-path"
@@ -34,7 +35,7 @@ import { findCachedResponse } from "@/lib/cached-responses"
import { isPdfFile, isTextFile } from "@/lib/pdf-utils"
import { type FileData, useFileProcessor } from "@/lib/use-file-processor"
import { useQuotaManager } from "@/lib/use-quota-manager"
import { formatXML, isMxCellXmlComplete, wrapWithMxFile } from "@/lib/utils"
import { formatXML } from "@/lib/utils"
import { ChatMessageDisplay } from "./chat-message-display"
import { DevXmlSimulator } from "./dev-xml-simulator"
@@ -214,9 +215,6 @@ export default function ChatPanel({
chartXMLRef.current = chartXML
}, [chartXML])
// Ref to hold stop function for use in onToolCall (avoids stale closure)
const stopRef = useRef<(() => void) | null>(null)
// Ref to track consecutive auto-retry count (reset on user action)
const autoRetryCountRef = useRef(0)
// Ref to track continuation retry count (for truncation handling)
@@ -239,6 +237,16 @@ export default function ChatPanel({
> | null>(null)
const LOCAL_STORAGE_DEBOUNCE_MS = 1000 // Save at most once per second
// Diagram tool handlers (display_diagram, edit_diagram, append_diagram)
const { handleToolCall } = useDiagramToolHandlers({
partialXmlRef,
editDiagramOriginalXmlRef,
chartXMLRef,
onDisplayChart,
onFetchChart,
onExport,
})
const {
messages,
sendMessage,
@@ -251,311 +259,8 @@ export default function ChatPanel({
transport: new DefaultChatTransport({
api: getApiEndpoint("/api/chat"),
}),
async onToolCall({ toolCall }) {
if (DEBUG) {
console.log(
`[onToolCall] Tool: ${toolCall.toolName}, CallId: ${toolCall.toolCallId}`,
)
}
if (toolCall.toolName === "display_diagram") {
const { xml } = toolCall.input as { xml: string }
// DEBUG: Log raw input to diagnose false truncation detection
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)
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! 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.",
)
}
}
} else if (toolCall.toolName === "edit_diagram") {
const { operations } = toolCall.input as {
operations: Array<{
type: "update" | "add" | "delete"
cell_id: string
new_xml?: string
}>
}
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,
)
}
} else if (toolCall.toolName === "append_diagram") {
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.`,
})
}
}
onToolCall: async ({ toolCall }) => {
await handleToolCall({ toolCall }, addToolOutput)
},
onError: (error) => {
// Handle server-side quota limit (429 response)
@@ -719,9 +424,6 @@ Continue from EXACTLY where you stopped.`,
},
})
// Update stopRef so onToolCall can access it
stopRef.current = stop
// Ref to track latest messages for unload persistence
const messagesRef = useRef(messages)
useEffect(() => {