import type { RefObject } from "react" import { useEffect, useRef } from "react" import { toast } from "sonner" import type { DiagramCommit, LoadMode } from "@/contexts/diagram-context" import { useDictionary } from "@/hooks/use-dictionary" import { countCells, diffDiagrams } from "@/lib/diagram-diff" import { clearHighlights, highlightCells } from "@/lib/drawio/editor-bridge" import { formatMessage } from "@/lib/i18n/utils" import { BLANK_MXFILE } from "@/packages/mcp-server/src/pages.ts" import { useVersionsStore } from "@/stores/versions-store" interface UseVersionsParams { loadDiagram: ( xml: string, skipValidation?: boolean, mode?: LoadMode, ) => string | null getVersionSvg: () => Promise setCommitHandler: ( handler: ((commit: DiagramCommit) => void) | null, ) => void /** Index of the user message whose turn is running */ currentTurnRef: RefObject /** The document on the canvas now */ chartXMLRef: RefObject isBusy: boolean /** Live busy state, for the restore toast's Undo that runs later */ isTurnActive: () => boolean /** Changes when another chat comes on screen */ getChatGeneration: () => number } /** * Turns every committed AI change into a version (card in the chat), marks * the changed shapes on the canvas, and restores or undoes versions. */ export function useVersions({ loadDiagram, getVersionSvg, setCommitHandler, currentTurnRef, chartXMLRef, isBusy, isTurnActive, getChatGeneration, }: UseVersionsParams) { const dict = useDictionary() // A restore commits too; that path must not add a version const restoringRef = useRef(false) // Registered once: re-registering on every render would leave a gap // (effect cleanups run before the streaming hook commits). The handler // only reads refs and stable functions. useEffect(() => { setCommitHandler(({ beforeXml, afterXml, toolCallId }) => { if (restoringRef.current) return const { summary, touchedIds } = diffDiagrams(beforeXml, afterXml) const before = countCells(beforeXml) const fromScratch = before.shapes + before.edges === 0 const id = useVersionsStore.getState().addVersion({ xml: afterXml, beforeXml, turnIndex: currentTurnRef.current, toolCallId, summary, fromScratch, }) // Mark what changed, unless the whole diagram is new if (!fromScratch) { setTimeout(() => { const marker = getComputedStyle(document.documentElement) .getPropertyValue("--marker") .trim() highlightCells(touchedIds, marker || "#ffd84d") }, 60) } // The picture is of the canvas: take it right away, before a // newer version can be on it; once one is, this one gets none // rather than the newer one's const isNewest = () => useVersionsStore.getState().versions.at(-1)?.id === id setTimeout(async () => { const svg = isNewest() ? await getVersionSvg() : null useVersionsStore .getState() .updateVersion(id, { svg: svg && isNewest() ? svg : "" }) }, 0) }) return () => setCommitHandler(null) }, []) const apply = (xml: string) => { restoringRef.current = true try { loadDiagram(xml, true, "commit") } finally { restoringRef.current = false } clearHighlights() } const restoreVersion = (versionId: string) => { if (isBusy) return const store = useVersionsStore.getState() const index = store.versions.findIndex((v) => v.id === versionId) if (index < 0) return // The canvas before the restore, to go back to it: a multi-page // document is loaded in full, which Ctrl+Z cannot undo const before = chartXMLRef.current const wasOnCanvas = store.onCanvasVersionId const wasUndone = store.undoneVersionId const generation = getChatGeneration() apply(store.versions[index].xml) store.setOnCanvas(versionId) toast.success( formatMessage(dict.versions.restoredVersion, { n: index + 1 }), { id: "version-restored", duration: 8000, action: { label: dict.versions.undoRestore, onClick: () => { // Another chat is on screen, or an answer is // changing the canvas: the old canvas is not theirs if (!before || isTurnActive()) return if (getChatGeneration() !== generation) return apply(before) const now = useVersionsStore.getState() if (wasUndone) now.setUndone(wasUndone) else if (wasOnCanvas) now.setOnCanvas(wasOnCanvas) }, }, }, ) } /** Undo the change of the latest version, or redo it when undone */ const undoVersion = (versionId: string) => { if (isBusy) return const store = useVersionsStore.getState() const version = store.versions.find((v) => v.id === versionId) if (!version) return if (store.undoneVersionId === versionId) { apply(version.xml) store.setOnCanvas(versionId) return } if (version.beforeXml === undefined) return apply(version.beforeXml || BLANK_MXFILE) store.setUndone(versionId) } return { restoreVersion, undoVersion } }