i/next-ai-draw-io
1
0
Fork 0
mirror of https://github.com/DayuanJiang/next-ai-draw-io.git synced 2026-10-12 20:49:53 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor(chat): drop redesign leftovers and old-format migrations

Browse Source 
- stores/versions-store.ts, hooks/use-session-manager.ts,
  lib/session-storage.ts: remove the diagramHistory conversion layer,
  the localStorage chat migration and the empty templates store
- stores/canvas-store.ts, lib/drawio/editor-bridge.ts: remove the
  pages list nothing read
- components/chat/chat-engine.tsx: use the context's chartXMLRef and
  latestSvgRef instead of local copies, subscribe to one settings field,
  stop prefetching quota limits (lib/use-quota-manager.tsx takes the
  values from the 429 body), readDiagramFile delegates to the MCP
  loader (which now strips XML prologs and accepts <svg:svg>)
- contexts/diagram-context.tsx: one requestTaggedExport path for
  thumbnails, validation, fetch and save; fetchDiagramXml replaces
  handleExport plus the exported resolver map
- hooks/use-validate-diagram.ts and app/api/validate-diagram: plain
  fetch returning JSON instead of a streamed object nobody consumed
- hooks/use-diagram-tool-handlers.ts: drop the DEBUG blocks; canvas
  vars inheritance lives in placeOnPage (lib/diagram-pages.ts)
- lobby-hero, icon-button, drawio-frame: remove unreachable branches
  and unused props
This commit is contained in:
dayuan.jiang
2026-10-12 08:17:11 +09:00
parent 3a9168e5b3
commit 06621764b2
27 changed files with 411 additions and 866 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/api/config/route.ts
-3
View File
@@ -3,8 +3,5 @@ import { NextResponse } from "next/server"
export async function GET() {
return NextResponse.json({
accessCodeRequired: !!process.env.ACCESS_CODE_LIST,
dailyRequestLimit: Number(process.env.DAILY_REQUEST_LIMIT) || 0,
dailyTokenLimit: Number(process.env.DAILY_TOKEN_LIMIT) || 0,
tpmLimit: Number(process.env.TPM_LIMIT) || 0,
})
}
app/api/validate-diagram/route.ts
+18 -30
View File
@@ -1,9 +1,9 @@
/**
* API endpoint for VLM-based diagram validation.
* Accepts a PNG image and streams validation results using useObject-compatible format.
* Accepts a PNG image and answers with the validation result as JSON.
*/
import { Output, streamText } from "ai"
import { generateText, Output } from "ai"
import { checkAccessCode, rejectCrossSite } from "@/lib/access-code"
import { getValidationModel } from "@/lib/ai-providers"
import {
@@ -35,13 +35,6 @@ const DEFAULT_VALID_RESULT: ValidationResult = {
suggestions: [],
}
/** A fixed result in the text format useObject reads */
function createStreamingResponse(result: ValidationResult): Response {
return new Response(JSON.stringify(result), {
headers: { "Content-Type": "text/plain; charset=utf-8" },
})
}
export async function POST(req: Request): Promise<Response> {
const crossSite = rejectCrossSite(req)
if (crossSite) return crossSite
@@ -53,7 +46,7 @@ export async function POST(req: Request): Promise<Response> {
// Check if VLM validation is enabled (default: true)
const enableValidation = process.env.ENABLE_VLM_VALIDATION !== "false"
if (!enableValidation) {
return createStreamingResponse(DEFAULT_VALID_RESULT)
return Response.json(DEFAULT_VALID_RESULT)
}
const body: ValidateDiagramRequest = await req.json()
@@ -123,7 +116,7 @@ export async function POST(req: Request): Promise<Response> {
error,
)
// Return valid if no vision model is configured
return createStreamingResponse(DEFAULT_VALID_RESULT)
return Response.json(DEFAULT_VALID_RESULT)
}
// Parse timeout with validation (minimum 1000ms, default 10000ms)
@@ -133,8 +126,7 @@ export async function POST(req: Request): Promise<Response> {
parseInt(process.env.VALIDATION_TIMEOUT || "10000", 10),
) || 10000
// Stream the VLM response for useObject consumption
const result = streamText({
const { output, totalUsage } = await generateText({
model,
output: Output.object({ schema: ValidationResultSchema }),
system: VALIDATION_SYSTEM_PROMPT,
@@ -155,23 +147,19 @@ export async function POST(req: Request): Promise<Response> {
],
maxOutputTokens: 1024,
abortSignal: AbortSignal.timeout(timeout),
onFinish: ({ output, totalUsage }) => {
if (countsQuota && totalUsage) {
recordTokenUsage(
userId,
(totalUsage.inputTokens || 0) +
(totalUsage.outputTokens || 0),
)
}
if (sessionId && output) {
console.log(
`[validate-diagram] Session ${sessionId}: valid=${output.valid}, issues=${output.issues?.length ?? 0}`,
)
}
},
})
return result.toTextStreamResponse()
if (countsQuota) {
await recordTokenUsage(
userId,
(totalUsage.inputTokens || 0) + (totalUsage.outputTokens || 0),
)
}
if (sessionId) {
console.log(
`[validate-diagram] Session ${sessionId}: valid=${output.valid}, issues=${output.issues?.length ?? 0}`,
)
}
return Response.json(output)
} catch (error) {
// Log with session context if available
const errorMessage =
@@ -179,6 +167,6 @@ export async function POST(req: Request): Promise<Response> {
console.error("[validate-diagram] Error:", errorMessage)
// On error, return valid to not block the user
return createStreamingResponse(DEFAULT_VALID_RESULT)
return Response.json(DEFAULT_VALID_RESULT)
}
}
components/canvas/canvas-stage.tsx
+1 -2
View File
@@ -136,8 +136,7 @@ export function CanvasStage({
key={src}
ref={drawioRef}
src={src}
configCss={DRAWIO_CSS}
editorCss={DRAWIO_CSS}
css={DRAWIO_CSS}
dark={isDark}
onLoad={onDrawioLoad}
onAutoSave={handleDiagramAutoSave}
components/canvas/drawio-frame.tsx
+11 -21
View File
@@ -30,10 +30,11 @@ export interface DrawioFrameHandle {
interface DrawioFrameProps {
src: string
/** CSS sent with the configure message (works cross-origin too) */
configCss: string
/** CSS injected once the editor is reachable directly (same origin) */
editorCss: string
/**
* Theme CSS: sent with the configure message (works cross-origin too)
* and, same origin, injected as a <style> so it updates without a reload
*/
css: string
/** draw.io's dark mode, switched in place when the editor is reachable */
dark: boolean
onLoad?: () => void
@@ -52,25 +53,14 @@ const STYLE_ID = "next-ai-drawio-theme"
*/
export const DrawioFrame = forwardRef<DrawioFrameHandle, DrawioFrameProps>(
function DrawioFrame(
{
src,
configCss,
editorCss,
dark,
onLoad,
onAutoSave,
onExport,
onEditorAccess,
},
{ src, css, dark, onLoad, onAutoSave, onExport, onEditorAccess },
ref,
) {
const iframeRef = useRef<HTMLIFrameElement>(null)
const editorUiRef = useRef<any>(null)
const attachedRef = useRef(false)
const configCssRef = useRef(configCss)
configCssRef.current = configCss
const editorCssRef = useRef(editorCss)
editorCssRef.current = editorCss
const cssRef = useRef(css)
cssRef.current = css
const darkRef = useRef(dark)
darkRef.current = dark
// Handlers change on every render; the message listener reads refs
@@ -111,7 +101,7 @@ export const DrawioFrame = forwardRef<DrawioFrameHandle, DrawioFrameProps>(
style.id = STYLE_ID
doc.head.appendChild(style)
}
style.textContent = editorCssRef.current
style.textContent = cssRef.current
} catch {
// Cross-origin: the configure message already applied it
}
@@ -119,7 +109,7 @@ export const DrawioFrame = forwardRef<DrawioFrameHandle, DrawioFrameProps>(
useEffect(() => {
applyThemeCss()
}, [editorCss])
}, [css])
// Same-origin only: draw.io's own dark mode, without a reload
const applyDarkMode = () => {
@@ -209,7 +199,7 @@ export const DrawioFrame = forwardRef<DrawioFrameHandle, DrawioFrameProps>(
send("configure", {
config: {
...DRAWIO_CONFIG,
css: configCssRef.current,
css: cssRef.current,
},
})
break
components/chat/chat-engine.tsx
+18 -106
View File
@@ -23,10 +23,7 @@ import { flushSync } from "react-dom"
import { toast } from "sonner"
import type { ValidationState } from "@/components/chat/ValidationCard"
import { useDiagram } from "@/contexts/diagram-context"
import {
keepFileVars,
useDiagramToolHandlers,
} from "@/hooks/use-diagram-tool-handlers"
import { useDiagramToolHandlers } from "@/hooks/use-diagram-tool-handlers"
import { useDictionary } from "@/hooks/use-dictionary"
import {
getSelectedAIConfig,
@@ -34,7 +31,6 @@ import {
useModelConfig,
} from "@/hooks/use-model-config"
import {
getSessionVersions,
type SessionData,
useSessionManager,
} from "@/hooks/use-session-manager"
@@ -121,31 +117,14 @@ function newLangfuseSessionId() {
return `session-${Date.now()}-${Math.random().toString(36).slice(2, 9)}`
}
const XML_PROLOG =
/^(?:\s|<\?[\s\S]*?\?>|<!--[\s\S]*?-->|<!DOCTYPE[^>[]*(?:\[[\s\S]*?\])?\s*>)*/i
/**
* Read the diagram document out of a .drawio, .xml or .drawio.svg file, with
* its pages inflated: previews and the model read it before draw.io saves
* it again. A broken or cut off file gives null.
*/
async function readDiagramFile(file: File): Promise<string | null> {
const raw = await file.text()
// What XML allows before the root element: the declaration, comments,
// a DOCTYPE (the parser below expects the root first)
let text = raw.replace(XML_PROLOG, "")
// Editable SVG: the diagram is stored in the root's `content` attribute.
// Parsed whole: its DOCTYPE may define entities the content uses.
if (/^<(?:[\w.-]+:)?svg[\s>]/i.test(text)) {
const doc = new DOMParser().parseFromString(raw, "image/svg+xml")
text = (doc.documentElement?.getAttribute("content") ?? "").replace(
XML_PROLOG,
"",
)
}
const loaded = parseDrawioFileContent(text)
if (!loaded.ok) return null
return withUniquePageIds(loaded.xml)
const loaded = parseDrawioFileContent(await file.text())
return loaded.ok ? withUniquePageIds(loaded.xml) : null
}
/**
@@ -267,11 +246,10 @@ export function ChatEngineProvider({
}) {
const {
loadDiagram: onDisplayChart,
handleExport: onExport,
exportResolversRef,
fetchDiagramXml,
chartXML,
chartXMLRef: liveChartXMLRef,
latestSvg,
chartXMLRef,
latestSvgRef,
clearDiagram,
getThumbnailSvg,
captureValidationPng,
@@ -284,26 +262,10 @@ export function ChatEngineProvider({
const searchParams = useSearchParams()
const urlSessionId = searchParams.get("session")
const settings = useSettingsStore()
const vlmValidationEnabled = useSettingsStore((s) => s.vlmValidationEnabled)
const openSettings = useUiStore((s) => s.openSettings)
const focusComposer = useUiStore((s) => s.focusComposer)
const onFetchChart = () => {
// Waits for the reply to its own export, by its tag
const tag = onExport()
return Promise.race([
new Promise<string>((resolve) => {
if (tag) exportResolversRef.current[tag] = resolve
}),
new Promise<string>((_, reject) => {
setTimeout(() => {
delete exportResolversRef.current[tag]
reject(new Error("Chart export timed out after 10 seconds"))
}, 10000)
}),
])
}
const { files, pdfData, handleFileChange, setFiles } = useFileProcessor({
tooLong: dict.errors.fileTooLong,
readFailed: dict.errors.fileReadFailed,
@@ -318,9 +280,6 @@ export function ChatEngineProvider({
const sessionManager = useSessionManager({ initialSessionId: urlSessionId })
const [input, setInputState] = useState("")
const [dailyRequestLimit, setDailyRequestLimit] = useState(0)
const [dailyTokenLimit, setDailyTokenLimit] = useState(0)
const [tpmLimit, setTpmLimit] = useState(0)
// Restore the unsent input from sessionStorage
useEffect(() => {
@@ -340,21 +299,7 @@ export function ChatEngineProvider({
else sessionStorage.removeItem(SESSION_STORAGE_INPUT_KEY)
}, [])
useEffect(() => {
fetch(getApiEndpoint("/api/config"))
.then((res) => res.json())
.then((data) => {
setDailyRequestLimit(data.dailyRequestLimit || 0)
setDailyTokenLimit(data.dailyTokenLimit || 0)
setTpmLimit(data.tpmLimit || 0)
})
.catch(() => {})
}, [])
const quotaManager = useQuotaManager({
dailyRequestLimit,
dailyTokenLimit,
tpmLimit,
onConfigModel: () => openSettings("models"),
})
@@ -392,20 +337,12 @@ export function ChatEngineProvider({
const hasRestoredRef = useRef(false)
const [isRestored, setIsRestored] = useState(false)
// Latest chartXML for callbacks (avoids stale closures)
const chartXMLRef = useRef(chartXML)
// Session loaded without a diagram (prevents thumbnail contamination)
const justLoadedSessionIdRef = useRef<string | null>(null)
useEffect(() => {
chartXMLRef.current = chartXML
if (chartXML) justLoadedSessionIdRef.current = null
}, [chartXML])
const latestSvgRef = useRef(latestSvg)
useEffect(() => {
latestSvgRef.current = latestSvg
}, [latestSvg])
const autoRetryCountRef = useRef(0)
const continuationRetryCountRef = useRef(0)
// Partial XML when output was truncated; non-empty means continuation mode
@@ -458,13 +395,13 @@ export function ChatEngineProvider({
processedToolCallsRef,
validationRetryCountRef,
// A preview undone just before the tool call is in this one already
chartXMLRef: liveChartXMLRef,
chartXMLRef,
turnPageIdRef,
onDisplayChart,
onFetchChart,
onFetchChart: fetchDiagramXml,
captureValidationPng,
validateDiagram: validateWithFallback,
enableVlmValidation: settings.vlmValidationEnabled,
enableVlmValidation: vlmValidationEnabled,
sessionId: langfuseSessionId,
watchStop: () => {
const stopsBefore = stopCountRef.current
@@ -832,13 +769,7 @@ export function ChatEngineProvider({
const messagesBeforeShownRef = useRef<UIMessage[] | null>(null)
const syncUIWithSession = useCallback(
(
data:
| (SessionData & {
diagramHistory?: { svg: string; xml: string }[]
})
| null,
) => {
(data: SessionData | null) => {
// An answer still coming in belongs to the chat that was on
// screen (browser back or forward): it stops, before this chat's
// messages are put in its place, so none of it lands here
@@ -877,24 +808,15 @@ export function ChatEngineProvider({
)
setMessages(data.messages as unknown as UIMessage[])
xmlSnapshotsRef.current = new Map(data.xmlSnapshots)
if (data.diagramXml) {
onDisplayChart(data.diagramXml, true)
chartXMLRef.current = data.diagramXml
} else {
clearDiagram()
chartXMLRef.current = ""
}
if (data.diagramXml) onDisplayChart(data.diagramXml, true)
else clearDiagram()
if (!isRealDiagram(data.diagramXml)) latestSvgRef.current = ""
useVersionsStore
.getState()
.setVersions(getSessionVersions(data))
useVersionsStore.getState().setVersions(data.versions ?? [])
} else {
loadedMessageIdsRef.current = new Set()
setMessages([])
xmlSnapshotsRef.current.clear()
clearDiagram()
chartXMLRef.current = ""
latestSvgRef.current = ""
useVersionsStore.getState().clear()
}
},
@@ -927,7 +849,6 @@ export function ChatEngineProvider({
if (isRealDiagram(currentDiagramXml) && options.withThumbnail) {
const freshThumb = await getThumbnailSvg()
if (freshThumb) {
latestSvgRef.current = freshThumb
thumbnailDataUrl = freshThumb
} else if (latestSvgRef.current) {
thumbnailDataUrl = latestSvgRef.current
@@ -1256,7 +1177,7 @@ export function ChatEngineProvider({
selectedCells: MessageSelection = [],
generation = getChatGeneration(),
): Promise<boolean> => {
const chartXml = formatXML(await onFetchChart())
const chartXml = formatXML(await fetchDiagramXml())
if (getChatGeneration() !== generation || leavingRef.current) {
return false
}
@@ -1339,12 +1260,8 @@ export function ChatEngineProvider({
})
if (prepared.ok) {
const pageId = viewedPageId()
const canvasXml = chartXMLRef.current || ""
onDisplayChart(
keepFileVars(
placeOnPage(prepared.xml, canvasXml, pageId),
canvasXml,
),
placeOnPage(prepared.xml, chartXMLRef.current, pageId),
true,
"commit",
{ toolCallId, pageId },
@@ -1473,11 +1390,6 @@ export function ChatEngineProvider({
}
}
const restoreDiagramFromSnapshot = (savedXml: string) => {
onDisplayChart(savedXml, true)
chartXMLRef.current = savedXml
}
// Snapshots of the messages after this one go with them
const dropLaterSnapshots = (messageIndex: number) => {
for (const key of xmlSnapshotsRef.current.keys()) {
@@ -1495,7 +1407,7 @@ export function ChatEngineProvider({
return false
}
const previousXml = getPreviousXml(index)
restoreDiagramFromSnapshot(savedXml)
onDisplayChart(savedXml, true)
dropLaterSnapshots(index)
useVersionsStore.getState().removeFromTurn(index)
@@ -1563,7 +1475,7 @@ export function ChatEngineProvider({
// Like submit: the diagram export comes first
preparingSendRef.current = true
try {
const chartXml = formatXML(await onFetchChart())
const chartXml = formatXML(await fetchDiagramXml())
// Another chat came on screen meanwhile (browser back/forward),
// or this one is being left
if (getChatGeneration() !== generation || leavingRef.current) {
components/chat/lobby-hero.tsx
+9 -22
View File
@@ -1,4 +1,4 @@
import { FileText, FileUp, History, Pencil } from "lucide-react"
import { FileUp, History, Pencil } from "lucide-react"
import type React from "react"
import { useDictionary } from "@/hooks/use-dictionary"
import { getAssetUrl } from "@/lib/base-path"
@@ -12,7 +12,7 @@ interface Example {
title: string
description: string
/** Picture of the result (or of the input, for "replicate" examples) */
image?: string
image: string
prompt: string
/** Bundled file sent with the prompt; matches the cached example */
file?: { path: string; name: string; type: string }
@@ -34,7 +34,6 @@ function ExampleTile({
title,
description,
picture,
plain,
testId,
}: {
onClick: () => void
@@ -42,8 +41,6 @@ function ExampleTile({
/** Shown as the tooltip: the title alone names the example */
description: string
picture: React.ReactNode
/** An icon on a tinted tile instead of a diagram on paper */
plain?: boolean
testId?: string
}) {
return (
@@ -54,12 +51,7 @@ function ExampleTile({
className="group/example flex min-w-0 flex-col gap-2 rounded-xl border border-border bg-card p-1.5 pb-2.5 text-left transition-colors hover:border-foreground/25"
data-testid={testId}
>
<span
className={cn(
"flex h-16 w-full items-center justify-center overflow-hidden rounded-lg",
plain ? "bg-muted" : "sheet-light bg-white",
)}
>
<span className="flex h-16 w-full items-center justify-center overflow-hidden rounded-lg sheet-light bg-white">
{picture}
</span>
<span className="px-1 text-[13px] leading-snug font-medium text-foreground">
@@ -220,18 +212,13 @@ export function LobbyHero({
description={example.description}
testId={`example-${example.key}`}
picture={
example.image ? (
// biome-ignore lint/performance/noImgElement: static preview
<img
src={getAssetUrl(example.image)}
alt=""
className="max-h-14 max-w-[88%] object-contain transition-transform duration-300 group-hover/example:scale-[1.06]"
/>
) : (
<FileText className="size-5 text-muted-foreground" />
)
// biome-ignore lint/performance/noImgElement: static preview
<img
src={getAssetUrl(example.image)}
alt=""
className="max-h-14 max-w-[88%] object-contain transition-transform duration-300 group-hover/example:scale-[1.06]"
/>
}
plain={!example.image}
/>
),
)}
components/workspace/icon-button.tsx
+1 -15
View File
@@ -13,9 +13,7 @@ interface IconButtonProps extends React.ComponentProps<"button"> {
label: string
/** Shown next to the label in the tooltip, e.g. "⌘B" */
shortcut?: string
active?: boolean
tooltipSide?: "top" | "bottom" | "left" | "right"
size?: "sm" | "md"
}
/** Square icon button with a tooltip; the tooltip text is also its label */
@@ -24,9 +22,7 @@ export const IconButton = forwardRef<HTMLButtonElement, IconButtonProps>(
{
label,
shortcut,
active = false,
tooltipSide = "bottom",
size = "md",
className,
children,
...props
@@ -40,17 +36,12 @@ export const IconButton = forwardRef<HTMLButtonElement, IconButtonProps>(
ref={ref}
type="button"
aria-label={label}
aria-pressed={active || undefined}
className={cn(
"inline-flex shrink-0 items-center justify-center rounded-lg text-muted-foreground transition-colors",
"inline-flex size-8 shrink-0 items-center justify-center rounded-lg text-muted-foreground transition-colors",
"hover:bg-accent hover:text-foreground",
"focus-visible:outline-2 focus-visible:outline-offset-1 focus-visible:outline-ring",
"disabled:pointer-events-none disabled:opacity-35",
"[&_svg]:size-4 [&_svg]:stroke-[1.75]",
size === "md" ? "size-8" : "size-7",
// A tinted tile: solid ink is for the main action
active &&
"bg-interactive-active text-foreground hover:bg-interactive-active [&_svg]:stroke-2",
className,
)}
{...props}
@@ -73,8 +64,3 @@ export const IconButton = forwardRef<HTMLButtonElement, IconButtonProps>(
)
},
)
/** Thin vertical divider between toolbar groups */
export function ToolbarDivider() {
return <span className="mx-1 h-5 w-px shrink-0 bg-border" aria-hidden />
}
contexts/diagram-context.tsx
+119 -144
View File
@@ -64,20 +64,16 @@ interface DiagramContextType {
chartXML: string
// chartXML right away, before the re-render (loadDiagram sets both)
chartXMLRef: React.MutableRefObject<string>
latestSvg: string
// The diagram's latest SVG export, a fallback thumbnail for saves
latestSvgRef: React.MutableRefObject<string>
loadDiagram: (
chart: string,
skipValidation?: boolean,
mode?: LoadMode,
meta?: { toolCallId?: string; pageId?: string | null },
) => string | null
// Returns the export's tag (empty when draw.io is not there yet)
handleExport: () => string
// Pending exports by tag; a plain export's resolver gets the XML of the
// page on screen
exportResolversRef: React.MutableRefObject<
Record<string, (data: string, xml?: string) => void>
>
/** The XML of the page on screen, exported from draw.io */
fetchDiagramXml: () => Promise<string>
drawioRef: React.MutableRefObject<DrawioFrameHandle | null>
handleDiagramExport: (data: DrawioExportEvent) => void
handleDiagramAutoSave: (data: { xml?: string }) => void
@@ -87,7 +83,7 @@ interface DiagramContextType {
format: ExportFormat,
sessionId?: string,
successMessage?: string,
) => void
) => Promise<void>
getThumbnailSvg: () => Promise<string | null>
getVersionSvg: () => Promise<string | null>
captureValidationPng: () => Promise<string | null>
@@ -108,11 +104,17 @@ const DiagramContext = createContext<DiagramContextType | undefined>(undefined)
// echoes the request back in the export event, so each result reaches its
// own caller. Tags end in a request number, so a late result never answers
// a newer request.
type ExportTag = "thumbnail" | "validation" | "version"
type ExportTag = "thumbnail" | "validation" | "version" | "fetch" | "save"
/** What draw.io answers: the data URL asked for, and the whole document */
interface ExportResult {
data: string
xml: string
}
export function DiagramProvider({ children }: { children: React.ReactNode }) {
const [chartXML, setChartXML] = useState<string>("")
const [latestSvg, setLatestSvg] = useState<string>("")
const latestSvgRef = useRef("")
const [isDrawioReady, setIsDrawioReady] = useState(false)
const hasCalledOnLoadRef = useRef(false)
const mountedRef = useRef(true)
@@ -125,7 +127,7 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
const drawioRef = useRef<DrawioFrameHandle | null>(null)
// Pending exports, keyed by their export tag
const exportResolversRef = useRef<
Record<string, (data: string, xml?: string) => void>
Record<string, (result: ExportResult) => void>
>({})
const exportSeqRef = useRef(0)
// Track latest chartXML for restoration after remount
@@ -184,26 +186,15 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
setChartXML(xml)
}
const handleExport = () => {
if (!drawioRef.current) return ""
// Export of the current state (for the chat and edit_diagram)
const tag = `fetch-${++exportSeqRef.current}`
drawioRef.current.exportDiagram({
format: "xmlsvg",
message: tag,
})
return tag
}
// Export with a tag in `message` (draw.io echoes it back in the export
// event) and wait for that result. Resolves to null on timeout, which is
// expected occasionally.
// (Reads refs only, so it keeps one identity)
const requestTaggedExport = useCallback(
(tag: ExportTag, format: "xmlsvg" | "svg" | "png", timeoutMs: number) =>
new Promise<string | null>((resolve) => {
new Promise<ExportResult | null>((resolve) => {
const id = `${tag}-${++exportSeqRef.current}`
const finish = (value: string | null) => {
const finish = (value: ExportResult | null) => {
clearTimeout(timer)
delete exportResolversRef.current[id]
resolve(value)
@@ -216,20 +207,18 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
)
// Get current diagram as SVG for thumbnail (used by session storage).
// One identity: the chat's auto-save depends on it, and each thumbnail
// renders this provider again (latestSvg), which would otherwise start
// the next save
// One identity: the chat's auto-save effect depends on it
const getThumbnailSvg = useCallback(async (): Promise<string | null> => {
if (!drawioRef.current) return null
// Don't export if diagram is empty
if (!isRealDiagram(chartXMLRef.current)) return null
// xmlsvg exports return an SVG data URL
const svgData = await requestTaggedExport("thumbnail", "xmlsvg", 3000)
if (svgData?.startsWith("data:image/svg")) {
// Update latestSvg so it's available for future saves
setLatestSvg(svgData)
return svgData
const result = await requestTaggedExport("thumbnail", "xmlsvg", 3000)
if (result?.data.startsWith("data:image/svg")) {
// Kept for later saves that take no fresh thumbnail
latestSvgRef.current = result.data
return result.data
}
return null
}, [requestTaggedExport])
@@ -237,8 +226,8 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
// Plain SVG (no embedded diagram data) for version cards
const getVersionSvg = async (): Promise<string | null> => {
if (!drawioRef.current) return null
const svgData = await requestTaggedExport("version", "svg", 4000)
return svgData?.startsWith("data:image/svg") ? svgData : null
const result = await requestTaggedExport("version", "svg", 4000)
return result?.data.startsWith("data:image/svg") ? result.data : null
}
// Capture current diagram as PNG for VLM validation
@@ -247,9 +236,25 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
// Don't export if diagram is empty
if (!isRealDiagram(chartXMLRef.current)) return null
const pngData = await requestTaggedExport("validation", "png", 5000)
const result = await requestTaggedExport("validation", "png", 5000)
// PNG data should be a base64 data URL
return pngData?.startsWith("data:image/png") ? pngData : null
return result?.data.startsWith("data:image/png") ? result.data : null
}
// The page the user is viewing (the first one with an external draw.io,
// which cannot tell us), for the chat and edit_diagram.
// Not written to chartXML: exports don't change the diagram, and the
// export's xml has compressed <diagram> payloads that would break
// edit_diagram/display_diagram. Autosave keeps chartXML up to date with
// the full uncompressed multi-page document (#879).
const fetchDiagramXml = async (): Promise<string> => {
const result = await requestTaggedExport("fetch", "xmlsvg", 10000)
if (!result) throw new Error("Chart export timed out after 10 seconds")
latestSvgRef.current = result.data
return extractDiagramXML(
result.data,
useCanvasStore.getState().currentPageId,
)
}
const loadDiagram = (
@@ -335,34 +340,11 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
resetPreview()
}
// Each export result goes to the caller that asked for it, by its tag
const handleDiagramExport = (data: DrawioExportEvent) => {
// Thumbnail, version, validation PNG and file save exports go only
// to their own caller
const tag = data.message?.message
if (/^(thumbnail|validation|version|save)-/.test(tag ?? "")) {
exportResolversRef.current[tag as string]?.(data.data, data.xml)
return
}
// Don't write chartXML here: exports don't change the diagram, and
// data.xml from xmlsvg exports has compressed <diagram> payloads that
// would break edit_diagram/display_diagram. Autosave keeps chartXML
// up to date with the full uncompressed multi-page document (#879).
// The chat gets the page the user is viewing (the first one with
// an external draw.io, which cannot tell us)
const extractedXML = extractDiagramXML(
data.data,
useCanvasStore.getState().currentPageId,
)
setLatestSvg(data.data)
// The chat's own export (onFetchChart), not another one in flight
const resolve =
tag !== undefined ? exportResolversRef.current[tag] : undefined
if (resolve) {
delete exportResolversRef.current[tag as string]
resolve(extractedXML)
}
if (!tag) return
exportResolversRef.current[tag]?.({ data: data.data, xml: data.xml })
}
// The frame registers this callback once per iframe mount, so it must
@@ -383,10 +365,10 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
const emptyDiagram = BLANK_MXFILE
// Skip validation for trusted internal template (loadDiagram also sets chartXML)
loadDiagram(emptyDiagram, true)
setLatestSvg("")
latestSvgRef.current = ""
}
const saveDiagramToFile = (
const saveDiagramToFile = async (
filename: string,
format: ExportFormat,
sessionId?: string,
@@ -401,85 +383,79 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
const drawioFormat =
format === "drawio" || format === "xmlsvg" ? "xmlsvg" : format
// Each save has its own tag, so two at once never swap results
const tag = `save-${++exportSeqRef.current}`
exportResolversRef.current[tag] = (
exportData: string,
fullDiagramXML?: string,
) => {
delete exportResolversRef.current[tag]
let fileContent: string | Blob
let mimeType: string
let extension: string
// Each save has its own tag, so two at once never swap results. A
// large PNG takes a while to render.
const result = await requestTaggedExport("save", drawioFormat, 30000)
if (!result) {
console.warn("Diagram export for saving timed out")
return
}
const { data: exportData, xml: fullDiagramXML } = result
let fileContent: string | Blob
let mimeType: string
let extension: string
if (format === "drawio") {
// Prefer the complete document from the export event so all pages are saved.
const xml = fullDiagramXML?.trim()
? fullDiagramXML
: extractDiagramXML(exportData)
fileContent =
normalizeToMxfile(xml, {
pageId: "page-1",
pageName: "Page-1",
}) ?? xml
mimeType = "application/xml"
extension = ".drawio"
} else if (format === "png") {
// PNG data comes as base64 data URL
fileContent = exportData
mimeType = "image/png"
extension = ".png"
} else if (format === "xmlsvg") {
// Editable SVG: pass data URL directly (like PNG)
fileContent = exportData
mimeType = "image/svg+xml"
extension = ".drawio.svg"
} else {
// SVG format (view-only)
fileContent = exportData
mimeType = "image/svg+xml"
extension = ".svg"
}
// Log save event to Langfuse (flags the trace)
logSaveToLangfuse(filename, format, sessionId)
// Handle download
let url: string
if (
typeof fileContent === "string" &&
fileContent.startsWith("data:")
) {
// Already a data URL (PNG)
url = fileContent
} else {
const blob = new Blob([fileContent], { type: mimeType })
url = URL.createObjectURL(blob)
}
const a = document.createElement("a")
a.href = url
a.download = `${filename}${extension}`
document.body.appendChild(a)
a.click()
document.body.removeChild(a)
// Show success toast after download is initiated
if (successMessage) {
toast.success(successMessage, { duration: 2500 })
}
// Delay URL revocation to ensure download completes
if (!url.startsWith("data:")) {
setTimeout(() => URL.revokeObjectURL(url), 100)
}
if (format === "drawio") {
// Prefer the complete document from the export event so all pages are saved.
const xml = fullDiagramXML?.trim()
? fullDiagramXML
: extractDiagramXML(exportData)
fileContent =
normalizeToMxfile(xml, {
pageId: "page-1",
pageName: "Page-1",
}) ?? xml
mimeType = "application/xml"
extension = ".drawio"
} else if (format === "png") {
// PNG data comes as base64 data URL
fileContent = exportData
mimeType = "image/png"
extension = ".png"
} else if (format === "xmlsvg") {
// Editable SVG: pass data URL directly (like PNG)
fileContent = exportData
mimeType = "image/svg+xml"
extension = ".drawio.svg"
} else {
// SVG format (view-only)
fileContent = exportData
mimeType = "image/svg+xml"
extension = ".svg"
}
// Export diagram - callback will be handled in handleDiagramExport
drawioRef.current.exportDiagram({
format: drawioFormat,
message: tag,
})
// Log save event to Langfuse (flags the trace)
logSaveToLangfuse(filename, format, sessionId)
// Handle download
let url: string
if (
typeof fileContent === "string" &&
fileContent.startsWith("data:")
) {
// Already a data URL (PNG)
url = fileContent
} else {
const blob = new Blob([fileContent], { type: mimeType })
url = URL.createObjectURL(blob)
}
const a = document.createElement("a")
a.href = url
a.download = `${filename}${extension}`
document.body.appendChild(a)
a.click()
document.body.removeChild(a)
// Show success toast after download is initiated
if (successMessage) {
toast.success(successMessage, { duration: 2500 })
}
// Delay URL revocation to ensure download completes
if (!url.startsWith("data:")) {
setTimeout(() => URL.revokeObjectURL(url), 100)
}
}
// Log save event to Langfuse (just flags the trace, doesn't send content)
@@ -510,10 +486,9 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
value={{
chartXML,
chartXMLRef,
latestSvg,
latestSvgRef,
loadDiagram,
handleExport,
exportResolversRef,
fetchDiagramXml,
drawioRef,
handleDiagramExport,
handleDiagramAutoSave,
hooks/use-diagram-tool-handlers.ts
+2 -89
View File
@@ -14,32 +14,10 @@ import { editDiagram } from "@/packages/mcp-server/src/edit-diagram.ts"
import { prepareNewDiagram } from "@/packages/mcp-server/src/new-diagram.ts"
import { hasCells } from "@/packages/mcp-server/src/pages.ts"
const DEBUG = process.env.NODE_ENV === "development"
// The page a drawing is wrapped in; on the canvas it replaces the page the
// model works on, so the canvas page keeps its own id and name
const NEW_PAGE = { pageId: "page-1", pageName: "Page-1" }
/**
* A new diagram written without file variables (%name% placeholders) keeps
* the canvas file's, as replacing the page in draw.io does
*/
export function keepFileVars(xml: string, canvasXml: string): string {
const varsOf = (doc: Document) =>
doc.documentElement?.nodeName === "mxfile"
? doc.documentElement.getAttribute("vars")
: null
const parser = new DOMParser()
const canvasVars = varsOf(parser.parseFromString(canvasXml, "text/xml"))
if (!canvasVars) return xml
const doc = parser.parseFromString(xml, "text/xml")
if (doc.documentElement?.nodeName !== "mxfile" || varsOf(doc) !== null) {
return xml
}
doc.documentElement.setAttribute("vars", canvasVars)
return new XMLSerializer().serializeToString(doc)
}
interface ToolCall {
toolCallId: string
toolName: string
@@ -157,12 +135,6 @@ export function useDiagramToolHandlers({
{ toolCall }: { toolCall: ToolCall },
addToolOutput: AddToolOutputFn,
) => {
if (DEBUG) {
console.log(
`[onToolCall] Tool: ${toolCall.toolName}, CallId: ${toolCall.toolCallId}`,
)
}
processedToolCallsRef.current.add(toolCall.toolCallId)
// Only display_diagram, edit_diagram and a completing append_diagram
// put their result on the canvas. Other tools (get_shape_library,
@@ -204,12 +176,9 @@ export function useDiagramToolHandlers({
// A drawn page goes onto the model's page of the canvas file, which
// keeps its other pages; a drawing that brings several pages replaces
// the file, and gets its variables
// the file. Either way it keeps the file's variables unless it sets its own.
const onAiPage = (drawnXml: string) =>
keepFileVars(
placeOnPage(drawnXml, chartXMLRef.current, aiPageId()),
chartXMLRef.current,
)
placeOnPage(drawnXml, chartXMLRef.current, aiPageId())
const takeOriginals = (): string | undefined => {
const [originalXml] = editDiagramOriginalXmlRef.current.values()
@@ -228,20 +197,8 @@ export function useDiagramToolHandlers({
) => {
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
@@ -298,11 +255,6 @@ NEXT STEP: Call append_diagram with the continuation XML.
// keeps the diagram from before this failed call
if (originalXml) onDisplayChart(originalXml, true, "revert")
// 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,
@@ -318,12 +270,6 @@ ${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 &&
@@ -348,12 +294,6 @@ ${finalXml}
if (stopped?.()) {
updateValidationState(toolCall.toolCallId, "skipped")
} else if (capturedPngData) {
if (DEBUG) {
console.log(
"[display_diagram] Captured PNG for validation",
)
}
// Each retry is a new tool call, so count attempts
// per user turn (the chat resets it when the user sends)
const attempt = validationRetryCountRef.current + 1
@@ -379,12 +319,6 @@ ${finalXml}
if (attempt < MAX_VALIDATION_RETRIES) {
const feedback =
formatValidationFeedback(result)
if (DEBUG) {
console.log(
`[display_diagram] Validation failed (attempt ${attempt}/${MAX_VALIDATION_RETRIES}):`,
result.issues,
)
}
// Notify UI of validation failure (include the image)
updateValidationState(
@@ -407,11 +341,6 @@ ${finalXml}
return
} else {
// Last attempt - accept the diagram with warning
if (DEBUG) {
console.log(
"[display_diagram] Max validation retries reached, accepting diagram",
)
}
// Notify UI that we're accepting with issues (include the image)
updateValidationState(
toolCall.toolCallId,
@@ -427,12 +356,6 @@ ${finalXml}
return
}
} else {
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
@@ -474,21 +397,11 @@ ${finalXml}
}
}
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.",
)
}
}
}
hooks/use-session-manager.ts
+3 -19
View File
@@ -12,7 +12,6 @@ import {
getAllSessionMetadata,
getSession,
isIndexedDBAvailable,
migrateFromLocalStorage,
readSessionCount,
renameSession as renameSessionInDB,
type SessionMetadata,
@@ -20,10 +19,7 @@ import {
saveSession,
} from "@/lib/session-storage"
import { STORAGE_KEYS } from "@/lib/storage"
import {
type DiagramVersion,
versionsFromLegacyHistory,
} from "@/stores/versions-store"
import type { DiagramVersion } from "@/stores/versions-store"
export interface SessionData {
messages: StoredMessage[]
@@ -35,13 +31,6 @@ export interface SessionData {
title?: string
}
/** Versions of a stored session, converting the old history format */
export function getSessionVersions(
session: Pick<ChatSession, "versions" | "diagramHistory">,
): DiagramVersion[] {
return session.versions ?? versionsFromLegacyHistory(session.diagramHistory)
}
// Taken right before a save's data is read: the chat on screen then, and
// the order of the reads
export interface SaveTicket {
@@ -192,9 +181,6 @@ export function useSessionManager(
setIsAvailable(true)
try {
// Run migration first (one-time conversion from localStorage)
await migrateFromLocalStorage()
// Load sessions list
const metadata = await getAllSessionMetadata()
setSessions(metadata)
@@ -334,7 +320,7 @@ export function useSessionManager(
xmlSnapshots: session.xmlSnapshots,
diagramXml: session.diagramXml,
thumbnailDataUrl: session.thumbnailDataUrl,
versions: getSessionVersions(session),
versions: session.versions ?? [],
}
},
[currentSessionId, changeChat],
@@ -417,9 +403,7 @@ export function useSessionManager(
thumbnailDataUrl:
data.thumbnailDataUrl ??
currentSession.thumbnailDataUrl,
versions:
data.versions ?? getSessionVersions(currentSession),
diagramHistory: undefined,
versions: data.versions ?? currentSession.versions ?? [],
updatedAt: Date.now(),
// Update title if it's still default and we have messages
title:
hooks/use-validate-diagram.ts
+42 -112
View File
@@ -1,10 +1,9 @@
"use client"
/**
* Hook for VLM-based diagram validation using AI SDK's useObject.
* Hook for VLM-based diagram validation: one request, one JSON result.
*/
import { experimental_useObject as useObject } from "@ai-sdk/react"
import { useCallback, useRef } from "react"
import { getSelectedAIConfig } from "@/hooks/use-model-config"
import { getApiEndpoint } from "@/lib/base-path"
@@ -13,8 +12,6 @@ import {
ValidationResultSchema,
} from "@/lib/validation-schema"
export type { ValidationResult }
// Default valid result for fallback cases
const DEFAULT_VALID_RESULT: ValidationResult = {
valid: true,
@@ -22,111 +19,42 @@ const DEFAULT_VALID_RESULT: ValidationResult = {
suggestions: [],
}
interface UseValidateDiagramOptions {
onSuccess?: (result: ValidationResult) => void
onError?: (error: Error) => void
}
// Track pending validation promises for imperative API
type PendingValidation = {
resolve: (result: ValidationResult) => void
reject: (error: Error) => void
}
export function useValidateDiagram(options: UseValidateDiagramOptions = {}) {
const { onSuccess, onError } = options
const pendingValidationRef = useRef<PendingValidation | null>(null)
const { object, submit, isLoading, error, stop } = useObject({
api: getApiEndpoint("/api/validate-diagram"),
schema: ValidationResultSchema,
// Resolved per request so a changed access code is picked up
headers: () => ({ "x-access-code": getSelectedAIConfig().accessCode }),
onFinish: ({
object,
error: finishError,
}: {
object: ValidationResult | undefined
error: Error | undefined
}) => {
if (finishError) {
console.error(
"[useValidateDiagram] Validation error:",
finishError,
)
onError?.(finishError)
pendingValidationRef.current?.reject(finishError)
pendingValidationRef.current = null
return
}
if (object) {
const result = object as ValidationResult
onSuccess?.(result)
pendingValidationRef.current?.resolve(result)
pendingValidationRef.current = null
}
},
onError: (err: Error) => {
console.error("[useValidateDiagram] Stream error:", err)
onError?.(err)
pendingValidationRef.current?.reject(err)
pendingValidationRef.current = null
},
})
export function useValidateDiagram() {
// The running check's controller; a new check or cancel() aborts it
const controllerRef = useRef<AbortController | null>(null)
/**
* Validate a diagram image.
* Returns a promise that resolves with the validation result.
*/
const validate = useCallback(
async (
imageData: string,
sessionId?: string,
): Promise<ValidationResult> => {
// Reject any pending validation to prevent promise leaks
if (pendingValidationRef.current) {
pendingValidationRef.current.reject(
new Error("Validation superseded by new request"),
)
pendingValidationRef.current = null
}
return new Promise((resolve, reject) => {
// Store the promise handlers
pendingValidationRef.current = { resolve, reject }
// Submit the validation request
submit({ imageData, sessionId })
})
},
[submit],
)
/**
* End a running check (the user pressed Stop): its promise rejects with
* an AbortError, so the tool handler can finish at once.
*/
const cancel = useCallback(() => {
const pending = pendingValidationRef.current
if (!pending) return
pendingValidationRef.current = null
stop()
pending.reject(new DOMException("Validation cancelled", "AbortError"))
}, [stop])
/**
* Validate with fallback - returns default valid result on error.
* Use this to avoid blocking the user on validation failures.
* A cancelled check is passed on as its AbortError.
* Validate a diagram image, with fallback: a failed check gives the
* default valid result so it never blocks the user. A cancelled check
* is passed on as its AbortError.
*/
const validateWithFallback = useCallback(
async (
imageData: string,
sessionId?: string,
): Promise<ValidationResult> => {
controllerRef.current?.abort()
const controller = new AbortController()
controllerRef.current = controller
try {
return await validate(imageData, sessionId)
const res = await fetch(
getApiEndpoint("/api/validate-diagram"),
{
method: "POST",
headers: {
"Content-Type": "application/json",
// Read per request so a changed access code is
// picked up
"x-access-code": getSelectedAIConfig().accessCode,
},
body: JSON.stringify({ imageData, sessionId }),
signal: controller.signal,
},
)
if (!res.ok) {
throw new Error(`Validation request failed: ${res.status}`)
}
return ValidationResultSchema.parse(await res.json())
} catch (error) {
if ((error as Error)?.name === "AbortError") throw error
console.warn(
@@ -134,21 +62,23 @@ export function useValidateDiagram(options: UseValidateDiagramOptions = {}) {
error,
)
return DEFAULT_VALID_RESULT
} finally {
if (controllerRef.current === controller) {
controllerRef.current = null
}
}
},
[validate],
[],
)
return {
// Validation functions
validate,
validateWithFallback,
stop,
cancel,
/**
* End a running check (the user pressed Stop): its promise rejects with
* an AbortError, so the tool handler can finish at once.
*/
const cancel = useCallback(() => {
controllerRef.current?.abort()
controllerRef.current = null
}, [])
// State
isValidating: isLoading,
partialResult: object as ValidationResult | undefined,
error,
}
return { validateWithFallback, cancel }
}
lib/diagram-diff.ts
+10 -15
View File
@@ -1,4 +1,9 @@
import { modelOfPage, pageElement } from "@/lib/diagram-pages"
import {
fileVars,
modelOfPage,
pageElement,
parseXml,
} from "@/lib/diagram-pages"
/** What changed on one page between two versions of a diagram */
export interface ChangeSummary {
@@ -32,18 +37,12 @@ interface CellInfo {
signature: string
}
function parse(xml: string): Document | null {
if (!xml?.trim()) return null
const doc = new DOMParser().parseFromString(xml, "text/xml")
return doc.querySelector("parsererror") ? null : doc
}
/**
* A page's mxGraphModel element, whatever wrapper the XML has: the page
* with this id, else the first
*/
function pageModel(xml: string, pageId?: string | null): Element | null {
const doc = parse(xml)
const doc = parseXml(xml)
if (!doc) return null
const diagram = pageElement(doc, pageId)
if (diagram) return modelOfPage(diagram)
@@ -298,14 +297,10 @@ export function isSameDocument(a: string, b: string): boolean {
model: modelOfPage(d),
}))
}
const docA = parse(a)
const docB = parse(b)
const docA = parseXml(a)
const docB = parseXml(b)
if (!docA || !docB) return a === b
const varsOf = (doc: Document) =>
doc.documentElement.nodeName === "mxfile"
? doc.documentElement.getAttribute("vars")
: null
if (!sameFileVars(varsOf(docA), varsOf(docB))) return false
if (!sameFileVars(fileVars(docA), fileVars(docB))) return false
const pagesA = pagesOf(docA)
const pagesB = pagesOf(docB)
if (pagesA.length !== pagesB.length) return false
lib/diagram-pages.ts
+35 -17
View File
@@ -12,12 +12,20 @@ import {
type PageSelector,
} from "@/packages/mcp-server/src/pages.ts"
function parse(xml: string): Document | null {
/** The parsed XML, or null when it is empty or does not parse */
export function parseXml(xml: string): Document | null {
if (!xml?.trim()) return null
const doc = new DOMParser().parseFromString(xml, "text/xml")
return doc.querySelector("parsererror") ? null : doc
}
/** The file variables (%name% placeholders) an <mxfile> sets, else null */
export function fileVars(doc: Document): string | null {
return doc.documentElement.nodeName === "mxfile"
? doc.documentElement.getAttribute("vars")
: null
}
/** The <diagram> the AI works on: the one with this id, else the first */
export function pageElement(
doc: Document,
@@ -40,7 +48,7 @@ export function modelOfPage(diagram: Element): Element | null {
if (model) return model
const inflated = decompressPageContent(diagram.textContent || "")
if (!inflated) return null
return parse(inflated)?.querySelector("mxGraphModel") ?? null
return parseXml(inflated)?.querySelector("mxGraphModel") ?? null
}
/**
@@ -52,7 +60,7 @@ export function pageSelectorFor(
pageId: string | null | undefined,
): PageSelector {
if (!pageId) return {}
const doc = parse(xml)
const doc = parseXml(xml)
const found = doc
? Array.from(doc.getElementsByTagName("diagram")).some(
(diagram) => diagram.getAttribute("id") === pageId,
@@ -69,7 +77,7 @@ export function pageModelXml(
xml: string,
pageId: string | null | undefined,
): string | null {
const doc = parse(xml)
const doc = parseXml(xml)
if (!doc) return null
const root = doc.documentElement
if (root.nodeName === "mxGraphModel") return xml
@@ -90,15 +98,15 @@ export function replacePageModel(
pageId: string | null | undefined,
modelXml: string,
): string {
const model = parse(modelXml)
const model = parseXml(modelXml)
if (model?.documentElement.nodeName !== "mxGraphModel") {
throw new Error("replacePageModel needs an <mxGraphModel>")
}
let doc = parse(fileXml)
if (doc?.documentElement.nodeName !== "mxfile") doc = parse(BLANK_MXFILE)
let doc = parseXml(fileXml)
if (doc?.documentElement.nodeName !== "mxfile") doc = parseXml(BLANK_MXFILE)
let diagram = pageElement(doc as Document, pageId)
if (!diagram) {
doc = parse(BLANK_MXFILE)
doc = parseXml(BLANK_MXFILE)
diagram = pageElement(doc as Document, null)
}
const target = diagram as Element
@@ -112,27 +120,37 @@ export function replacePageModel(
/**
* Where a drawn diagram goes. One page's worth (bare cells wrapped by
* prepareNewDiagram, or a one-page file) replaces the AI's page of the
* canvas file; file variables the drawn file sets replace the canvas
* file's. A file with several pages replaces the whole document: the model
* wrote the pages on purpose.
* canvas file. A file with several pages replaces the whole document: the
* model wrote the pages on purpose. Either way, file variables the drawn
* file sets replace the canvas file's, and a drawn file without any keeps
* the canvas file's (as replacing the page in draw.io does).
*/
export function placeOnPage(
drawnXml: string,
canvasXml: string,
pageId: string | null | undefined,
): string {
const doc = parse(drawnXml)
const doc = parseXml(drawnXml)
const drawnFile =
doc?.documentElement.nodeName === "mxfile" ? doc.documentElement : null
if (drawnFile && drawnFile.getElementsByTagName("diagram").length > 1) {
return drawnXml
const severalPages =
(drawnFile?.getElementsByTagName("diagram").length ?? 0) > 1
const model = severalPages ? null : pageModelXml(drawnXml, null)
if (!model) {
// The drawn document as it is, with the canvas file's variables
// when it sets none itself
const canvas = parseXml(canvasXml)
const canvasVars = canvas ? fileVars(canvas) : null
if (!doc || !drawnFile || fileVars(doc) !== null || !canvasVars) {
return drawnXml
}
drawnFile.setAttribute("vars", canvasVars)
return new XMLSerializer().serializeToString(doc)
}
const model = pageModelXml(drawnXml, null)
if (!model) return drawnXml
const placed = replacePageModel(canvasXml, pageId, model)
const vars = drawnFile?.getAttribute("vars")
if (vars === null || vars === undefined) return placed
const result = parse(placed) as Document
const result = parseXml(placed) as Document
result.documentElement.setAttribute("vars", vars)
return new XMLSerializer().serializeToString(result)
}
lib/drawio/editor-bridge.ts
-13
View File
@@ -600,18 +600,6 @@ function readSelectionRect() {
}
}
// ---------------------------------------------------------------------------
// Pages
// ---------------------------------------------------------------------------
function readPages() {
const pages = (ui?.pages ?? []) as any[]
return pages.map((page, index) => ({
id: String(page.getId?.() ?? index),
name: String(page.getName?.() ?? `Page-${index + 1}`),
}))
}
// ---------------------------------------------------------------------------
// Keeping the store in sync
// ---------------------------------------------------------------------------
@@ -624,7 +612,6 @@ function syncAll() {
if (!ui) return
const g = graph()
useCanvasStore.getState().set({
pages: readPages(),
currentPageId: ui.currentPage?.getId
? String(ui.currentPage.getId())
: null,
lib/session-storage.ts
+5 -111
View File
@@ -2,13 +2,11 @@ import { type DBSchema, type IDBPDatabase, openDB } from "idb"
import { nanoid } from "nanoid"
import { toast } from "sonner"
import type { DiagramVersion } from "@/stores/versions-store"
import type { Template } from "./template-storage"
// Constants
const DB_NAME = "next-ai-drawio"
const DB_VERSION = 2
const STORE_NAME = "sessions"
const MIGRATION_FLAG = "next-ai-drawio-migrated-to-idb"
const MAX_SESSIONS = 50
// Types
@@ -21,8 +19,6 @@ export interface ChatSession {
xmlSnapshots: [number, string][]
diagramXml: string
thumbnailDataUrl?: string // Small PNG preview of the diagram
// Old version list (before AI versions existed); read once, then replaced
diagramHistory?: { svg: string; xml: string }[]
// Diagram states produced by AI changes, shown as cards in the chat
versions?: DiagramVersion[]
}
@@ -51,16 +47,6 @@ interface ChatSessionDB extends DBSchema {
value: ChatSession
indexes: { "by-updated": number }
}
templates: {
key: string
value: Template
indexes: {
"by-updated": number
"by-pinned": number
"by-run-count": number
"by-last-used": number
}
}
}
// Database singleton
@@ -77,24 +63,6 @@ async function getDB(): Promise<IDBPDatabase<ChatSessionDB>> {
})
store.createIndex("by-updated", "updatedAt")
}
// Version 2: templates store (added by template-storage.ts)
// Note: We also need to include this here to ensure the upgrade
// callback properly handles all migrations when opening from this file
if (oldVersion < 2) {
// Check if templates store already exists (created by template-storage.ts)
if (!db.objectStoreNames.contains("templates")) {
const templateStore = db.createObjectStore(
"templates",
{
keyPath: "id",
},
)
templateStore.createIndex("by-updated", "updatedAt")
templateStore.createIndex("by-pinned", "pinned")
templateStore.createIndex("by-run-count", "runCount")
templateStore.createIndex("by-last-used", "lastUsedAt")
}
}
},
blocked() {
// An older tab keeps the DB open, so the upgrade has to wait
@@ -307,16 +275,13 @@ export function sanitizeMessage(message: unknown): StoredMessage | null {
const role = msg.role as string
if (!["user", "assistant", "system"].includes(role)) return null
// Extract parts, removing streaming state artifacts
let parts: Array<{ type: string; [key: string]: unknown }> = []
if (Array.isArray(msg.parts)) {
parts = msg.parts.map((part: unknown) => {
if (!part || typeof part !== "object") return { type: "unknown" }
const p = part as Record<string, unknown>
// Remove streaming-related fields
const { isStreaming, streamingState, ...cleanPart } = p
return cleanPart as { type: string; [key: string]: unknown }
})
parts = msg.parts.map((part: unknown) =>
part && typeof part === "object"
? (part as { type: string; [key: string]: unknown })
: { type: "unknown" },
)
}
return {
@@ -332,74 +297,3 @@ export function sanitizeMessages(messages: unknown[]): StoredMessage[] {
.map(sanitizeMessage)
.filter((m): m is StoredMessage => m !== null)
}
// Migration from localStorage
export async function migrateFromLocalStorage(): Promise<string | null> {
if (typeof window === "undefined") return null
if (!isIndexedDBAvailable()) return null
// Check if already migrated
if (localStorage.getItem(MIGRATION_FLAG)) return null
try {
const savedMessages = localStorage.getItem("next-ai-draw-io-messages")
const savedSnapshots = localStorage.getItem(
"next-ai-draw-io-xml-snapshots",
)
const savedXml = localStorage.getItem("next-ai-draw-io-diagram-xml")
let newSessionId: string | null = null
let migrationSucceeded = false
if (savedMessages) {
const messages = JSON.parse(savedMessages)
if (Array.isArray(messages) && messages.length > 0) {
const sanitized = sanitizeMessages(messages)
const session: ChatSession = {
id: nanoid(),
title: extractTitle(sanitized),
createdAt: Date.now(),
updatedAt: Date.now(),
messages: sanitized,
xmlSnapshots: savedSnapshots
? JSON.parse(savedSnapshots)
: [],
diagramXml: savedXml || "",
}
const saved = await saveSession(session)
if (saved) {
// Verify the session was actually written
const verified = await getSession(session.id)
if (verified) {
newSessionId = session.id
migrationSucceeded = true
}
}
} else {
// Empty array or invalid data - nothing to migrate, mark as success
migrationSucceeded = true
}
} else {
// No data to migrate - mark as success
migrationSucceeded = true
}
// Only clean up old data if migration succeeded
if (migrationSucceeded) {
localStorage.setItem(MIGRATION_FLAG, "true")
localStorage.removeItem("next-ai-draw-io-messages")
localStorage.removeItem("next-ai-draw-io-xml-snapshots")
localStorage.removeItem("next-ai-draw-io-diagram-xml")
} else {
console.warn(
"Migration to IndexedDB failed - keeping localStorage data for retry",
)
}
return newSessionId
} catch (error) {
console.error("Migration failed:", error)
// Don't mark as migrated - allow retry on next load
return null
}
}
lib/use-quota-manager.tsx
+17 -24
View File
@@ -7,34 +7,30 @@ import { useDictionary } from "@/hooks/use-dictionary"
import { formatMessage } from "@/lib/i18n/utils"
export interface QuotaConfig {
dailyRequestLimit: number
dailyTokenLimit: number
tpmLimit: number
onConfigModel?: () => void
}
/**
* Hook for displaying quota limit toasts.
* Server-side handles actual quota enforcement via DynamoDB.
* This hook only provides UI feedback when limits are exceeded.
* Server-side handles actual quota enforcement via DynamoDB, and its 429
* responses carry the used and limit values shown here.
*/
export function useQuotaManager(config: QuotaConfig): {
showQuotaLimitToast: (used?: number, limit?: number) => void
showTokenLimitToast: (used?: number, limit?: number) => void
showTPMLimitToast: (limit?: number) => void
showQuotaLimitToast: (used: number, limit: number) => void
showTokenLimitToast: (used: number, limit: number) => void
showTPMLimitToast: (limit: number) => void
} {
const { dailyRequestLimit, dailyTokenLimit, tpmLimit, onConfigModel } =
config
const { onConfigModel } = config
const dict = useDictionary()
// Show quota limit toast (request-based)
const showQuotaLimitToast = useCallback(
(used?: number, limit?: number) => {
(used: number, limit: number) => {
toast.custom(
(t) => (
<QuotaLimitToast
used={used ?? dailyRequestLimit}
limit={limit ?? dailyRequestLimit}
used={used}
limit={limit}
onDismiss={() => toast.dismiss(t)}
onConfigModel={onConfigModel}
/>
@@ -42,18 +38,18 @@ export function useQuotaManager(config: QuotaConfig): {
{ duration: 15000 },
)
},
[dailyRequestLimit, onConfigModel],
[onConfigModel],
)
// Show token limit toast
const showTokenLimitToast = useCallback(
(used?: number, limit?: number) => {
(used: number, limit: number) => {
toast.custom(
(t) => (
<QuotaLimitToast
type="token"
used={used ?? dailyTokenLimit}
limit={limit ?? dailyTokenLimit}
used={used}
limit={limit}
onDismiss={() => toast.dismiss(t)}
onConfigModel={onConfigModel}
/>
@@ -61,24 +57,21 @@ export function useQuotaManager(config: QuotaConfig): {
{ duration: 15000 },
)
},
[dailyTokenLimit, onConfigModel],
[onConfigModel],
)
// Show TPM limit toast
const showTPMLimitToast = useCallback(
(limit?: number) => {
const effectiveLimit = limit ?? tpmLimit
(limit: number) => {
const limitDisplay =
effectiveLimit >= 1000
? `${effectiveLimit / 1000}k`
: String(effectiveLimit)
limit >= 1000 ? `${limit / 1000}k` : String(limit)
const message = formatMessage(dict.quota.tpmMessageDetailed, {
limit: limitDisplay,
seconds: 60,
})
toast.error(message, { duration: 8000 })
},
[tpmLimit, dict],
[dict],
)
return {
packages/mcp-server/src/load-diagram.ts
+12 -4
View File
@@ -46,7 +46,8 @@ export function decompressPageContent(compressed: string): string | null {
/**
* Read the diagram XML a .drawio.svg (Editable SVG) carries in its root
* element's content attribute. Returns undefined when the text is not an
* SVG, and "" when the SVG has no embedded diagram.
* SVG, and "" when the SVG has no embedded diagram. The root may carry a
* namespace prefix (<svg:svg>).
*
* Uses getAttributeNode().value rather than getAttribute(): linkedom's
* getAttribute re-escapes <, > and & in XML documents, which would turn
@@ -55,10 +56,16 @@ export function decompressPageContent(compressed: string): string | null {
export function extractSvgEmbeddedXml(text: string): string | undefined {
const doc = new DOMParser().parseFromString(text, "text/xml")
const root = doc.documentElement
if (root?.tagName !== "svg") return undefined
// (linkedom keeps the prefix in localName too)
if (root?.tagName.split(":").pop() !== "svg") return undefined
return (root.getAttributeNode("content")?.value ?? "").trim()
}
// What XML allows before the root element: the declaration, comments, a
// DOCTYPE (the root checks below expect the root first)
const XML_PROLOG =
/^(?:\s|<\?[\s\S]*?\?>|<!--[\s\S]*?-->|<!DOCTYPE[^>[]*(?:\[[\s\S]*?\])?\s*>)*/i
/**
* Parse the content of a .drawio file into the canonical session shape:
* an <mxfile> whose every page holds plain <mxGraphModel> XML. Accepts a
@@ -66,7 +73,7 @@ export function extractSvgEmbeddedXml(text: string): string | undefined {
* the diagram embedded, and decompresses any compressed pages.
*/
export function parseDrawioFileContent(content: string): LoadResult {
let trimmed = content.trim()
let trimmed = content.replace(XML_PROLOG, "").trim()
if (!trimmed) return { ok: false, error: "File is empty." }
if (isMxGraphModel(trimmed)) {
@@ -78,7 +85,8 @@ export function parseDrawioFileContent(content: string): LoadResult {
trimmed = normalized
}
if (!isMxFile(trimmed)) {
const embedded = extractSvgEmbeddedXml(trimmed)
// Parsed whole: the SVG's DOCTYPE may define entities its content uses
const embedded = extractSvgEmbeddedXml(content)
if (embedded === undefined) {
return {
ok: false,
packages/mcp-server/tests/load-diagram.test.ts
+34
View File
@@ -190,12 +190,46 @@ const SVG_PROLOG = `<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE svg PUBLIC "-//W3C//DTD SVG 1.1//EN" "http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd">
`
describe("parseDrawioFileContent with a prolog before the root", () => {
it("accepts a declaration, a comment and a DOCTYPE before <mxfile>", () => {
const prolog = `<?xml version="1.0" encoding="UTF-8"?>\n<!-- Architecture, v2 -->\n<!DOCTYPE mxfile>\n`
expect(parseDrawioFileContent(prolog + PLAIN_MXFILE)).toEqual({
ok: true,
xml: PLAIN_MXFILE,
hadCompressedPages: false,
})
})
it("accepts a comment before a bare <mxGraphModel>", () => {
const r = parseDrawioFileContent(`<!-- exported -->\n${MODEL_XML}`)
expect(r.ok).toBe(true)
if (r.ok) expect(r.xml).toContain('value="Hello"')
})
it("treats a file with nothing after its prolog as empty", () => {
expect(parseDrawioFileContent(`<?xml version="1.0"?>\n`)).toEqual({
ok: false,
error: "File is empty.",
})
})
})
describe("extractSvgEmbeddedXml", () => {
it("returns undefined for non-SVG text", () => {
expect(extractSvgEmbeddedXml(PLAIN_MXFILE)).toBeUndefined()
expect(extractSvgEmbeddedXml("hello")).toBeUndefined()
})
it("accepts a root element with a namespace prefix", () => {
const svg = `<svg:svg xmlns:svg="http://www.w3.org/2000/svg" content="${escapeAttr(PLAIN_MXFILE)}"></svg:svg>`
expect(extractSvgEmbeddedXml(svg)).toBe(PLAIN_MXFILE)
expect(parseDrawioFileContent(svg)).toEqual({
ok: true,
xml: PLAIN_MXFILE,
hadCompressedPages: false,
})
})
it("returns an empty string for an SVG without content", () => {
expect(extractSvgEmbeddedXml("<svg><rect/></svg>")).toBe("")
})
stores/canvas-store.ts
-7
View File
@@ -14,15 +14,9 @@ interface CanvasRect {
height: number
}
interface CanvasPage {
id: string
name: string
}
interface CanvasState {
/** The page can call the draw.io editor directly (same-origin iframe) */
hasEditor: boolean
pages: CanvasPage[]
currentPageId: string | null
selection: SelectedCell[]
selectionRect: CanvasRect | null
@@ -35,7 +29,6 @@ interface CanvasState {
export const useCanvasStore = create<CanvasState>((set) => ({
hasEditor: false,
pages: [],
currentPageId: null,
selection: [],
selectionRect: null,
stores/versions-store.ts
+1 -19
View File
@@ -1,7 +1,6 @@
import { nanoid } from "nanoid"
import { create } from "zustand"
import type { ChangeSummary } from "@/lib/diagram-diff"
import { EMPTY_SUMMARY } from "@/lib/diagram-diff"
/** One diagram state produced by an AI change, shown as a card in the chat */
export interface DiagramVersion {
@@ -71,10 +70,8 @@ export const useVersionsStore = create<VersionsState>((set) => ({
setVersions: (versions) =>
set({
// A missing thumbnail could not be made: no "rendering" forever.
// Versions saved before they had numbers count from 1.
versions: versions.map((v, i) => ({
versions: versions.map((v) => ({
...v,
number: v.number ?? i + 1,
svg: v.svg ?? "",
})),
undoneVersionId: versions.find((v) => v.undone)?.id ?? null,
@@ -141,18 +138,3 @@ export const useVersionsStore = create<VersionsState>((set) => ({
clear: () =>
set({ versions: [], undoneVersionId: null, onCanvasVersionId: null }),
}))
/** Convert the old `diagramHistory` list ({svg, xml}) saved by earlier versions */
export function versionsFromLegacyHistory(
history: { svg: string; xml: string }[] | undefined,
): DiagramVersion[] {
if (!history?.length) return []
return history.slice(-MAX_VERSIONS).map((entry, index) => ({
id: `legacy-${index}`,
number: index + 1,
xml: entry.xml,
svg: entry.svg,
turnIndex: -1,
summary: { ...EMPTY_SUMMARY },
}))
}
tests/e2e/settings-providers.spec.ts
+1 -8
View File
@@ -219,14 +219,7 @@ test("an access code error opens the general tab with the code field", async ({
page,
}) => {
await page.route("**/api/config", (route) =>
route.fulfill({
json: {
accessCodeRequired: true,
dailyRequestLimit: 0,
dailyTokenLimit: 0,
tpmLimit: 0,
},
}),
route.fulfill({ json: { accessCodeRequired: true } }),
)
await page.route("**/api/chat", (route) =>
route.fulfill({
tests/e2e/workspace.spec.ts
+18 -5
View File
@@ -2987,9 +2987,25 @@ test.describe("Edge cases", () => {
).toBeVisible()
})
test("compare closes when another chat with the same old version ids comes on screen", async ({
test("compare closes when another chat with the same version ids comes on screen", async ({
page,
}) => {
// Both chats' versions share their ids
const version = (id: string, number: number) => ({
id,
number,
xml: `<mxGraphModel><root/></mxGraphModel>`,
svg: "",
turnIndex: 0,
summary: {
shapesAdded: 0,
shapesRemoved: 0,
shapesChanged: 0,
edgesAdded: 0,
edgesRemoved: 0,
edgesChanged: 0,
},
})
const legacy = (id: string, label: string) => ({
id,
title: `Chat ${id}`,
@@ -3004,10 +3020,7 @@ test.describe("Edge cases", () => {
],
xmlSnapshots: [],
diagramXml: `<mxfile><diagram name="P" id="p"><mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/>${boxCell("a", label, 40)}</root></mxGraphModel></diagram></mxfile>`,
diagramHistory: [
{ svg: "", xml: `<mxGraphModel><root/></mxGraphModel>` },
{ svg: "", xml: `<mxGraphModel><root/></mxGraphModel>` },
],
versions: [version("v-1", 1), version("v-2", 2)],
})
await openApp(page)
await page.evaluate(
tests/unit/diagram-context.test.tsx
+18 -22
View File
@@ -18,17 +18,18 @@ function setup() {
exportDiagram: (r: any) => requests.push(r),
load: vi.fn(),
} as any
// draw.io's reply to a request: it echoes the request in `message`
// draw.io's reply to a request: it echoes the request in `message`.
// Async, so what the caller does with the result runs before it returns
const reply = (request: { message: string }, data: string, xml = "") =>
act(() =>
act(async () => {
result.current.handleDiagramExport({
event: "export",
data,
xml,
format: "xmlsvg",
message: request,
} as any),
)
} as any)
})
return { result, requests, reply }
}
@@ -53,25 +54,20 @@ afterEach(() => {
})
describe("exports in flight at the same time", () => {
it("give the chat's export only its own reply", () => {
it("give the chat's export only its own reply", async () => {
const { result, requests, reply } = setup()
// An edit's history export is still on its way when the chat exports
act(() => {
result.current.handleExport()
const older = result.current.fetchDiagramXml()
const current = result.current.fetchDiagramXml()
let currentDone = false
current.then(() => {
currentDone = true
})
let tag = ""
const got: string[] = []
act(() => {
tag = result.current.handleExport()
result.current.exportResolversRef.current[tag] = (xml) =>
got.push(xml)
})
reply(requests[0], svgOf("older"))
expect(got).toEqual([])
reply(requests[1], svgOf("current"))
expect(got).toHaveLength(1)
expect(got[0]).toContain('value="current"')
expect(result.current.exportResolversRef.current[tag]).toBeUndefined()
await reply(requests[0], svgOf("older"))
expect(await older).toContain('value="older"')
expect(currentDone).toBe(false)
await reply(requests[1], svgOf("current"))
expect(await current).toContain('value="current"')
})
it("save each file with its own result", async () => {
@@ -101,8 +97,8 @@ describe("exports in flight at the same time", () => {
result.current.saveDiagramToFile("pic", "png")
})
// The PNG answers first
reply(requests[1], "data:image/png;base64,iVBORw0KGgo=")
reply(requests[0], svgOf("doc"), twoPages)
await reply(requests[1], "data:image/png;base64,iVBORw0KGgo=")
await reply(requests[0], svgOf("doc"), twoPages)
expect(saved.map((s) => s.name)).toEqual(["pic.png", "doc.drawio"])
expect(saved[0].href).toMatch(/^data:image\/png/)
tests/unit/diagram-pages.test.ts
+18 -1
View File
@@ -113,7 +113,24 @@ describe("placeOnPage", () => {
it("lets a drawing with several pages replace the document", () => {
const drawn = `<mxfile>${page("x", box("X"))}${page("y", box("Y"))}</mxfile>`
expect(placeOnPage(drawn, TWO_PAGES, "b")).toBe(drawn)
expect(pagesOf(placeOnPage(drawn, TWO_PAGES, "b"))).toEqual([
{ id: "x", name: "Page x", cells: ["X"] },
{ id: "y", name: "Page y", cells: ["Y"] },
])
// A canvas without variables leaves it as it is
const plainCanvas = `<mxfile>${page("a", box("A"))}</mxfile>`
expect(placeOnPage(drawn, plainCanvas, "a")).toBe(drawn)
})
it("gives a drawing with several pages the canvas's variables unless it has its own", () => {
const drawn = `<mxfile>${page("x", box("X"))}${page("y", box("Y"))}</mxfile>`
expect(placeOnPage(drawn, TWO_PAGES, "b")).toContain(
'vars="{&quot;x&quot;:&quot;1&quot;}"',
)
const withVars = `<mxfile vars="{&quot;team&quot;:&quot;New&quot;}">${page("x", box("X"))}${page("y", box("Y"))}</mxfile>`
const result = placeOnPage(withVars, TWO_PAGES, "b")
expect(result).toContain('vars="{&quot;team&quot;:&quot;New&quot;}"')
expect(result).not.toContain("&quot;x&quot;")
})
it("draws on the blank file when the canvas is empty", () => {
tests/unit/redesign-helpers.test.ts
+2 -23
View File
@@ -7,11 +7,7 @@ import jaDict from "@/lib/i18n/dictionaries/ja.json"
import zhDict from "@/lib/i18n/dictionaries/zh.json"
import zhHantDict from "@/lib/i18n/dictionaries/zh-Hant.json"
import { describeChanges, describeTotals } from "@/lib/version-text"
import {
MAX_VERSIONS,
useVersionsStore,
versionsFromLegacyHistory,
} from "@/stores/versions-store"
import { MAX_VERSIONS, useVersionsStore } from "@/stores/versions-store"
describe("withPageDefaults", () => {
it("adds page and grid when the model has none", () => {
@@ -157,7 +153,7 @@ describe("versions store", () => {
expect(versions.at(-1)?.number).toBe(MAX_VERSIONS + 5)
})
it("numbers on from the last version kept, and numbers saved ones", () => {
it("numbers on from the last version kept", () => {
add(0)
add(2)
add(4)
@@ -166,14 +162,6 @@ describe("versions store", () => {
expect(
useVersionsStore.getState().versions.map((v) => v.number),
).toEqual([1, 2, 3])
const { versions } = useVersionsStore.getState()
// Saved before versions had numbers
useVersionsStore
.getState()
.setVersions(versions.map(({ number, ...v }) => v as any))
expect(
useVersionsStore.getState().versions.map((v) => v.number),
).toEqual([1, 2, 3])
})
it("drops versions from a turn that is re-run, none left on the canvas", () => {
@@ -199,13 +187,4 @@ describe("versions store", () => {
expect(state.undoneVersionId).toBeNull()
expect(state.versions[0].undone).toBeUndefined()
})
it("converts the old history list", () => {
const versions = versionsFromLegacyHistory([
{ svg: "data:image/svg+xml;base64,AA", xml: "<a/>" },
])
expect(versions).toHaveLength(1)
expect(versions[0]).toMatchObject({ xml: "<a/>", turnIndex: -1 })
expect(versionsFromLegacyHistory(undefined)).toEqual([])
})
})
tests/unit/use-session-manager.test.tsx
-1
View File
@@ -17,7 +17,6 @@ vi.mock("@/lib/session-storage", async (importActual) => {
createEmptySession: actual.createEmptySession,
extractTitle: actual.extractTitle,
isIndexedDBAvailable: () => true,
migrateFromLocalStorage: async () => null,
readSessionCount: async () => stored.size,
enforceSessionLimit: async () => {},
getSession: async (id: string) => {
tests/unit/validate-diagram-route.test.ts
+16 -33
View File
@@ -1,5 +1,4 @@
// @vitest-environment node
import { simulateReadableStream } from "ai"
import { MockLanguageModelV3 } from "ai/test"
import { afterEach, describe, expect, it, vi } from "vitest"
import { POST as validateDiagram } from "@/app/api/validate-diagram/route"
@@ -16,36 +15,19 @@ const RESULT = {
suggestions: ["Move box B to the right"],
}
// A vision model that answers with the JSON in a few text chunks
// A vision model that answers with the result as JSON text
vi.mock("@/lib/ai-providers", () => ({
getValidationModel: () =>
new MockLanguageModelV3({
doStream: (async () => {
const json = JSON.stringify(RESULT)
return {
stream: simulateReadableStream({
chunks: [
{ type: "text-start", id: "t" },
...[json.slice(0, 20), json.slice(20)].map(
(delta) => ({
type: "text-delta",
id: "t",
delta,
}),
),
{ type: "text-end", id: "t" },
{
type: "finish",
finishReason: { unified: "stop", raw: "stop" },
usage: {
inputTokens: { total: 1 },
outputTokens: { total: 1 },
},
},
],
}),
}
}) as any,
doGenerate: (async () => ({
content: [{ type: "text", text: JSON.stringify(RESULT) }],
finishReason: { unified: "stop", raw: "stop" },
usage: {
inputTokens: { total: 1 },
outputTokens: { total: 1 },
},
warnings: [],
})) as any,
}),
}))
@@ -110,24 +92,25 @@ describe("the quota", () => {
// get its check, and the check does not count as a chat
quota.enabled = true
const res = await validateDiagram(post())
expect(JSON.parse(await res.text())).toEqual(RESULT)
expect(await res.json()).toEqual(RESULT)
expect(quota.checks).toHaveLength(1)
expect(quota.checks[0].increment).toBe(0)
expect(quota.checks[0].limits.requests).toBe(0)
await vi.waitFor(() => expect(quota.recorded).toEqual([2]))
expect(quota.recorded).toEqual([2])
})
})
describe("POST /api/validate-diagram", () => {
it("streams the model's result as JSON text for useObject", async () => {
it("answers the model's result as JSON", async () => {
const res = await validateDiagram(post())
expect(JSON.parse(await res.text())).toEqual(RESULT)
expect(res.headers.get("content-type")).toContain("application/json")
expect(await res.json()).toEqual(RESULT)
})
it("answers valid when the check is turned off", async () => {
process.env.ENABLE_VLM_VALIDATION = "false"
const res = await validateDiagram(post())
expect(JSON.parse(await res.text())).toEqual({
expect(await res.json()).toEqual({
valid: true,
issues: [],
suggestions: [],