mirror of
https://github.com/DayuanJiang/next-ai-draw-io.git
synced 2026-10-04 08:47:45 +08:00
fix(chat): keep saved diagrams and pages when restoring, editing and retrying
- Restored sessions no longer replay the last display_diagram over the saved diagram - Failed or stopped edit_diagram restores the canvas - Message snapshots keep the full multi-page document - "Improve with suggestions" uses the normal send path (headers, xml, retry counters) - Editing a message keeps its file/URL sections; cached example edits work - New chat's first autosave no longer resets the UI - Validation retries counted per user turn; validate-diagram sends the access code - Cached examples only match the example files on an empty canvas - Template sends keep attachments and wait for extraction
This commit is contained in:
dayuan.jiang
@@ -129,12 +129,14 @@ const getMessageTextContent = (message: UIMessage): string => {
|
||||
.join("\n")
|
||||
}
|
||||
|
||||
// Matches the [PDF: ...], [File: ...] and [URL: ...] sections appended to the user's text
|
||||
export const APPENDED_FILE_SECTIONS_PATTERN =
|
||||
/\n\n\[(PDF|File|URL):\s*[^\]]+\]\n[\s\S]*$/
|
||||
|
||||
// Get only the user's original text, excluding appended file content
|
||||
const getUserOriginalText = (message: UIMessage): string => {
|
||||
const fullText = getMessageTextContent(message)
|
||||
// Strip out [PDF: ...], [File: ...], and [URL: ...] sections that were appended
|
||||
const filePattern = /\n\n\[(PDF|File|URL):\s*[^\]]+\]\n[\s\S]*$/
|
||||
return fullText.replace(filePattern, "").trim()
|
||||
return fullText.replace(APPENDED_FILE_SECTIONS_PATTERN, "").trim()
|
||||
}
|
||||
|
||||
interface SessionMetadata {
|
||||
@@ -458,6 +460,11 @@ export function ChatMessageDisplay({
|
||||
messages.length > 0 ? [messages[messages.length - 1]] : []
|
||||
|
||||
messagesToProcess.forEach((message) => {
|
||||
// Messages restored from a saved session were applied before it was
|
||||
// saved; the saved diagram is authoritative, so don't replay them
|
||||
const isRestoredMessage =
|
||||
loadedMessageIdsRef?.current.has(message.id) ?? false
|
||||
|
||||
if (message.parts) {
|
||||
message.parts.forEach((part) => {
|
||||
if (part.type?.startsWith("tool-")) {
|
||||
@@ -475,6 +482,8 @@ export function ChatMessageDisplay({
|
||||
})
|
||||
}
|
||||
|
||||
if (isRestoredMessage) return
|
||||
|
||||
if (
|
||||
part.type === "tool-display_diagram" &&
|
||||
input?.xml
|
||||
@@ -541,6 +550,32 @@ export function ChatMessageDisplay({
|
||||
part.type === "tool-edit_diagram" &&
|
||||
input?.operations
|
||||
) {
|
||||
// Failed or stopped: drop the queued preview. If the original
|
||||
// XML is still stored, the tool handler never ran (user pressed
|
||||
// stop), so undo the streamed preview here.
|
||||
if (state === "output-error") {
|
||||
if (
|
||||
pendingEditRef.current?.toolCallId ===
|
||||
toolCallId &&
|
||||
editDebounceTimeoutRef.current
|
||||
) {
|
||||
clearTimeout(editDebounceTimeoutRef.current)
|
||||
editDebounceTimeoutRef.current = null
|
||||
pendingEditRef.current = null
|
||||
}
|
||||
const originalXml =
|
||||
editDiagramOriginalXmlRef.current.get(
|
||||
toolCallId,
|
||||
)
|
||||
if (originalXml) {
|
||||
editDiagramOriginalXmlRef.current.delete(
|
||||
toolCallId,
|
||||
)
|
||||
onDisplayChart(originalXml, true)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
const completeOps = getCompleteOperations(
|
||||
input.operations as DiagramOperation[],
|
||||
)
|
||||
@@ -610,9 +645,10 @@ export function ChatMessageDisplay({
|
||||
origXml,
|
||||
pending.operations,
|
||||
)
|
||||
handleDisplayChart(
|
||||
// Load the full document so other pages stay intact
|
||||
onDisplayChart(
|
||||
editedXml,
|
||||
false,
|
||||
true,
|
||||
)
|
||||
lastProcessedXmlRef.current.set(
|
||||
pending.toolCallId +
|
||||
|
||||
+103
-50
@@ -32,6 +32,7 @@ import { useSessionManager } from "@/hooks/use-session-manager"
|
||||
import { useValidateDiagram } from "@/hooks/use-validate-diagram"
|
||||
import { getApiEndpoint } from "@/lib/base-path"
|
||||
import { findCachedResponse } from "@/lib/cached-responses"
|
||||
import { isMinimalDiagram } from "@/lib/chat-helpers"
|
||||
import type { DrawioTheme } from "@/lib/drawio-themes"
|
||||
import { formatMessage } from "@/lib/i18n/utils"
|
||||
import { isPdfFile, isTextFile } from "@/lib/pdf-utils"
|
||||
@@ -40,9 +41,12 @@ import { STORAGE_KEYS } from "@/lib/storage"
|
||||
import type { UrlData } from "@/lib/url-utils"
|
||||
import { type FileData, useFileProcessor } from "@/lib/use-file-processor"
|
||||
import { useQuotaManager } from "@/lib/use-quota-manager"
|
||||
import { cn, formatXML, isRealDiagram } from "@/lib/utils"
|
||||
import { cn, formatXML, isRealDiagram, wrapWithMxFile } from "@/lib/utils"
|
||||
import type { ValidationState } from "./chat/ValidationCard"
|
||||
import { ChatMessageDisplay } from "./chat-message-display"
|
||||
import {
|
||||
APPENDED_FILE_SECTIONS_PATTERN,
|
||||
ChatMessageDisplay,
|
||||
} from "./chat-message-display"
|
||||
import { DevXmlSimulator } from "./dev-xml-simulator"
|
||||
|
||||
// localStorage keys for persistence
|
||||
@@ -107,6 +111,18 @@ function hasToolErrors(messages: ChatMessage[]): boolean {
|
||||
return lastToolPart?.state === TOOL_ERROR_STATE
|
||||
}
|
||||
|
||||
/**
|
||||
* 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.
|
||||
*/
|
||||
function getFirstPageXml(xml: string): 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
|
||||
}
|
||||
|
||||
export default function ChatPanel({
|
||||
isVisible,
|
||||
onToggleVisibility,
|
||||
@@ -336,19 +352,8 @@ export default function ChatPanel({
|
||||
localStorage.setItem(STORAGE_KEYS.maxOutputTokens, digitsOnly)
|
||||
}, [])
|
||||
|
||||
// Ref to store the sendMessage function for use in callbacks
|
||||
const sendMessageRef = useRef<typeof sendMessage | null>(null)
|
||||
|
||||
// Callback to improve diagram with validation suggestions
|
||||
const handleImproveWithSuggestions = useCallback((feedback: string) => {
|
||||
if (sendMessageRef.current) {
|
||||
// Send the feedback as a new user message to trigger regeneration
|
||||
sendMessageRef.current({
|
||||
role: "user",
|
||||
parts: [{ type: "text", text: feedback }],
|
||||
})
|
||||
}
|
||||
}, [])
|
||||
// Failed VLM validations in the current user turn (reset on user action)
|
||||
const validationRetryCountRef = useRef(0)
|
||||
|
||||
// VLM validation hook using AI SDK's useObject
|
||||
const { validateWithFallback } = useValidateDiagram()
|
||||
@@ -357,6 +362,7 @@ export default function ChatPanel({
|
||||
const { handleToolCall } = useDiagramToolHandlers({
|
||||
partialXmlRef,
|
||||
editDiagramOriginalXmlRef,
|
||||
validationRetryCountRef,
|
||||
chartXMLRef,
|
||||
onDisplayChart,
|
||||
onFetchChart,
|
||||
@@ -518,11 +524,6 @@ export default function ChatPanel({
|
||||
},
|
||||
})
|
||||
|
||||
// Store sendMessage in ref for use in callbacks (like handleImproveWithSuggestions)
|
||||
useEffect(() => {
|
||||
sendMessageRef.current = sendMessage
|
||||
}, [sendMessage])
|
||||
|
||||
// Ref to track latest messages for unload persistence
|
||||
const messagesRef = useRef(messages)
|
||||
useEffect(() => {
|
||||
@@ -531,6 +532,9 @@ export default function ChatPanel({
|
||||
|
||||
// Track last synced session ID to detect external changes (e.g., URL back/forward)
|
||||
const lastSyncedSessionIdRef = useRef<string | null>(null)
|
||||
// Messages array from our latest save. A session holding this exact array was
|
||||
// created by our own save, so it must not be treated as an external switch.
|
||||
const lastSavedMessagesRef = useRef<unknown[] | null>(null)
|
||||
|
||||
// Helper: Sync UI state with session data (eliminates duplication)
|
||||
// Track message IDs that are being loaded from session (to skip animations/scroll)
|
||||
@@ -597,8 +601,10 @@ export default function ChatPanel({
|
||||
thumbnailDataUrl = latestSvgRef.current
|
||||
}
|
||||
}
|
||||
const messages = sanitizeMessages(messagesRef.current)
|
||||
lastSavedMessagesRef.current = messages
|
||||
return {
|
||||
messages: sanitizeMessages(messagesRef.current),
|
||||
messages,
|
||||
xmlSnapshots: Array.from(xmlSnapshotsRef.current.entries()),
|
||||
diagramXml: currentDiagramXml,
|
||||
thumbnailDataUrl,
|
||||
@@ -651,8 +657,13 @@ export default function ChatPanel({
|
||||
// Skip if session ID hasn't changed (our own saves don't change the ID)
|
||||
if (newSessionId === lastSyncedSessionIdRef.current) return
|
||||
|
||||
// Our own save created this session; the UI already shows its content
|
||||
const isOwnNewSession =
|
||||
newSession?.messages === lastSavedMessagesRef.current
|
||||
|
||||
// Update last synced ID
|
||||
lastSyncedSessionIdRef.current = newSessionId
|
||||
if (isOwnNewSession) return
|
||||
|
||||
// Sync UI with new session
|
||||
if (newSession) {
|
||||
@@ -793,12 +804,23 @@ export default function ChatPanel({
|
||||
const onFormSubmit = async (e: React.FormEvent<HTMLFormElement>) => {
|
||||
e.preventDefault()
|
||||
const isProcessing = status === "streaming" || status === "submitted"
|
||||
if (input.trim() && !isProcessing) {
|
||||
// Check if input matches a cached example (only when no messages yet)
|
||||
if (messages.length === 0) {
|
||||
// Attachments still extracting have no text yet. Template sends call
|
||||
// requestSubmit() and skip the disabled send button, so check here too.
|
||||
const isExtracting =
|
||||
files.some((f) => pdfData.get(f)?.isExtracting) ||
|
||||
Array.from(urlData.values()).some((d) => d.isExtracting)
|
||||
if (input.trim() && !isProcessing && !isExtracting) {
|
||||
// Check if input matches a cached example (only when no messages
|
||||
// yet and the canvas is empty, same rule as the server)
|
||||
if (
|
||||
messages.length === 0 &&
|
||||
isMinimalDiagram(chartXMLRef.current || "")
|
||||
) {
|
||||
// Pass the file name so a user's own file never matches an example
|
||||
const cached = findCachedResponse(
|
||||
input.trim(),
|
||||
files.length > 0,
|
||||
files.length === 1 ? files[0].name : undefined,
|
||||
)
|
||||
if (cached) {
|
||||
// Add user message and fake assistant response to messages
|
||||
@@ -834,6 +856,11 @@ export default function ChatPanel({
|
||||
],
|
||||
},
|
||||
] as any)
|
||||
// Snapshot the canvas before the example so editing this message works
|
||||
xmlSnapshotsRef.current.set(
|
||||
0,
|
||||
chartXMLRef.current || wrapWithMxFile(""),
|
||||
)
|
||||
setInput("")
|
||||
sessionStorage.removeItem(SESSION_STORAGE_INPUT_KEY)
|
||||
setFiles([])
|
||||
@@ -843,9 +870,6 @@ export default function ChatPanel({
|
||||
}
|
||||
|
||||
try {
|
||||
let chartXml = await onFetchChart()
|
||||
chartXml = formatXML(chartXml)
|
||||
|
||||
// Build user text by concatenating input with pre-extracted text
|
||||
// (Backend only reads first text part, so we must combine them)
|
||||
const parts: any[] = []
|
||||
@@ -860,20 +884,7 @@ export default function ChatPanel({
|
||||
// Add the combined text as the first part
|
||||
parts.unshift({ type: "text", text: userText })
|
||||
|
||||
// Get previous XML from the last snapshot (before this message)
|
||||
const snapshotKeys = Array.from(
|
||||
xmlSnapshotsRef.current.keys(),
|
||||
).sort((a, b) => b - a)
|
||||
const previousXml =
|
||||
snapshotKeys.length > 0
|
||||
? xmlSnapshotsRef.current.get(snapshotKeys[0]) || ""
|
||||
: ""
|
||||
|
||||
// Save XML snapshot for this message (will be at index = current messages.length)
|
||||
const messageIndex = messages.length
|
||||
xmlSnapshotsRef.current.set(messageIndex, chartXml)
|
||||
|
||||
sendChatMessage(parts, chartXml, previousXml, sessionId)
|
||||
await sendWithCurrentDiagram(parts)
|
||||
|
||||
// Token count is tracked in onFinish with actual server usage
|
||||
setInput("")
|
||||
@@ -882,10 +893,37 @@ export default function ChatPanel({
|
||||
setUrlData(new Map())
|
||||
} catch (error) {
|
||||
console.error("Error fetching chart data:", error)
|
||||
toast.error(dict.errors.failedToExport)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Export the current diagram, snapshot it for this message, and send
|
||||
const sendWithCurrentDiagram = async (parts: any[]) => {
|
||||
const chartXml = formatXML(await onFetchChart())
|
||||
const previousXml = getPreviousXml(messages.length)
|
||||
|
||||
// Snapshot the full multi-page document (kept fresh by autosave) so
|
||||
// regenerate/edit can restore every page; the model gets page 1 only
|
||||
xmlSnapshotsRef.current.set(
|
||||
messages.length,
|
||||
chartXMLRef.current || chartXml,
|
||||
)
|
||||
|
||||
sendChatMessage(parts, chartXml, previousXml, sessionId)
|
||||
}
|
||||
|
||||
// Send VLM validation feedback as a new user message through the normal send path
|
||||
const handleImproveWithSuggestions = async (feedback: string) => {
|
||||
if (status === "streaming" || status === "submitted") return
|
||||
try {
|
||||
await sendWithCurrentDiagram([{ type: "text", text: feedback }])
|
||||
} catch (error) {
|
||||
console.error("Error fetching chart data:", error)
|
||||
toast.error(dict.errors.failedToExport)
|
||||
}
|
||||
}
|
||||
|
||||
// Handle session switching from history dropdown
|
||||
const handleSelectSession = useCallback(
|
||||
async (sessionId: string) => {
|
||||
@@ -989,10 +1027,9 @@ export default function ChatPanel({
|
||||
// Handle sending a template directly (called from TemplatePanel)
|
||||
const handleSendTemplate = useCallback(
|
||||
async (template: { prompt: string }) => {
|
||||
// Keep attachments: they are sent along with the template prompt
|
||||
flushSync(() => {
|
||||
setInput(template.prompt)
|
||||
setFiles([])
|
||||
setUrlData(new Map())
|
||||
})
|
||||
|
||||
const formElement = document.getElementById(
|
||||
@@ -1002,7 +1039,7 @@ export default function ChatPanel({
|
||||
formElement.requestSubmit()
|
||||
}
|
||||
},
|
||||
[setInput, setFiles, setUrlData],
|
||||
[setInput],
|
||||
)
|
||||
|
||||
const handleInputChange = (
|
||||
@@ -1017,13 +1054,15 @@ export default function ChatPanel({
|
||||
}
|
||||
|
||||
// Helper functions for message actions (regenerate/edit)
|
||||
// Extract previous XML snapshot before a given message index
|
||||
// Extract previous XML snapshot (first page, as sent to the model) before a given message index
|
||||
const getPreviousXml = (beforeIndex: number): string => {
|
||||
const snapshotKeys = Array.from(xmlSnapshotsRef.current.keys())
|
||||
.filter((k) => k < beforeIndex)
|
||||
.sort((a, b) => b - a)
|
||||
return snapshotKeys.length > 0
|
||||
? xmlSnapshotsRef.current.get(snapshotKeys[0]) || ""
|
||||
? getFirstPageXml(
|
||||
xmlSnapshotsRef.current.get(snapshotKeys[0]) || "",
|
||||
)
|
||||
: ""
|
||||
}
|
||||
|
||||
@@ -1075,6 +1114,7 @@ export default function ChatPanel({
|
||||
// Reset all retry/continuation state on user-initiated message
|
||||
autoRetryCountRef.current = 0
|
||||
continuationRetryCountRef.current = 0
|
||||
validationRetryCountRef.current = 0
|
||||
partialXmlRef.current = ""
|
||||
|
||||
const config = getSelectedAIConfig()
|
||||
@@ -1223,7 +1263,12 @@ export default function ChatPanel({
|
||||
})
|
||||
|
||||
// Now send the message after state is guaranteed to be updated
|
||||
sendChatMessage(userParts, savedXml, previousXml, sessionId)
|
||||
sendChatMessage(
|
||||
userParts,
|
||||
getFirstPageXml(savedXml),
|
||||
previousXml,
|
||||
sessionId,
|
||||
)
|
||||
}
|
||||
|
||||
const handleEditMessage = async (messageIndex: number, newText: string) => {
|
||||
@@ -1250,10 +1295,13 @@ export default function ChatPanel({
|
||||
// Clean up snapshots for messages after the user message (they will be removed)
|
||||
cleanupSnapshotsAfter(messageIndex)
|
||||
|
||||
// Create new parts with updated text
|
||||
// Create new parts with updated text. The edit box only shows the typed
|
||||
// text, so keep the appended PDF/file/URL content
|
||||
const newParts = message.parts?.map((part: any) => {
|
||||
if (part.type === "text") {
|
||||
return { ...part, text: newText }
|
||||
const appended =
|
||||
part.text.match(APPENDED_FILE_SECTIONS_PATTERN)?.[0] ?? ""
|
||||
return { ...part, text: newText + appended }
|
||||
}
|
||||
return part
|
||||
}) || [{ type: "text", text: newText }]
|
||||
@@ -1266,7 +1314,12 @@ export default function ChatPanel({
|
||||
})
|
||||
|
||||
// Now send the edited message after state is guaranteed to be updated
|
||||
sendChatMessage(newParts, savedXml, previousXml, sessionId)
|
||||
sendChatMessage(
|
||||
newParts,
|
||||
getFirstPageXml(savedXml),
|
||||
previousXml,
|
||||
sessionId,
|
||||
)
|
||||
}
|
||||
|
||||
// Collapsed view (desktop only)
|
||||
|
||||
@@ -1,5 +1,4 @@
|
||||
import type { MutableRefObject } from "react"
|
||||
import { useRef } from "react"
|
||||
import type { DiagramOperation } from "@/components/chat/types"
|
||||
import type {
|
||||
ValidationState,
|
||||
@@ -48,6 +47,8 @@ type ValidateDiagramFn = (
|
||||
interface UseDiagramToolHandlersParams {
|
||||
partialXmlRef: MutableRefObject<string>
|
||||
editDiagramOriginalXmlRef: MutableRefObject<Map<string, string>>
|
||||
// Failed VLM validations in the current user turn (reset on each user message)
|
||||
validationRetryCountRef: MutableRefObject<number>
|
||||
chartXMLRef: MutableRefObject<string>
|
||||
onDisplayChart: (xml: string, skipValidation?: boolean) => string | null
|
||||
onFetchChart: (saveToHistory?: boolean) => Promise<string>
|
||||
@@ -72,6 +73,7 @@ interface UseDiagramToolHandlersParams {
|
||||
export function useDiagramToolHandlers({
|
||||
partialXmlRef,
|
||||
editDiagramOriginalXmlRef,
|
||||
validationRetryCountRef,
|
||||
chartXMLRef,
|
||||
onDisplayChart,
|
||||
onFetchChart,
|
||||
@@ -82,9 +84,6 @@ export function useDiagramToolHandlers({
|
||||
sessionId,
|
||||
onValidationStateChange,
|
||||
}: UseDiagramToolHandlersParams) {
|
||||
// Track validation retry count per tool call
|
||||
const validationRetryCountRef = useRef<Map<string, number>>(new Map())
|
||||
|
||||
// Helper to update validation state
|
||||
const updateValidationState = (
|
||||
toolCallId: string,
|
||||
@@ -232,17 +231,15 @@ ${finalXml}
|
||||
)
|
||||
}
|
||||
|
||||
const retryCount =
|
||||
validationRetryCountRef.current.get(
|
||||
toolCall.toolCallId,
|
||||
) || 0
|
||||
// Each retry is a new tool call, so count attempts per user turn
|
||||
const attempt = validationRetryCountRef.current + 1
|
||||
|
||||
// Notify UI that we're validating (include the image)
|
||||
updateValidationState(
|
||||
toolCall.toolCallId,
|
||||
"validating",
|
||||
{
|
||||
attempt: retryCount + 1,
|
||||
attempt,
|
||||
maxAttempts: MAX_VALIDATION_RETRIES,
|
||||
imageData: capturedPngData,
|
||||
},
|
||||
@@ -254,17 +251,14 @@ ${finalXml}
|
||||
)
|
||||
|
||||
if (!result.valid) {
|
||||
if (retryCount < MAX_VALIDATION_RETRIES) {
|
||||
validationRetryCountRef.current.set(
|
||||
toolCall.toolCallId,
|
||||
retryCount + 1,
|
||||
)
|
||||
if (attempt < MAX_VALIDATION_RETRIES) {
|
||||
validationRetryCountRef.current = attempt
|
||||
|
||||
const feedback =
|
||||
formatValidationFeedback(result)
|
||||
if (DEBUG) {
|
||||
console.log(
|
||||
`[display_diagram] Validation failed (attempt ${retryCount + 1}/${MAX_VALIDATION_RETRIES}):`,
|
||||
`[display_diagram] Validation failed (attempt ${attempt}/${MAX_VALIDATION_RETRIES}):`,
|
||||
result.issues,
|
||||
)
|
||||
}
|
||||
@@ -274,7 +268,7 @@ ${finalXml}
|
||||
toolCall.toolCallId,
|
||||
"failed",
|
||||
{
|
||||
attempt: retryCount + 1,
|
||||
attempt,
|
||||
maxAttempts: MAX_VALIDATION_RETRIES,
|
||||
result,
|
||||
imageData: capturedPngData,
|
||||
@@ -285,19 +279,17 @@ ${finalXml}
|
||||
tool: "display_diagram",
|
||||
toolCallId: toolCall.toolCallId,
|
||||
state: "output-error",
|
||||
errorText: `[Validation attempt ${retryCount + 1}/${MAX_VALIDATION_RETRIES}]\n${feedback}`,
|
||||
errorText: `[Validation attempt ${attempt}/${MAX_VALIDATION_RETRIES}]\n${feedback}`,
|
||||
})
|
||||
return
|
||||
} else {
|
||||
// Max retries reached - accept the diagram with warning
|
||||
// Last attempt - accept the diagram with warning
|
||||
if (DEBUG) {
|
||||
console.log(
|
||||
"[display_diagram] Max validation retries reached, accepting diagram",
|
||||
)
|
||||
}
|
||||
validationRetryCountRef.current.delete(
|
||||
toolCall.toolCallId,
|
||||
)
|
||||
validationRetryCountRef.current = 0
|
||||
|
||||
// Notify UI that we're accepting with issues (include the image)
|
||||
updateValidationState(
|
||||
@@ -314,10 +306,8 @@ ${finalXml}
|
||||
return
|
||||
}
|
||||
} else {
|
||||
// Validation passed - clean up retry count
|
||||
validationRetryCountRef.current.delete(
|
||||
toolCall.toolCallId,
|
||||
)
|
||||
// Validation passed - reset retry count
|
||||
validationRetryCountRef.current = 0
|
||||
if (DEBUG) {
|
||||
console.log(
|
||||
"[display_diagram] Validation passed!",
|
||||
@@ -382,12 +372,17 @@ ${finalXml}
|
||||
}
|
||||
|
||||
let currentXml = ""
|
||||
// Use the original XML captured during streaming (shared with chat-message-display)
|
||||
// This ensures we apply operations to the same base XML that streaming used
|
||||
const originalXml = editDiagramOriginalXmlRef.current.get(
|
||||
toolCall.toolCallId,
|
||||
)
|
||||
// On failure, undo the streaming preview so the canvas matches the XML
|
||||
// reported back to the model
|
||||
const restoreOriginal = () => {
|
||||
if (originalXml) onDisplayChart(originalXml, true)
|
||||
}
|
||||
try {
|
||||
// Use the original XML captured during streaming (shared with chat-message-display)
|
||||
// This ensures we apply operations to the same base XML that streaming used
|
||||
const originalXml = editDiagramOriginalXmlRef.current.get(
|
||||
toolCall.toolCallId,
|
||||
)
|
||||
if (originalXml) {
|
||||
currentXml = originalXml
|
||||
} else {
|
||||
@@ -416,6 +411,7 @@ ${finalXml}
|
||||
)
|
||||
.join("\n")
|
||||
|
||||
restoreOriginal()
|
||||
addToolOutput({
|
||||
tool: "edit_diagram",
|
||||
toolCallId: toolCall.toolCallId,
|
||||
@@ -441,6 +437,7 @@ Please check the cell IDs and retry.`,
|
||||
"[edit_diagram] Validation error:",
|
||||
validationError,
|
||||
)
|
||||
restoreOriginal()
|
||||
addToolOutput({
|
||||
tool: "edit_diagram",
|
||||
toolCallId: toolCall.toolCallId,
|
||||
@@ -472,6 +469,7 @@ Please fix the operations to avoid structural issues.`,
|
||||
const errorMessage =
|
||||
error instanceof Error ? error.message : String(error)
|
||||
|
||||
restoreOriginal()
|
||||
addToolOutput({
|
||||
tool: "edit_diagram",
|
||||
toolCallId: toolCall.toolCallId,
|
||||
@@ -496,6 +494,19 @@ Please check cell IDs and retry, or use display_diagram to regenerate.`,
|
||||
) => {
|
||||
const { xml } = toolCall.input as { xml: string }
|
||||
|
||||
// Nothing to continue: loading the fragment alone would replace the whole diagram
|
||||
if (!partialXmlRef.current) {
|
||||
addToolOutput({
|
||||
tool: "append_diagram",
|
||||
toolCallId: toolCall.toolCallId,
|
||||
state: "output-error",
|
||||
errorText: `ERROR: There is no truncated diagram to continue, so append_diagram cannot be used now.
|
||||
|
||||
Use display_diagram to create the complete diagram, or edit_diagram to change the current one.`,
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// Detect if LLM incorrectly started fresh instead of continuing
|
||||
// LLM should only output bare mxCells now, so wrapper tags indicate error
|
||||
const trimmed = xml.trim()
|
||||
|
||||
@@ -6,6 +6,7 @@
|
||||
|
||||
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"
|
||||
import {
|
||||
type ValidationResult,
|
||||
@@ -39,6 +40,8 @@ export function useValidateDiagram(options: UseValidateDiagramOptions = {}) {
|
||||
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,
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
export interface CachedResponse {
|
||||
promptText: string
|
||||
hasImage: boolean
|
||||
// Name of the bundled example file the prompt is sent with
|
||||
fileName?: string
|
||||
xml: string
|
||||
}
|
||||
|
||||
@@ -254,6 +256,7 @@ export const CACHED_EXAMPLE_RESPONSES: CachedResponse[] = [
|
||||
{
|
||||
promptText: "Replicate this in aws style",
|
||||
hasImage: true,
|
||||
fileName: "architecture.png",
|
||||
xml: `<mxCell id="2" value="AWS" style="sketch=0;outlineConnect=0;gradientColor=none;html=1;whiteSpace=wrap;fontSize=12;fontStyle=0;container=1;pointerEvents=0;collapsible=0;recursiveResize=0;shape=mxgraph.aws4.group;grIcon=mxgraph.aws4.group_aws_cloud;strokeColor=#232F3E;fillColor=none;verticalAlign=top;align=left;spacingLeft=30;fontColor=#232F3E;dashed=0;rounded=1;arcSize=5;" vertex="1" parent="1">
|
||||
<mxGeometry x="340" y="40" width="880" height="520" as="geometry"/>
|
||||
</mxCell>
|
||||
@@ -318,6 +321,7 @@ export const CACHED_EXAMPLE_RESPONSES: CachedResponse[] = [
|
||||
{
|
||||
promptText: "Replicate this flowchart.",
|
||||
hasImage: true,
|
||||
fileName: "example.png",
|
||||
xml: `<mxCell id="2" value="Lamp doesn't work" style="rounded=1;whiteSpace=wrap;html=1;fillColor=#ffcccc;strokeColor=#000000;strokeWidth=2;fontSize=18;fontStyle=0;" vertex="1" parent="1">
|
||||
<mxGeometry x="140" y="40" width="180" height="60" as="geometry"/>
|
||||
</mxCell>
|
||||
@@ -379,6 +383,7 @@ export const CACHED_EXAMPLE_RESPONSES: CachedResponse[] = [
|
||||
{
|
||||
promptText: "Summarize this paper as a diagram",
|
||||
hasImage: true,
|
||||
fileName: "chain-of-thought.txt",
|
||||
xml: `<mxCell id="title_bg" parent="1"
|
||||
style="rounded=1;whiteSpace=wrap;html=1;fillColor=#1a237e;strokeColor=none;arcSize=8;"
|
||||
value="" vertex="1">
|
||||
@@ -879,14 +884,19 @@ export const CACHED_EXAMPLE_RESPONSES: CachedResponse[] = [
|
||||
},
|
||||
]
|
||||
|
||||
// Examples that come with a file only match when that exact example file is
|
||||
// attached, so a user's own file with the same prompt still goes to the model.
|
||||
// Callers that can't tell file names (the server) only get text-only examples.
|
||||
export function findCachedResponse(
|
||||
promptText: string,
|
||||
hasImage: boolean,
|
||||
fileName?: string,
|
||||
): CachedResponse | undefined {
|
||||
return CACHED_EXAMPLE_RESPONSES.find(
|
||||
(c) =>
|
||||
c.promptText === promptText &&
|
||||
c.hasImage === hasImage &&
|
||||
(!c.fileName || c.fileName === fileName) &&
|
||||
c.xml !== "",
|
||||
)
|
||||
}
|
||||
|
||||
@@ -14,12 +14,35 @@ describe("findCachedResponse", () => {
|
||||
expect(result?.xml).toContain("Transformer Architecture")
|
||||
})
|
||||
|
||||
it("returns cached response for exact match with image", () => {
|
||||
const result = findCachedResponse("Replicate this in aws style", true)
|
||||
it("returns cached response for exact match with the example file", () => {
|
||||
const result = findCachedResponse(
|
||||
"Replicate this in aws style",
|
||||
true,
|
||||
"architecture.png",
|
||||
)
|
||||
expect(result).toBeDefined()
|
||||
expect(result?.xml).toContain("AWS")
|
||||
})
|
||||
|
||||
it("returns undefined when the user attached their own file", () => {
|
||||
expect(
|
||||
findCachedResponse("Replicate this flowchart.", true, "mine.png"),
|
||||
).toBeUndefined()
|
||||
expect(
|
||||
findCachedResponse(
|
||||
"Summarize this paper as a diagram",
|
||||
true,
|
||||
"thesis.pdf",
|
||||
),
|
||||
).toBeUndefined()
|
||||
})
|
||||
|
||||
it("returns undefined for file examples when the file name is unknown", () => {
|
||||
// The server only knows whether a file is attached, not which one
|
||||
const result = findCachedResponse("Replicate this in aws style", true)
|
||||
expect(result).toBeUndefined()
|
||||
})
|
||||
|
||||
it("returns undefined for non-matching prompt", () => {
|
||||
const result = findCachedResponse(
|
||||
"random prompt that doesn't exist",
|
||||
|
||||
Reference in New Issue
Block a user