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

fix(chat): the AI reads and draws the page the user is viewing (#974)

Browse Source 
* fix(chat): the AI reads and draws the page the user is viewing

The model always got the first page and display_diagram replaced the whole
file with one page, so drawing on page 3 of a multi-page file deleted the
other pages (#241), and edits went to page 1 while the user looked at
another page.

Now a turn records the page the user is viewing when the message is sent.
That page's model goes to the model (also as previousXml, on regenerate and
in edit_diagram error messages), selected shapes are sent on any page,
display_diagram and append_diagram replace only that page (a drawing that
brings several pages still replaces the file), edit_diagram targets it,
and the streaming preview draws on it. Compact cells and default styles
take the layer and the edges of that page, in the preview and in
editDiagram. The editor bridge replaces the current page in place when the
other pages are unchanged, so Ctrl+Z undoes an AI change on a multi-page
file in one step; when a full load is needed, the user's page is shown
again afterwards. Version summaries and the compare dialog count the page
the change was made on.

* test: read the failed edit's output without unsafe optional chaining
This commit is contained in:
Dayuan Jiang
2026-10-10 23:30:51 +09:00
committed by GitHub
parent cd5352ca88
commit d1ee1594f1
21 changed files with 907 additions and 120 deletions
Download Patch File Download Diff File Expand all files Collapse all files
components/chat/chat-engine.tsx
+32 -17
View File
@@ -42,6 +42,7 @@ import { useValidateDiagram } from "@/hooks/use-validate-diagram"
import { getApiEndpoint } from "@/lib/base-path"
import { findCachedResponse } from "@/lib/cached-responses"
import { buildChatHeaders } from "@/lib/chat-request"
import { pageModelXml, placeOnPage } from "@/lib/diagram-pages"
import { EMPTY_DRAWIO_DOCUMENT } from "@/lib/drawio/drawio-config"
import { formatMessage } from "@/lib/i18n/utils"
import { isPdfFile, isTextFile } from "@/lib/pdf-utils"
@@ -90,15 +91,20 @@ function hasToolErrors(messages: UIMessage[]): boolean {
}
/**
* Snapshots keep the full multi-page document, but the model only sees and
* edits the first page, so give it the first page's mxGraphModel.
* Older snapshots already hold a single mxGraphModel and are returned as is.
* Snapshots keep the full multi-page document, but the model sees and
* edits one page, the one the user is viewing (the first when unknown), so
* give it that page's mxGraphModel. Older snapshots already hold a single
* mxGraphModel and are returned as is.
*/
function getFirstPageXml(xml: string): string {
function getPageXml(xml: string, pageId: string | null): string {
if (!xml.includes("<mxfile")) return xml
const doc = new DOMParser().parseFromString(xml, "text/xml")
const model = doc.querySelector("diagram")?.querySelector("mxGraphModel")
return model ? formatXML(new XMLSerializer().serializeToString(model)) : xml
const model = pageModelXml(xml, pageId)
return model ? formatXML(model) : xml
}
/** The page the user is viewing; null with an external draw.io */
function viewedPageId(): string | null {
return useCanvasStore.getState().currentPageId
}
// Shapes sent with a user message (also kept in its metadata)
@@ -374,6 +380,9 @@ export function ChatEngineProvider({
// XML snapshot taken before each user message (keyed by message index)
const xmlSnapshotsRef = useRef<Map<number, string>>(new Map())
// The page the running turn's model reads and writes: the one the user
// viewed when the message was sent (null: the first page)
const turnPageIdRef = useRef<string | null>(null)
// Index of the user message whose turn is running (versions belong to it)
const currentTurnRef = useRef(0)
// The chat (its generation) the running turn was sent in: a reply that
@@ -450,6 +459,7 @@ export function ChatEngineProvider({
validationRetryCountRef,
// A preview undone just before the tool call is in this one already
chartXMLRef: liveChartXMLRef,
turnPageIdRef,
onDisplayChart,
onFetchChart,
captureValidationPng,
@@ -748,6 +758,7 @@ export function ChatEngineProvider({
processedToolCallsRef,
editDiagramOriginalXmlRef,
loadedMessageIdsRef,
turnPageIdRef,
})
const { restoreVersion, undoVersion } = useVersions({
@@ -762,14 +773,10 @@ export function ChatEngineProvider({
// ---------------------------------------------------------------------
const selection = useCanvasStore((s) => s.selection)
// The model sees and edits the first page only
const onFirstPage = useCanvasStore(
(s) => s.pages.length === 0 || s.pages[0]?.id === s.currentPageId,
)
const [dismissedSelectionKey, setDismissedSelectionKey] = useState("")
const selectionKey = selection.map((c) => c.id).join(",")
const chatSelection =
!onFirstPage || selectionKey === dismissedSelectionKey ? [] : selection
selectionKey === dismissedSelectionKey ? [] : selection
// Once nothing is selected, picking the same shapes again attaches them
useEffect(() => {
if (selection.length === 0) setDismissedSelectionKey("")
@@ -1178,8 +1185,9 @@ export function ChatEngineProvider({
.filter((k) => k < beforeIndex)
.sort((a, b) => b - a)
return snapshotKeys.length > 0
? getFirstPageXml(
? getPageXml(
xmlSnapshotsRef.current.get(snapshotKeys[0]) || "",
viewedPageId(),
)
: ""
}
@@ -1198,6 +1206,7 @@ export function ChatEngineProvider({
continuationOriginalRef.current = null
stoppedRef.current = false
currentTurnRef.current = turnIndex
turnPageIdRef.current = viewedPageId()
turnChatRef.current = getChatGeneration()
// Busy from now on, before the next render says so
busyRef.current = true
@@ -1254,7 +1263,8 @@ export function ChatEngineProvider({
const turnIndex = messagesRef.current.length
const previousXml = getPreviousXml(turnIndex)
// Snapshot the full multi-page document (kept fresh by autosave) so
// regenerate/edit can restore every page; the model gets page 1 only
// regenerate/edit can restore every page; the model gets the page
// the user is viewing
xmlSnapshotsRef.current.set(turnIndex, chartXMLRef.current || chartXml)
sendChatMessage(parts, chartXml, previousXml, turnIndex, selectedCells)
return true
@@ -1328,11 +1338,16 @@ export function ChatEngineProvider({
pageName: "Page-1",
})
if (prepared.ok) {
const pageId = viewedPageId()
const canvasXml = chartXMLRef.current || ""
onDisplayChart(
keepFileVars(prepared.xml, chartXMLRef.current || ""),
keepFileVars(
placeOnPage(prepared.xml, canvasXml, pageId),
canvasXml,
),
true,
"commit",
{ toolCallId },
{ toolCallId, pageId },
)
}
clearComposer()
@@ -1492,7 +1507,7 @@ export function ChatEngineProvider({
})
sendChatMessage(
parts,
getFirstPageXml(savedXml),
getPageXml(savedXml, viewedPageId()),
previousXml,
index,
selectionOf(messages[index]),
components/chat/compare-dialog.tsx
+3 -3
View File
@@ -56,11 +56,11 @@ export function CompareDialog() {
() => !version || isSameDocument(version.xml, chartXML),
[version, chartXML],
)
// What the canvas has changed since this version on the first page
// (null: nothing there)
// What the canvas has changed since this version on the page it was
// made on (null: nothing there)
const changes = useMemo(() => {
if (!version || same) return null
const { summary } = diffDiagrams(version.xml, chartXML)
const { summary } = diffDiagrams(version.xml, chartXML, version.pageId)
return Object.values(summary).some((n) => n > 0) ? summary : null
}, [version, chartXML, same])
components/chat/use-diagram-streaming.ts
+23 -5
View File
@@ -5,6 +5,7 @@ import type { RefObject } from "react"
import { useCallback, useEffect, useRef } from "react"
import type { DiagramOperation, ToolPartLike } from "@/components/chat/types"
import { useDiagram } from "@/contexts/diagram-context"
import { pageSelectorFor } from "@/lib/diagram-pages"
import {
convertToLegalXml,
extractCompleteMxCells,
@@ -15,6 +16,7 @@ import {
expandCompactCells,
} from "@/packages/mcp-server/src/compact-cells.ts"
import { applyDiagramOperations } from "@/packages/mcp-server/src/diagram-operations.ts"
import { targetPageXml } from "@/packages/mcp-server/src/edit-diagram.ts"
import { BLANK_MXFILE } from "@/packages/mcp-server/src/pages.ts"
import {
addDefaultStyles,
@@ -44,6 +46,8 @@ interface UseDiagramStreamingParams {
processedToolCallsRef: RefObject<Set<string>>
editDiagramOriginalXmlRef: RefObject<Map<string, string>>
loadedMessageIdsRef: RefObject<Set<string>>
// The page the model draws on in this turn; null: the first page
turnPageIdRef?: RefObject<string | null>
}
/**
@@ -60,6 +64,7 @@ export function useDiagramStreaming({
processedToolCallsRef,
editDiagramOriginalXmlRef,
loadedMessageIdsRef,
turnPageIdRef,
}: UseDiagramStreamingParams) {
const { chartXML, chartXMLRef, loadDiagram } = useDiagram()
const previousXML = useRef<string>("")
@@ -102,10 +107,15 @@ export function useDiagramStreaming({
if (testDoc.querySelector("parsererror")) return
try {
// Replace the first page's cells so other pages stay intact.
// An empty canvas gets a default mxfile to put the cells in.
// Replace the model's page's cells so other pages stay
// intact. An empty canvas gets a default mxfile to put the
// cells in.
const baseXML = chartXML || BLANK_MXFILE
const replacedXML = replaceNodes(baseXML, convertedXml)
const replacedXML = replaceNodes(
baseXML,
convertedXml,
turnPageIdRef?.current,
)
previousXML.current = convertedXml
loadDiagram(replacedXML, true, "preview")
} catch (error) {
@@ -220,8 +230,15 @@ export function useDiagramStreaming({
if (lastProcessedXmlRef.current.get(opsKey) === ops) continue
lastProcessedXmlRef.current.set(opsKey, ops)
try {
const edges = edgeIdsOf(originalXml)
const layer = defaultLayerOf(originalXml)
// The model's page: the edges on it and its first layer,
// as editDiagram reads them
const selector = pageSelectorFor(
originalXml,
turnPageIdRef?.current,
)
const page = targetPageXml(originalXml, selector)
const edges = edgeIdsOf(page)
const layer = defaultLayerOf(page)
const { result } = applyDiagramOperations(
originalXml,
completeOps.map((op) =>
@@ -235,6 +252,7 @@ export function useDiagramStreaming({
}
: op,
),
selector,
)
if (lastProcessedXmlRef.current.get(resultKey) === result) {
continue
components/chat/use-versions.ts
+3 -1
View File
@@ -46,15 +46,17 @@ export function useVersions({
// (effect cleanups run before the streaming hook commits). The handler
// only reads refs and stable functions.
useEffect(() => {
setCommitHandler(({ beforeXml, afterXml, toolCallId }) => {
setCommitHandler(({ beforeXml, afterXml, toolCallId, pageId }) => {
if (restoringRef.current) return
const { summary, touchedIds, fromScratch } = diffDiagrams(
beforeXml,
afterXml,
pageId,
)
const id = useVersionsStore.getState().addVersion({
xml: afterXml,
beforeXml,
...(pageId && { pageId }),
turnIndex: currentTurnRef.current,
toolCallId,
summary,