"use client" import { useChat } from "@ai-sdk/react" import { DefaultChatTransport } from "ai" import { CheckCircle, PanelRightClose, PanelRightOpen, Settings, } from "lucide-react" import Image from "next/image" import Link from "next/link" import type React from "react" import { useEffect, useRef, useState } from "react" import { flushSync } from "react-dom" import { FaGithub } from "react-icons/fa" import { Toaster } from "sonner" import { ButtonWithTooltip } from "@/components/button-with-tooltip" import { ChatInput } from "@/components/chat-input" import { SettingsDialog, STORAGE_ACCESS_CODE_KEY, } from "@/components/settings-dialog" import { useDiagram } from "@/contexts/diagram-context" import { formatXML, validateMxCellStructure } from "@/lib/utils" import { ChatMessageDisplay } from "./chat-message-display" interface ChatPanelProps { isVisible: boolean onToggleVisibility: () => void drawioUi: "min" | "sketch" onToggleDrawioUi: () => void isMobile?: boolean onCloseProtectionChange?: (enabled: boolean) => void } export default function ChatPanel({ isVisible, onToggleVisibility, drawioUi, onToggleDrawioUi, isMobile = false, onCloseProtectionChange, }: ChatPanelProps) { const { loadDiagram: onDisplayChart, handleExport: onExport, handleExportWithoutHistory, resolverRef, chartXML, clearDiagram, } = useDiagram() const onFetchChart = (saveToHistory = true) => { return Promise.race([ new Promise((resolve) => { if (resolverRef && "current" in resolverRef) { resolverRef.current = resolve } if (saveToHistory) { onExport() } else { handleExportWithoutHistory() } }), new Promise((_, reject) => setTimeout( () => reject( new Error( "Chart export timed out after 10 seconds", ), ), 10000, ), ), ]) } const [files, setFiles] = useState([]) const [showHistory, setShowHistory] = useState(false) const [showSettingsDialog, setShowSettingsDialog] = useState(false) const [accessCodeRequired, setAccessCodeRequired] = useState(false) const [input, setInput] = useState("") // Check if access code is required on mount useEffect(() => { fetch("/api/config") .then((res) => res.json()) .then((data) => setAccessCodeRequired(data.accessCodeRequired)) .catch(() => setAccessCodeRequired(false)) }, []) // Generate a unique session ID for Langfuse tracing const [sessionId, setSessionId] = useState( () => `session-${Date.now()}-${Math.random().toString(36).slice(2, 9)}`, ) // Store XML snapshots for each user message (keyed by message index) const xmlSnapshotsRef = useRef>(new Map()) // Ref to track latest chartXML for use in callbacks (avoids stale closure) const chartXMLRef = useRef(chartXML) useEffect(() => { chartXMLRef.current = chartXML }, [chartXML]) const { messages, sendMessage, addToolResult, status, error, setMessages } = useChat({ transport: new DefaultChatTransport({ api: "/api/chat", }), async onToolCall({ toolCall }) { if (toolCall.toolName === "display_diagram") { const { xml } = toolCall.input as { xml: string } const validationError = validateMxCellStructure(xml) if (validationError) { addToolResult({ tool: "display_diagram", toolCallId: toolCall.toolCallId, output: validationError, }) } else { addToolResult({ tool: "display_diagram", toolCallId: toolCall.toolCallId, output: "Successfully displayed the diagram.", }) } } else if (toolCall.toolName === "edit_diagram") { const { edits } = toolCall.input as { edits: Array<{ search: string; replace: string }> } let currentXml = "" try { console.log("[edit_diagram] Starting...") // Use chartXML from ref directly - more reliable than export // especially on Vercel where DrawIO iframe may have latency issues // Using ref to avoid stale closure in callback const cachedXML = chartXMLRef.current if (cachedXML) { currentXml = cachedXML console.log( "[edit_diagram] Using cached chartXML, length:", currentXml.length, ) } else { // Fallback to export only if no cached XML console.log( "[edit_diagram] No cached XML, fetching from DrawIO...", ) currentXml = await onFetchChart(false) console.log( "[edit_diagram] Got XML from export, length:", currentXml.length, ) } const { replaceXMLParts } = await import("@/lib/utils") const editedXml = replaceXMLParts(currentXml, edits) onDisplayChart(editedXml) addToolResult({ tool: "edit_diagram", toolCallId: toolCall.toolCallId, output: `Successfully applied ${edits.length} edit(s) to the diagram.`, }) console.log("[edit_diagram] Success") } catch (error) { console.error("[edit_diagram] Failed:", error) const errorMessage = error instanceof Error ? error.message : String(error) addToolResult({ tool: "edit_diagram", toolCallId: toolCall.toolCallId, output: `Edit failed: ${errorMessage} Current diagram XML: \`\`\`xml ${currentXml || "No XML available"} \`\`\` Please retry with an adjusted search pattern or use display_diagram if retries are exhausted.`, }) } } }, onError: (error) => { // Silence access code error in console since it's handled by UI if (!error.message.includes("Invalid or missing access code")) { console.error("Chat error:", error) } // Add system message for error so it can be cleared setMessages((currentMessages) => { const errorMessage = { id: `error-${Date.now()}`, role: "system" as const, content: error.message, parts: [{ type: "text" as const, text: error.message }], } return [...currentMessages, errorMessage] }) if (error.message.includes("Invalid or missing access code")) { // Show settings button and open dialog to help user fix it setAccessCodeRequired(true) setShowSettingsDialog(true) } }, }) const messagesEndRef = useRef(null) useEffect(() => { if (messagesEndRef.current) { messagesEndRef.current.scrollIntoView({ behavior: "smooth" }) } }, [messages]) const onFormSubmit = async (e: React.FormEvent) => { e.preventDefault() const isProcessing = status === "streaming" || status === "submitted" if (input.trim() && !isProcessing) { try { let chartXml = await onFetchChart() chartXml = formatXML(chartXml) // Update ref directly to avoid race condition with React's async state update // This ensures edit_diagram has the correct XML before AI responds chartXMLRef.current = chartXml const parts: any[] = [{ type: "text", text: input }] if (files.length > 0) { for (const file of files) { const reader = new FileReader() const dataUrl = await new Promise((resolve) => { reader.onload = () => resolve(reader.result as string) reader.readAsDataURL(file) }) parts.push({ type: "file", url: dataUrl, mediaType: file.type, }) } } // Save XML snapshot for this message (will be at index = current messages.length) const messageIndex = messages.length xmlSnapshotsRef.current.set(messageIndex, chartXml) const accessCode = localStorage.getItem(STORAGE_ACCESS_CODE_KEY) || "" sendMessage( { parts }, { body: { xml: chartXml, sessionId, }, headers: { "x-access-code": accessCode, }, }, ) setInput("") setFiles([]) } catch (error) { console.error("Error fetching chart data:", error) } } } const handleInputChange = ( e: React.ChangeEvent, ) => { setInput(e.target.value) } const handleFileChange = (newFiles: File[]) => { setFiles(newFiles) } const handleRegenerate = async (messageIndex: number) => { const isProcessing = status === "streaming" || status === "submitted" if (isProcessing) return // Find the user message before this assistant message let userMessageIndex = messageIndex - 1 while ( userMessageIndex >= 0 && messages[userMessageIndex].role !== "user" ) { userMessageIndex-- } if (userMessageIndex < 0) return const userMessage = messages[userMessageIndex] const userParts = userMessage.parts // Get the text from the user message const textPart = userParts?.find((p: any) => p.type === "text") if (!textPart) return // Get the saved XML snapshot for this user message const savedXml = xmlSnapshotsRef.current.get(userMessageIndex) if (!savedXml) { console.error( "No saved XML snapshot for message index:", userMessageIndex, ) return } // Restore the diagram to the saved state onDisplayChart(savedXml) // Update ref directly to ensure edit_diagram has the correct XML chartXMLRef.current = savedXml // Clean up snapshots for messages after the user message (they will be removed) for (const key of xmlSnapshotsRef.current.keys()) { if (key > userMessageIndex) { xmlSnapshotsRef.current.delete(key) } } // Remove the user message AND assistant message onwards (sendMessage will re-add the user message) // Use flushSync to ensure state update is processed synchronously before sending const newMessages = messages.slice(0, userMessageIndex) flushSync(() => { setMessages(newMessages) }) // Now send the message after state is guaranteed to be updated sendMessage( { parts: userParts }, { body: { xml: savedXml, sessionId, }, }, ) } const handleEditMessage = async (messageIndex: number, newText: string) => { const isProcessing = status === "streaming" || status === "submitted" if (isProcessing) return const message = messages[messageIndex] if (!message || message.role !== "user") return // Get the saved XML snapshot for this user message const savedXml = xmlSnapshotsRef.current.get(messageIndex) if (!savedXml) { console.error( "No saved XML snapshot for message index:", messageIndex, ) return } // Restore the diagram to the saved state onDisplayChart(savedXml) // Update ref directly to ensure edit_diagram has the correct XML chartXMLRef.current = savedXml // Clean up snapshots for messages after the user message (they will be removed) for (const key of xmlSnapshotsRef.current.keys()) { if (key > messageIndex) { xmlSnapshotsRef.current.delete(key) } } // Create new parts with updated text const newParts = message.parts?.map((part: any) => { if (part.type === "text") { return { ...part, text: newText } } return part }) || [{ type: "text", text: newText }] // Remove the user message AND assistant message onwards (sendMessage will re-add the user message) // Use flushSync to ensure state update is processed synchronously before sending const newMessages = messages.slice(0, messageIndex) flushSync(() => { setMessages(newMessages) }) // Now send the edited message after state is guaranteed to be updated sendMessage( { parts: newParts }, { body: { xml: savedXml, sessionId, }, }, ) } // Collapsed view (desktop only) if (!isVisible && !isMobile) { return (
AI Chat
) } // Full view return (
{/* Header */}
Next AI Drawio

Next AI Drawio

{!isMobile && ( About )} {!isMobile && ( )}
setShowSettingsDialog(true)} className="hover:bg-accent" > {!isMobile && ( )}
{/* Messages */}
{/* Input */}
{ setMessages([]) clearDiagram() setSessionId( `session-${Date.now()}-${Math.random() .toString(36) .slice(2, 9)}`, ) xmlSnapshotsRef.current.clear() }} files={files} onFileChange={handleFileChange} showHistory={showHistory} onToggleHistory={setShowHistory} sessionId={sessionId} error={error} drawioUi={drawioUi} onToggleDrawioUi={onToggleDrawioUi} />
) }