mirror of
https://github.com/DayuanJiang/next-ai-draw-io.git
synced 2026-10-12 04:29:51 +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
@@ -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)
|
||||
})
|
||||
|
||||
Reference in New Issue
Block a user