fix(chat): undo rejected or stopped previews, and more fixes from the third review

- A call the server runs (get_shape_library) still reaches the browser's
  tool handler, and it dropped the stored diagram of an earlier broken
  edit, whose preview then stayed. Only the tools that draw take it now.
- A display_diagram whose final XML fails the checks loads the diagram
  from before its preview again, as a failed edit does.
- After Stop, a tool result that arrives later (a screenshot check still
  running) no longer sends a new request; Stop also skips calls the tool
  handler already took.
- New chat and opening another chat kept nothing of a diagram drawn
  without messages when it could not be saved; now they stay on it.
- The settings dialog drops a model list or test result whose provider
  credentials changed meanwhile, also in another tab.
- A saved provider this version does not know crashed the whole page on
  load; it is skipped.
- The input emptied a moment after the message showed in the chat, so it
  briefly appeared twice (seen as a flaky e2e test).
- The desktop app's preset switch on the same port refetches the server
  models instead of reloading the page, which lost unsent attachments.
This commit is contained in:
dayuan.jiang
2026-10-05 13:04:12 +09:00
parent d5f31cb253
commit 49468f9478
8 changed files with 425 additions and 30 deletions
+39 -14
View File
@@ -281,6 +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())
// Set by Stop until the user sends the next message
const stoppedRef = useRef(false)
// 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
@@ -458,6 +461,10 @@ export default function ChatPanel({
// diagram preview draws on each update, so this also limits redraws
experimental_throttle: 150,
sendAutomaticallyWhen: ({ messages }) => {
// The user stopped: a tool result that arrives later (a VLM
// check still running) must not start a new request
if (stoppedRef.current) return false
const isInContinuationMode = partialXmlRef.current.length > 0
const shouldRetry = hasToolErrors(messages)
@@ -869,12 +876,12 @@ export default function ChatPanel({
// Add the combined text as the first part
parts.unshift({ type: "text", text: userText })
await sendWithCurrentDiagram(parts)
setInput("")
sessionStorage.removeItem(SESSION_STORAGE_INPUT_KEY)
setFiles([])
setUrlData(new Map())
await sendWithCurrentDiagram(parts, () => {
setInput("")
sessionStorage.removeItem(SESSION_STORAGE_INPUT_KEY)
setFiles([])
setUrlData(new Map())
})
} catch (error) {
console.error("Error fetching chart data:", error)
toast.error(dict.errors.failedToExport)
@@ -882,8 +889,13 @@ export default function ChatPanel({
}
}
// Export the current diagram, snapshot it for this message, and send
const sendWithCurrentDiagram = async (parts: any[]) => {
// Export the current diagram, snapshot it for this message, and send.
// onSent runs right after sending, so the input empties as the message
// shows in the chat
const sendWithCurrentDiagram = async (
parts: any[],
onSent?: () => void,
) => {
const chartXml = formatXML(await onFetchChart())
const previousXml = getPreviousXml(messages.length)
@@ -895,6 +907,7 @@ export default function ChatPanel({
)
sendChatMessage(parts, chartXml, previousXml, sessionId)
onSent?.()
}
// Send VLM validation feedback as a new user message through the normal send path
@@ -913,12 +926,15 @@ export default function ChatPanel({
async (sessionId: string) => {
if (!sessionManager.isAvailable) return
// Save current session before switching
if (messages.length > 0) {
// Save current session before switching (also a diagram drawn
// without messages); if that failed (storage full), stay on it
if (messages.length > 0 || isRealDiagram(chartXMLRef.current)) {
const sessionData = await buildSessionData({
withThumbnail: true,
})
await sessionManager.saveCurrentSession(sessionData)
if (!(await sessionManager.saveCurrentSession(sessionData))) {
return
}
}
// Switch to selected session
@@ -961,8 +977,12 @@ export default function ChatPanel({
)
const handleNewChat = useCallback(async () => {
// Save current session before creating new one
if (sessionManager.isAvailable && messages.length > 0) {
// Save current session before creating new one (also a diagram
// drawn without messages)
if (
sessionManager.isAvailable &&
(messages.length > 0 || isRealDiagram(chartXMLRef.current))
) {
const sessionData = await buildSessionData({ withThumbnail: true })
// Not saved (storage full): keep the chat on screen
if (!(await sessionManager.saveCurrentSession(sessionData))) return
@@ -1068,11 +1088,15 @@ export default function ChatPanel({
// Handle stop button click
const handleStop = useCallback(() => {
stoppedRef.current = true
const lastMessage = messages[messages.length - 1]
// Calls the tool handler already took can still show as streaming:
// the messages update at most every 150 ms (useChat throttle)
const toolParts = lastMessage?.parts?.filter(
(part: any) =>
part.type?.startsWith("tool-") &&
part.state === "input-streaming",
part.state === "input-streaming" &&
!processedToolCallsRef.current.has(part.toolCallId),
)
toolParts?.forEach((part: any) => {
@@ -1101,6 +1125,7 @@ export default function ChatPanel({
continuationRetryCountRef.current = 0
validationRetryCountRef.current = 0
partialXmlRef.current = ""
stoppedRef.current = false
const config = getSelectedAIConfig()
+23 -8
View File
@@ -157,9 +157,6 @@ export function ModelConfigDialog({
id: string
value: string
} | null>(null)
// Bumped on every credential edit so a running test can tell that its
// results belong to the old credentials
const credentialsVersionRef = useRef(0)
// Models fetched from the provider, per provider config
const [fetchedModels, setFetchedModels] = useState<
Record<string, ListedModel[]>
@@ -192,6 +189,20 @@ export function ModelConfigDialog({
selectedProviderIdRef.current = selectedProviderId
const configRef = useRef(config)
configRef.current = config
// A model list or test result belongs to the credentials it was asked
// with; they can change meanwhile, here or in another tab
const credentialsOf = (providerId: string) => {
const p = configRef.current.providers.find((x) => x.id === providerId)
return JSON.stringify([
p?.provider,
p?.apiKey,
p?.baseUrl,
p?.awsAccessKeyId,
p?.awsSecretAccessKey,
p?.awsRegion,
p?.vertexApiKey,
])
}
// Discard an unfinished model ID edit when the dialog closes
useEffect(() => {
@@ -217,6 +228,7 @@ export function ModelConfigDialog({
const handleFetchModels = async () => {
if (!selectedProvider) return
const providerId = selectedProvider.id
const askedWith = credentialsOf(providerId)
setFetchingModels(true)
setFetchModelsError("")
try {
@@ -237,6 +249,7 @@ export function ModelConfigDialog({
},
)
const data = await response.json().catch(() => ({}))
if (credentialsOf(providerId) !== askedWith) return
// The picker and the error belong to the provider shown
const stillShown = selectedProviderIdRef.current === providerId
if (Array.isArray(data.models)) {
@@ -253,7 +266,10 @@ export function ModelConfigDialog({
)
}
} catch {
if (selectedProviderIdRef.current === providerId) {
if (
selectedProviderIdRef.current === providerId &&
credentialsOf(providerId) === askedWith
) {
setFetchModelsError(dict.errors.networkError)
}
} finally {
@@ -313,7 +329,6 @@ export function ModelConfigDialog({
"vertexApiKey",
]
if (credentialFields.includes(field)) {
credentialsVersionRef.current++
setValidationStatus("idle")
setValidatingModelIds(new Set())
setFetchedModels(({ [selectedProviderId]: _, ...rest }) => rest)
@@ -398,7 +413,7 @@ export function ModelConfigDialog({
let allValid = true
let errorCount = 0
let idChanged = false
const credentialsVersion = credentialsVersionRef.current
const askedWith = credentialsOf(selectedProviderId)
// For EdgeOne, construct baseUrl from current origin
const baseUrl = isEdgeOne
@@ -474,7 +489,7 @@ export function ModelConfigDialog({
}
}
// Credentials changed during the test: drop the result
if (credentialsVersionRef.current !== credentialsVersion) return
if (credentialsOf(selectedProviderId) !== askedWith) return
// So did this model's id: the result is for the old one
const current = configRef.current.providers
.find((p) => p.id === selectedProviderId)
@@ -500,7 +515,7 @@ export function ModelConfigDialog({
})
}),
)
if (credentialsVersionRef.current !== credentialsVersion) return
if (credentialsOf(selectedProviderId) !== askedWith) return
// A model whose id changed was not tested
if (allValid && !idChanged) {
+18 -5
View File
@@ -121,20 +121,28 @@ export function useDiagramToolHandlers({
)
}
processedToolCallsRef.current.add(toolCall.toolCallId)
// Only these two put their result on the canvas. Other tools
// (get_shape_library, which the server runs, still arrives here)
// leave the stored originals for the preview code to undo.
const drawsDiagram =
toolCall.toolName === "display_diagram" ||
toolCall.toolName === "edit_diagram"
// 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)
if (drawsDiagram) {
for (const id of editDiagramOriginalXmlRef.current.keys()) {
processedToolCallsRef.current.add(id)
}
editDiagramOriginalXmlRef.current.clear()
}
processedToolCallsRef.current.add(toolCall.toolCallId)
editDiagramOriginalXmlRef.current.clear()
if (toolCall.toolName === "display_diagram") {
await handleDisplayDiagram(toolCall, addToolOutput)
await handleDisplayDiagram(toolCall, addToolOutput, originalXml)
} else if (toolCall.toolName === "edit_diagram") {
await handleEditDiagram(toolCall, addToolOutput, originalXml)
} else if (toolCall.toolName === "append_diagram") {
@@ -142,9 +150,11 @@ export function useDiagramToolHandlers({
}
}
// originalXml: the diagram before the streamed previews, if any were drawn
const handleDisplayDiagram = async (
toolCall: ToolCall,
addToolOutput: AddToolOutputFn,
originalXml: string | undefined,
) => {
const { xml } = toolCall.input as { xml: string }
@@ -202,6 +212,9 @@ NEXT STEP: Call append_diagram with the continuation XML.
if (validationError) {
console.warn("[display_diagram] Validation error:", validationError)
// Undo the streamed preview, as a failed edit does: the canvas
// keeps the diagram from before this failed call
if (originalXml) onDisplayChart(originalXml, true)
// Return error to model - sendAutomaticallyWhen will trigger retry
if (DEBUG) {
console.log(
+23 -3
View File
@@ -13,6 +13,7 @@ import {
flattenModels,
type ModelConfig,
type MultiModelConfig,
PROVIDER_INFO,
type ProviderConfig,
type ProviderName,
} from "@/lib/types/model-config"
@@ -73,7 +74,21 @@ function loadConfig(): MultiModelConfig {
const stored = localStorage.getItem(STORAGE_KEYS.modelConfigs)
if (stored) {
try {
return JSON.parse(stored) as MultiModelConfig
const config = JSON.parse(stored) as MultiModelConfig
// A provider this version does not know (saved by another
// version, or edited by hand) would break every list of models
const known = config.providers.filter((p) =>
Object.hasOwn(PROVIDER_INFO, p.provider),
)
if (known.length < config.providers.length) {
console.warn(
"Skipped saved providers this version does not know:",
config.providers
.filter((p) => !known.includes(p))
.map((p) => p.provider),
)
}
return { ...config, providers: known }
} catch {
console.error("Failed to parse model config")
}
@@ -163,10 +178,15 @@ export function useModelConfig(): UseModelConfigReturn {
return () => window.removeEventListener("storage", handleStorage)
}, [])
// Load server models on mount (if any)
// Load server models on mount (if any), and again when the desktop app
// restarted its server for another preset
useEffect(() => {
if (typeof window === "undefined") return
loadServerModels()
return window.electronAPI?.onServerRestarted?.(loadServerModels)
}, [])
function loadServerModels() {
fetch(getApiEndpoint("/api/server-models"))
.then((res) => {
if (!res.ok) {
@@ -214,7 +234,7 @@ export function useModelConfig(): UseModelConfigReturn {
console.error("Error while loading server models:", error)
setServerLoaded(true)
})
}, [])
}
// Save config whenever it changes (after initial load)
useEffect(() => {
+126
View File
@@ -559,3 +559,129 @@ test("stopping while drawing undoes the half drawn diagram", async ({
await expect(canvas.getByText("Beta", { exact: true })).toHaveCount(0)
await expect(canvas.getByText("Alpha", { exact: true })).toBeVisible()
})
test("a broken edit's preview is undone after a shape library call", async ({
page: p,
}) => {
// The server runs get_shape_library, but its call still reaches the
// browser's tool handler, before the UI shows the broken edit's error
const library = toolCallEvents("s1", "get_shape_library", {
library: "aws4",
})
const canvas = await chunkedReplies(p, [
[drawReply("d1", cell("a", "Alpha", 40))],
[
sse([{ type: "start" }, { type: "start-step" }]),
sse([editStart("e1"), ...editDeltas("e1")]),
220,
`${sse([
{
type: "tool-input-error",
toolCallId: "e1",
toolName: "edit_diagram",
input: "{broken",
errorText: "JSON parsing failed",
},
{ type: "finish-step" },
{ type: "start-step" },
library.start,
...library.deltas,
library.done,
{
type: "tool-output-available",
toolCallId: "s1",
output: "AWS shapes",
},
{ 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")
// The preview shows Gamma only for a moment; afterwards it must be gone
await expect(p.getByText("Get Shape Library").first()).toBeVisible({
timeout: 15000,
})
await p.waitForTimeout(2000)
await expect(canvas.getByText("Gamma", { exact: true })).toHaveCount(0)
await expect(canvas.getByText("Alpha", { exact: true })).toBeVisible()
})
test("a drawing rejected by the checks undoes its preview", async ({
page: p,
}) => {
// A linked shape that uses the root cell id "1": the finished XML is
// rejected, while the preview still draws the other shape
const redraw = toolCallEvents("d2", "display_diagram", {
xml:
cell("b", "Beta", 40) +
`<UserObject id="1" label="Bad" link="https://example.com"><mxCell vertex="1" parent="1"><mxGeometry x="220" y="40" width="120" height="60" as="geometry"/></mxCell></UserObject>`,
})
const canvas = await chunkedReplies(p, [
[drawReply("d1", cell("a", "Alpha", 40))],
[
sse([{ type: "start" }, redraw.start, ...redraw.deltas]),
`${sse([redraw.done, { type: "finish" }])}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,
})
// The tool card shows the rejection
await expect(p.locator('text="Error"').first()).toBeVisible({
timeout: 15000,
})
await p.waitForTimeout(1000)
await expect(canvas.getByText("Beta", { exact: true })).toHaveCount(0)
await expect(canvas.getByText("Alpha", { exact: true })).toBeVisible()
})
test("stopping during the screenshot check starts no new request", async ({
page: p,
}) => {
await p.addInitScript(() => {
localStorage.setItem("next-ai-draw-io-vlm-validation-enabled", "true")
})
let chatRequests = 0
await p.route("**/api/chat", async (route) => {
chatRequests++
await route.fulfill({
status: 200,
contentType: "text/event-stream",
body: drawReply(`d${chatRequests}`, cell("a", "Alpha", 40)),
})
})
// The check answers late, and finds a problem
let checking = false
await p.route("**/api/validate-diagram", async (route) => {
checking = true
await new Promise((r) => setTimeout(r, 3000))
await route.fulfill({
status: 200,
contentType: "text/plain",
body: JSON.stringify({
valid: false,
issues: [
{
type: "overlap",
severity: "critical",
description: "Boxes overlap",
},
],
suggestions: ["Move them apart"],
}),
})
})
await p.goto("/", { waitUntil: "networkidle" })
await getIframe(p).waitFor({ state: "visible", timeout: 30000 })
await sendMessage(p, "Draw a box")
await expect.poll(() => checking, { timeout: 15000 }).toBe(true)
await p.getByRole("button", { name: "Stop generation" }).click()
await p.waitForTimeout(5000)
expect(chatRequests).toBe(1)
})
+89
View File
@@ -1,3 +1,4 @@
import type { Page } from "@playwright/test"
import { SINGLE_BOX_XML } from "./fixtures/diagrams"
import {
expect,
@@ -85,6 +86,94 @@ test.describe("History and Session Restore", () => {
).toBeVisible()
})
// A diagram drawn by hand, without chat messages: loaded into draw.io
// directly, then moved with an arrow key, which draw.io reports as an
// edit like any manual change
async function drawByHand(page: Page, label: string) {
const shape = getIframeContent(page).getByText(label, { exact: true })
// draw.io may still be starting (after a new chat): send until shown
for (let i = 0; i < 10 && (await shape.count()) === 0; i++) {
await page.evaluate((label) => {
const xml = `<mxfile><diagram id="p" name="Page-1"><mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/><mxCell id="h" value="${label}" style="rounded=1;" vertex="1" parent="1"><mxGeometry x="60" y="60" width="140" height="60" as="geometry"/></mxCell></root></mxGraphModel></diagram></mxfile>`
document
.querySelector("iframe")
?.contentWindow?.postMessage(
JSON.stringify({ action: "load", xml, autosave: 1 }),
"*",
)
}, label)
await page.waitForTimeout(1000)
}
await shape.click({ timeout: 10000 })
await page.keyboard.press("ArrowRight")
// The app saves the edit about a second later
await page.waitForTimeout(2000)
}
const storageFull = (page: Page) =>
page.evaluate(() => {
IDBObjectStore.prototype.put = () => {
throw new DOMException("Storage is full", "QuotaExceededError")
}
})
test("new chat keeps a diagram without messages that could not be saved", async ({
page,
}) => {
await page.goto("/", { waitUntil: "networkidle" })
await getIframe(page).waitFor({ state: "visible", timeout: 30000 })
await drawByHand(page, "Hand drawn")
await expect(
getIframeContent(page).getByText("Hand drawn", { exact: true }),
).toBeVisible({ timeout: 10000 })
await storageFull(page)
await page.locator('[data-testid="new-chat-button"]').click()
await expect(
page.getByText(/Could not save this chat/).first(),
).toBeVisible({ timeout: 5000 })
await expect(
getIframeContent(page).getByText("Hand drawn", { exact: true }),
).toBeVisible()
})
test("opening another chat keeps a diagram 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.")
await page.waitForTimeout(1500)
await page.locator('[data-testid="new-chat-button"]').click()
// The empty chat lists the first one; draw something by hand
const firstChat = page.getByRole("button", {
name: /Create a test diagram/,
})
await expect(firstChat).toBeVisible({ timeout: 10000 })
await drawByHand(page, "Hand drawn")
await expect(
getIframeContent(page).getByText("Hand drawn", { exact: true }),
).toBeVisible({ timeout: 10000 })
await storageFull(page)
await firstChat.click()
await expect(
page.getByText(/Could not save this chat/).first(),
).toBeVisible({ timeout: 5000 })
await page.waitForTimeout(1000)
await expect(
getIframeContent(page).getByText("Hand drawn", { 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 })
+40
View File
@@ -232,3 +232,43 @@ test("a test result does not count for a model id changed meanwhile", async ({
await page.waitForTimeout(500)
await expect(dialog.locator('[title="1.0 s"]')).toHaveCount(0)
})
test("a model list fetched with an old API key is dropped", async ({
page,
}) => {
const release = await holdRoute(page, "**/api/provider-models", {
models: [{ id: "model-of-old-key", tools: true }],
})
const dialog = await openQwenSettings(page, TWO_PROVIDERS)
await dialog
.getByRole("button", { name: "Fetch models from the provider" })
.click()
// The user corrects the key while the list is loading
await dialog.locator("#api-key").fill("new-key")
release()
await page.waitForTimeout(500)
await expect(page.getByText("model-of-old-key")).toHaveCount(0)
})
test("a test result for an old API key is dropped", 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()
// Another tab saves a different key for this provider
await page.evaluate(() => {
const key = "next-ai-draw-io-model-configs"
const config = JSON.parse(localStorage.getItem(key) ?? "{}")
config.providers[0].apiKey = "key-from-another-tab"
const value = JSON.stringify(config)
localStorage.setItem(key, value)
window.dispatchEvent(
new StorageEvent("storage", { key, newValue: value }),
)
})
release()
await page.waitForTimeout(500)
await expect(dialog.locator('[title="1.0 s"]')).toHaveCount(0)
})
+67
View File
@@ -106,6 +106,26 @@ describe("useModelConfig server model selection", () => {
)
})
it("skips a saved provider this version does not know", async () => {
// Saved by another version, or edited by hand: it used to crash the
// whole page on load
storeConfig({
...USER_CONFIG,
providers: [
...USER_CONFIG.providers,
{
id: "p9",
provider: "not-a-provider" as any,
apiKey: "k",
models: [{ id: "m9", modelId: "x" }],
},
],
})
const { result } = await renderLoaded()
expect(result.current.config.providers.map((p) => p.id)).toEqual(["p1"])
expect(result.current.models.map((m) => m.id)).toContain("m1")
})
it("keeps a selected user model", async () => {
storeConfig({ ...USER_CONFIG, selectedModelId: "m1" })
const { result } = await renderLoaded()
@@ -127,6 +147,53 @@ describe("useModelConfig server model selection", () => {
})
})
describe("useModelConfig in the desktop app", () => {
it("reloads the server models after a preset switch restarts the server", async () => {
// The new preset offers other models; the saved one is gone
let restarted: (() => void) | undefined
;(window as any).electronAPI = {
onServerRestarted: (callback: () => void) => {
restarted = callback
return () => {
restarted = undefined
}
},
}
try {
storeConfig({
...USER_CONFIG,
selectedModelId: "server:openai-main:gpt-4o-mini",
})
const { result } = await renderLoaded()
await waitFor(() => expect(restarted).toBeDefined())
const nextModels: FlattenedServerModel[] = [
{
id: "server:claude:claude-sonnet-5-5",
modelId: "claude-sonnet-5-5",
provider: "anthropic",
providerLabel: "Claude",
isDefault: true,
},
]
vi.stubGlobal(
"fetch",
vi.fn(async () => ({
ok: true,
json: async () => ({ models: nextModels }),
})),
)
act(() => restarted?.())
await waitFor(() =>
expect(result.current.selectedModelId).toBe(
"server:claude:claude-sonnet-5-5",
),
)
} finally {
delete (window as any).electronAPI
}
})
})
describe("useModelConfig across tabs", () => {
it("reloads the config when another tab saves it", async () => {
storeConfig({ ...USER_CONFIG, selectedModelId: "m1" })