diff --git a/components/chat-message-display.tsx b/components/chat-message-display.tsx index ada6ae5..8f8fd99 100644 --- a/components/chat-message-display.tsx +++ b/components/chat-message-display.tsx @@ -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 + diff --git a/components/chat-panel.tsx b/components/chat-panel.tsx index 66f44b6..6f3a158 100644 --- a/components/chat-panel.tsx +++ b/components/chat-panel.tsx @@ -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("(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(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(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) => { 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) diff --git a/hooks/use-diagram-tool-handlers.ts b/hooks/use-diagram-tool-handlers.ts index d1970d8..a435078 100644 --- a/hooks/use-diagram-tool-handlers.ts +++ b/hooks/use-diagram-tool-handlers.ts @@ -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 editDiagramOriginalXmlRef: MutableRefObject> + // Failed VLM validations in the current user turn (reset on each user message) + validationRetryCountRef: MutableRefObject chartXMLRef: MutableRefObject onDisplayChart: (xml: string, skipValidation?: boolean) => string | null onFetchChart: (saveToHistory?: boolean) => Promise @@ -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>(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() diff --git a/hooks/use-validate-diagram.ts b/hooks/use-validate-diagram.ts index b9cea01..c4593c3 100644 --- a/hooks/use-validate-diagram.ts +++ b/hooks/use-validate-diagram.ts @@ -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, diff --git a/lib/cached-responses.ts b/lib/cached-responses.ts index 8b8375f..c2b474f 100644 --- a/lib/cached-responses.ts +++ b/lib/cached-responses.ts @@ -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: ` @@ -318,6 +321,7 @@ export const CACHED_EXAMPLE_RESPONSES: CachedResponse[] = [ { promptText: "Replicate this flowchart.", hasImage: true, + fileName: "example.png", xml: ` @@ -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: ` @@ -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 !== "", ) } diff --git a/tests/unit/cached-responses.test.ts b/tests/unit/cached-responses.test.ts index 2e0f873..7eb6b84 100644 --- a/tests/unit/cached-responses.test.ts +++ b/tests/unit/cached-responses.test.ts @@ -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",