From a20d14ef9d6fba23feef40d60df08a71ec7c87da Mon Sep 17 00:00:00 2001 From: "dayuan.jiang" Date: Tue, 2 Dec 2025 19:11:23 +0900 Subject: [PATCH] feat: improve diagram edit tracking and cache handling - Track lastGeneratedXml to detect user modifications - Only send XML context when needed (saves tokens) - Fix cached response streaming to include tool output - Add debug logging for model messages and steps - Enable multi-step tool execution with maxSteps: 5 --- app/api/chat/route.ts | 47 +++++++++++++++++++++++------ components/chat-message-display.tsx | 24 +++++++++------ components/chat-panel.tsx | 30 +++++++++++++----- contexts/diagram-context.tsx | 38 ++++++++++++++++++++++- lib/utils.ts | 9 ++++++ 5 files changed, 120 insertions(+), 28 deletions(-) diff --git a/app/api/chat/route.ts b/app/api/chat/route.ts index a4aeaf9..842c6d8 100644 --- a/app/api/chat/route.ts +++ b/app/api/chat/route.ts @@ -19,8 +19,12 @@ function createCachedStreamResponse(xml: string): Response { execute: async ({ writer }) => { writer.write({ type: 'start' }); writer.write({ type: 'tool-input-start', toolCallId, toolName: 'display_diagram' }); - writer.write({ type: 'tool-input-delta', toolCallId, inputTextDelta: xml }); + // Stream the XML as JSON input so it matches the tool schema exactly + writer.write({ type: 'tool-input-delta', toolCallId, inputTextDelta: JSON.stringify({ xml }) }); + // Input must match the tool schema (only xml field, no extra fields like fromCache) writer.write({ type: 'tool-input-available', toolCallId, toolName: 'display_diagram', input: { xml } }); + // Include tool output so the message is complete for follow-up conversations + writer.write({ type: 'tool-output-available', toolCallId, output: 'Successfully displayed the diagram.' }); writer.write({ type: 'finish' }); }, }); @@ -30,7 +34,7 @@ function createCachedStreamResponse(xml: string): Response { export async function POST(req: Request) { try { - const { messages, xml } = await req.json(); + const { messages, xml, lastGeneratedXml } = await req.json(); // === CACHE CHECK START === const isFirstMessage = messages.length === 1; @@ -74,6 +78,7 @@ parameters: { IMPORTANT: Choose the right tool: - Use display_diagram for: Creating new diagrams, major restructuring, or when the current diagram XML is empty - Use edit_diagram for: Small modifications, adding/removing elements, changing text/colors, repositioning items +- When using edit_diagram: If the current diagram XML is provided in the user message context, use it as the source of truth for constructing search patterns. If no XML is provided, you can use your memory of the diagram structure. Core capabilities: - Generate valid, well-formed XML strings for draw.io diagrams @@ -120,19 +125,34 @@ When using edit_diagram tool: // Extract file parts (images) from the last message const fileParts = lastMessage.parts?.filter((part: any) => part.type === 'file') || []; - const formattedTextContent = ` -Current diagram XML: -"""xml -${xml || ''} -""" -User input: + // Check diagram state + const hasDiagram = xml && !isMinimalDiagram(xml); + const noHistory = !lastGeneratedXml || lastGeneratedXml.trim() === ''; + const userModified = hasDiagram && lastGeneratedXml && xml !== lastGeneratedXml; + + // Build context based on diagram state + let diagramContext = ''; + if (hasDiagram && noHistory) { + // No history (e.g., cached response) - include XML directly + diagramContext = `\n\n[Current diagram XML - use this as source of truth for edits:]\n\`\`\`xml\n${xml}\n\`\`\``; + } else if (userModified) { + // User modified - include XML + diagramContext = `\n\n[User modified the diagram. Current XML:]\n\`\`\`xml\n${xml}\n\`\`\``; + } + // If unchanged and has history, agent can use memory (no XML sent = save tokens) + + const formattedTextContent = `User input: """md ${lastMessageText} -"""`; +"""${diagramContext}`; // Convert UIMessages to ModelMessages and add system message const modelMessages = convertToModelMessages(messages); + // Debug: Log the full structure of model messages to diagnose Bedrock API errors + console.log('[Debug] Model messages structure:'); + console.log('[Debug] Full messages JSON:', JSON.stringify(modelMessages, null, 2)); + // Log messages with empty content for debugging (helps identify root cause) const emptyMessages = modelMessages.filter((msg: any) => !msg.content || !Array.isArray(msg.content) || msg.content.length === 0 @@ -217,7 +237,13 @@ ${lastMessageText} messages: [systemMessageWithCache, ...enhancedMessages], ...(providerOptions && { providerOptions }), ...(headers && { headers }), - onFinish: ({ usage, providerMetadata }) => { + onStepFinish: ({ stepType, toolCalls, toolResults }) => { + console.log('[Step] Type:', stepType); + console.log('[Step] Tool calls:', toolCalls?.map(t => t.toolName)); + console.log('[Step] Tool results:', toolResults?.length); + }, + onFinish: ({ usage, providerMetadata, steps }) => { + console.log('[Finish] Total steps:', steps?.length); console.log('[Cache] Usage:', JSON.stringify({ inputTokens: usage?.inputTokens, outputTokens: usage?.outputTokens, @@ -261,6 +287,7 @@ IMPORTANT: Keep edits concise: }, }, temperature: 0, + maxSteps: 5, // Allow model to continue after server-side tool execution }); // Error handler function to provide detailed error messages diff --git a/components/chat-message-display.tsx b/components/chat-message-display.tsx index ad55dd2..5a01ed2 100644 --- a/components/chat-message-display.tsx +++ b/components/chat-message-display.tsx @@ -157,19 +157,23 @@ export function ChatMessageDisplay({
) : state === "output-available" ? (
- {output || (toolName === "display_diagram" - ? "Diagram generated" - : toolName === "edit_diagram" - ? "Diagram edited" - : "Tool executed")} + {typeof output === "object" && output !== null + ? (output as any).message || JSON.stringify(output) + : output || (toolName === "display_diagram" + ? "Diagram generated" + : toolName === "edit_diagram" + ? "Diagram edited" + : "Tool executed")}
) : state === "output-error" ? (
- {output || (toolName === "display_diagram" - ? "Error generating diagram" - : toolName === "edit_diagram" - ? "Error editing diagram" - : "Tool error")} + {typeof output === "object" && output !== null + ? (output as any).message || JSON.stringify(output) + : output || (toolName === "display_diagram" + ? "Error generating diagram" + : toolName === "edit_diagram" + ? "Error editing diagram" + : "Tool error")}
) : null}
diff --git a/components/chat-panel.tsx b/components/chat-panel.tsx index 2eb8f30..780b43d 100644 --- a/components/chat-panel.tsx +++ b/components/chat-panel.tsx @@ -33,6 +33,8 @@ export default function ChatPanel({ isVisible, onToggleVisibility }: ChatPanelPr resolverRef, chartXML, clearDiagram, + getLastAgentGeneratedXml, + markAgentDiagramPending, } = useDiagram(); const onFetchChart = () => { @@ -73,11 +75,18 @@ export default function ChatPanel({ isVisible, onToggleVisibility }: ChatPanelPr }), async onToolCall({ toolCall }) { if (toolCall.toolName === "display_diagram") { - // Diagram is handled streamingly in the ChatMessageDisplay component + // Check if this is a cached response by looking at the toolCallId prefix + const isCached = toolCall.toolCallId.startsWith('cached-'); + + // Only mark as pending if agent actually generated it (not cached) + // This ensures lastAgentGeneratedXml stays empty for cached responses + if (!isCached) { + markAgentDiagramPending(); + } + addToolResult({ - tool: "display_diagram", toolCallId: toolCall.toolCallId, - output: "Successfully displayed the diagram.", + result: "Successfully displayed the diagram.", }); } else if (toolCall.toolName === "edit_diagram") { const { edits } = toolCall.input as { @@ -96,10 +105,12 @@ export default function ChatPanel({ isVisible, onToggleVisibility }: ChatPanelPr // Load the edited diagram onDisplayChart(editedXml); + // Mark that an agent diagram is pending - the next export will update lastAgentGeneratedXml + markAgentDiagramPending(); + addToolResult({ - tool: "edit_diagram", toolCallId: toolCall.toolCallId, - output: `Successfully applied ${edits.length} edit(s) to the diagram.`, + result: `Successfully applied ${edits.length} edit(s) to the diagram.`, }); } catch (error) { console.error("Edit diagram failed:", error); @@ -108,9 +119,8 @@ export default function ChatPanel({ isVisible, onToggleVisibility }: ChatPanelPr // Provide detailed error with current diagram XML addToolResult({ - tool: "edit_diagram", toolCallId: toolCall.toolCallId, - output: `Edit failed: ${errorMessage} + result: `Edit failed: ${errorMessage} Current diagram XML: \`\`\`xml @@ -171,11 +181,17 @@ Please retry with an adjusted search pattern or use display_diagram if retries a } } + const lastGenXml = getLastAgentGeneratedXml(); + console.log('[ChatPanel] Sending message with xml length:', chartXml.length); + console.log('[ChatPanel] lastGeneratedXml length:', lastGenXml.length); + console.log('[ChatPanel] Are they equal:', chartXml === lastGenXml); + sendMessage( { parts }, { body: { xml: chartXml, + lastGeneratedXml: lastGenXml, }, } ); diff --git a/contexts/diagram-context.tsx b/contexts/diagram-context.tsx index a47f8ae..4fe8394 100644 --- a/contexts/diagram-context.tsx +++ b/contexts/diagram-context.tsx @@ -2,12 +2,16 @@ import React, { createContext, useContext, useRef, useState } from "react"; import type { DrawIoEmbedRef } from "react-drawio"; -import { extractDiagramXML } from "../lib/utils"; +import { extractDiagramXML, formatXML } from "../lib/utils"; interface DiagramContextType { chartXML: string; latestSvg: string; diagramHistory: { svg: string; xml: string }[]; + lastAgentGeneratedXml: string; + getLastAgentGeneratedXml: () => string; + setLastAgentGeneratedXml: (xml: string) => void; + markAgentDiagramPending: () => void; loadDiagram: (chart: string) => void; handleExport: () => void; resolverRef: React.Ref<((value: string) => void) | null>; @@ -24,9 +28,26 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) { const [diagramHistory, setDiagramHistory] = useState< { svg: string; xml: string }[] >([]); + const [lastAgentGeneratedXml, setLastAgentGeneratedXmlState] = useState(""); + const lastAgentGeneratedXmlRef = useRef(""); + const agentDiagramPendingRef = useRef(false); const drawioRef = useRef(null); const resolverRef = useRef<((value: string) => void) | null>(null); + // Wrapper to keep ref and state in sync + const setLastAgentGeneratedXml = (xml: string) => { + lastAgentGeneratedXmlRef.current = xml; + setLastAgentGeneratedXmlState(xml); + }; + + // Getter that returns the ref value (always up-to-date, even in async contexts) + const getLastAgentGeneratedXml = () => lastAgentGeneratedXmlRef.current; + + const markAgentDiagramPending = () => { + console.log('[DiagramContext] markAgentDiagramPending called'); + agentDiagramPendingRef.current = true; + }; + const handleExport = () => { if (drawioRef.current) { drawioRef.current.exportDiagram({ @@ -54,6 +75,16 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) { xml: extractedXML, }, ]); + + // If agent just generated a diagram, update lastAgentGeneratedXml with the exported XML + // This ensures we compare apples-to-apples (both formatted the same way) + if (agentDiagramPendingRef.current) { + const formatted = formatXML(extractedXML); + console.log('[DiagramContext] Setting lastAgentGeneratedXml from export, length:', formatted.length); + setLastAgentGeneratedXml(formatted); + agentDiagramPendingRef.current = false; + } + if (resolverRef.current) { resolverRef.current(extractedXML); resolverRef.current = null; @@ -66,6 +97,7 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) { setChartXML(emptyDiagram); setLatestSvg(""); setDiagramHistory([]); + setLastAgentGeneratedXml(""); }; return ( @@ -74,6 +106,10 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) { chartXML, latestSvg, diagramHistory, + lastAgentGeneratedXml, + getLastAgentGeneratedXml, + setLastAgentGeneratedXml, + markAgentDiagramPending, loadDiagram, handleExport, resolverRef, diff --git a/lib/utils.ts b/lib/utils.ts index a74eecf..9c7ed99 100644 --- a/lib/utils.ts +++ b/lib/utils.ts @@ -190,9 +190,15 @@ export function replaceXMLParts( let result = formatXML(xmlContent); let lastProcessedIndex = 0; + console.log('[replaceXMLParts] Input XML length:', xmlContent.length); + console.log('[replaceXMLParts] Formatted XML length:', result.length); + console.log('[replaceXMLParts] Number of edits:', searchReplacePairs.length); + for (const { search, replace } of searchReplacePairs) { // Also format the search content for consistency const formattedSearch = formatXML(search); + console.log('[replaceXMLParts] Search pattern (first 200):', search.substring(0, 200)); + console.log('[replaceXMLParts] Formatted search (first 200):', formattedSearch.substring(0, 200)); const searchLines = formattedSearch.split('\n'); // Split into lines for exact line matching @@ -276,6 +282,9 @@ export function replaceXMLParts( } if (!matchFound) { + console.log('[replaceXMLParts] SEARCH FAILED!'); + console.log('[replaceXMLParts] Current XML content:\n', result); + console.log('[replaceXMLParts] Search pattern:\n', formattedSearch); throw new Error(`Search pattern not found in the diagram. The pattern may not exist in the current structure.`); }