mirror of
https://github.com/DayuanJiang/next-ai-draw-io.git
synced 2026-10-08 02:37:46 +08:00
fix(chat): keep the canvas after an unrelated error, and more review fixes
Found by the second PR review: - After a streamed edit, an older render of the stream stored the edit's original diagram again, and the next failed request (no quota, a lost connection) put that old diagram back on the canvas. The tool handler now marks its call as handled, so the preview code leaves it alone. - An edit applied before the UI showed an earlier broken edit's error was erased when that error undid its preview, or was built on that preview. The handler now starts from the diagram before all unhandled previews, and reads the diagram state that updates at once. - A failed or stopped display_diagram left its half drawn diagram on the canvas. Its preview is undone now, like an edit's. - "New chat" cleared a chat that could not be saved (storage full). - The settings dialog showed a model list, a fetch error or a test result on the provider that was opened after the request started, and marked a model id changed during the test as tested. - A tool call with broken JSON was shown as cut off by the output limit. - History entries and session thumbnails could pair with a later diagram when draw.io answered an export late. - A server model saved before non-ASCII provider names got into the id was reset to the default model.
This commit is contained in:
dayuan.jiang
@@ -453,71 +453,81 @@ export function ChatMessageDisplay({
|
||||
if (isRestoredMessage) return
|
||||
|
||||
if (
|
||||
part.type === "tool-display_diagram" &&
|
||||
input?.xml
|
||||
part.type !== "tool-display_diagram" &&
|
||||
part.type !== "tool-edit_diagram"
|
||||
) {
|
||||
const xml = input.xml as string
|
||||
return
|
||||
}
|
||||
|
||||
// Failed or stopped: if the original XML is still
|
||||
// stored, the tool handler never ran (invalid JSON,
|
||||
// or the user pressed stop), so undo the streamed
|
||||
// preview here. Invalid JSON leaves no input, so
|
||||
// check this first.
|
||||
if (state === "output-error") {
|
||||
const originalXml =
|
||||
editDiagramOriginalXmlRef.current.get(
|
||||
toolCallId,
|
||||
)
|
||||
if (originalXml) {
|
||||
editDiagramOriginalXmlRef.current.delete(
|
||||
toolCallId,
|
||||
)
|
||||
onDisplayChart(originalXml, true)
|
||||
baseXml = originalXml
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// Input complete, or the tool handler, a stop or an
|
||||
// error took the call already: the tool handler loads
|
||||
// the checked diagram (with the original XML). The
|
||||
// messages update at most every 150 ms (useChat
|
||||
// throttle in chat-panel), so they can still show the
|
||||
// call streaming after that.
|
||||
if (
|
||||
state !== "input-streaming" ||
|
||||
processedToolCalls.current.has(toolCallId)
|
||||
) {
|
||||
processedToolCalls.current.add(toolCallId)
|
||||
lastProcessedXmlRef.current.delete(toolCallId)
|
||||
lastProcessedXmlRef.current.delete(
|
||||
`${toolCallId}-opCount`,
|
||||
)
|
||||
return
|
||||
}
|
||||
|
||||
if (part.type === "tool-display_diagram") {
|
||||
const xml = input?.xml as string | undefined
|
||||
// Skip if XML hasn't changed since last processing
|
||||
const lastXml =
|
||||
lastProcessedXmlRef.current.get(toolCallId)
|
||||
if (lastXml === xml) {
|
||||
return // Skip redundant processing
|
||||
}
|
||||
|
||||
// Messages update at most every 150 ms while
|
||||
// streaming (useChat throttle in chat-panel)
|
||||
if (state === "input-streaming") {
|
||||
handleDisplayChart(xml)
|
||||
lastProcessedXmlRef.current.set(toolCallId, xml)
|
||||
} else if (
|
||||
!processedToolCalls.current.has(toolCallId)
|
||||
if (
|
||||
!xml ||
|
||||
lastProcessedXmlRef.current.get(toolCallId) ===
|
||||
xml
|
||||
) {
|
||||
// Input complete: the tool handler loads the
|
||||
// validated diagram
|
||||
processedToolCalls.current.add(toolCallId)
|
||||
lastProcessedXmlRef.current.delete(toolCallId)
|
||||
return
|
||||
}
|
||||
// Keep the diagram from before the preview, to
|
||||
// undo it on a stop or an error
|
||||
if (
|
||||
!editDiagramOriginalXmlRef.current.has(
|
||||
toolCallId,
|
||||
)
|
||||
) {
|
||||
editDiagramOriginalXmlRef.current.set(
|
||||
toolCallId,
|
||||
baseXml || BLANK_MXFILE,
|
||||
)
|
||||
}
|
||||
handleDisplayChart(xml)
|
||||
lastProcessedXmlRef.current.set(toolCallId, xml)
|
||||
return
|
||||
}
|
||||
|
||||
// Handle edit_diagram streaming - apply operations incrementally for preview
|
||||
// Uses shared editDiagramOriginalXmlRef to coordinate with tool handler
|
||||
if (part.type === "tool-edit_diagram") {
|
||||
// Failed or stopped: if the original XML is still
|
||||
// stored, the tool handler never ran (invalid
|
||||
// JSON, or the user pressed stop), so undo the
|
||||
// streamed preview here. Invalid JSON leaves no
|
||||
// operations in the input, so check this first.
|
||||
if (state === "output-error") {
|
||||
const originalXml =
|
||||
editDiagramOriginalXmlRef.current.get(
|
||||
toolCallId,
|
||||
)
|
||||
if (originalXml) {
|
||||
editDiagramOriginalXmlRef.current.delete(
|
||||
toolCallId,
|
||||
)
|
||||
onDisplayChart(originalXml, true)
|
||||
baseXml = originalXml
|
||||
}
|
||||
return
|
||||
}
|
||||
if (!input?.operations) return
|
||||
|
||||
if (state !== "input-streaming") {
|
||||
// Input complete: the tool handler applies the
|
||||
// checked edit (it reads the original XML too)
|
||||
if (
|
||||
!processedToolCalls.current.has(toolCallId)
|
||||
) {
|
||||
lastProcessedXmlRef.current.delete(
|
||||
toolCallId + "-opCount",
|
||||
)
|
||||
processedToolCalls.current.add(toolCallId)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
const completeOps = getCompleteOperations(
|
||||
input.operations as DiagramOperation[],
|
||||
)
|
||||
|
||||
@@ -112,6 +112,7 @@ export default function ChatPanel({
|
||||
handleExportWithoutHistory,
|
||||
resolverRef,
|
||||
chartXML,
|
||||
chartXMLRef: liveChartXMLRef,
|
||||
latestSvg,
|
||||
clearDiagram,
|
||||
getThumbnailSvg,
|
||||
@@ -280,8 +281,9 @@ export default function ChatPanel({
|
||||
// Persist processed tool call IDs so collapsing the chat doesn't replay old tool outputs
|
||||
const processedToolCallsRef = useRef<Set<string>>(new Set())
|
||||
|
||||
// Store original XML for edit_diagram streaming - shared between streaming preview and tool handler
|
||||
// Key: toolCallId, Value: original XML before any operations applied
|
||||
// Store original XML for display_diagram and edit_diagram streaming -
|
||||
// shared between streaming preview and tool handler
|
||||
// Key: toolCallId, Value: XML before the call's preview was drawn
|
||||
const editDiagramOriginalXmlRef = useRef<Map<string, string>>(new Map())
|
||||
|
||||
// Debounce timeout for localStorage writes (prevents blocking during streaming)
|
||||
@@ -336,8 +338,10 @@ export default function ChatPanel({
|
||||
const { handleToolCall } = useDiagramToolHandlers({
|
||||
partialXmlRef,
|
||||
editDiagramOriginalXmlRef,
|
||||
processedToolCallsRef,
|
||||
validationRetryCountRef,
|
||||
chartXMLRef,
|
||||
// A preview undone just before the tool call is in this one already
|
||||
chartXMLRef: liveChartXMLRef,
|
||||
onDisplayChart,
|
||||
onFetchChart,
|
||||
onExport,
|
||||
@@ -363,10 +367,22 @@ export default function ChatPanel({
|
||||
onToolCall: async ({ toolCall }) => {
|
||||
await handleToolCall({ toolCall }, addToolOutput)
|
||||
},
|
||||
onFinish: ({ message, isAbort, isError }) => {
|
||||
// Stopped or failed: tool calls still streaming never reach the
|
||||
// tool handler. Mark them handled so a later render of the
|
||||
// stream does not draw their preview again.
|
||||
if (!isAbort && !isError) return
|
||||
for (const part of message.parts as any[]) {
|
||||
if (part.state === "input-streaming" && part.toolCallId) {
|
||||
processedToolCallsRef.current.add(part.toolCallId)
|
||||
}
|
||||
}
|
||||
},
|
||||
onError: (error) => {
|
||||
// An edit still streaming when the request failed never reaches
|
||||
// the tool handler: undo its preview. The first stored original
|
||||
// is the diagram before any of them.
|
||||
// A diagram still streaming when the request failed never
|
||||
// reaches the tool handler: undo its preview. Only previews not
|
||||
// handled yet are stored, and the first one holds the diagram
|
||||
// before any of them.
|
||||
const [originalXml] = editDiagramOriginalXmlRef.current.values()
|
||||
if (originalXml) onDisplayChart(originalXml, true)
|
||||
editDiagramOriginalXmlRef.current.clear()
|
||||
@@ -948,7 +964,8 @@ export default function ChatPanel({
|
||||
// Save current session before creating new one
|
||||
if (sessionManager.isAvailable && messages.length > 0) {
|
||||
const sessionData = await buildSessionData({ withThumbnail: true })
|
||||
await sessionManager.saveCurrentSession(sessionData)
|
||||
// Not saved (storage full): keep the chat on screen
|
||||
if (!(await sessionManager.saveCurrentSession(sessionData))) return
|
||||
// Refresh sessions list to ensure dropdown shows the saved session
|
||||
await sessionManager.refreshSessions()
|
||||
}
|
||||
|
||||
@@ -93,12 +93,14 @@ export function ToolCallCard({
|
||||
}
|
||||
|
||||
// Incomplete XML means the output hit the length limit, unless the user
|
||||
// stopped the generation themselves
|
||||
// stopped the generation themselves. Without an input the JSON was
|
||||
// broken (the server repairs JSON cut short by the limit).
|
||||
const isTruncated =
|
||||
state === "output-error" &&
|
||||
errorText !== "Stopped by user" &&
|
||||
(toolName === "display_diagram" || toolName === "append_diagram") &&
|
||||
!isMxCellXmlComplete(input?.xml)
|
||||
typeof input?.xml === "string" &&
|
||||
!isMxCellXmlComplete(input.xml)
|
||||
|
||||
const handleCopy = () => {
|
||||
let textToCopy = ""
|
||||
|
||||
@@ -186,6 +186,12 @@ export function ModelConfigDialog({
|
||||
const selectedProvider = config.providers.find(
|
||||
(p) => p.id === selectedProviderId,
|
||||
)
|
||||
// For requests that finish after the user switched provider or edited
|
||||
// a model id
|
||||
const selectedProviderIdRef = useRef(selectedProviderId)
|
||||
selectedProviderIdRef.current = selectedProviderId
|
||||
const configRef = useRef(config)
|
||||
configRef.current = config
|
||||
|
||||
// Discard an unfinished model ID edit when the dialog closes
|
||||
useEffect(() => {
|
||||
@@ -231,13 +237,15 @@ export function ModelConfigDialog({
|
||||
},
|
||||
)
|
||||
const data = await response.json().catch(() => ({}))
|
||||
// The picker and the error belong to the provider shown
|
||||
const stillShown = selectedProviderIdRef.current === providerId
|
||||
if (Array.isArray(data.models)) {
|
||||
setFetchedModels((current) => ({
|
||||
...current,
|
||||
[providerId]: data.models,
|
||||
}))
|
||||
setModelPickerOpen(true)
|
||||
} else {
|
||||
if (stillShown) setModelPickerOpen(true)
|
||||
} else if (stillShown) {
|
||||
const hints = dict.errors.llm as Record<string, string>
|
||||
setFetchModelsError(
|
||||
[hints[data.code], data.error].filter(Boolean).join(" ") ||
|
||||
@@ -245,7 +253,9 @@ export function ModelConfigDialog({
|
||||
)
|
||||
}
|
||||
} catch {
|
||||
setFetchModelsError(dict.errors.networkError)
|
||||
if (selectedProviderIdRef.current === providerId) {
|
||||
setFetchModelsError(dict.errors.networkError)
|
||||
}
|
||||
} finally {
|
||||
setFetchingModels(false)
|
||||
}
|
||||
@@ -387,6 +397,7 @@ export function ModelConfigDialog({
|
||||
|
||||
let allValid = true
|
||||
let errorCount = 0
|
||||
let idChanged = false
|
||||
const credentialsVersion = credentialsVersionRef.current
|
||||
|
||||
// For EdgeOne, construct baseUrl from current origin
|
||||
@@ -464,6 +475,19 @@ export function ModelConfigDialog({
|
||||
}
|
||||
// Credentials changed during the test: drop the result
|
||||
if (credentialsVersionRef.current !== credentialsVersion) return
|
||||
// So did this model's id: the result is for the old one
|
||||
const current = configRef.current.providers
|
||||
.find((p) => p.id === selectedProviderId)
|
||||
?.models.find((m) => m.id === model.id)
|
||||
if (current?.modelId !== model.modelId) {
|
||||
idChanged = true
|
||||
setValidatingModelIds((prev) => {
|
||||
const next = new Set(prev)
|
||||
next.delete(model.id)
|
||||
return next
|
||||
})
|
||||
return
|
||||
}
|
||||
if (update.validated === false) {
|
||||
allValid = false
|
||||
errorCount++
|
||||
@@ -478,9 +502,16 @@ export function ModelConfigDialog({
|
||||
)
|
||||
if (credentialsVersionRef.current !== credentialsVersion) return
|
||||
|
||||
if (allValid) {
|
||||
setValidationStatus("success")
|
||||
// A model whose id changed was not tested
|
||||
if (allValid && !idChanged) {
|
||||
updateProvider(selectedProviderId, { validated: true })
|
||||
}
|
||||
// The status line is about the provider shown now
|
||||
if (selectedProviderIdRef.current !== selectedProviderId) return
|
||||
if (idChanged) {
|
||||
setValidationStatus("idle")
|
||||
} else if (allValid) {
|
||||
setValidationStatus("success")
|
||||
// Reset to idle after showing success briefly (with cleanup)
|
||||
if (validationResetTimeoutRef.current) {
|
||||
clearTimeout(validationResetTimeoutRef.current)
|
||||
|
||||
@@ -15,6 +15,8 @@ import { extractDiagramXML, isRealDiagram } from "../lib/utils"
|
||||
|
||||
interface DiagramContextType {
|
||||
chartXML: string
|
||||
// chartXML right away, before the re-render (loadDiagram sets both)
|
||||
chartXMLRef: React.MutableRefObject<string>
|
||||
latestSvg: string
|
||||
diagramHistory: { svg: string; xml: string }[]
|
||||
setDiagramHistory: (history: { svg: string; xml: string }[]) => void
|
||||
@@ -43,10 +45,12 @@ interface DiagramContextType {
|
||||
|
||||
const DiagramContext = createContext<DiagramContextType | undefined>(undefined)
|
||||
|
||||
// Exports for thumbnails, validation PNGs and file saves carry a tag in the
|
||||
// request's `message` field. draw.io echoes the request back in the export
|
||||
// event, so each result reaches its own caller; untagged exports (chat-panel's
|
||||
// onFetchChart) resolve resolverRef.
|
||||
// Exports for thumbnails, validation PNGs, history entries and file saves
|
||||
// carry a tag in the request's `message` field. draw.io echoes the request
|
||||
// back in the export event, so each result reaches its own caller; untagged
|
||||
// exports (chat-panel's onFetchChart) resolve resolverRef. Thumbnail,
|
||||
// validation and history tags end in a request number, so a late result
|
||||
// never answers a newer request.
|
||||
type ExportTag = "thumbnail" | "validation"
|
||||
|
||||
export function DiagramProvider({ children }: { children: React.ReactNode }) {
|
||||
@@ -61,11 +65,12 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
|
||||
const drawioRef = useRef<DrawIoEmbedRef | null>(null)
|
||||
const resolverRef = useRef<((value: string) => void) | null>(null)
|
||||
// Pending thumbnail and validation PNG exports, keyed by their export tag
|
||||
const taggedResolversRef = useRef<
|
||||
Partial<Record<ExportTag, (value: string) => void>>
|
||||
>({})
|
||||
// Track if we're expecting an export for history (user-initiated)
|
||||
const expectHistoryExportRef = useRef<boolean>(false)
|
||||
const taggedResolversRef = useRef<Record<string, (value: string) => void>>(
|
||||
{},
|
||||
)
|
||||
// Pending history exports: the document each one was asked for
|
||||
const historyXmlRef = useRef(new Map<string, string>())
|
||||
const exportSeqRef = useRef(0)
|
||||
// Track latest chartXML for restoration after remount
|
||||
const chartXMLRef = useRef<string>("")
|
||||
|
||||
@@ -100,10 +105,13 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
|
||||
|
||||
const handleExport = () => {
|
||||
if (drawioRef.current) {
|
||||
// Mark that this export should be saved to history
|
||||
expectHistoryExportRef.current = true
|
||||
// Save this export to history, with the document shown now:
|
||||
// chartXML can change before the result comes back
|
||||
const tag = `history-${++exportSeqRef.current}`
|
||||
historyXmlRef.current.set(tag, chartXMLRef.current)
|
||||
drawioRef.current.exportDiagram({
|
||||
format: "xmlsvg",
|
||||
message: tag,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -126,16 +134,15 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
|
||||
timeoutMs: number,
|
||||
) =>
|
||||
new Promise<string | null>((resolve) => {
|
||||
const id = `${tag}-${++exportSeqRef.current}`
|
||||
const finish = (value: string | null) => {
|
||||
clearTimeout(timer)
|
||||
if (taggedResolversRef.current[tag] === finish) {
|
||||
delete taggedResolversRef.current[tag]
|
||||
}
|
||||
delete taggedResolversRef.current[id]
|
||||
resolve(value)
|
||||
}
|
||||
const timer = setTimeout(() => finish(null), timeoutMs)
|
||||
taggedResolversRef.current[tag] = finish
|
||||
drawioRef.current?.exportDiagram({ format, message: tag })
|
||||
taggedResolversRef.current[id] = finish
|
||||
drawioRef.current?.exportDiagram({ format, message: id })
|
||||
})
|
||||
|
||||
// Get current diagram as SVG for thumbnail (used by session storage)
|
||||
@@ -209,8 +216,8 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
|
||||
// Tagged exports (thumbnail, validation PNG, file save) go only to
|
||||
// their own caller, so they never take the result meant for resolverRef
|
||||
const tag = data.message?.message
|
||||
if (tag === "thumbnail" || tag === "validation") {
|
||||
taggedResolversRef.current[tag]?.(data.data)
|
||||
if (/^(thumbnail|validation)-/.test(tag ?? "")) {
|
||||
taggedResolversRef.current[tag as string]?.(data.data)
|
||||
return
|
||||
}
|
||||
if (tag === "save") {
|
||||
@@ -229,10 +236,13 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
|
||||
// Only add to history if this was a user-initiated export
|
||||
// Limit to 20 entries to prevent memory leaks during long sessions
|
||||
const MAX_HISTORY_SIZE = 20
|
||||
if (expectHistoryExportRef.current) {
|
||||
const askedXml =
|
||||
tag !== undefined ? historyXmlRef.current.get(tag) : undefined
|
||||
if (askedXml !== undefined) {
|
||||
historyXmlRef.current.delete(tag as string)
|
||||
// Store the full multi-page document (extractedXML is only the
|
||||
// first page), so restoring a version keeps every page
|
||||
const historyXml = chartXMLRef.current || extractedXML
|
||||
const historyXml = askedXml || extractedXML
|
||||
setDiagramHistory((prev) => {
|
||||
const newHistory = [
|
||||
...prev,
|
||||
@@ -244,7 +254,6 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
|
||||
// Keep only the last MAX_HISTORY_SIZE entries (circular buffer)
|
||||
return newHistory.slice(-MAX_HISTORY_SIZE)
|
||||
})
|
||||
expectHistoryExportRef.current = false
|
||||
}
|
||||
|
||||
if (resolverRef.current) {
|
||||
@@ -391,6 +400,7 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
|
||||
<DiagramContext.Provider
|
||||
value={{
|
||||
chartXML,
|
||||
chartXMLRef,
|
||||
latestSvg,
|
||||
diagramHistory,
|
||||
setDiagramHistory,
|
||||
|
||||
@@ -52,6 +52,8 @@ type ValidateDiagramFn = (
|
||||
interface UseDiagramToolHandlersParams {
|
||||
partialXmlRef: MutableRefObject<string>
|
||||
editDiagramOriginalXmlRef: MutableRefObject<Map<string, string>>
|
||||
// Tool calls the streaming preview must leave alone (shared with it)
|
||||
processedToolCallsRef: MutableRefObject<Set<string>>
|
||||
// Failed VLM validations in the current user turn (reset on each user message)
|
||||
validationRetryCountRef: MutableRefObject<number>
|
||||
chartXMLRef: MutableRefObject<string>
|
||||
@@ -78,6 +80,7 @@ interface UseDiagramToolHandlersParams {
|
||||
export function useDiagramToolHandlers({
|
||||
partialXmlRef,
|
||||
editDiagramOriginalXmlRef,
|
||||
processedToolCallsRef,
|
||||
validationRetryCountRef,
|
||||
chartXMLRef,
|
||||
onDisplayChart,
|
||||
@@ -118,10 +121,22 @@ export function useDiagramToolHandlers({
|
||||
)
|
||||
}
|
||||
|
||||
// Stored originals belong to previews not handled yet: this call's,
|
||||
// and those of earlier calls with invalid input, which never get
|
||||
// here. The first is the diagram before all of them. This call's
|
||||
// result replaces those previews, so the preview code must neither
|
||||
// draw them again nor undo them later.
|
||||
const [originalXml] = editDiagramOriginalXmlRef.current.values()
|
||||
for (const id of editDiagramOriginalXmlRef.current.keys()) {
|
||||
processedToolCallsRef.current.add(id)
|
||||
}
|
||||
processedToolCallsRef.current.add(toolCall.toolCallId)
|
||||
editDiagramOriginalXmlRef.current.clear()
|
||||
|
||||
if (toolCall.toolName === "display_diagram") {
|
||||
await handleDisplayDiagram(toolCall, addToolOutput)
|
||||
} else if (toolCall.toolName === "edit_diagram") {
|
||||
await handleEditDiagram(toolCall, addToolOutput)
|
||||
await handleEditDiagram(toolCall, addToolOutput, originalXml)
|
||||
} else if (toolCall.toolName === "append_diagram") {
|
||||
handleAppendDiagram(toolCall, addToolOutput)
|
||||
}
|
||||
@@ -369,20 +384,18 @@ ${finalXml}
|
||||
}
|
||||
}
|
||||
|
||||
// originalXml: the diagram before the streamed previews, if any were drawn.
|
||||
// Operations apply to it, the same base XML that streaming used.
|
||||
const handleEditDiagram = async (
|
||||
toolCall: ToolCall,
|
||||
addToolOutput: AddToolOutputFn,
|
||||
originalXml: string | undefined,
|
||||
) => {
|
||||
const { operations } = toolCall.input as {
|
||||
operations: DiagramOperation[]
|
||||
}
|
||||
|
||||
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 = () => {
|
||||
@@ -423,8 +436,6 @@ ${currentXml}
|
||||
|
||||
Please check the cell IDs and retry.`,
|
||||
})
|
||||
// Clean up the shared original XML ref
|
||||
editDiagramOriginalXmlRef.current.delete(toolCall.toolCallId)
|
||||
return
|
||||
}
|
||||
|
||||
@@ -435,8 +446,6 @@ Please check the cell IDs and retry.`,
|
||||
toolCallId: toolCall.toolCallId,
|
||||
output: `Successfully applied ${outcome.applied} operation(s) to the diagram.`,
|
||||
})
|
||||
// Clean up the shared original XML ref
|
||||
editDiagramOriginalXmlRef.current.delete(toolCall.toolCallId)
|
||||
} catch (error) {
|
||||
console.error("[edit_diagram] Failed:", error)
|
||||
|
||||
@@ -457,8 +466,6 @@ ${currentXml || "No XML available"}
|
||||
|
||||
Please check cell IDs and retry, or use display_diagram to regenerate.`,
|
||||
})
|
||||
// Clean up the shared original XML ref even on error
|
||||
editDiagramOriginalXmlRef.current.delete(toolCall.toolCallId)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -192,7 +192,19 @@ export function useModelConfig(): UseModelConfigReturn {
|
||||
id?.startsWith("server:") &&
|
||||
!raw.some((m) => m.id === id)
|
||||
if (id && !isStale) return prev
|
||||
const fallback = defaultServerModelId(raw)
|
||||
// Saved before non-ASCII characters in provider names
|
||||
// got into the id: they were dropped from it
|
||||
const renamed = raw.filter(
|
||||
(m) =>
|
||||
`server:${m.providerLabel
|
||||
.toLowerCase()
|
||||
.replace(/[^a-z0-9]+/g, "-")
|
||||
.replace(/^-|-$/g, "")}:${m.modelId}` === id,
|
||||
)
|
||||
const fallback =
|
||||
renamed.length === 1
|
||||
? renamed[0].id
|
||||
: defaultServerModelId(raw)
|
||||
return fallback === id
|
||||
? prev
|
||||
: { ...prev, selectedModelId: fallback }
|
||||
|
||||
@@ -38,10 +38,11 @@ export interface UseSessionManagerReturn {
|
||||
switchSession: (id: string) => Promise<SessionData | null>
|
||||
deleteSession: (id: string) => Promise<{ wasCurrentSession: boolean }>
|
||||
// forSessionId: optional session ID to verify save targets correct session (prevents stale debounce writes)
|
||||
// Resolves to false when the save failed (the user was told)
|
||||
saveCurrentSession: (
|
||||
data: SessionData,
|
||||
forSessionId?: string | null,
|
||||
) => Promise<void>
|
||||
) => Promise<boolean>
|
||||
refreshSessions: () => Promise<void>
|
||||
clearCurrentSession: () => void
|
||||
}
|
||||
@@ -247,17 +248,17 @@ export function useSessionManager(
|
||||
async (
|
||||
data: SessionData,
|
||||
forSessionId?: string | null,
|
||||
): Promise<void> => {
|
||||
): Promise<boolean> => {
|
||||
// If forSessionId is provided, verify it matches current session
|
||||
// This prevents stale debounced saves from overwriting a newly switched session
|
||||
if (
|
||||
forSessionId !== undefined &&
|
||||
forSessionId !== currentSessionId
|
||||
) {
|
||||
return
|
||||
return true
|
||||
}
|
||||
// Nothing can be stored without IndexedDB
|
||||
if (!isIndexedDBAvailable()) return
|
||||
if (!isIndexedDBAvailable()) return true
|
||||
|
||||
if (!currentSession) {
|
||||
// Create a new session if none exists
|
||||
@@ -274,13 +275,13 @@ export function useSessionManager(
|
||||
// up in the URL and point to nothing after a reload)
|
||||
if (!(await saveSession(newSession))) {
|
||||
notifySaveFailed(dict.errors.sessionSaveFailed)
|
||||
return
|
||||
return false
|
||||
}
|
||||
await enforceSessionLimit()
|
||||
setCurrentSession(newSession)
|
||||
setCurrentSessionId(newSession.id)
|
||||
await refreshSessions()
|
||||
return
|
||||
return true
|
||||
}
|
||||
|
||||
// Update existing session
|
||||
@@ -304,7 +305,7 @@ export function useSessionManager(
|
||||
|
||||
if (!(await saveSession(updatedSession))) {
|
||||
notifySaveFailed(dict.errors.sessionSaveFailed)
|
||||
return
|
||||
return false
|
||||
}
|
||||
setCurrentSession(updatedSession)
|
||||
|
||||
@@ -325,6 +326,7 @@ export function useSessionManager(
|
||||
: s,
|
||||
),
|
||||
)
|
||||
return true
|
||||
},
|
||||
[currentSession, currentSessionId, refreshSessions, dict],
|
||||
)
|
||||
|
||||
@@ -67,24 +67,32 @@ const editDeltas = (id: string) =>
|
||||
/**
|
||||
* Answer each chat request with the next reply. Each string in a reply is
|
||||
* one network chunk, sent 300 ms apart, so the throttled UI renders between
|
||||
* chunks like with a real model.
|
||||
* chunks like with a real model. A number in a reply sets the wait before
|
||||
* the next chunk instead.
|
||||
*/
|
||||
async function chunkedReplies(p: Page, replies: string[][]) {
|
||||
async function chunkedReplies(p: Page, replies: (string | number)[][]) {
|
||||
await p.addInitScript((replies) => {
|
||||
const realFetch = window.fetch
|
||||
let n = 0
|
||||
window.fetch = async (input, init) => {
|
||||
const url =
|
||||
typeof input === "string" ? input : (input as Request).url
|
||||
// Next.js passes URL objects for its own requests
|
||||
const url = input instanceof Request ? input.url : String(input)
|
||||
if (!url.endsWith("/api/chat")) return realFetch(input, init)
|
||||
const chunks = replies[n++] ?? [
|
||||
'data: {"type":"start"}\n\ndata: {"type":"finish"}\n\ndata: [DONE]\n\n',
|
||||
]
|
||||
const body = new ReadableStream({
|
||||
async start(controller) {
|
||||
for (const chunk of chunks) {
|
||||
for (const [i, chunk] of chunks.entries()) {
|
||||
if (typeof chunk === "number") continue
|
||||
controller.enqueue(new TextEncoder().encode(chunk))
|
||||
await new Promise((r) => setTimeout(r, 300))
|
||||
const next = chunks[i + 1]
|
||||
await new Promise((r) =>
|
||||
setTimeout(
|
||||
r,
|
||||
typeof next === "number" ? next : 300,
|
||||
),
|
||||
)
|
||||
}
|
||||
controller.close()
|
||||
},
|
||||
@@ -381,3 +389,173 @@ test("a request that fails during an edit undoes its preview", async ({
|
||||
await expect(canvas.getByText("Gamma", { exact: true })).toHaveCount(0)
|
||||
await expect(canvas.getByText("Alpha", { exact: true })).toBeVisible()
|
||||
})
|
||||
|
||||
/** A streamed tool call split into its start, input deltas and finished input */
|
||||
function toolCallEvents(id: string, toolName: string, input: unknown) {
|
||||
const deltas = (JSON.stringify(input).match(/[\s\S]{1,40}/g) ?? []).map(
|
||||
(d) => ({
|
||||
type: "tool-input-delta",
|
||||
toolCallId: id,
|
||||
inputTextDelta: d,
|
||||
}),
|
||||
)
|
||||
return {
|
||||
start: { type: "tool-input-start", toolCallId: id, toolName },
|
||||
deltas,
|
||||
done: { type: "tool-input-available", toolCallId: id, toolName, input },
|
||||
}
|
||||
}
|
||||
const drawReply = (id: string, xml: string) => {
|
||||
const call = toolCallEvents(id, "display_diagram", { xml })
|
||||
return `${sse([{ type: "start" }, call.start, ...call.deltas, call.done, { type: "finish" }])}data: [DONE]\n\n`
|
||||
}
|
||||
// SSE comments keep a stream open without sending anything
|
||||
const KEEP_OPEN = Array(20).fill(":\n\n")
|
||||
|
||||
test("an error after a finished edit keeps the current diagram", async ({
|
||||
page: p,
|
||||
}) => {
|
||||
const edit = toolCallEvents("e1", "edit_diagram", EDIT_GAMMA)
|
||||
const half = Math.ceil(edit.deltas.length / 2)
|
||||
const canvas = await chunkedReplies(p, [
|
||||
[drawReply("d1", cell("a", "Alpha", 40))],
|
||||
[
|
||||
sse([
|
||||
{ type: "start" },
|
||||
{ type: "start-step" },
|
||||
edit.start,
|
||||
...edit.deltas.slice(0, half),
|
||||
]),
|
||||
// The last input and the finished call arrive together, so the
|
||||
// tool handler runs before the UI shows the call as finished
|
||||
`${sse([...edit.deltas.slice(half), edit.done, { type: "finish-step" }, { type: "finish" }])}data: [DONE]\n\n`,
|
||||
],
|
||||
[drawReply("d2", cell("b", "Beta", 40))],
|
||||
[
|
||||
`${sse([{ type: "start" }, { type: "error", errorText: "Upstream down" }])}data: [DONE]\n\n`,
|
||||
],
|
||||
])
|
||||
await sendMessage(p, "Draw a box")
|
||||
await waitForCompleteCount(p, 1)
|
||||
await sendMessage(p, "Add a box")
|
||||
await waitForCompleteCount(p, 2)
|
||||
await expect(canvas.getByText("Gamma", { exact: true })).toBeVisible({
|
||||
timeout: 15000,
|
||||
})
|
||||
await sendMessage(p, "Start over")
|
||||
await waitForCompleteCount(p, 3)
|
||||
await expect(canvas.getByText("Beta", { exact: true })).toBeVisible({
|
||||
timeout: 15000,
|
||||
})
|
||||
await sendMessage(p, "Once more")
|
||||
await expect(p.getByText("Upstream down").first()).toBeVisible({
|
||||
timeout: 15000,
|
||||
})
|
||||
await p.waitForTimeout(1000)
|
||||
await expect(canvas.getByText("Beta", { exact: true })).toBeVisible()
|
||||
await expect(canvas.getByText("Alpha", { exact: true })).toHaveCount(0)
|
||||
})
|
||||
|
||||
// The broken call's error and the whole next edit arrive together. 220 ms
|
||||
// after the preview: the UI (throttled to 150 ms) shows the error only after
|
||||
// the next edit was applied. 600 ms: the UI undoes the preview first.
|
||||
for (const gap of [220, 600]) {
|
||||
test(`an edit that arrives with a broken edit's error keeps its change (${gap} ms)`, async ({
|
||||
page: p,
|
||||
}) => {
|
||||
const delta = toolCallEvents("e2", "edit_diagram", {
|
||||
operations: [
|
||||
{
|
||||
operation: "add",
|
||||
cell_id: "d",
|
||||
new_xml: cell("d", "Delta", 220),
|
||||
},
|
||||
],
|
||||
})
|
||||
const canvas = await chunkedReplies(p, [
|
||||
[drawReply("d1", cell("a", "Alpha", 40))],
|
||||
[
|
||||
sse([{ type: "start" }, { type: "start-step" }]),
|
||||
sse([editStart("e1"), ...editDeltas("e1")]),
|
||||
gap,
|
||||
`${sse([
|
||||
{
|
||||
type: "tool-input-error",
|
||||
toolCallId: "e1",
|
||||
toolName: "edit_diagram",
|
||||
input: "{broken",
|
||||
errorText: "JSON parsing failed",
|
||||
},
|
||||
{ type: "finish-step" },
|
||||
{ type: "start-step" },
|
||||
delta.start,
|
||||
...delta.deltas,
|
||||
delta.done,
|
||||
{ type: "finish-step" },
|
||||
{ type: "finish" },
|
||||
])}data: [DONE]\n\n`,
|
||||
],
|
||||
])
|
||||
await sendMessage(p, "Draw a box")
|
||||
await waitForCompleteCount(p, 1)
|
||||
await sendMessage(p, "Add another box")
|
||||
await expect(canvas.getByText("Delta", { exact: true })).toBeVisible({
|
||||
timeout: 15000,
|
||||
})
|
||||
await p.waitForTimeout(1000)
|
||||
await expect(canvas.getByText("Delta", { exact: true })).toBeVisible()
|
||||
await expect(canvas.getByText("Gamma", { exact: true })).toHaveCount(0)
|
||||
await expect(canvas.getByText("Alpha", { exact: true })).toBeVisible()
|
||||
await expect(p.getByText(/No changes were made/)).toHaveCount(0)
|
||||
})
|
||||
}
|
||||
|
||||
test("a request that fails while drawing undoes the half drawn diagram", async ({
|
||||
page: p,
|
||||
}) => {
|
||||
const redraw = toolCallEvents("d2", "display_diagram", {
|
||||
xml: cell("b", "Beta", 40) + cell("c", "Gamma", 220),
|
||||
})
|
||||
const canvas = await chunkedReplies(p, [
|
||||
[drawReply("d1", cell("a", "Alpha", 40))],
|
||||
[
|
||||
sse([{ type: "start" }, redraw.start, ...redraw.deltas]),
|
||||
`${sse([{ type: "error", errorText: "Upstream connection lost" }])}data: [DONE]\n\n`,
|
||||
],
|
||||
])
|
||||
await sendMessage(p, "Draw a box")
|
||||
await waitForCompleteCount(p, 1)
|
||||
await sendMessage(p, "Draw it again")
|
||||
await expect(canvas.getByText("Beta", { exact: true })).toBeVisible({
|
||||
timeout: 15000,
|
||||
})
|
||||
await expect(p.getByText("Upstream connection lost").first()).toBeVisible({
|
||||
timeout: 15000,
|
||||
})
|
||||
await expect(canvas.getByText("Beta", { exact: true })).toHaveCount(0)
|
||||
await expect(canvas.getByText("Alpha", { exact: true })).toBeVisible()
|
||||
})
|
||||
|
||||
test("stopping while drawing undoes the half drawn diagram", async ({
|
||||
page: p,
|
||||
}) => {
|
||||
const redraw = toolCallEvents("d2", "display_diagram", {
|
||||
xml: cell("b", "Beta", 40) + cell("c", "Gamma", 220),
|
||||
})
|
||||
const canvas = await chunkedReplies(p, [
|
||||
[drawReply("d1", cell("a", "Alpha", 40))],
|
||||
[
|
||||
sse([{ type: "start" }, redraw.start, ...redraw.deltas]),
|
||||
...KEEP_OPEN,
|
||||
],
|
||||
])
|
||||
await sendMessage(p, "Draw a box")
|
||||
await waitForCompleteCount(p, 1)
|
||||
await sendMessage(p, "Draw it again")
|
||||
await expect(canvas.getByText("Beta", { exact: true })).toBeVisible({
|
||||
timeout: 15000,
|
||||
})
|
||||
await p.getByRole("button", { name: "Stop generation" }).click()
|
||||
await expect(canvas.getByText("Beta", { exact: true })).toHaveCount(0)
|
||||
await expect(canvas.getByText("Alpha", { exact: true })).toBeVisible()
|
||||
})
|
||||
|
||||
@@ -133,4 +133,45 @@ test.describe("Error Handling", () => {
|
||||
timeout: 15000,
|
||||
})
|
||||
})
|
||||
|
||||
test("shows the error of a tool call with broken input", async ({
|
||||
page,
|
||||
}) => {
|
||||
// Invalid JSON that the server could not repair: not a length limit
|
||||
const toolCallId = `call_${Date.now()}`
|
||||
const events = [
|
||||
{ type: "start", messageId: `msg_${Date.now()}` },
|
||||
{
|
||||
type: "tool-input-start",
|
||||
toolCallId,
|
||||
toolName: "display_diagram",
|
||||
},
|
||||
{
|
||||
type: "tool-input-error",
|
||||
toolCallId,
|
||||
toolName: "display_diagram",
|
||||
input: '{"xml": "<mxCell value="a"/>"}',
|
||||
errorText: "Invalid input for tool display_diagram",
|
||||
},
|
||||
{ type: "finish" },
|
||||
]
|
||||
await page.route("**/api/chat", async (route) => {
|
||||
await route.fulfill({
|
||||
status: 200,
|
||||
contentType: "text/event-stream",
|
||||
body:
|
||||
events
|
||||
.map((e) => `data: ${JSON.stringify(e)}\n\n`)
|
||||
.join("") + "data: [DONE]\n\n",
|
||||
})
|
||||
})
|
||||
await page.goto("/", { waitUntil: "networkidle" })
|
||||
await getIframe(page).waitFor({ state: "visible", timeout: 30000 })
|
||||
await sendMessage(page, "Draw something")
|
||||
|
||||
await expect(
|
||||
page.getByText("Invalid input for tool display_diagram").first(),
|
||||
).toBeVisible({ timeout: 15000 })
|
||||
await expect(page.locator('text="Truncated"')).toHaveCount(0)
|
||||
})
|
||||
})
|
||||
|
||||
@@ -48,6 +48,43 @@ test.describe("History and Session Restore", () => {
|
||||
})
|
||||
})
|
||||
|
||||
test("new chat keeps a conversation that could not be saved", async ({
|
||||
page,
|
||||
}) => {
|
||||
await page.route("**/api/chat", async (route) => {
|
||||
await route.fulfill({
|
||||
status: 200,
|
||||
contentType: "text/event-stream",
|
||||
body: createMockSSEResponse(
|
||||
SINGLE_BOX_XML,
|
||||
"Created your test diagram.",
|
||||
),
|
||||
})
|
||||
})
|
||||
await page.goto("/", { waitUntil: "networkidle" })
|
||||
await getIframe(page).waitFor({ state: "visible", timeout: 30000 })
|
||||
await sendMessage(page, "Create a test diagram")
|
||||
await waitForText(page, "Created your test diagram.")
|
||||
|
||||
// Browser storage is full from now on
|
||||
await page.evaluate(() => {
|
||||
IDBObjectStore.prototype.put = () => {
|
||||
throw new DOMException("Storage is full", "QuotaExceededError")
|
||||
}
|
||||
})
|
||||
await page.locator('[data-testid="new-chat-button"]').click()
|
||||
|
||||
await expect(
|
||||
page.getByText(/Could not save this chat/).first(),
|
||||
).toBeVisible({ timeout: 5000 })
|
||||
await page.waitForTimeout(1000)
|
||||
// Still the conversation and its diagram, not the empty chat's examples
|
||||
await expect(page.getByText("Paper to Diagram")).toHaveCount(0)
|
||||
await expect(
|
||||
getIframeContent(page).getByText("Test Box", { exact: true }),
|
||||
).toBeVisible()
|
||||
})
|
||||
|
||||
test("chat history sidebar shows past conversations", async ({ page }) => {
|
||||
await page.goto("/", { waitUntil: "networkidle" })
|
||||
await getIframe(page).waitFor({ state: "visible", timeout: 30000 })
|
||||
|
||||
@@ -157,3 +157,78 @@ test("editing a model id clears the old test warning", async ({ page }) => {
|
||||
await expect(dialog.getByText(warning)).toHaveCount(0)
|
||||
await expect(dialog.getByText("may not be able to draw")).toBeVisible()
|
||||
})
|
||||
|
||||
/** Hold requests to an endpoint until release() answers them with json */
|
||||
async function holdRoute(page: Page, url: string, json: object) {
|
||||
let release!: () => void
|
||||
const released = new Promise<void>((r) => {
|
||||
release = r
|
||||
})
|
||||
await page.route(url, async (route) => {
|
||||
await released
|
||||
await route.fulfill({ status: 200, json })
|
||||
})
|
||||
return release
|
||||
}
|
||||
|
||||
const TWO_PROVIDERS = {
|
||||
version: 1,
|
||||
providers: [
|
||||
{
|
||||
...CONFIG.providers[0],
|
||||
models: [{ id: "m1", modelId: "qwen-max" }],
|
||||
},
|
||||
{ id: "p2", provider: "glm", apiKey: "k", models: [] },
|
||||
],
|
||||
}
|
||||
|
||||
test("a model list that arrives after switching provider stays with its provider", async ({
|
||||
page,
|
||||
}) => {
|
||||
const release = await holdRoute(page, "**/api/provider-models", {
|
||||
code: "invalid_api_key",
|
||||
error: "Incorrect API key",
|
||||
})
|
||||
const dialog = await openQwenSettings(page, TWO_PROVIDERS)
|
||||
await dialog
|
||||
.getByRole("button", { name: "Fetch models from the provider" })
|
||||
.click()
|
||||
await dialog.getByText("GLM (Zhipu)").first().click()
|
||||
release()
|
||||
await page.waitForTimeout(500)
|
||||
await expect(dialog.getByText("Incorrect API key")).toHaveCount(0)
|
||||
})
|
||||
|
||||
test("a test result that arrives after switching provider stays with its provider", async ({
|
||||
page,
|
||||
}) => {
|
||||
const release = await holdRoute(page, "**/api/validate-model", {
|
||||
valid: false,
|
||||
error: "Model not found",
|
||||
})
|
||||
const dialog = await openQwenSettings(page, TWO_PROVIDERS)
|
||||
await dialog.getByRole("button", { name: "Test", exact: true }).click()
|
||||
await dialog.getByText("GLM (Zhipu)").first().click()
|
||||
release()
|
||||
await page.waitForTimeout(500)
|
||||
await expect(dialog.getByText(/model\(s\) failed validation/)).toHaveCount(
|
||||
0,
|
||||
)
|
||||
})
|
||||
|
||||
test("a test result does not count for a model id changed meanwhile", async ({
|
||||
page,
|
||||
}) => {
|
||||
const release = await holdRoute(page, "**/api/validate-model", {
|
||||
valid: true,
|
||||
responseTime: 1000,
|
||||
})
|
||||
const dialog = await openQwenSettings(page, TWO_PROVIDERS)
|
||||
await dialog.getByRole("button", { name: "Test", exact: true }).click()
|
||||
const input = dialog.locator('input[title="qwen-max"]')
|
||||
await input.fill("qwen-plus")
|
||||
await input.blur()
|
||||
release()
|
||||
await page.waitForTimeout(500)
|
||||
await expect(dialog.locator('[title="1.0 s"]')).toHaveCount(0)
|
||||
})
|
||||
|
||||
@@ -70,6 +70,31 @@ describe("useModelConfig server model selection", () => {
|
||||
expect(result.current.selectedModelId).toBe("server:openai-main:gpt-4o")
|
||||
})
|
||||
|
||||
it("keeps a server model saved under its old id", async () => {
|
||||
// Before non-ASCII provider names got their own slug, "主力 OpenAI"
|
||||
// became "openai"
|
||||
const renamed: FlattenedServerModel = {
|
||||
id: "server:4e3b-529b-openai:gpt-4o-mini",
|
||||
modelId: "gpt-4o-mini",
|
||||
provider: "openai",
|
||||
providerLabel: "主力 OpenAI",
|
||||
isDefault: false,
|
||||
}
|
||||
vi.stubGlobal(
|
||||
"fetch",
|
||||
vi.fn(async () => ({
|
||||
ok: true,
|
||||
json: async () => ({ models: [...SERVER_MODELS, renamed] }),
|
||||
})),
|
||||
)
|
||||
storeConfig({
|
||||
...USER_CONFIG,
|
||||
selectedModelId: "server:openai:gpt-4o-mini",
|
||||
})
|
||||
const { result } = await renderLoaded()
|
||||
expect(result.current.selectedModelId).toBe(renamed.id)
|
||||
})
|
||||
|
||||
it("keeps a saved server model that still exists", async () => {
|
||||
storeConfig({
|
||||
...USER_CONFIG,
|
||||
|
||||
Reference in New Issue
Block a user